=== PAGE: https://www.union.ai/docs/v2/flyte ===

# Documentation

Welcome to the documentation.

## Subpages

- **Flyte OSS**
- **Tutorials**
- **Integrations**
- **Reference**
- **Community**
- **Platform deployment**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide ===

# Flyte OSS

Flyte is a free and open source platform that provides a full suite of features for orchestrating AI workflows.
Flyte enables AI development teams to rapidly ship high-quality code to production by offering optimized performance, unparalleled resource efficiency, and a delightful workflow authoring experience.
You deploy and manage Flyte yourself, on your own cloud infrastructure.

> [!NOTE]
> These are the Flyte **2.0** docs.
> To switch to [version 1.0](https://www.union.ai/docs/v1/flyte) or to the commercial product, [**Union.ai**](https://www.union.ai/docs/v2/union), use the selectors above.

## Basics

Learn the basics of Flyte, covering all the core concepts around tasks, apps, and agents.

### **Get started**

What Flyte 2 is, how to install it, the core concepts, and the ways to run your code.

### **Tasks**

Configure, build, and deploy the durable batch workloads that everything else is made of.

### **Apps**

Long-running services for dashboards, REST APIs, and model endpoints.

### **Agents**

Durable, self-healing agents built from tasks and apps, with sandboxing and MCP.

## Advanced guides

Organize your codebase, optimize performance for production, and migrate from
other workflow orchestrators.

### **Project patterns**

Patterns for BYO images, monorepos with uv, CI/CD, and multi-team resource management.

### **Scale your runs**

Tune task overhead, batching, reusable containers, and fanout to scale your workflows.

### **Advanced project: LLM reporting agent**

An advanced guide for building an LLM reporting agent on Flyte.

### **Migration**

Port a Flyte 1 codebase to Flyte 2, or map Airflow concepts to their Flyte 2 equivalents.

## Subpages
- **Get started**
- **Tasks**
- **Apps**
- **Agents**
- **Project patterns**
- **Scale your runs**
- **Advanced project: LLM reporting agent**
- **Migration**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started ===

# Get started

This guide covers how to build AI applications, data pipelines, and ML workflows using the Flyte 2 SDK.

Programs written using the Flyte 2 SDK can run on either a Union.ai or Flyte OSS back-end. This guide applies to both.

## Pure Python, no DSL

Flyte lets you write workflows in standard Python: no domain-specific language, no special syntax, no restrictions.

Your "workflow" is simply a task that calls other tasks:

```python
@env.task()
async def my_workflow(data: list[str]) -> list[str]:
    results = []
    for item in data:
        if should_process(item):
            result = await process_item(item)
            results.append(result)
    return results
```

You can use everything Python offers:

- **Loops and conditionals**: standard `for`, `while`, `if-elif-else`
- **Error handling**: `try/except` blocks work as expected
- **Async/await**: native Python concurrency model
- **Any library**: import and use whatever you need

This means no learning curve beyond Python itself, and no fighting a DSL when your requirements don't fit its constraints.

## Durability

Every task execution in Flyte is automatically persisted. Inputs, outputs, and intermediate results are stored in an object store, giving you:

- **Full observability**: see exactly what data flowed through each step
- **Audit trail**: track what ran, when, and with what parameters
- **Data lineage**: trace outputs back to their inputs

This persistence happens automatically. You don't need to add logging or manually save state. Flyte handles it.

## Reproducibility

Flyte ensures that runs can be reproduced exactly:

- **Deterministic execution**: same inputs produce same outputs
- **Caching**: task results are cached and reused when inputs match
- **Versioned containers**: code runs in the same environment every time

Caching is configurable per task:

```python
@env.task(cache="auto")
async def expensive_computation(data: str) -> str:
    # This result will be cached and reused for identical inputs
    ...
```

When you rerun a workflow, Flyte serves cached results for unchanged tasks rather than recomputing them.

## Recoverability

When something fails, Flyte doesn't make you start over. Failed workflows can resume from where they left off:

- **Completed tasks are preserved**: successful outputs remain cached
- **Retry from failure point**: no need to re-execute what already succeeded
- **Fine-grained checkpoints**: the `@flyte.trace` decorator creates checkpoints within tasks

This reduces wasted compute and speeds up debugging. When a task fails after hours of prior computation, you fix the issue and continue, not restart.

## Built for scale

Flyte handles the hard parts of distributed execution:

- **Parallel execution**: express parallelism with `asyncio.gather()`, Flyte handles the rest
- **Dynamic workflows**: construct workflows based on runtime data, not just static definitions
- **Fast scheduling**: reusable containers achieve millisecond-level task startup
- **Resource management**: specify CPU, memory, and GPU requirements per task

## What this means in practice

Consider a data pipeline that processes thousands of files, trains a model, and deploys it:

- If file processing fails on item 847, you fix the issue and resume from item 847
- If training succeeds, but deployment fails, you redeploy without retraining
- If you rerun next week with the same data, cached results skip redundant computation
- If you need to audit what happened, every step is recorded

Flyte gives you the flexibility of Python scripts with the reliability of a production system.

## Start here

### **Get started > Quickstart**

Install the SDK and run your first workflow locally in a few minutes.

### **Get started > Core concepts**

The building blocks of every Flyte program: TaskEnvironments, tasks, runs, actions, and apps.

### **Get started > Run modes**

Run the same task code locally, on a devbox, or on a remote cluster.

## Subpages
- **Get started > Quickstart**
- **Get started > Core concepts**
- **Get started > Run modes**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/quickstart ===

# Quickstart

Let's get you up and running with your first workflow on your local machine.

<!--
Hidden until the in-browser demo is repaired. The linked demo
(https://flyte2intro.apps.demo.hosted.unionai.cloud/) is currently unreachable.
To restore: re-add a note callout reading "Want to try Flyte without installing
anything? Try Flyte 2 in your browser" linking to the demo URL once it is live.
-->

## What you'll need

- Python 3.10+ in a virtual environment

## Install the SDK

Install the `flyte` package:

```bash
pip install 'flyte[tui]'
```

> **📝 Note**
>
> We also install the `tui` extra to enable the terminal user interface.

Verify it worked:

```bash
flyte --version
```

Output:

```bash
Flyte SDK version: 2.*.*
```

> **📝 Note**
>
> If you have [`uv`](https://docs.astral.sh/uv/) installed, you can run the `flyte` CLI directly with `uvx`, without installing the package into your environment:
>
> ```bash
> uvx flyte --version
> uvx flyte get run
> ```

## Configure

Create a config file for local execution. Runs will be persisted locally in a SQLite database.

```bash
flyte create config --local-persistence
```

This creates `.flyte/config.yaml` in your current directory.

See [Setting up a configuration file](./run-modes/running-devbox#configure) for more options when connecting to a cluster.

> **📝 Note**
>
> Run `flyte get config` to check which configuration is currently active.

## Write your first workflow

> [!TIP] Author workflows with an AI assistant
> [`flyte-agent-plugins`](https://github.com/flyteorg/flyte-agent-plugins) — a
> portable agent harness plugin for Claude Code, Codex, OpenCode, and other
> harnesses — adds skills that scaffold projects and generate tasks, workflows,
> apps, and tests for you, plus MCP servers that ground the agent in the Flyte SDK
> and docs. See [Flyte agent plugins](../../api-reference/agent-plugins) to get started.

Create `hello.py`:

```python
# hello.py

import flyte

# The `hello_env` TaskEnvironment is assigned to the variable `env`.
# It is then used in the `@env.task` decorator to define tasks.
# The environment groups configuration for all tasks defined within it.
env = flyte.TaskEnvironment(name="hello_env")

# We use the `@env.task` decorator to define a task called `fn`.
@env.task
def fn(x: int) -> int: # Type annotations are required
    slope, intercept = 2, 5
    return slope * x + intercept

# We also use the `@env.task` decorator to define another task called `main`.
# This is the entrypoint task of the workflow.
# It calls the `fn` task defined above multiple times using `flyte.map`.
@env.task
def main(x_list: list[int] = list(range(10))) -> float:
    y_list = list(flyte.map(fn, x_list)) # flyte.map is like Python map, but runs in parallel.
    y_mean = sum(y_list) / len(y_list)
    return y_mean
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/getting-started/hello.py*

Here's what's happening:

- **`TaskEnvironment`** specifies configuration for your tasks (container image, resources, etc.)
- **`@env.task`** turns Python functions into tasks that run remotely
- Both tasks share the same `env`, so they'll have identical configurations

## Run it

Create a project directory and place your files there:

CODE5

> [!WARNING]
> Do not run `flyte run` from your home directory. Flyte packages the current directory when running remotely, so running from `$HOME` would attempt to bundle your entire home folder. Always work from a dedicated project directory.

Run the workflow:

CODE6

This executes the workflow locally on your machine.

## See the results

You can see the run in the TUI by running:

CODE7

The TUI will open into the explorer view

![Explorer View](../../_static/images/user-guide/quickstart/explorer-tui.png)

To navigate to the run details, double-click it or press `Enter` to view the run details.

![Run Details View](../../_static/images/user-guide/quickstart/run-tui.png)

## Next steps

Now that you've run your first workflow:

- [**Core concepts**](./core-concepts/_index): Understand the core concepts of Flyte programming
- [**Run locally**](./run-modes/running-locally): Learn about the TUI, caching, and other features that work locally
- [**Run on the devbox**](./run-modes/running-devbox): Learn about the devbox cluster and how to run workflows on it


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts ===

# Core concepts

Now that you've completed the **Get started > Quickstart**, let's explore Flyte's core concepts through working examples.

By the end of this section, you'll understand:

- **TaskEnvironment**: The container configuration that defines where and how your code runs
- **Tasks**: Python functions that execute remotely in containers
- **Runs and Actions**: How Flyte tracks and manages your executions
- **Apps**: Long-running services for APIs, dashboards, and inference endpoints

Each concept is introduced with a practical example you can run yourself.

## How Flyte works

When you run code with Flyte, here's what happens:

1. You define a **TaskEnvironment** that specifies the container image and resources
2. You decorate Python functions with `@env.task` to create **tasks**
3. When you execute a task, Flyte creates a **run** that tracks the execution
4. Each task execution within a run is an **action**

Let's explore each of these in detail.

## Subpages
- **Get started > Core concepts > TaskEnvironment**
- **Get started > Core concepts > Tasks**
- **Get started > Core concepts > Runs and actions**
- **Get started > Core concepts > Where your data lives** - A developer's map of what Flyte stores in the control plane database versus the data plane object store, and what "metadata," "literals," and "raw data" actually mean.
- **Get started > Core concepts > Apps**
- **Get started > Core concepts > Projects and domains**
- **Get started > Core concepts > Key capabilities**
- **Get started > Core concepts > Basic project: RAG**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/task-environment ===

# TaskEnvironment

A `TaskEnvironment` defines the hardware and software environment where your tasks run. Think of it as the container configuration for your code.

## A minimal example

Here's the simplest possible TaskEnvironment:

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
def hello() -> str:
    return "Hello from Flyte!"
```

With just a `name`, you get Flyte's default container image and resource allocation. This is enough for simple tasks that only need Python and the Flyte SDK.

## What TaskEnvironment controls

A TaskEnvironment specifies two things:

**Hardware environment** - The compute resources allocated to each task:

- CPU cores
- Memory
- GPU type and count

**Software environment** - The container image your code runs in:

- Base image (Python version, OS)
- Installed packages and dependencies
- Environment variables

## Configuring resources

Use the `limits` parameter to specify compute resources:

```python
env = flyte.TaskEnvironment(
    name="compute_heavy",
    limits=flyte.Resources(cpu="4", mem="16Gi"),
)
```

For GPU workloads:

```python
env = flyte.TaskEnvironment(
    name="gpu_training",
    limits=flyte.Resources(cpu="8", mem="32Gi", gpu="1"),
    accelerator=flyte.GPUAccelerator.NVIDIA_A10G,
)
```

## Configuring container images

For tasks that need additional Python packages, specify a custom image:

```python
image = flyte.Image.from_debian_base().with_pip_packages("pandas", "scikit-learn")

env = flyte.TaskEnvironment(
    name="ml_env",
    image=image,
)
```

The image doesn't hard-code a container registry. For remote runs you set that once in your config (`image.registry`), so it stays out of your code. See [Container images](../../tasks/task-configuration/container-images) for detailed image configuration options.

## Multiple tasks, one environment

All tasks decorated with the same `@env.task` share that environment's configuration:

```python
env = flyte.TaskEnvironment(
    name="data_processing",
    limits=flyte.Resources(cpu="2", mem="8Gi"),
)

@env.task
def load_data(path: str) -> dict:
    # Runs with 2 CPU, 8Gi memory
    ...

@env.task
def transform_data(data: dict) -> dict:
    # Also runs with 2 CPU, 8Gi memory
    ...
```

This is useful when multiple tasks have similar requirements.

## Multiple environments

When tasks have different requirements, create separate environments:

```python
light_env = flyte.TaskEnvironment(
    name="light",
    limits=flyte.Resources(cpu="1", mem="2Gi"),
)

heavy_env = flyte.TaskEnvironment(
    name="heavy",
    limits=flyte.Resources(cpu="8", mem="32Gi"),
)

@light_env.task
def preprocess(data: str) -> str:
    # Light processing
    ...

@heavy_env.task
def train_model(data: str) -> dict:
    # Resource-intensive training
    ...
```

## Next steps

Now that you understand TaskEnvironments, let's look at how to define [tasks](./tasks) that run inside them.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/tasks ===

# Tasks

A task is a Python function that runs remotely in a container. You create tasks by decorating functions with `@env.task`.

> **📝 Note**
>
> In Flyte 1, tasks and workflows were defined with separate `@task`, `@workflow`, and `@dynamic` decorators. Flyte 2 uses a single `@env.task` decorator off a `flyte.TaskEnvironment`: everything is a task.

## Defining a task

Here's a simple task:

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
def greet(name: str) -> str:
    return f"Hello, {name}!"
```

The `@env.task` decorator tells Flyte to run this function in a container configured by `env`.

## Type hints are required

Flyte uses type hints to understand your data and serialize it between tasks:

```python
@env.task
def process_numbers(values: list[int]) -> int:
    return sum(values)
```

Supported types include:

- Primitives: `int`, `float`, `str`, `bool`
- Collections: `list`, `dict`, `tuple`
- DataFrames: `pandas.DataFrame`, `polars.DataFrame`
- Files: `flyte.File`, `flyte.Directory`
- Custom: dataclasses, Pydantic models

See [Data classes and structures](../../tasks/task-programming/dataclasses-and-structures) for complex types.

## Tasks calling tasks

In Flyte 2, tasks can call other tasks directly. The called task runs in its own container:

```python
@env.task
def fetch_data(url: str) -> dict:
    # Runs in container 1
    ...

@env.task
def process_data(url: str) -> str:
    data = fetch_data(url)  # Calls fetch_data, runs in container 2
    return transform(data)
```

This is how you build workflows in Flyte 2. There's no separate `@workflow` decorator - just tasks calling tasks.

## The top-level task

The task you execute directly is the "top-level" or "driver" task. It orchestrates other tasks:

```python
@env.task
def step_one(x: int) -> int:
    return x * 2

@env.task
def step_two(x: int) -> int:
    return x + 10

@env.task
def pipeline(x: int) -> int:
    a = step_one(x)   # Run step_one
    b = step_two(a)   # Run step_two with result
    return b
```

When you run `pipeline`, it becomes the top-level task and orchestrates `step_one` and `step_two`.

## Running tasks locally

For quick testing, you can call a task like a regular function:

```python
# Direct call - runs locally, not in a container
result = greet("World")
print(result)  # "Hello, World!"
```

This bypasses Flyte entirely and is useful for debugging logic. However, local calls don't track data, use remote resources, or benefit from Flyte's features.

## Running tasks remotely

To run a task on your Flyte backend:

```python
import flyte

flyte.init_from_config()
result = flyte.run(greet, name="World")
print(result)  # "Hello, World!"
```

Or from the command line:

```bash
flyte run my_script.py greet --name World
```

This sends your code to the Flyte backend, runs it in a container, and returns the result.

## Next steps

Now that you can define and run tasks, let's understand how Flyte tracks executions with [runs and actions](./runs-and-actions).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/runs-and-actions ===

# Runs and actions

When you execute a task on Flyte, the system creates a **run** to track it. Each individual task execution within that run is an **action**. Understanding this hierarchy helps you navigate the UI and debug your workflows.

## What is a run?

A **run** is the execution of a task that you directly initiate, plus all its descendant task executions, considered as a single unit.

When you execute:

```bash
flyte run my_script.py pipeline --x 5
```

Flyte creates a run for `pipeline`. If `pipeline` calls other tasks, those executions are part of the same run.

## What is an action?

An **action** is the execution of a single task, considered independently. A run consists of one or more actions.

Consider this workflow:

```python
@env.task
def step_one(x: int) -> int:
    return x * 2

@env.task
def step_two(x: int) -> int:
    return x + 10

@env.task
def pipeline(x: int) -> int:
    a = step_one(x)
    b = step_two(a)
    return b
```

When you run `pipeline(5)`:

- **1 run** is created for the entire execution
- **3 actions** are created: one for `pipeline`, one for `step_one`, one for `step_two`

## Runs vs actions in practice

| Concept | What it represents | In the UI |
|---------|-------------------|-----------|
| **Run** | Complete execution initiated by user | Runs list, top-level view |
| **Action** | Single task execution | Individual task details, logs |

For details on how to run tasks locally and remotely, see [Tasks](./tasks#running-tasks-locally).

## Viewing runs in the UI

After running a task remotely, click the URL in the output to see your run in the UI:

```bash
flyte run my_script.py pipeline --x 5
```

Output:

```bash
abc123xyz
https://my-instance.example.com/v2/runs/project/my-project/domain/development/abc123xyz
Run 'a0' completed successfully.
```

In the UI, you can:

- See the overall run status and duration
- Navigate to individual actions
- View inputs and outputs for each task
- Access logs for debugging
- See the execution graph

## Understanding the execution graph

The UI shows how tasks relate to each other:

```
pipeline (action)
├── step_one (action)
└── step_two (action)
```

Each box is an action. Arrows show data flow between tasks. This visualization helps you understand complex workflows and identify bottlenecks.

## Checking run status

From the command line:

```bash
flyte get run <run-id>
```

From Python:

```python
import flyte

flyte.init_from_config()
run = flyte.run(pipeline, x=5)

# The run object has status information
print(run.status)
```

## Next steps

You now understand tasks and how Flyte tracks their execution. Next, let's learn about [apps](./introducing-apps) - Flyte's approach to long-running services.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/where-data-lives ===

A developer's map of what Flyte stores in the control plane database versus the data plane object store, and what "metadata," "literals," and "raw data" actually mean.

# Where your data lives

When you run a Flyte task, your data ends up in two stores: a **database** in the control plane and an **object-store bucket** in the data plane.

## The two stores

| | Control plane database | Data plane object store |
|---|---|---|
| **Backing tech** | Postgres (plus a few internal coordination stores) | S3, GCS, or ABS bucket |
| **What's in it** | Every record Flyte uses to *describe* your runs, plus pointers to where each run's inputs and outputs live | Every run's inputs and outputs, and all bulk/offloaded content |
| **Lifetime** | Durable; long-lived history | Durable, but you can apply lifecycle/retention rules |

The database is the **source of truth for what executed**. The bucket is **where your runs' actual input and output values live**.

## What goes in the database

The control plane database holds everything Flyte needs to enumerate, schedule, and replay your work. Specifically:

- **Registrations**: every task you've deployed, every trigger you've registered, every project and domain. A task's definition includes its *default* input values, which are stored inline as part of the registration.
- **Execution records**: every run, every action (task / trace / condition) inside that run, attempts, phases, timing, error messages, parent/child relationships.
- **Schedules and triggers**: `Cron`, event triggers, and their revision history.
- **Pointers to runtime inputs and outputs**: the database stores the *URI* of each run's `inputs.pb` / `outputs.pb`, not the values themselves. (One exception: an awaited *condition* / approval action stores the value that satisfies it inline.)
- **Caches**: the cache key → output-URI mapping for `@env.task(cache=...)`.

The values your tasks actually pass at runtime, even a bare `int`, do **not** live in the database. They are written to `inputs.pb` / `outputs.pb` in the bucket, and the database keeps only the pointer. See the next section.

(Internally, Flyte uses several backing databases: Postgres for registrations and run history, separate stores for in-flight action coordination and caches. For developer purposes the only thing that matters is that they're all small-record, structured stores; none of them hold bulk content.)

## What goes in the bucket

Every run's inputs and outputs are written to the bucket as `inputs.pb` / `outputs.pb`, and the database stores a **pointer** (URI) to them. Within those files, small scalar values are inlined directly while large values are offloaded to separate objects and referenced by URI. The bucket holds:

- **Task inputs**, serialized as `inputs.pb` per run.
- **Task outputs**, serialized as `outputs.pb` per attempt.
- **Offloaded values**: `flyte.io.File`, `flyte.io.Dir`, `flyte.io.DataFrame`, pickled objects, models, anything large.
- **Decks**: the HTML reports your task renders.
- **Trace checkpoints**: used by `@flyte.trace` to resume partial work.
- **Fast-registered code bundles**: what `flyte deploy` and `flyte run --copy-style all` upload so the cluster can run your local Python.
- **Image-build contexts**: when Flyte builds a container image from an `Image` definition that requires a build context.

The layout under your bucket is `<project>/<domain>/...`, with the bulk of execution artifacts under per-run, per-action subprefixes (`<run-name>/<action>/...` for outputs / Decks / checkpoints) and sibling prefixes for offloaded inputs and SDK uploads (code bundles, image-build contexts). You don't typically need to know the exact paths; you do need to know that **everything above lives behind one configured bucket prefix**.

## What "literal" and "raw data" mean

Every value a task takes in or returns is, in Flyte's data model, a **literal**: a typed, serialized representation of that value. Literals are how data flows between tasks, and each one is stored in one of two ways:

- **Inline**: small values (primitives like `int` / `float` / `str` / `bool`, collections, and JSON-serializable dataclasses and Pydantic models) are serialized *by value* directly into the run's `inputs.pb` / `outputs.pb`.
- **By reference**: large values (`flyte.io.File`, `flyte.io.Dir`, `flyte.io.DataFrame`, large model objects, larger pickled objects) are offloaded to their own objects in the bucket; the literal recorded in `inputs.pb` / `outputs.pb` then holds only a *URI pointer* to them.

**Raw data** is that offloaded content itself: the file, directory, DataFrame, or model bytes a by-reference literal points at, plus other offloaded artifacts such as checkpoints. It's exactly the offloaded values listed under **Get started > Core concepts > Where your data lives > What goes in the bucket** above, and it's what `raw_data_path` (below) relocates. Because it's carried as a URI rather than copied inline, raw data is often described as being "passed by reference" (as opposed to inline literals, which are "passed by value").

(In the deployment and architecture guides you may also see execution data split into *literal data* and *raw data*. There, *literal data* means specifically the inline values above and *raw data* the offloaded ones, the same distinction, named as two categories.)

In short: every input and output is a **literal**; a literal is stored either **inline** (the value lives in `inputs.pb` / `outputs.pb`) or **by reference** to **raw data** offloaded elsewhere in the bucket.

## What "metadata" means

The word "metadata" appears in several places and means a different thing each time. The two senses that matter for developers:

### 1. "Metadata" as in the control plane database (Flyte's usage)

When Flyte documentation says **"metadata is preserved"** or **"metadata lives in the control plane,"** it means the database records above: registrations (including task default values), run history, and status. It does **not** mean "the contents of the bucket."

This is the sense most relevant to you: the database is durable, and losing the bucket does not lose your execution history. It loses the *large values* those history records pointed at.

### 2. "Metadata bucket" (a deployment/ops term you may see)

The Helm chart and some operational guides refer to a **"metadata bucket"** or `metadataContainer`. **This is a legacy name.** The bucket it refers to does *not* hold the database-style metadata above. It holds `inputs.pb`, `outputs.pb`, Decks, checkpoints, code bundles, and offloaded data. In other words, it holds exactly the "bucket" contents listed in the previous section.

If you see "metadata bucket" in an ops context, read it as **"the data plane object-store bucket."** The naming is unfortunate; the contents are what you'd expect from a data bucket.

You can largely ignore other appearances of the word in API surfaces (`TaskMetadata`, `ActionMetadata`, and `metadata_path` on `RunContext`, which is a local scratch directory used only by `from_local()` execution). Those are small property bags or local scratch paths and don't change where your data is stored.

## Per-run customization: `raw_data_path`

By default, offloaded values (`File`, `Dir`, `DataFrame`, checkpoints) land alongside everything else under the deployment's configured bucket prefix. You can route them to a different prefix, including a different bucket entirely, for a single run:

```python
import flyte

flyte.init_from_config()

run = flyte.with_runcontext(
    raw_data_path="s3://my-other-bucket/some/prefix",
).run(my_task, x=1)
```

This is the supported way to send a sensitive run to an isolated bucket, point at a bucket with different lifecycle rules, or otherwise route offloaded data per run. The `inputs.pb` / `outputs.pb` themselves still land in the deployment's bucket; only the *raw* offloaded contents move.

See [Run context](../../tasks/task-deployment/run-context) for the full set of `with_runcontext` options.

## What happens if the bucket is purged

If a retention rule deletes objects out of the bucket, the database records that pointed at them are **not** deleted, but their pointers now dangle. Concretely:

- Execution history, status, timing, structure: **still visible** in the UI. They come from the database.
- Input/output **previews, Deck views, artifact payloads**: show "not found" if the underlying bytes were purged.
- **Cache hits** for purged outputs: the cached pointer is dead, the task re-executes.
- **Trace resumption**: not possible if the checkpoint blob is gone.
- **Re-running an old execution**: fails if any input it needs has been purged.

This is the trade-off behind retention policies: you save storage cost at the price of being able to inspect or re-run old executions whose offloaded values have aged out. New executions are unaffected.

Lifecycle / retention rules should be scoped to the offloaded-data prefixes, **not** applied bucket-wide: `inputs.pb` and `outputs.pb` are needed for in-flight executions to complete, so purging them mid-run breaks things.

## The short version

- **Database** = the system of record. Holds registrations (including task default values), run history, schedules, and pointers to each run's inputs/outputs.
- **Bucket** = the object-store bucket. Holds every run's `inputs.pb`/`outputs.pb`, Decks, checkpoints, code bundles, and offloaded `File` / `Dir` / `DataFrame` contents.
- **Values** = every task input/output is a **literal**, stored either *inline* in `inputs.pb`/`outputs.pb` or *by reference* to **raw data** offloaded in the bucket.
- **"Metadata" in docs** usually means database-side records. **"Metadata bucket" in Helm/ops** is legacy naming for the data plane bucket. It does *not* hold database metadata.
- **`flyte.with_runcontext(raw_data_path=...)`** is your knob to send offloaded data elsewhere per run.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/introducing-apps ===

# Apps

Now that you understand tasks, let's learn about apps - Flyte's way of running long-lived services.

## Tasks vs apps

You've already learned about **tasks**: Python functions that run to completion in containers. Tasks are great for data processing, training, and batch operations.

**Apps** are different. An app is a long-running service that stays active and handles requests over time. Apps are ideal for:

- REST APIs and webhooks
- Model inference endpoints
- Interactive dashboards
- Real-time data services

| Aspect | Task | App |
|--------|------|-----|
| Lifecycle | Runs once, then exits | Stays running indefinitely |
| Invocation | Called with inputs, returns outputs | Receives HTTP requests |
| Use case | Batch processing, training | APIs, inference, dashboards |
| Durability | Inputs/outputs stored, can resume | Stateless request handling |

## AppEnvironment

Just as tasks use `TaskEnvironment`, apps use `AppEnvironment` to configure their runtime.

An `AppEnvironment` specifies:

- **Hardware**: CPU, memory, GPU allocation
- **Software**: Container image with dependencies
- **App-specific settings**: Ports, scaling, authentication

Here's a simple example:

```python
import flyte
from flyte.app.extras import FastAPIAppEnvironment

env = FastAPIAppEnvironment(
    name="my-app",
    image=flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn"),
    limits=flyte.Resources(cpu="1", mem="2Gi"),
)
```

## A hello world app

Let's create a minimal FastAPI app to see how this works.

First, create `hello_app.py`:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "uvicorn",
# ]
# ///

"""A simple "Hello World" FastAPI app example for serving."""

from fastapi import FastAPI
import pathlib
import flyte
from flyte.app.extras import FastAPIAppEnvironment

# Define a simple FastAPI application
app = FastAPI(
    title="Hello World API",
    description="A simple FastAPI application",
    version="1.0.0",
)

# Create an AppEnvironment for the FastAPI app
env = FastAPIAppEnvironment(
    name="hello-app",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)

# Define API endpoints
@app.get("/")
async def root():
    return {"message": "Hello, World!"}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# Serving this script will deploy and serve the app on your Union/Flyte instance.
if __name__ == "__main__":
    # Initialize Flyte from a config file.
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)

    # Serve the app remotely.
    app_instance = flyte.serve(env)

    # Print the app URL.
    print(app_instance.url)
    print("App 'hello-app' is now serving.")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/getting-started/serving/hello_app.py*

### Understanding the code

- **`FastAPI()`** creates the web application with its endpoints
- **`FastAPIAppEnvironment`** configures the container and resources
- **`@app.get("/")`** defines an HTTP endpoint that returns a greeting
- **`flyte.serve()`** deploys and starts the app on your Flyte backend

### Serving the app

With your config file in place, serve the app:

```bash
flyte serve hello_app.py env
```

Or run the Python file directly (which calls `flyte.serve()` in the main block):

```bash
python hello_app.py
```

You'll see output like:

```output
https://my-instance.flyte.com/v2/domain/development/project/my-project/apps/hello-app
App 'hello-app' is now serving.
```

Click the link to view your app in the UI. You can find the app URL there, or visit `/docs` for FastAPI's interactive API documentation.

## When to use apps vs tasks

Use **tasks** when:

- Processing takes seconds to hours
- You need durability (inputs/outputs tracked)
- Work is triggered by events or schedules
- Results need to be cached or resumed

Use **apps** when:

- Responses must be fast (milliseconds)
- You're serving an API or dashboard
- Users interact in real-time
- You need a persistent endpoint

## Common patterns

**Model serving with FastAPI**: Train a model with a Flyte pipeline, then serve predictions from it. During local development, the app loads the model from a local file. When deployed remotely, Flyte's `Parameter` system automatically resolves the model from the latest training run output. See [FastAPI app](../../apps/native-app-integrations/fastapi-app) for the full example.

**Agent UI with Gradio**: Build an interactive UI that kicks off agent runs using `flyte.with_runcontext()`. A single `RUN_MODE` environment variable controls the deployment progression: fully local (rapid iteration), local UI with remote task execution (cluster compute), or fully remote (production). See [Build apps](../../apps/build-apps/_index) for details.

## Next steps

You now understand the core building blocks of Flyte:

- **TaskEnvironment** and **AppEnvironment** configure where code runs
- **Tasks** are functions that execute and complete
- **Apps** are long-running services
- **Runs** and **Actions** track executions

Before diving deeper, check out [Key capabilities](./key-capabilities) for an overview of what Flyte can do: from parallelism and caching to LLM serving and error recovery.

Then head to [Basic project](./basic-project) to build a RAG application with an embedding pipeline and a Streamlit app.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/projects-and-domains ===

# Projects and domains

Flyte organizes work into a hierarchy of **organization**, **projects**, and **domains**.

- **Organization**: Your Flyte instance, typically representing a company or department. Set up during onboarding and mapped to your endpoint URL (e.g., `my-org.my-company.com`). You do not create or manage organizations directly. The organization is normally determined automatically from your endpoint URL, but you can override it with the `--org` flag on any CLI command (e.g., `flyte --org my-org get project`). This is only relevant if you have a multi-organization installation.
- **Project**: A logical grouping of related workflows, tasks, launch plans, and executions. Projects are the primary unit you create and manage.
- **Domain**: An environment classification within each project. Three fixed domains exist: `development`, `staging`, and `production`. Domains cannot be created or deleted.

Every project contains all three domains, creating **project-domain pairs** like `my-project/development`, `my-project/staging`, and `my-project/production`. Workflows, executions, and data are scoped to a specific project-domain pair.

## How projects and domains are used

When you run or deploy workflows, you target a project and domain:

- **CLI**: Use `--project` and `--domain` flags with `flyte run` or `flyte deploy`, or set defaults in your [configuration file](../run-modes/running-devbox#configure).

- **Python SDK**: Specify `project` and `domain` in `flyte.init` or `flyte.init_from_config`.

Projects and domains also determine data isolation. Storage and cache are isolated per project-domain pair.

## Managing projects via CLI

### Create a project

```shell
flyte create project --id my-project --name "My Project"
```

The `--id` is a unique identifier used in CLI commands and configuration (immutable once set). The `--name` is a human-readable display name.

You can also add a description and labels:

```shell
flyte create project \
    --id my-project \
    --name "My Project" \
    --description "ML platform workflows" \
    -l team=ml-platform \
    -l env=prod
```

Labels are specified as `-l key=value` and can be repeated.

### List projects

List all active projects:

```shell
flyte get project
```

Get details of a specific project:

```shell
flyte get project my-project
```

List archived projects:

```shell
flyte get project --archived
```

### Update a project

Update the name, description, or labels of a project:

```shell
flyte update project my-project --description "Updated description"
flyte update project my-project --name "New Display Name"
flyte update project my-project -l team=ml -l env=staging
```

> [!NOTE]
> Setting labels replaces all existing labels on the project.

### Archive a project

Archiving a project hides it from default listings but does not delete its data:

```shell
flyte update project my-project --archive
```

### Unarchive a project

Restore an archived project to active status:

```shell
flyte update project my-project --unarchive
```

## Listing projects programmatically

You can list and retrieve projects from Python using `flyte.remote.Project`:

```python
import flyte

flyte.init_from_config()

# Get a specific project
project = flyte.remote.Project.get(name="my-project", org="my-org")

# List all projects
for project in flyte.remote.Project.listall():
    print(project.to_dict())

# List with filtering and sorting
for project in flyte.remote.Project.listall(sort_by=("created_at", "desc")):
    print(project.to_dict())
```

Both `get()` and `listall()` support async execution via `.aio()`:

```python
project = await flyte.remote.Project.get.aio(name="my-project", org="my-org")
```

> [!NOTE]
> The Python SDK provides read-only access to projects. To create or modify projects, use the `flyte` CLI or the UI.

## Managing projects via the UI

When you log in to your Flyte instance, you land on the **Projects** page, which lists all projects in your organization. By default, the domain is set to `development`. You can change the active domain using the selector in the top left.

A **Recently viewed** list on the left sidebar provides quick access to your most commonly used projects.

From the project list you can:

* **Open a project**: Select a project from the list to navigate to it.
* **Create a project**: Click **+ New project** in the top right. In the dialog, specify a name and description. The project will be created across all three domains.
* **Archive a project**: Click the three-dot menu on a project's entry and select **Archive project**.

## Domains

Domains provide environment separation within each project. The three domains are:

| Domain | Purpose |
|--------|---------|
| `development` | For iterating on workflows during active development. |
| `staging` | For testing workflows before promoting to production. |
| `production` | For production workloads. |

Domains are predefined and cannot be created, renamed, or deleted.

### Targeting a domain

Set the default domain in your configuration file:

```yaml
task:
  domain: development
```

Or override per command:

```shell
flyte run --domain staging hello.py main
```

When using `flyte deploy`, the domain determines where the deployed workflows will execute:

```shell
flyte deploy --project my-project --domain production workflows
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/key-capabilities ===

# Key capabilities

Now that you understand the core concepts -- `TaskEnvironment`, tasks, runs, and apps -- here's an overview of what Flyte can do. Each capability is covered in detail later in the documentation.

## Environment and resources

Configure how and where your code runs.

- **Multiple environments**: Create separate configurations for different use cases (dev, prod, GPU vs CPU)
  → [Multiple environments](../../tasks/task-configuration/multiple-environments)

- **Resource specification**: Request specific CPU, memory, GPU, and storage for your tasks
  → [Resources](../../tasks/task-configuration/resources)

## Deployment

Get your code running remotely.

- **Code packaging**: Your local code is automatically bundled and deployed to remote execution
  → [Packaging](../../tasks/task-deployment/packaging)

- **Local testing**: Test tasks locally before deploying with `flyte run --local`
  → [How task run works](../../tasks/task-deployment/how-task-run-works)

## Data handling

Pass data efficiently between tasks.

- **Files and directories**: Pass large files and directories between tasks using `flyte.io.File` and `flyte.io.Dir`
  → [Files and directories](../../tasks/task-programming/files-and-directories)

- **DataFrames**: Work with pandas, Polars, and other DataFrame types natively
  → [DataFrames](../../tasks/task-programming/dataframes)

## Parallelism and composition

Scale out and compose workflows.

- **Fanout parallelism**: Process items in parallel using `flyte.map` or `asyncio.gather`
  → [Fanout](../../tasks/task-programming/fanout)

- **Remote tasks**: Call previously deployed tasks from within your workflows
  → [Remote tasks](../../tasks/task-programming/remote-tasks)

## Security and automation

Manage credentials and automate execution.

- **Secrets**: Inject API keys, passwords, and other credentials securely into tasks
  → [Secrets](../../tasks/task-configuration/secrets)

- **Triggers**: Schedule tasks on a cron schedule or trigger them from external events
  → [Triggers](../../tasks/task-configuration/triggers)

- **Webhooks**: Build APIs that trigger task execution from external systems
  → [Hybrid graphs](../../apps/build-apps/hybrid-graphs)

## Durability and reliability

Handle failures and avoid redundant work.

- **Error handling**: Catch failures and retry with different resources (e.g., more memory)
  → [Error handling](../../tasks/task-programming/error-handling)

- **Retries and timeouts**: Configure automatic retries and execution time limits
  → [Retries and timeouts](../../tasks/task-configuration/retries-and-timeouts)

- **Caching**: Add `cache="auto"` to any task and Flyte stores its outputs keyed on task name and inputs. Same inputs means instant results with no recomputation. This speeds up your development loop: skip re-downloading data, avoid replaying earlier steps in agentic chains, or bypass any expensive computation while you iterate.
  → [Caching](../../tasks/task-configuration/caching)

  ```python
  @env.task(cache="auto")
  async def load_data(data_dir: str = "./data") -> str:
      """Downloads once, then returns instantly on subsequent runs."""
      # ... expensive download ...
      return data_dir
  ```

- **Traces**: Use `@flyte.trace` to get visibility into the internal steps of a task without the overhead of making each step a separate task. Traced functions show up as child nodes under their parent task, each with their own timing, inputs, and outputs. This is particularly useful for AI agents where you want to see which tools were called.
  → [Traces](../../tasks/task-programming/traces)

  ```python
  @flyte.trace
  async def search(query: str) -> str:
      """Shows up as a child node under the parent task."""
      return await do_search(query)

  @env.task
  async def agent(request: str) -> str:
      results = await search(request)    # Traced
      answer = await summarize(results)   # Also traced if decorated
      return answer
  ```

- **Reports**: Add `report=True` to a task and it can generate an HTML report (charts, tables, images) saved alongside the task output. Combined with caching and persisted inputs/outputs, reports act as lightweight experiment tracking: each run produces a self-contained HTML file you can compare across runs and share with your team.
  → [Reports](../../tasks/task-programming/reports)

  ```python
  import flyte.report

  @env.task(report=True)
  async def evaluate(model_file: File, test_data: str) -> str:
      # ... run evaluation ...
      await flyte.report.replace.aio(
          f"<h2>Training Report</h2>"
          f"<h3>Test Results</h3>"
          f"<p>Accuracy: {accuracy:.4f}</p>"
      )
      await flyte.report.flush.aio()
      return f"Accuracy: {accuracy:.4f}"
  ```

## Apps and serving

Deploy long-running services.

- **FastAPI apps**: Deploy REST APIs and webhooks
  → [FastAPI app](../../apps/native-app-integrations/fastapi-app)

- **LLM serving**: Serve large language models with vLLM or SGLang
  → [vLLM app](../../apps/native-app-integrations/vllm-app), [SGLang app](../../apps/native-app-integrations/sglang-app)

- **Autoscaling**: Scale apps up and down based on traffic, including scale-to-zero
  → [Autoscaling apps](../../apps/configure-apps/auto-scaling-apps)

- **Streamlit dashboards**: Deploy interactive data dashboards
  → [Streamlit app](../../apps/native-app-integrations/streamlit-app)

## Notebooks

Work interactively.

- **Jupyter support**: Author and run workflows directly from Jupyter notebooks, and fetch workflow metadata (inputs, outputs, logs)
  → [Notebooks](../../tasks/task-programming/notebooks)

## Next steps

Ready to put it all together? Head to [Basic project](./basic-project) to build an end-to-end RAG pipeline with embeddings and a Streamlit app.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/core-concepts/basic-project ===

# Basic project: RAG

This example demonstrates a two-stage RAG (Retrieval-Augmented Generation) pattern:
an offline embedding pipeline that processes and stores quotes, followed by an online
serving application that enables semantic search.

## Concepts covered

- `TaskEnvironment` for defining task execution environments
- `Dir` artifacts for passing directories between tasks
- `AppEnvironment` for serving applications
- `Parameter` and `RunOutput` for connecting apps to task outputs
- Semantic search with sentence-transformers and ChromaDB

## Part 1: The embedding pipeline

The embedding pipeline fetches quotes from a public API, creates vector embeddings
using sentence-transformers, and stores them in a ChromaDB database.

### Setting up the environment

The `TaskEnvironment` defines the execution environment for all tasks in the pipeline.
It specifies the container image, required packages, and resource allocations:

```python
# Define the embedding environment
embedding_env = flyte.TaskEnvironment(
    name="quote-embedding",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "sentence-transformers>=2.2.0",
        "chromadb>=0.4.0",
        "requests>=2.31.0",
    ),
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    cache="auto",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/embed.py*

The environment uses:

- `Image.from_debian_base()` to create a container with Python 3.12
- `with_pip_packages()` to install sentence-transformers and ChromaDB
- `Resources` to request 2 CPUs and 4GB of memory
- `cache="auto"` to enable automatic caching of task outputs

### Fetching data

The `fetch_quotes` task retrieves quotes from a public API:

```python
@embedding_env.task
async def fetch_quotes() -> list[dict]:
    """
    Fetch quotes from a public quotes API.

    Returns:
        List of quote dictionaries with 'quote' and 'author' fields.
    """
    import requests

    print("Fetching quotes from API...")
    response = requests.get("https://dummyjson.com/quotes?limit=100")
    response.raise_for_status()

    data = response.json()
    quotes = data.get("quotes", [])

    print(f"Fetched {len(quotes)} quotes")
    return quotes
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/embed.py*

This task demonstrates:

- Async task definition with `async def`
- Returning structured data (`list[dict]`) from a task
- Using the `@embedding_env.task` decorator to associate the task with its environment

### Creating embeddings

The `embed_quotes` task creates vector embeddings and stores them in ChromaDB:

```python
@embedding_env.task
async def embed_quotes(quotes: list[dict]) -> Dir:
    """
    Create embeddings for quotes and store them in ChromaDB.

    Args:
        quotes: List of quote dictionaries with 'quote' and 'author' fields.

    Returns:
        Directory containing the ChromaDB database.
    """
    import chromadb
    from sentence_transformers import SentenceTransformer

    print("Loading embedding model...")
    model = SentenceTransformer("all-MiniLM-L6-v2")

    # Create ChromaDB in a temporary directory
    db_dir = tempfile.mkdtemp()
    print(f"Creating ChromaDB at {db_dir}...")

    client = chromadb.PersistentClient(path=db_dir)
    collection = client.create_collection(
        name="quotes",
        metadata={"hnsw:space": "cosine"},
    )

    # Prepare data for insertion
    texts = [q["quote"] for q in quotes]
    ids = [str(q["id"]) for q in quotes]
    metadatas = [{"author": q["author"], "quote": q["quote"]} for q in quotes]

    print(f"Embedding {len(texts)} quotes...")
    embeddings = model.encode(texts, show_progress_bar=True)

    # Add to collection
    collection.add(
        ids=ids,
        embeddings=embeddings.tolist(),
        metadatas=metadatas,
        documents=texts,
    )

    print(f"Stored {len(quotes)} quotes in ChromaDB")
    return await Dir.from_local(db_dir)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/embed.py*

Key points:

- Uses the `all-MiniLM-L6-v2` model from sentence-transformers (runs on CPU)
- Creates a persistent ChromaDB database with cosine similarity
- Returns a `Dir` artifact that captures the entire database directory
- The `await Dir.from_local()` call uploads the directory to artifact storage

### Orchestrating the pipeline

The main pipeline task composes the individual tasks:

```python
@embedding_env.task
async def embedding_pipeline() -> Dir:
    """
    Main pipeline that fetches quotes and creates embeddings.

    Returns:
        Directory containing the ChromaDB database with quote embeddings.
    """
    print("Starting embedding pipeline...")

    # Fetch quotes from API
    quotes = await fetch_quotes()

    # Create embeddings and store in ChromaDB
    db_dir = await embed_quotes(quotes)

    print("Embedding pipeline complete!")
    return db_dir
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/embed.py*

### Running the pipeline

To run the embedding pipeline:

```python
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(embedding_pipeline)
    print(f"Embedding run URL: {run.url}")
    run.wait()
    print(f"Embedding complete! Database directory: {run.outputs()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/embed.py*

```bash
uv run embed.py
```

The pipeline will:

1. Fetch 100 quotes from the API
2. Create embeddings using sentence-transformers
3. Store everything in a ChromaDB database
4. Return the database as a `Dir` artifact

## Part 2: The serving application

The serving application provides a Streamlit web interface for searching quotes
using the embeddings created by the pipeline.

### App environment configuration

The `AppEnvironment` defines how the application runs:

```python
# Define the app environment
env = AppEnvironment(
    name="quote-search-app",
    description="Semantic search over quotes using embeddings",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit>=1.41.0",
        "sentence-transformers>=2.2.0",
        "chromadb>=0.4.0",
    ),
    args=["streamlit", "run", "app.py", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    parameters=[
        Parameter(
            name="quotes_db",
            value=RunOutput(task_name="quote-embedding.embedding_pipeline", type="directory"),
            download=True,
            env_var="QUOTES_DB_PATH",
        ),
    ],
    include=["app.py"],
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/serve.py*

Key configuration:

- `args` specifies the command to run the Streamlit app
- `port=8080` exposes the application on port 8080
- `parameters` defines inputs to the app:
  - `RunOutput` connects to the embedding pipeline's output
  - `download=True` downloads the directory to local storage
  - `env_var="QUOTES_DB_PATH"` makes the path available to the app
- `include=["app.py"]` bundles the Streamlit app with the deployment

### The Streamlit application

The app loads the ChromaDB database using the path from the environment variable:

```python
# Load the database
@st.cache_resource
def load_db():
    db_path = os.environ.get("QUOTES_DB_PATH")
    if not db_path:
        st.error("QUOTES_DB_PATH environment variable not set")
        st.stop()

    client = chromadb.PersistentClient(path=db_path)
    collection = client.get_collection("quotes")
    model = SentenceTransformer("all-MiniLM-L6-v2")
    return collection, model

collection, model = load_db()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/app.py*

The search interface provides a text input and result count slider:

```python
# Search interface
query = st.text_input("Enter your search query:", placeholder="e.g., love, wisdom, success")
top_k = st.slider("Number of results:", min_value=1, max_value=20, value=5)

col1, col2 = st.columns([1, 1])
with col1:
    search_button = st.button("Search", type="primary", use_container_width=True)
with col2:
    random_button = st.button("Random Quote", use_container_width=True)

st.divider()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/app.py*

When the user searches, the app encodes the query and finds similar quotes:

```python
if search_button and query:
    # Encode query and search
    query_embedding = model.encode([query])[0].tolist()
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
    )

    if results["documents"] and results["documents"][0]:
        for i, (doc, metadata, distance) in enumerate(
            zip(results["documents"][0], results["metadatas"][0], results["distances"][0])
        ):
            similarity = 1 - distance  # Convert distance to similarity
            st.markdown(f'**{i+1}.** "{doc}"')
            st.caption(f"— {metadata['author']} | Similarity: {similarity:.2%}")
            st.write("")
    else:
        st.info("No results found.")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/app.py*

The app also includes a random quote feature:

```python
elif random_button:
    # Get a random quote from the collection
    all_data = collection.get(limit=100)
    if all_data["documents"]:
        idx = random.randint(0, len(all_data["documents"]) - 1)
        quote = all_data["documents"][idx]
        author = all_data["metadatas"][idx]["author"]
        st.markdown(f'**"{quote}"**')
        st.caption(f"— {author}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/app.py*

### Deploying the app

To deploy the quote search application:

```python
if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy the quote search app
    print("Deploying quote search app...")
    deployment = flyte.serve(env)
    print(f"App deployed at: {deployment.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/basic-project/serve.py*

```bash
uv run serve.py
```

The app will be deployed and automatically connected to the embedding pipeline's
output through the `RunOutput` parameter.

## Key takeaways

1. **Two-stage RAG pattern**: Separate offline embedding creation from online serving
   for better resource utilization and cost efficiency.

2. **Dir artifacts**: Use `Dir` to pass entire directories (like databases) between
   tasks and to serving applications.

3. **RunOutput**: Connect applications to task outputs declaratively, enabling
   automatic data flow between pipelines and apps.

4. **CPU-friendly embeddings**: The `all-MiniLM-L6-v2` model runs efficiently on CPU,
   making this pattern accessible without GPU resources.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/run-modes ===

# Run modes

Flyte OSS supports three execution modes, letting you choose the right trade-off between speed and fidelity at each stage of development:

### **Get started > Run modes > Run locally in Python**

Run tasks and apps directly in your local Python process with no Kubernetes cluster or Docker required. Ideal for rapid iteration and debugging.

### **Get started > Run modes > Run locally on the devbox**

Run tasks and apps in a lightweight Flyte cluster using Docker. Get the full Flyte UI and backend experience on your machine.

### **Get started > Run modes > Run on a remote cluster**

Run tasks and apps on a remote cluster with full production capabilities including GPUs, distributed compute, and cloud-scale resources.

| Aspect | Local (`--local`) | Devbox | Remote |
|--------|-------------------|--------|--------|
| **⚡️ Execution** | In-process Python | Containerized, local Docker | Containerized, on your cluster |
| **🐳 Docker required** | No | Yes | Yes (local image build) |
| **💻 Flyte UI** | No (TUI only) | Yes (`localhost:30080`) | Yes |
| **📦 Container images** | Ignored | Built locally | Built locally, pushed to a registry |
| **🔀 Parallelism** | Sequential | Cluster-level | Cluster-level |
| **⭐️ Best for** | Fast iteration, debugging | Testing container builds, full Flyte features | Production, GPUs, scale |

The same task code runs unchanged across all three modes. Start with local execution for fast feedback, move to the Devbox to validate containerized execution, then deploy to your Flyte cluster for production.

## Subpages
- **Get started > Run modes > Run locally in Python**
- **Get started > Run modes > Run locally on the devbox**
- **Get started > Run modes > Run on a remote cluster**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/run-modes/running-locally ===

# Run locally in Python

Flyte runs locally with no cluster or Docker needed. Install the SDK, write tasks, and run them on your machine. When you're ready to scale, drop the `--local` flag and the same code runs on a remote cluster with GPUs.

## Getting started

If you haven't already, install the SDK and configure local persistence as described in the [Quickstart](../quickstart).

## Running tasks locally

The `--local` flag tells Flyte to execute a task in your local Python environment rather than on a remote cluster. Add `--tui` to launch the interactive Terminal UI for real-time monitoring.

Basic local execution:

```bash
flyte run --local hello.py main
```

With the interactive TUI:

```bash
flyte run --local --tui hello.py main
```

You can also run tasks programmatically using the Python SDK with `flyte.run()`. See [Run and deploy tasks](../../tasks/task-deployment/_index) for details.

## Two ways to run a task

There are two distinct ways to run a task, and it's easy to confuse them when you're new. They differ in *who* calls `flyte.run()`.

### As a script: `python hello.py`

You call `flyte.run()` yourself, from inside an `if __name__ == "__main__":` block, and execute the file with plain Python:

```python
# hello.py
import flyte

env = flyte.TaskEnvironment(name="hello_env")

@env.task
def main(x_list: list[int] = list(range(10))) -> float:
    return sum(x_list) / len(x_list)

if __name__ == "__main__":
    flyte.init_from_config()       # load your local config (.flyte/config.yaml)
    run = flyte.run(main)          # call the task; pass inputs as keyword args
    print(run.name)
    run.wait()
```

Then run the file as an ordinary Python script:

```bash
python hello.py
```

Because *your code* decides which task runs and with what inputs, you pass inputs directly as arguments to `flyte.run()`, for example `flyte.run(main, x_list=[1, 2, 3])`. Use `flyte.run.aio(...)` from within async code.

### Via the CLI: `flyte run`

The `flyte run` CLI does the calling for you. You don't need a `__main__` block; instead you name the file and the task on the command line, and the CLI invokes it:

```bash
flyte run --local hello.py main
```

The task's **parameters become CLI options**. Each task input maps to a `--<name>` flag (run `flyte run --local hello.py main --help` to see them, with their defaults). For example, to override `x_list`:

```bash
flyte run --local hello.py main --x-list '[1, 2, 3]'
```

> [!NOTE]
> A common first-run trip-up: invoking `flyte run` against a task whose inputs have **no defaults** without supplying them. The CLI then can't construct the input and you'll see a confusing type-converter error rather than a "missing argument" message. If you hit one, check `--help` and pass the required `--<name>` values (or give the parameters defaults in the task signature, as `main` does above).

## Terminal UI

The TUI is an interactive split-screen dashboard. Task tree on the left, details and logs on the right.

```bash
flyte run --local hello.py main
```

![TUI agent run](../../../_static/images/user-guide/quickstart/run-tui.png)

What you see:

- **Task tree** with live status: `●` running, `✓` done, `✗` failed
- **Cache indicators**: `$` cache hit, `~` cache enabled but missed
- **Live logs**: `print()` output streams in real time
- **Details panel**: inputs, outputs, timing, report paths
- **Traced sub-tasks**: child nodes for `@flyte.trace` decorated functions

**Keyboard shortcuts:**

| Key | Action |
|-----|--------|
| `q` | Quit |
| `d` | Details tab |
| `l` | Logs tab |

### Exploring past runs

If you created a config file via `flyte create config --local-persistence`, Flyte
persists the inputs and outputs of every task run locally, so you can always go back and inspect what a task received and produced. Launch the TUI on its own to browse past runs, compare inputs and outputs, and review reports:

```bash
flyte start tui
```

---

## What works locally

Most Flyte features work in both local and remote execution. The table below summarizes how each feature behaves locally.

| Feature | Local behavior | Details |
|---------|---------------|---------|
| **Caching** | Outputs stored in local SQLite, keyed on task name and inputs. Same inputs = instant results. | [Caching](../../tasks/task-configuration/caching) |
| **Tracing** | `@flyte.trace` functions appear as child nodes in the TUI with their own timing, inputs, and outputs. | [Traces](../../tasks/task-programming/traces) |
| **Reports** | HTML files saved locally. TUI shows the file path. | [Reports](../../tasks/task-programming/reports) |
| **Serving** | Run apps locally with `python serve.py` or `flyte.with_servecontext(mode="local")`. | [Serve and deploy apps](../../apps/serve-and-deploy-apps/_index) |
| **Plugins** | Same decorators and APIs as remote. Secrets come from environment variables. | [Integrations](../../../api-reference/integrations/_index) |
| **Secrets** | Read from `.env` files or environment variables. No `flyte create secret` needed. | [Secrets](../../tasks/task-configuration/secrets) |

---

## Local to devbox/remote

The same code runs in both environments. Here's what changes:

| Aspect | Local | Devbox/Remote |
|--------|-------|--------|
| **Run pipeline** | `flyte run --local` | `flyte run` |
| **TUI** | `--tui` flag | Dashboard in Flyte UI |
| **Caching** | Local SQLite | Cluster-wide distributed cache |
| **Reports** | Local HTML files | Rendered in the Flyte UI |
| **Serving** | `python serve.py` | `flyte deploy serve.py env` |
| **Secrets** | `.env` / environment variables | `flyte create secret` / `flyte.Secret` |
| **Compute** | Your CPU/GPU | `Resources(cpu=2, memory="4Gi", gpu=1)` |

The [`TaskEnvironment`](../core-concepts/task-environment) is the bridge. Locally, image and resource settings are ignored. On the cluster, Flyte builds containers and allocates compute from the same definition.

---

## Next steps

- [**Run on the devbox**](./running-devbox): Run a full local Flyte cluster with Docker to test containerized execution before deploying remotely.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/run-modes/running-devbox ===

<!-- markdownlint-disable-next-line MD024 -- same heading as the union variant block; only one renders per variant -->
# Run locally on the devbox

The Flyte devbox is a lightweight local cluster that runs on your machine with Docker. It gives you a full Flyte environment (including the UI, scheduler, and object store) so you can test remote execution without deploying to a real cluster.

## What you'll need

- Python 3.10+ in a virtual environment
- [Docker](https://docs.docker.com/get-docker/) installed and running
- [kubectl](https://kubernetes.io/docs/tasks/tools/)

## Install the SDK

If you haven't already, install the `flyte` package:

```bash
pip install flyte
```

## Start the devbox

Launch the local cluster:

### CPU

```bash
flyte start devbox
```

### GPU

```bash
flyte start devbox --gpu
```

> **📝 Note**
>
> The `--gpu` flag requires an NVIDIA-enabled host. It currently *does not* support Apple Silicon or AMD GPUs.

![Devbox start](../../../_static/images/user-guide/run-modes/flyte-start-devbox.png)

This pulls the necessary containers and starts a local Flyte instance. Once ready, the Flyte UI is available at `http://localhost:30080`.

> **📝 Note**
>
> The first start may take a few minutes while Docker images are downloaded.

## Configure

Create a config file that points to the devbox:

```bash
flyte create config \
    --endpoint localhost:30080 \
    --project flytesnacks \
    --domain development \
    --builder local \
    --insecure
```

This creates `.flyte/config.yaml` configured to talk to your local devbox cluster.

## Run a workflow on the devbox

Using the same `hello.py` from the [Quickstart](../quickstart):

```python
# hello.py

import flyte

# The `hello_env` TaskEnvironment is assigned to the variable `env`.
# It is then used in the `@env.task` decorator to define tasks.
# The environment groups configuration for all tasks defined within it.
env = flyte.TaskEnvironment(name="hello_env")

# We use the `@env.task` decorator to define a task called `fn`.
@env.task
def fn(x: int) -> int: # Type annotations are required
    slope, intercept = 2, 5
    return slope * x + intercept

# We also use the `@env.task` decorator to define another task called `main`.
# This is the entrypoint task of the workflow.
# It calls the `fn` task defined above multiple times using `flyte.map`.
@env.task
def main(x_list: list[int] = list(range(10))) -> float:
    y_list = list(flyte.map(fn, x_list)) # flyte.map is like Python map, but runs in parallel.
    y_mean = sum(y_list) / len(y_list)
    return y_mean
CODE4bash
flyte run hello.py main
```

Without the `--local` flag, the workflow runs on the devbox cluster rather than in your local Python process. Tasks execute inside containers, just like they would on a remote cluster.

## View results in the UI

Open `http://localhost:30080` to see your workflow execution in the Flyte UI. You can inspect task inputs, outputs, logs, and execution timelines.

![Devbox UI](../../../_static/images/user-guide/run-modes/flyte-ui-devbox.png)

## Stop the devbox

When you're done, shut down the cluster:

CODE5

## Inline configuration

Skip the config file entirely by passing parameters directly.

### Programmatic

Use `flyte.init`:

CODE6

### CLI

Some parameters go after `flyte`, others after the subcommand:

CODE7

See the [CLI reference](../../../api-reference/flyte-cli) for details.

## Delete the devbox

CODE8

## Using a CUDA-enabled GPU host

If you started the devbox with `flyte start devbox --gpu`, you can use GPUs in your workflows.

CODE9

## Next steps

With your environment fully configured, you're ready to build:

- [**Core concepts**](../core-concepts/_index): Understand `TaskEnvironment`s, tasks, runs, and actions through working examples.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/get-started/run-modes/running-remote ===

# Run on a remote cluster

This guide covers setting up your local development environment and configuring the `flyte` CLI and SDK to connect to your Flyte instance.

<!--
Hidden until the in-browser demo is repaired. The linked demo
(https://flyte2intro.apps.demo.hosted.unionai.cloud/) is currently unreachable.
To restore: re-add a note callout reading "Want to try Flyte without installing
anything? Try Flyte 2 in your browser" linking to the demo URL once it is live.
-->

## Prerequisites

- **Python 3.10+**
- **`uv`**: A fast Python package installer. See the [`uv` installation guide](https://docs.astral.sh/uv/getting-started/installation/).
- Access to a Flyte instance (URL and a project where you can run workflows)

> [!NOTE]
> Don't have a Flyte cluster yet? See [Platform deployment](../../../oss-deployment/_index)
> to stand one up, or use the [Devbox](./running-devbox) to run a local cluster in Docker.

## Install the flyte package

Create a virtual environment and install the `flyte` package:

```bash
uv venv
source .venv/bin/activate
uv pip install flyte
```

> [!NOTE]
> On Windows, use `.venv\Scripts\activate` instead.

Verify installation:

```bash
flyte --version
```

## Configuration file

As we did in [Quickstart](../quickstart), use `flyte create config` to create a configuration file:

```bash
flyte create config \
    --endpoint my-org.my-company.com \
    --domain development \
    --project my-project \
    --builder local \
    --registry ghcr.io/my-org
```

This creates `./.flyte/config.yaml`:

```yaml
admin:
  endpoint: dns:///my-org.my-company.com
image:
  builder: local
  registry: ghcr.io/my-org
task:
  org: my-org
  domain: development
  project: my-project
```

> [!NOTE]
> The registry (`--registry`, the `image.registry` config entry, or the `FLYTE_IMAGE_REGISTRY` environment variable) sets where **locally-built** images are pushed. It applies whenever the builder is `local`: always on Flyte OSS, and on Union if you opt out of remote builds. With `--builder remote` (the Union default) images are built on the cluster, so no registry is required, which is why the Union example above omits it. Setting the registry from config requires flyte 2.5.9 or later.

### Set up local Docker

The `--builder local` setting means container images are
[built locally](../../tasks/task-configuration/container-images) on your machine and pushed to a
container registry that your Flyte cluster can pull from. You'll need Docker running and
logged into that registry, for example:

```bash
docker login ghcr.io
```

Because you set `image.registry` in your config above, your `Image` definitions don't
need a registry; the local build pushes there automatically. (Set `registry=` on an
individual `Image` only to override it.) See
[Image building](../../tasks/task-configuration/container-images#image-building) for details.

<details>
<summary>Full example with all options</summary>

Create a custom config file with all available options:

```bash
flyte create config \
    --endpoint my-org.my-company.com \
    --org my-org \
    --domain development \
    --project my-project \
    --builder local \
    --registry ghcr.io/my-org \
    --insecure \
    --output my-config.yaml \
    --force
```

See the [CLI reference](../../../api-reference/flyte-cli#flyte-create-config) for all parameters.

</details>

<details>
<summary>Config properties explained</summary>

**`admin`**: Connection details for your Flyte instance.

- `endpoint`: URL with `dns:///` prefix. If your UI is at `https://my-org.my-company.com`, use `dns:///my-org.my-company.com`.
- `insecure`: Set to `true` only for local instances without TLS.

**`image`**: Docker image building configuration.

- `builder`: How container images are built.
  - `remote` (Union): Images built on Union's infrastructure.
  - `local` (Flyte OSS): Images built on your machine. Requires Docker. See [Image building](../../tasks/task-configuration/container-images#image-building).
- `registry`: Optional registry prefix to use for image builds. This is helpful when you want the SDK to push or pull images from a custom registry without changing your code. You can also set it with the `FLYTE_IMAGE_REGISTRY` environment variable.

**`task`**: Default settings for task execution.

- `org`: Organization name (usually matches the first part of your endpoint URL).
- `domain`: Environment separation (`development`, `staging`, `production`).
- `project`: Default project for deployments. Must already exist on your instance. See [Projects and domains](../core-concepts/projects-and-domains) for how to create projects.

</details>

## Using the configuration

You can reference your config file explicitly or let the SDK find it automatically.

### Explicit configuration

### Programmatic

Initialize with `flyte.init_from_config`:

```python
flyte.init_from_config("my-config.yaml")
run = flyte.run(main)
```

### CLI

Use `--config` or `-c`:

```bash
flyte --config my-config.yaml run hello.py main
flyte -c my-config.yaml run hello.py main
```

<details>
<summary>Configuration precedence</summary>

Without an explicit path, the SDK searches these locations in order:

1. `./config.yaml`
2. `./.flyte/config.yaml`
3. `UCTL_CONFIG` environment variable
4. `FLYTECTL_CONFIG` environment variable
5. `~/.union/config.yaml`
6. `~/.flyte/config.yaml`

</details>

### Programmatic

```python
flyte.init_from_config()
```

### CLI

```bash
flyte run hello.py main
```

### Check current configuration

```bash
flyte get config
```

Output:

```bash
CLIConfig(
    Config(
        platform=PlatformConfig(endpoint='dns:///my-org.my-company.com', scopes=[]),
        task=TaskConfig(org='my-org', project='my-project', domain='development'),
        source=PosixPath('/Users/me/.flyte/config.yaml')
    ),
    ...
)
```

## Inline configuration

Skip the config file entirely by passing parameters directly.

### Programmatic

Use `flyte.init`:

```python
flyte.init(
    endpoint="dns:///my-org.my-company.com",
    org="my-org",
    project="my-project",
    domain="development",
)
```

### CLI

Some parameters go after `flyte`, others after the subcommand:

```bash
flyte \
    --endpoint my-org.my-company.com \
    --org my-org \
    run \
    --domain development \
    --project my-project \
    hello.py \
    main
```

See the [CLI reference](../../../api-reference/flyte-cli) for details.

See related methods:

* `flyte.init_from_api_key`
* `flyte.init_from_config`
* `flyte.init_in_cluster`
* `flyte.init_passthrough`

## Next steps

With your environment fully configured, you're ready to build:

- [**Core concepts**](../core-concepts/_index): Understand `TaskEnvironment`s, tasks, runs, and actions through working examples.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks ===

# Tasks

A task is a Python function that runs remotely in a container. Tasks are the unit of work in Flyte: they are versioned, cached, retried, and recorded, so a run you start today can be reproduced and inspected later.

Every task belongs to a `TaskEnvironment`, which declares the container image, the resources, and the secrets the task needs. You define the environment once and reuse it across the tasks that share it.

```python
env = flyte.TaskEnvironment(name="etl", image=flyte.Image.from_debian_base())

@env.task
def extract(url: str) -> str:
    ...
```

Tasks compose. Calling one task from another builds the graph as your code executes, so fanout, branching, and error handling are ordinary Python rather than a separate DSL.

The three sections below follow the order you meet them in: describe the environment a task runs in, write the task logic, then get it onto a cluster.

### **Tasks > Configure tasks**

Define `TaskEnvironment`s for container images, resources, secrets, caching, retries, and more; use triggers for schedules.

### **Tasks > Build tasks**

Compose tasks with fanout, parallelism, error handling, traces, files, and DataFrames.

### **Tasks > Run and deploy tasks**

Use `flyte run` for iteration or `flyte deploy` to register a stable task version.

## Related

Tasks are also the substrate for the other two building blocks. An app serves a task's results over HTTP; an agent drives tasks in a loop.

### **Apps**

Long-running services for dashboards, APIs, and model endpoints.

### **Agents**

Durable, self-healing agents built from tasks and apps.

## Subpages
- **Tasks > Configure tasks**
- **Tasks > Build tasks**
- **Tasks > Run and deploy tasks**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration ===

# Configure tasks

As we saw in **Get started > Quickstart**, you can run any Python function as a task in Flyte just by decorating it with `@env.task`.

This allows you to run your Python code in a distributed manner, with each function running in its own container.
Flyte manages the spinning up of the containers, the execution of the code, and the passing of data between the tasks.

The simplest possible case is a `TaskEnvironment` with only a `name` parameter, and an `env.task` decorator, with no parameters:

```
env = flyte.TaskEnvironment(name="my_env")

@env.task
async def my_task(name:str) -> str:
    return f"Hello {name}!"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/task_config.py*

> [!NOTE]
> Notice how the `TaskEnvironment` is assigned to the variable `env` and then that variable is
> used in the `@env.task`. This is what connects the `TaskEnvironment` to the task definition.
>
> In the following we will often use `@env.task` generically to refer to the decorator,
> but it is important to remember that it is actually a decorator attached to a specific
> `TaskEnvironment` object, and the `env` part can be any variable name you like.

This will run your task in the default container environment with default settings.

But, of course, one of the key advantages of Flyte is the ability to control the software environment, hardware environment, and other execution parameters for each task, right in your Python code.

## Task configuration levels

Task configuration is done at three levels. From most general to most specific, they are:

* The `TaskEnvironment` level: setting parameters when defining the `TaskEnvironment` object.
* The `@env.task` decorator level: Setting parameters in the `@env.task` decorator when defining a task function.
* The task invocation level: Using the **Tasks > Configure tasks > Overrides** method when invoking task execution.

Each level has its own set of parameters, and some parameters are shared across levels.
For shared parameters, the more specific level will override the more general one.

### Example

Here is an example of how these levels work together, showing each level with all available parameters:

```
# Level 1: TaskEnvironment - Base configuration
env_2 = flyte.TaskEnvironment(
    name="data_processing_env",
    image=flyte.Image.from_debian_base(),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    env_vars={"MY_VAR": "value"},
    # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY"),
    cache="disable",
    # pod_template=my_pod_template,
    # reusable=flyte.ReusePolicy(replicas=2, idle_ttl=300),
    depends_on=[another_env],
    description="Data processing task environment",
    # plugin_config=my_plugin_config
)

# Level 2: Decorator - Override some environment settings
@env_2.task(
    short_name="process",
    # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_2"),
    cache="auto",
    # pod_template=my_pod_template,
    report=True,
    max_inline_io_bytes=100 * 1024,
    retries=3,
    timeout=60,
    docs="This task processes data and generates a report."
)
async def process_data(data_path: str) -> str:
    return f"Processed {data_path}"

@env_2.task
async def invoke_process_data() -> str:
    result = await process_data.override(
        resources=flyte.Resources(cpu=4, memory="2Gi"),
        env_vars={"MY_VAR": "new_value"},
        # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_3"),
        cache="auto",
        max_inline_io_bytes=100 * 1024,
        retries=3,
        timeout=60
    )("input.csv")
    return result
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/task_config.py*

## Task configuration parameters

Each parameter is documented in detail on its dedicated page in this section.
For the complete parameter interaction matrix showing which parameters can be set at which level, and for full type signatures and constraints, see the **Flyte SDK > flyte > TaskEnvironment**.

| Parameter | Set at | Details |
|-----------|--------|---------|
| **name** | `TaskEnvironment` only | **Tasks > Configure tasks > Additional task settings** &bull; **Flyte SDK > flyte > TaskEnvironment** |
| **image** | `TaskEnvironment` only | **Tasks > Configure tasks > Container images** &bull; **Flyte SDK > flyte > Image** |
| **depends_on** | `TaskEnvironment` only | **Tasks > Configure tasks > Multiple environments** |
| **description** | `TaskEnvironment` only | **Tasks > Configure tasks > Additional task settings** |
| **plugin_config** | `TaskEnvironment` only | **Tasks > Configure tasks > Task plugins** |
| **resources** | `TaskEnvironment`, `override`\* | **Tasks > Configure tasks > Resources** &bull; **Flyte SDK > flyte > Resources** |
| **env_vars** | `TaskEnvironment`, `override`\* | **Tasks > Configure tasks > Additional task settings > Environment variables** |
| **secrets** | `TaskEnvironment`, `override`\* | **Tasks > Configure tasks > Secrets** &bull; **Flyte SDK > flyte > Secret** |
| **cache** | All three levels | **Tasks > Configure tasks > Caching** &bull; **Flyte SDK > flyte > Cache** |
| **pod_template** | All three levels | **Tasks > Configure tasks > Pod templates** &bull; **Flyte SDK > flyte > PodTemplate** |
| **reusable** | `TaskEnvironment`, `override` | **Tasks > Configure tasks > Reusable containers** &bull; **Flyte SDK > flyte > ReusePolicy** |
| **interruptible** | All three levels | **Tasks > Configure tasks > Interruptible tasks** |
| **short_name** | `@env.task`, `override` | **Tasks > Configure tasks > Additional task settings** |
| **retries** | `@env.task`, `override` | **Tasks > Configure tasks > Retries and timeouts** &bull; **Flyte SDK > flyte > RetryStrategy** |
| **timeout** | `@env.task`, `override` | **Tasks > Configure tasks > Retries and timeouts** &bull; **Flyte SDK > flyte > Timeout** |
| **max_inline_io_bytes** | `@env.task`, `override` | **Tasks > Configure tasks > Additional task settings > Inline I/O threshold** |
| **links** | `@env.task`, `override` | **Tasks > Configure tasks > Additional task settings > Naming and metadata > `links`** |
| **report** | `@env.task` only | **Tasks > Configure tasks > Additional task settings > Naming and metadata > `report`** |
| **triggers** | `@env.task` only | **Tasks > Configure tasks > Triggers** &bull; **Flyte SDK > flyte > Trigger** |
| **docs** | `@env.task` only | **Tasks > Configure tasks > Additional task settings > Naming and metadata > `docs`** |

\*When `reusable` is set, `resources`, `env_vars`, and `secrets` can only be overridden via `task.override()` with `reusable="off"` in the same call.

## Subpages
- **Tasks > Configure tasks > Container images**
- **Tasks > Configure tasks > Resources**
- **Tasks > Configure tasks > Secrets**
- **Tasks > Configure tasks > Caching**
- **Tasks > Configure tasks > Reusable containers**
- **Tasks > Configure tasks > Pod templates**
- **Tasks > Configure tasks > Multiple environments**
- **Tasks > Configure tasks > Retries and timeouts**
- **Tasks > Configure tasks > Triggers**
- **Tasks > Configure tasks > Interruptible tasks**
- **Tasks > Configure tasks > Task plugins**
- **Tasks > Configure tasks > Additional task settings**
- **Tasks > Configure tasks > Logging**
- **Tasks > Configure tasks > Overrides**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/container-images ===

# Container images

The `image` parameter of the [`TaskEnvironment`](../../../api-reference/flyte-sdk/flyte/taskenvironment) is used to specify a container image.
Every task defined using that `TaskEnvironment` will run in a container based on that image.

If a `TaskEnvironment` does not specify an `image`, it will use the default Flyte image ([`ghcr.io/flyteorg/flyte:py{python-version}-v{flyte_version}`](https://github.com/orgs/flyteorg/packages/container/package/flyte)).

> **📝 Note**
>
> In Flyte 1 the container image was defined with `ImageSpec` (the `flytekit.ImageSpec` API). Flyte 2 uses `flyte.Image`, described below.

## Specifying your own image directly

You can directly reference an image by URL in the `image` parameter, like this:

```python
env = flyte.TaskEnvironment(
    name="my_task_env",
    image="docker.io/myorg/myimage:mytag"
)
```

This works well if you have a pre-built image available in a public registry like Docker Hub or in a private registry that your Union/Flyte instance can access.

## Specifying your own image with the `flyte.Image` object

You can also construct an image programmatically using the `flyte.Image` object.

The `flyte.Image` object provides a fluent interface for building container images: start with a `from_*` base constructor, then customize with `with_*` methods. Each method returns a new immutable `Image`.

For a complete list of all available methods and their parameters, see the [`Image` API reference](../../../api-reference/flyte-sdk/flyte/image).

Here are some examples of the most common patterns for building images with `flyte.Image`.

## Example: Defining a custom image with `Image.from_debian_base`

The `[[Image.from_debian_base()]]` provides the default Flyte image as the base.
This image is itself based on the official Python Docker image (specifically `python:{version}-slim-bookworm`) with the addition of the Flyte SDK pre-installed.
Starting there, you can layer additional features onto your image.
For example:

```python
import flyte
import numpy as np

# Define the task environment
env = flyte.TaskEnvironment(
    name="my_env",
    image = (
        flyte.Image.from_debian_base(
            name="my-image",
            python_version=(3, 13)
            # registry="registry.example.com/my-org" # Only needed for local builds
        )
        .with_apt_packages("libopenblas-dev")
        .with_pip_packages("numpy")
        .with_env_vars({"OMP_NUM_THREADS": "4"})
    )
)

@env.task
def main(x_list: list[int]) -> float:
    arr = np.array(x_list)
    return float(np.mean(arr))

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, x_list=list(range(10)))
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/container-images/from_debian_base.py*

> [!NOTE]
> A registry is only needed when the image is **built locally** (it's where the built image is pushed); it isn't required when using the Union backend `ImageBuilder`, which builds on the cluster.
> The easiest way to set it is once in your config, `image.registry` (or the `FLYTE_IMAGE_REGISTRY` environment variable), so you don't have to repeat it in every `Image`. Set `registry=` on an `Image` only to override. See **Tasks > Configure tasks > Container images > Image building**.

> [!NOTE]
> Images built with `[[Image.from_debian_base()]]` do not include CA certificates by default, which can cause TLS
> validation errors and block access to HTTPS-based storage such as Amazon S3. Libraries like Polars (e.g., `polars.scan_parquet()`) are particularly affected.
> **Solution:** Add `"ca-certificates"` using `.with_apt_packages()` in your image definition.

## Example: Defining an image based on uv script metadata

Another common technique for defining an image is to use [`uv` inline script metadata](https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies) to specify your dependencies right in your Python file and then use the `flyte.Image.from_uv_script()` method to create a `flyte.Image` object.
The `from_uv_script` method starts with the default Flyte image and adds the dependencies specified in the `uv` metadata.
For example:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "numpy"
# ]
# main = "main"
# params = "x_list=[1,2,3,4,5,6,7,8,9,10]"
# ///

import flyte
import numpy as np

env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_uv_script(
            __file__,
            name="my-image"
            # registry="registry.example.com/my-org" # Only needed for local builds
        )
)

@env.task
def main(x_list: list[int]) -> float:
    arr = np.array(x_list)
    return float(np.mean(arr))

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, x_list=list(range(10)))
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/container-images/from_uv_script.py*

The advantage of this approach is that the dependencies used when running a script locally and when running it on the Flyte/Union backend are always the same (as long as you use `uv` to run your scripts locally).
This means you can develop and test your scripts in a consistent environment, reducing the chances of encountering issues when deploying to the backend.

In the above example you can see how to use `flyte.init_from_config()` for remote runs and `flyte.init()` for local runs.
Uncomment the `flyte.init()` line (and comment out `flyte.init_from_config()`) to enable local runs.
Do the opposite to enable remote runs.

> [!NOTE]
> When using `uv` metadata in this way, be sure to include the `flyte` package in your `uv` script dependencies.
> This will ensure that `flyte` is installed when running the script locally using `uv run`.
> When running on the Flyte/Union backend, the `flyte` package from the uv script dependencies will overwrite the one included automatically from the default Flyte image.

## Customizing an image with `with_*` methods

Images from `from_debian_base()` and `from_uv_script()` are *extendable*: you can layer additional customizations on top using the `with_*` methods. (Images from `from_base()` and `from_dockerfile()` are **not** extendable &mdash; calling a `with_*` method on one raises an error; customize those at their source instead.)
Each method returns a new `flyte.Image`, so you can chain them together in a fluent style:

```python
import flyte
from flyte import Image

image = (
    Image.from_debian_base()
    .with_apt_packages("git", "vim")
    .with_pip_packages("pandas", "numpy")
    .with_env_vars({"MY_ENV_VAR": "my_value"})
    .with_commands(["echo 'building image'"])
)

env = flyte.TaskEnvironment(name="my_env", image=image)
```

The available customization methods are:

| Method | Description |
|---|---|
| `flyte.Image.with_pip_packages()` | Install one or more packages with `pip` (supports `index_url`, `extra_index_urls`, `pre` for pre-releases, and `secret_mounts` for private indexes). |
| `flyte.Image.with_apt_packages()` | Install one or more system packages with `apt`. |
| `flyte.Image.with_requirements()` | Install Python dependencies from a `requirements.txt` file. |
| `flyte.Image.with_uv_project()` | Install dependencies from a `pyproject.toml` + `uv.lock` pair. |
| `flyte.Image.with_poetry_project()` | Install dependencies from a `pyproject.toml` + `poetry.lock` pair. |
| `flyte.Image.with_pixi_project()` | Install dependencies from a [pixi](https://pixi.sh) project (conda and PyPI packages) defined in a `pixi.toml`, or a `pyproject.toml` with a `[tool.pixi]` section. |
| `flyte.Image.with_source_file()` | Copy a single local file into the image. |
| `flyte.Image.with_source_folder()` | Copy a local directory into the image. |
| `flyte.Image.with_commands()` | Run additional shell commands during the build (do not prefix them with `RUN`). |
| `flyte.Image.with_env_vars()` | Set environment variables in the image. |
| `flyte.Image.with_workdir()` | Set the working directory in the image. |
| `flyte.Image.with_dockerignore()` | Point at a `.dockerignore` file to exclude paths from the build context. |

For the full signature of each method, see the [`Image` API reference](../../../api-reference/flyte-sdk/flyte/image).

> [!NOTE]
> The `with_*` methods that install Python dependencies (`with_pip_packages`, `with_requirements`, `with_uv_project`, `with_poetry_project`) cannot be combined with a conda-based image.

### Installing dependencies from a `uv` or Poetry project

If your project already declares its dependencies in a `pyproject.toml`, you can install them directly into the image rather than listing packages individually.
Use `flyte.Image.with_uv_project()` for a `uv.lock` or `flyte.Image.with_poetry_project()` for a `poetry.lock`:

```python
from pathlib import Path

from flyte import Image

image = (
    Image.from_debian_base(install_flyte=False)
    .with_apt_packages("git")
    .with_uv_project(
        pyproject_file=Path("pyproject.toml"),
        uvlock=Path("uv.lock"),
    )
)
```

By default only the dependencies are installed. To also install the project itself as a package, pass `project_install_mode="install_project"`.

### Installing dependencies from a pixi project

If your dependencies are managed with [pixi](https://pixi.sh) (conda and PyPI packages together), use `[[Image.with_pixi_project()]]` to build the image from a pixi manifest.
The manifest is resolved and installed with `pixi install` at build time, and the resulting pixi environment becomes the image's runtime environment:

```python
from flyte import Image

image = (
    Image.from_debian_base()
    .with_pixi_project("pixi.toml")
)
```

The manifest argument can be a `pixi.toml` file, a `pyproject.toml` with a `[tool.pixi]` section, or the project directory containing either.
When a `pixi.lock` sits next to the manifest, the build uses `pixi install --locked` so it reproduces the lock exactly.
By default only the manifest and lock file are copied into the image; pass `project_install_mode="install_project"` to copy the whole project directory (use this when the manifest installs the project itself, for example a `pyproject.toml` that declares the project as an editable dependency).

Three things to keep in mind:

* **`flyte` must be present in the pixi environment.** After this layer the pixi environment replaces the image's virtualenv as the runtime, so tasks run only if `flyte` is installed in it. Declare `flyte` in the manifest (for example under `[pypi-dependencies]`), or add `.with_pip_packages("flyte")` after the pixi layer (it installs into the pixi environment). The environment must also provide `python`.
* **`platforms` must cover every build architecture.** A multi-architecture image (`linux/amd64` plus `linux/arm64`) needs `platforms = ["linux-64", "linux-aarch64"]` in the manifest, or `pixi install` fails for the missing architecture at build time.
* **GPU-less builders with a CUDA manifest.** If the manifest declares a CUDA `[system-requirements]` and image builds run on machines without a GPU, set `.with_env_vars({"CONDA_OVERRIDE_CUDA": "<version>"})` before the pixi layer so install-time validation of the `__cuda` virtual package succeeds.

For the full parameter list (`environment`, `extra_args`, `secret_mounts`, `project_install_mode`), see the [`Image` API reference](../../../api-reference/flyte-sdk/flyte/image#with_pixi_project).

### Copying local files into the image

Use `flyte.Image.with_source_file()` to copy a single file, or `flyte.Image.with_source_folder()` to copy a directory, into the image:

```python
from pathlib import Path

from flyte import Image

image = (
    Image.from_debian_base()
    .with_source_file(Path("config.yaml"), dst="/app/config.yaml")
    .with_source_folder(Path("./src"), dst="/app/src")
)
```

## More ways to define a base image

Beyond `from_debian_base` and `from_uv_script`, `flyte.Image` provides several other base constructors.

### Starting from an existing image

If you already have a pre-built image in a registry, use `flyte.Image.from_base()` to use it as the starting point:

```python
from flyte import Image

image = Image.from_base("ghcr.io/my-org/my-base-image:latest")
```

> [!NOTE]
> An image from `from_base()` is **not extendable** by default &mdash; chaining `with_*` methods onto it raises an error. To add layers, first clone it with `extendable=True`, then chain as usual:
>
> ```python
> image = Image.from_base("ghcr.io/my-org/my-base-image:latest").clone(extendable=True).with_pip_packages("pandas")
> ```
>
> Otherwise, bake any extra dependencies into the pre-built image itself, or start from `from_debian_base()` / `from_uv_script()`.

You can also use `flyte.Image.clone()` to reuse an existing image definition while overriding its registry, name, Python version, or extendability:

```python
from flyte import Image

image = Image.from_base("ghcr.io/my-org/my-base-image:latest").clone(name="my-flyte-image")
```

### Building from a Dockerfile

If you need full control over the build, use `flyte.Image.from_dockerfile()` to build the image from a Dockerfile you provide:

```python
from pathlib import Path

from flyte import Image

image = Image.from_dockerfile(
    file=Path("Dockerfile").absolute(),
    registry="ghcr.io/my-org",
    name="my-image",
)
```

> [!NOTE]
> Because Flyte does not parse the Dockerfile, you cannot layer additional `with_*` methods on top of a `from_dockerfile()` image — put all of your build logic into the Dockerfile itself.
> Use an absolute `Path` for the Dockerfile; the build context is the directory containing it.

### Referencing an image defined in configuration

Use `flyte.Image.from_ref_name()` to reference an image by a name defined in your configuration, rather than hard-coding it in your task code.
This lets you swap images without changing code:

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_ref_name("custom-image"),
)
```

The named references are supplied at initialization — either in your `config.yaml`:

```yaml
image:
  image_refs:
    custom-image: ghcr.io/flyteorg/flyte:py{python-version}-v{flyte_version}
```

or through `flyte.init_from_config()`:

```python
flyte.init_from_config(images=("custom-image=ghcr.io/flyteorg/flyte:py{python-version}-v{flyte_version}",))
```

Calling `flyte.Image.from_ref_name()` with no argument references the image named `default`.

## Using different images for different tasks

A single project can use multiple images: each `flyte.TaskEnvironment` specifies its own image, so tasks in different environments run in different containers.
This is useful when, for example, a data-preparation task needs only a lightweight image while a training task needs a large GPU image.

```python
import flyte
from flyte import Image

prep_image = Image.from_debian_base().with_pip_packages("pandas", "pyarrow")
train_image = Image.from_debian_base().with_pip_packages("torch")

prep_env = flyte.TaskEnvironment(name="prep", image=prep_image)
train_env = flyte.TaskEnvironment(name="train", image=train_image, depends_on=[prep_env])

@prep_env.task
async def prepare(data: str) -> str:
    return data

@train_env.task
async def train(data: str) -> str:
    return data
```

> [!NOTE]
> For patterns where each team fully owns and builds its own image (rather than layering with `flyte.Image`), see [Bring your own image](../../project-patterns/bring-your-own-image).

## Image building

There are two ways that the image can be built:

* If you are running a Flyte OSS instance then the image is built locally on your machine and pushed to a container registry that your cluster can pull from.
* If you are running a Union instance, the image can be built locally, as with Flyte OSS, or using the Union `ImageBuilder`, which runs remotely on Union's infrastructure (no registry required).

**Setting the registry for local builds.** Rather than repeat a registry in every `Image` definition, set it once, globally, in any of these ways:

* **Config file**: add a `registry` key under `image:` in your `config.yaml`:

  ```yaml
  image:
    registry: ghcr.io/my-org
  ```

* **Environment variable**: set `FLYTE_IMAGE_REGISTRY=ghcr.io/my-org`.
* **CLI**: pass `--registry` when generating the config: `flyte create config --registry ghcr.io/my-org`. (The `--registry` flag requires flyte 2.5.9 or later; the `image.registry` config key and `FLYTE_IMAGE_REGISTRY` variable also require 2.5.9.)

Any of these sets the base registry for all image builds, so your `Image` definitions can omit `registry=` entirely. Set `registry=` on an individual `Image` only to override the global value.

### Configuring the `builder`

[Earlier](../../get-started/run-modes/running-devbox#configure), we discussed the `image.builder` property in the `config.yaml`.

For Flyte OSS instances, this property must be set to `local`.

For Union instances, this property can be set to `remote` to use the Union `ImageBuilder`, or `local` to build the image locally on your machine.

### Local image building

When `image.builder` in the `config.yaml` is set to `local`, `flyte.run()` does the following:

* Builds the Docker image using your local Docker installation, installing the dependencies specified in the `uv` inline script metadata.
* Pushes the image to the configured registry (`image.registry`, or a per-`Image` `registry=`).
* Deploys your code to the backend.
* Kicks off the execution of your workflow
* Before the task that uses your custom image is executed, the backend pulls the image from the registry to set up the container.

> [!NOTE]
> Above, we used `registry="ghcr.io/my_gh_org"`.
>
> Be sure to change `ghcr.io/my_gh_org` to the URL of your actual container registry.

You must ensure that:

* Docker is running on your local machine.
* You have successfully run `docker login` to that registry from your local machine (For example GitHub uses the syntax `echo $GITHUB_TOKEN | docker login ghcr.io -u USERNAME --password-stdin`)
* Your Union/Flyte installation has read access to that registry.

> [!NOTE]
> If you are using the GitHub container registry (`ghcr.io`)
> note that images pushed there are private by default.
> You may need to go to the image URI, click **Package Settings**, and change the visibility to public in order to access the image.
>
> Other registries (such as Docker Hub) require that you pre-create the image repository before pushing the image.
> In that case you can set it to public when you create it.
>
> Public images are on the public internet and should only be used for testing purposes.
> Do not place proprietary code in public images.

### Remote `ImageBuilder`

`ImageBuilder` is a service provided by Union that builds container images on Union's infrastructure and provides an internal container registry for storing the built images.

When `image.builder` in the `config.yaml` is set to `remote` (and you are running Union.ai), `flyte.run()` does the following:

* Builds the Docker image on your Union instance with `ImageBuilder`.
* Pushes the image to a registry
  * If you did not specify a `registry` in the `Image` definition, it pushes to the internal registry in your Union instance.
  * If you did specify a `registry`, it pushes to that registry. Be sure to also set the `registry_secret` parameter in the `Image` definition to enable `ImageBuilder` to authenticate to that registry (see **Tasks > Configure tasks > Container images > Image building > Remote `ImageBuilder` > ImageBuilder with external registries**).
* Deploys your code to the backend.
* Kicks off the execution of your workflow.
* Before the task that uses your custom image is executed, the backend pulls the image from the registry to set up the container.

There is no set up of Docker nor any other local configuration required on your part.

> [!NOTE]
> The Flyte SDK checks whether the image builder is enabled for your cluster by verifying that the `image_build` task is deployed in the `system` project within the `production` domain.
> If you are using custom roles and policies, ensure that users are granted the `view_flyte_inventory` action for the `production/system` project-domain pair.
> See the [V1 user management documentation](https://www.union.ai/docs/v1/union/user-guide/administration/user-management) for more details on creating and assigning custom roles and policies (V2 user management currently works identically to V1).

#### ImageBuilder with external registries

If you are want to push the images built by `ImageBuilder` to an external registry, you can do this by setting the `registry` parameter in the `Image` object.
You will also need to set the `registry_secret` parameter to provide the secret needed to push and pull images to the private registry.
For example:

```python
# Add registry credentials so the Union remote builder can pull the base image
# and push the resulting image to your private registry.
image=flyte.Image.from_debian_base(
    name="my-image",
    base_image="registry.example.com/my-org/my-private-image:latest",
    registry="registry.example.com/my-org"
    registry_secret="my-secret"
)

# Reference the same secret in the TaskEnvironment so Flyte can pull the image at runtime.
env = flyte.TaskEnvironment(
    name="my_task_env",
    image=image,
    secrets="my-secret"
)
```

The value of the `registry_secret` parameter must be the name of a Flyte secret of type `image_pull` that contains the credentials needed to access the private registry. It must match the name specified in the `secrets` parameter of the `TaskEnvironment` so that Flyte can use it to pull the image at runtime.

To create an `image_pull` secret for the remote builder and the task environment, run the following command:

```bash
flyte create secret --type image_pull my-secret --from-file ~/.docker/config.json
```

The format of this secret matches the standard Kubernetes [image pull secret](https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/#log-in-to-docker-hub), and should look like this:

```json
{
  "auths": {
    "registry.example.com": {
      "auth": "base64-encoded-auth"
    }
  }
}
```

> [!NOTE]
> The `auth` field contains the base64-encoded credentials for your registry (username and password or token).

### Install private PyPI packages

To install Python packages from a private PyPI index (for example, from GitHub), you can mount a secret to the image layer.
This allows your build to authenticate securely during dependency installation.
For example:

```python
private_package = "git+https://$GITHUB_PAT@github.com/pingsutw/flytex.git@2e20a2acebfc3877d84af643fdd768edea41d533"
image = (
    Image.from_debian_base()
    .with_apt_packages("git")
    .with_pip_packages(private_package, pre=True, secret_mounts=Secret("GITHUB_PAT"))
)
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/resources ===

# Resources

Task resources specify the computational limits and requests (CPU, memory, GPU, storage) that will be allocated to each task's container during execution.

To specify resource requirements for your task, instantiate a `Resources` object with the desired parameters and assign it to either
the `resources` parameter of the `TaskEnvironment` or the `resources` parameter of the `override` function (for invocation overrides).

Every task defined using that `TaskEnvironment` will run with the specified resources.
If a specific task has its own `resources` defined in the decorator, it will override the environment's resources for that task only.

If neither `TaskEnvironment` nor the task decorator specifies `resources`, the default resource allocation will be used.

## Resources data class

For the full class definition, parameter types, and accepted formats, see the [`Resources` API reference](../../../api-reference/flyte-sdk/flyte/resources).

The main parameters are:

- **`cpu`**: CPU allocation, as a number, string (`"500m"`), or `(request, limit)` tuple.
- **`memory`**: Memory with Kubernetes units, such as `"4Gi"`, or a `(request, limit)` tuple. Leave headroom below a node's total RAM: its *allocatable* memory is smaller (the kubelet reserves overhead for the OS and system daemons), so a request near a node's nominal capacity can leave the pod stuck `Pending`.
- **`gpu`**: GPU or other accelerator allocation, as `"A100:2"`, an integer count, or `flyte.GPU()`, `flyte.TPU()`, `flyte.AMD_GPU()`, or `flyte.Device()` for advanced config. See **Tasks > Configure tasks > Resources > Accelerators** below.
- **`disk`**: Ephemeral storage, such as `"10Gi"`.
- **`shm`**: Shared memory, such as `"1Gi"` or `"auto"`.

## Ephemeral storage

The `disk` parameter requests *ephemeral storage* — node-local scratch disk for the task's container. Set it as a string with Kubernetes units, for example `"50Gi"`:

```python
env = flyte.TaskEnvironment(
    name="etl_env",
    resources=flyte.Resources(cpu=2, memory="4Gi", disk="50Gi"),
)
```

Under the hood, `disk` maps to the Kubernetes [`ephemeral-storage`](https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#local-ephemeral-storage) resource on the task's container.

**What it covers.** Ephemeral storage is the local disk a task writes to while it runs: the container's writable filesystem and any temporary files your code creates on the local filesystem during execution (downloaded datasets, intermediate outputs, model checkpoints staged before offload). It is distinct from the offloaded storage backing `flyte.io.File` and `flyte.io.Dir`, which lives in the blob store rather than on the node.

**Lifecycle.** Ephemeral storage is tied to the task's pod: it is provisioned when the task starts and reclaimed when the pod terminates. Nothing written to it survives beyond the task run, so use it for scratch work — persist anything you need to keep to a `flyte.io.File` or `flyte.io.Dir` in object storage.

**Single value, not a request/limit range.** Unlike `cpu` and `memory`, `disk` takes a single string, not a `(request, limit)` tuple.

**Default behavior.** Flyte does not set an `ephemeral-storage` request or limit when `disk` is unset. (A cluster-level Kubernetes `LimitRange`, if configured, may still inject a default.) The pod can still write to node-local disk, but it may be evicted if the node comes under storage pressure. Tasks doing heavy local data processing should set `disk` explicitly.

## Examples

### Usage in TaskEnvironment

Here's a complete example of defining a TaskEnvironment with resource specifications for a machine learning training workload:

```
import flyte

# Define a TaskEnvironment for ML training tasks
env = flyte.TaskEnvironment(
    name="ml-training",
    resources=flyte.Resources(
        cpu=("2", "4"),        # Request 2 cores, allow up to 4 cores for scaling
        memory=("2Gi", "12Gi"), # Request 2 GiB, allow up to 12 GiB for large datasets
        disk="50Gi",           # 50 GiB ephemeral storage for checkpoints
        shm="8Gi"              # 8 GiB shared memory for efficient data loading
    )
)

# Use the environment for tasks
@env.task
async def train_model(dataset_path: str) -> str:
    # This task will run with flexible resource allocation
    return "model trained"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/resources/resources.py*

### Usage in a task-specific override

```
# Demonstrate resource override at task invocation level
@env.task
async def heavy_training_task() -> str:
    return "heavy model trained with overridden resources"

@env.task
async def main():
    # Task using environment-level resources
    result = await train_model("data.csv")
    print(result)

    # Task with overridden resources at invocation time
    result = await heavy_training_task.override(
        resources=flyte.Resources(
            cpu="4",
            memory="24Gi",
            disk="100Gi",
            shm="16Gi"
        )
    )()
    print(result)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/resources/resources.py*

For complete format specifications for each resource type (CPU, memory, GPU/TPU/Device, disk, shared memory), including accepted string formats, request/limit ranges, GPU partitioning, and supported accelerator types, see the [`Resources` API reference](../../../api-reference/flyte-sdk/flyte/resources).

## Accelerators

To run a task on a hardware accelerator — an NVIDIA or AMD GPU, a Google TPU, an AWS Trainium/Inferentia (Neuron) chip, or an Intel Habana Gaudi device — set the `gpu` parameter of `Resources`. Despite its name, `gpu` selects any supported accelerator type.

You can request an accelerator in three ways:

- **Type-and-count string** — `"<device>:<count>"`, for example `"T4:1"` or `"A100:8"`. This is the most common form.
- **Integer count** — `gpu=2` requests that many GPUs of whatever type is available.
- **Device object** — `flyte.GPU()`, `flyte.TPU()`, `flyte.AMD_GPU()`, or `flyte.Device()` for advanced configuration such as GPU partitioning (MIG) and TPU slice topologies.

The `gpu` value can be set on the `TaskEnvironment` (applying to every task in it) or on a per-task `override`, exactly like the other resource fields shown above.

### NVIDIA GPUs

Request an NVIDIA GPU by type and count:

```python
import flyte

env = flyte.TaskEnvironment(
    name="nvidia",
    resources=flyte.Resources(gpu="T4:1"),  # one NVIDIA T4
)
```

Supported NVIDIA types include `T4`, `L4`, `L40s`, `A10`, `A10G`, `A100`, `A100 80G`, `B200`, `H100`, `H200`, and `V100`.

#### GPU partitioning (MIG)

To request a Multi-Instance GPU (MIG) partition, use `flyte.GPU()` with a `partition`:

```python
resources=flyte.Resources(gpu=flyte.GPU(device="A100", quantity=1, partition="1g.5gb"))
```

Partitioning is available on `A100`, `A100 80G`, `H100`, and `H200`.

### Google TPUs

Use `flyte.TPU()` with the device type and, optionally, a slice topology:

```python
resources=flyte.Resources(gpu=flyte.TPU(device="V5P", partition="2x2x1"))
```

Supported TPU device types are `V5P` and `V6E`.

### AMD GPUs

Request an AMD GPU by type and count, or with `flyte.AMD_GPU()`:

```python
resources=flyte.Resources(gpu="MI300X:1")
```

Supported AMD types include `MI100`, `MI210`, `MI250`, `MI250X`, `MI300A`, `MI300X`, `MI325X`, `MI350X`, and `MI355X`.

### AWS Trainium and Inferentia (Neuron)

Request AWS Neuron accelerators — Trainium (`Trn`) or Inferentia (`Inf`) — by type and count:

```python
resources=flyte.Resources(gpu="Trn1:1")
```

Supported types include `Trn1`, `Trn1n`, `Trn2`, `Trn2u`, `Inf1`, and `Inf2`.

### Intel Habana Gaudi

Request an Intel Habana Gaudi accelerator by type and count:

```python
resources=flyte.Resources(gpu="Gaudi1:1")
```

> [!NOTE]
> Which accelerator types are actually available depends on your deployment and the node pools configured in your cluster. Requesting a type that no node provides will leave the task's pod `Pending`.

For the full list of accepted accelerator strings and device-configuration options, see the [`Resources` API reference](../../../api-reference/flyte-sdk/flyte/resources).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/secrets ===

# Secrets

Flyte secrets enable you to securely store and manage sensitive information, such as API keys, passwords, and other credentials.
Secrets reside in a secret store on the data plane of your Union/Flyte backend.
You can create, list, and delete secrets in the store using the Flyte CLI or SDK.
Secrets in the store can be accessed and used within your workflow tasks, without exposing any cleartext values in your code.

## Creating a literal string secret

You can create a secret using the [`flyte create secret`](../../../api-reference/flyte-cli#flyte-create-secret) command like this:

```bash
flyte create secret MY_SECRET_KEY --value my_secret_value
```

This will create a secret called `MY_SECRET_KEY` with the value `my_secret_value`.
This secret will be scoped to your entire organization.
It will be available across all projects and domains in your organization.
See the **Tasks > Configure tasks > Secrets > Scoping secrets** section below for more details.
See **Tasks > Configure tasks > Secrets > Using a literal string secret** for how to access the secret in your task code.

## Creating a file secret

You can also create a secret by specifying a local file:

```bash
flyte create secret MY_SECRET_KEY --from-file /local/path/to/my_secret_file
```

In this case, when accessing the secret in your task code, you will need to **Tasks > Configure tasks > Secrets > Using a file secret**.

## Scoping secrets

When you create a secret without specifying a project or domain, as we did above, the secret is scoped to the organization level.
This means that the secret will be available across all projects and domains in the organization.

You can optionally specify either or both of the `--project` and `--domain` flags to restrict the scope of the secret to:

* A specific project (across all domains)
* A specific domain (across all projects)
* A specific project and a specific domain.

For example, to create a secret that it is only available in `my_project/development`, you would execute the following command:

```bash
flyte create secret MY_SECRET_KEY --value my_secret_value --project my_project --domain development
```

## Listing secrets

You can list existing secrets with the [`flyte get secret`](../../../api-reference/flyte-cli#flyte-get-secret) command.
For example, the following command will list all secrets in the organization:

```bash
flyte get secret
```

Specifying either or both of the `--project` and `--domain` flags will list the secrets that are **only** available in that project and/or domain.

For example, to list the secrets that are only available in `my_project` and domain `development`, you would run:

```bash
flyte get secret --project my_project --domain development
```

## Deleting secrets

To delete a secret, use the [`flyte delete secret`](../../../api-reference/flyte-cli#flyte-delete-secret) command:

```bash
flyte delete secret MY_SECRET_KEY
```

## Using a literal string secret

To use a literal string secret, specify it in the `TaskEnvironment` along with the name of the environment variable into which it will be injected.
You can then access it using `os.getenv()` in your task code.
For example:

```
env_1 = flyte.TaskEnvironment(
    name="env_1",
    secrets=[
        flyte.Secret(key="my_secret", as_env_var="MY_SECRET_ENV_VAR"),
    ]
)

@env_1.task
def task_1():
    my_secret_value = os.getenv("MY_SECRET_ENV_VAR")
    print(f"My secret value is: {my_secret_value}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/secrets/secrets.py*

## Using a file secret

To use a file secret, specify it in the `TaskEnvironment` along with the `mount="/etc/flyte/secrets"` argument (with that precise value).

The file will be mounted at `/etc/flyte/secrets/<SECRET_KEY>`.

For example:

```
env_2 = flyte.TaskEnvironment(
    name="env_2",
    secrets=[
        flyte.Secret(key="my_secret", mount="/etc/flyte/secrets"),
    ]
)

@env_2.task
def task_2():
    with open("/etc/flyte/secrets/my_secret", "r") as f:
        my_secret_file_content = f.read()
    print(f"My secret file content is: {my_secret_file_content}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/secrets/secrets.py*

> [!NOTE]
> Currently, to access a file secret you must specify a `mount` parameter value of `"/etc/flyte/secrets"`.
> This fixed path is the directory in which the secret file will be placed.
> The name of the secret file will be equal to the key of the secret.

## Overriding secrets at invocation time

The secrets above are declared when the task is defined, on the `TaskEnvironment`.
You can also override which secrets are injected for a **single invocation** of a task using `task.override(secrets=...)` &mdash; useful when the same task needs different credentials depending on how it's called.
See [Overriding secrets](./overrides#overriding-secrets) for details and an example.

> [!NOTE]
> A `TaskEnvironment` can only access a secret if the scope of the secret includes the project and domain where the `TaskEnvironment` is deployed.

> [!WARNING]
> Do not return secret values from tasks. Returned values are stored in plaintext in your data plane's object store and shown in the UI and to downstream tasks, defeating the secret store's protections.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/caching ===

# Caching

Flyte 2 provides intelligent **task output caching** that automatically avoids redundant computation by reusing previously computed task results.

> [!NOTE]
> Caching works at the task level and caches complete task outputs.
> For function-level checkpointing and resumption *within tasks*, see [Traces](../task-programming/traces).

## Overview

By default, caching is disabled.

If caching is enabled for a task, then Flyte determines a **cache key** for the task.
The key is composed of the following:

* Final inputs: The set of inputs after removing any specified in the `ignored_inputs`.
* Task name: The fully-qualified name of the task.
* Interface hash: A hash of the task's input and output types.
* Cache version: The cache version string.

If the cache behavior is set to `"auto"`, the cache version is automatically generated using a hash of the task's source code (or according to the custom policy if one is specified).
If the cache behavior is set to `"override"`, the cache version can be specified explicitly using the `version_override` parameter.

When the task runs, Flyte checks if a cache entry exists for the key.
If found, the cached result is returned immediately instead of re-executing the task.

## Basic caching usage

Flyte 2 supports three main cache behaviors:

### `"auto"` - automatic versioning

```
@env.task(cache=flyte.Cache(behavior="auto"))
async def auto_versioned_task(data: str) -> str:
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

With `behavior="auto"`, the cache version is automatically generated based on the function's source code.
If you change the function implementation, the cache is automatically invalidated.

- **When to use**: Development and most production scenarios.
- **Cache invalidation**: Automatic when function code changes.
- **Benefits**: Zero-maintenance caching that "just works".

You can also use the direct string shorthand:

```
@env.task(cache="auto")
async def auto_versioned_task_2(data: str) -> str:
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

### `"override"`

With `behavior="override"`, you can specify a custom cache key in the `version_override` parameter.
Since the cache key is fixed as part of the code, it can be manually changed when you need to invalidate the cache.

```
@env.task(cache=flyte.Cache(behavior="override", version_override="v1.2"))
async def manually_versioned_task(data: str) -> str:
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

- **When to use**: When you need explicit control over cache invalidation.
- **Cache invalidation**: Manual, by changing `version_override`.
- **Benefits**: Stable caching across code changes that don't affect logic.

### `"disable"` - No caching

To explicitly disable caching, use the `"disable"` behavior.
**This is the default behavior.**

```
@env.task(cache=flyte.Cache(behavior="disable"))
async def always_fresh_task(data: str) -> str:
    return get_current_timestamp() + await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

- **When to use**: Non-deterministic functions, side effects, or always-fresh data.
- **Cache invalidation**: N/A - never cached.
- **Benefits**: Ensures execution every time.

You can also use the direct string shorthand:

```
@env.task(cache="disable")
async def always_fresh_task_2(data: str) -> str:
    return get_current_timestamp() + await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

## Advanced caching configuration

### Ignoring specific inputs

Sometimes you want to cache based on some inputs but not others:

```
@env.task(cache=flyte.Cache(behavior="auto", ignored_inputs=("debug_flag",)))
async def selective_caching(data: str, debug_flag: bool) -> str:
    if debug_flag:
        print(f"Debug: transforming {data}")
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

**This is useful for**:

- Debug flags that don't affect computation
- Logging levels or output formats
- Metadata that doesn't impact results

### Cache serialization

Cache serialization ensures that only one instance of a task runs at a time for identical inputs:

```
@env.task(cache=flyte.Cache(behavior="auto", serialize=True))
async def expensive_model_training(data: str) -> str:
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

**When to use serialization**:

- Very expensive computations (model training, large data processing)
- Shared resources that shouldn't be accessed concurrently
- Operations where multiple parallel executions provide no benefit

**How it works**:

1. First execution acquires a reservation and runs normally.
2. Concurrent executions with identical inputs wait for the first to complete.
3. Once complete, all waiting executions receive the cached result.
4. If the running execution fails, another waiting execution takes over.

### Salt for cache key variation

Use `salt` to vary cache keys without changing function logic:

```
@env.task(cache=flyte.Cache(behavior="auto", salt="experiment_2024_q4"))
async def experimental_analysis(data: str) -> str:
    return await transform_data(data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

**`salt` is useful for**:

- A/B testing with identical code.
- Temporary cache namespaces for experiments.
- Environment-specific cache isolation.

## Content-based caching for DataFrames, files, and directories

When a task input is a DataFrame (`pandas`, `polars`, or `flyte.io.DataFrame`), a `flyte.io.File`, or a `flyte.io.Dir`, the value is passed *by reference* - the cache key is derived from the data's storage location, not its contents. As a result, a downstream task keyed on such an input does **not** get a cache hit when the underlying data is identical but lives at a new path (the common case, since each run writes to a fresh location).

To cache on **content** instead, attach a hash of the data at the point where it is produced. Flyte then uses that content hash when computing the cache key of any downstream consuming task, so identical content produces a cache hit regardless of where it is stored.

> [!NOTE]
> Caching applies only on a remote cluster - local execution does not produce cache hits across runs.

### DataFrames

For a raw `pandas` or `polars` DataFrame, supply a content hash with `flyte.io.HashFunction.from_fn` in a `typing.Annotated` return type. Define the hash function once and reuse the annotated alias:

```
def hash_pandas_dataframe(df: pd.DataFrame) -> str:
    # Content-based hash using pandas' built-in row hashing.
    return str(pd.util.hash_pandas_object(df).sum())

# Reusable type alias: a pandas DataFrame whose cache key is its content hash.
HashedPandasDataFrame = Annotated[pd.DataFrame, HashFunction.from_fn(hash_pandas_dataframe)]

@env.task
async def produce_pandas() -> HashedPandasDataFrame:
    # The HashFunction in the return annotation tells Flyte to compute a
    # content hash for this output.
    return pd.DataFrame(SAMPLE_DATA)

@env.task(cache=Cache(behavior="override", version_override="v1"))
async def consume_pandas(df: pd.DataFrame) -> int:
    # Cached on the input's content hash: identical content -> cache hit.
    return int(df["value"].sum())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/content_caching.py*

The producer's return annotation tells Flyte to compute the content hash; the consumer is cached on it. The same pattern works for `polars`:

```
def hash_polars_dataframe(df: pl.DataFrame) -> str:
    return str(df.hash_rows().sum())

HashedPolarsDataFrame = Annotated[pl.DataFrame, HashFunction.from_fn(hash_polars_dataframe)]

@env.task
async def produce_polars() -> HashedPolarsDataFrame:
    return pl.DataFrame(SAMPLE_DATA)

@env.task(cache=Cache(behavior="override", version_override="v1"))
async def consume_polars(df: pl.DataFrame) -> int:
    return int(df["value"].sum())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/content_caching.py*

For `flyte.io.DataFrame`, pass the `HashFunction` to `DataFrame.from_local` via the `hash_method` parameter instead of annotating the return type:

```
@env.task
async def produce_flyte_dataframe() -> DataFrame:
    df = pd.DataFrame(SAMPLE_DATA)
    # For flyte.io.DataFrame, pass the HashFunction to `from_local` instead of
    # annotating the return type.
    hash_method = HashFunction.from_fn(hash_pandas_dataframe)
    return await DataFrame.from_local(df, hash_method=hash_method)

@env.task(cache=Cache(behavior="override", version_override="v1"))
async def consume_flyte_dataframe(df: DataFrame) -> int:
    pdf = await df.open(pd.DataFrame).all()
    return int(pdf["value"].sum())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/content_caching.py*

### Files

A `flyte.io.File` accepts a `hash_method` on `File.from_local` (and on `File.new_remote`). Pass a `HashFunction` that hashes the file's bytes, and the file is cached on its content rather than its remote path:

```
def hash_bytes(data: bytes) -> str:
    import hashlib

    return hashlib.sha256(data).hexdigest()

@env.task
async def produce_file() -> File:
    import aiofiles

    async with aiofiles.open("/tmp/data.csv", "w") as fh:
        await fh.write("id,value\n1,100\n2,200\n")
    # Pass a HashFunction (over the uploaded bytes) to `from_local` - the same
    # mechanism works for `File.new_remote(...)`. The File is then cached on its
    # content rather than its remote path.
    return await File.from_local("/tmp/data.csv", hash_method=HashFunction.from_fn(hash_bytes))

@env.task(cache=Cache(behavior="override", version_override="v1"))
async def consume_file(f: File) -> str:
    async with f.open("rb") as fh:
        return hash_bytes(bytes(await fh.read()))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/content_caching.py*

### Directories

`flyte.io.Dir.from_local` does **not** take a `HashFunction` callable. Instead, compute a content key yourself and pass it as the precomputed `dir_cache_key`:

```
@env.task
async def produce_dir() -> Dir:
    import os

    os.makedirs("/tmp/data_dir", exist_ok=True)
    with open("/tmp/data_dir/part.csv", "w") as fh:
        fh.write("id,value\n1,100\n")

    # `Dir.from_local` does not take a HashFunction callable. Instead, compute a
    # content key yourself and pass it as the precomputed `dir_cache_key`.
    content_key = hash_bytes(b"id,value\n1,100\n")
    return await Dir.from_local("/tmp/data_dir/", dir_cache_key=content_key)

@env.task(cache=Cache(behavior="override", version_override="v1"))
async def consume_dir(d: Dir) -> int:
    count = 0
    async for _ in d.walk():
        count += 1
    return count
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/content_caching.py*

## Cache policies

For details on implementing custom cache policies, see the [`CachePolicy` protocol](../../../api-reference/flyte-sdk/flyte/cachepolicy) and [`Cache` class](../../../api-reference/flyte-sdk/flyte/cache) API references.

For `behavior="auto"`, Flyte uses cache policies to generate version hashes.

### Function body policy (default)

The default `FunctionBodyPolicy` generates cache versions from the function's source code:

CODE10

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

### Custom cache policies

You can implement custom cache policies by following the `CachePolicy` protocol:

CODE11

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

## Caching configuration at different levels

You can configure caching at three levels: `TaskEnvironment` definition, `@env.task` decorator, and task invocation.

### `TaskEnvironment` level

You can configure caching at the `TaskEnvironment` level.
This will set the default cache behavior for all tasks defined using that environment.
For example:

CODE12

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

### `@env.task` decorator level

By setting the cache parameter in the `@env.task` decorator, you can override the environment's default cache behavior for specific tasks:

CODE13

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

### `task.override` level

By setting the cache parameter in the `task.override` method, you can override the cache behavior for specific task invocations:

CODE14

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/caching/caching.py*

## Runtime cache control

You can also force cache invalidation for a specific run:

CODE15

## Project and domain cache isolation

Caches are automatically isolated by:

- **Project**: Tasks in different projects have separate cache namespaces.
- **Domain**: Development, staging, and production domains maintain separate caches.

## Local development caching

When running locally, Flyte maintains a local cache:

CODE16

Local cache behavior:

- Stored in `~/.flyte/local-cache/` directory
- No project/domain isolation (since running locally)
- Disabled by setting `FLYTE_LOCAL_CACHE_ENABLED=false`


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/reusable-containers ===

# Reusable containers

By default, each task execution in Flyte and Union runs in a fresh container instance that is created just for that execution and then discarded.
With reusable containers, the same container can be reused across multiple executions and tasks.
This approach reduces start up overhead and improves resource efficiency.

> [!NOTE]
> The reusable container feature is only available when running your Flyte code on a Union backend.
> See [one of the Union.ai product variants of this page](https://www.union.ai/docs/v2/union/user-guide/reusable-containers) for details.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/pod-templates ===

# Pod templates

Flyte is built on Kubernetes and uses its powerful container orchestration capabilities. A Kubernetes [pod](https://kubernetes.io/docs/concepts/workloads/pods/) is a group of one or more containers that share storage and network resources. While Flyte automatically runs your task code in a container, pod templates let you customize the entire pod specification for advanced use cases.

The `pod_template` parameter in `TaskEnvironment` allows you to:

- **Add sidecar containers**: Run metrics exporters, service proxies, or specialized services alongside your task
- **Mount volumes**: Attach persistent storage or cloud storage like GCS or S3
- **Configure metadata**: Set custom labels and annotations for monitoring, routing, or cluster policies
- **Manage resources**: Configure resource requests, limits, and affinities
- **Inject configuration**: Add secrets, environment variables, or config maps
- **Access private registries**: Specify image pull secrets

## How it works

When you define a pod template:

1. **Primary container**: Flyte automatically injects your task code into the container specified by `primary_container_name` (default: `"primary"`)
2. **Automatic monitoring**: Flyte watches the primary container and exits the entire pod when it completes
3. **Image handling**: The image for your task environment is built automatically by Flyte; images for sidecar containers must be provided by you
4. **Local execution**: When running locally, only the task code executes; additional containers are not started

## Requirements

To use pod templates, install the Kubernetes Python client:

```bash
pip install kubernetes
```

Or add it to your image dependencies:

```python
image = flyte.Image.from_debian_base().with_pip_packages("kubernetes")
```

## Basic usage

Here's a complete example showing how to configure labels, annotations, environment variables, and image pull secrets:

```
# /// script
# requires-python = "==3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "kubernetes"
# ]
# ///

import flyte
from kubernetes.client import (
    V1Container,
    V1EnvVar,
    V1LocalObjectReference,
    V1PodSpec,
)

# Create a custom pod template
pod_template = flyte.PodTemplate(
    primary_container_name="primary",           # Name of the main container
    labels={"lKeyA": "lValA"},                 # Custom pod labels
    annotations={"aKeyA": "aValA"},            # Custom pod annotations
    pod_spec=V1PodSpec(                        # Kubernetes pod specification
        containers=[
            V1Container(
                name="primary",
                env=[V1EnvVar(name="hello", value="world")]  # Environment variables
            )
        ],
        image_pull_secrets=[                   # Access to private registries
            V1LocalObjectReference(name="regcred-test")
        ],
    ),
)

# Use the pod template in a TaskEnvironment
env = flyte.TaskEnvironment(
    name="hello_world",
    pod_template=pod_template,                 # Apply the custom pod template
    image=flyte.Image.from_uv_script(__file__, name="flyte", pre=True),
)

@env.task
async def say_hello(data: str) -> str:
    return f"Hello {data}"

@env.task
async def say_hello_nested(data: str = "default string") -> str:
    return await say_hello(data=data)

if __name__ == "__main__":
    flyte.init_from_config()
    result = flyte.run(say_hello_nested, data="hello world")
    print(result.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/pod-templates/pod_template.py*

## PodTemplate parameters

The `PodTemplate` class provides the following parameters:

| Parameter | Type | Description |
|-----------|------|-------------|
| `primary_container_name` | `str` | Name of the container where task code runs (default: `"primary"`). Must match a container in `pod_spec`. |
| `pod_spec` | `V1PodSpec` | Kubernetes pod specification for configuring containers, volumes, security contexts, and more. |
| `labels` | `dict[str, str]` | Pod labels for organization and selection by Kubernetes selectors. |
| `annotations` | `dict[str, str]` | Pod annotations for metadata and integrations (doesn't affect scheduling). |

## Volume mounts

Pod templates are commonly used to mount volumes for persistent storage or cloud storage access:

```python
from kubernetes.client import (
    V1Container,
    V1PodSpec,
    V1Volume,
    V1VolumeMount,
    V1CSIVolumeSource,
)
import flyte

pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="primary",
                volume_mounts=[
                    V1VolumeMount(
                        name="data-volume",
                        mount_path="/mnt/data",
                        read_only=False,
                    )
                ],
            )
        ],
        volumes=[
            V1Volume(
                name="data-volume",
                csi=V1CSIVolumeSource(
                    driver="your-csi-driver",
                    volume_attributes={"key": "value"},
                ),
            )
        ],
    ),
)

env = flyte.TaskEnvironment(
    name="volume-example",
    pod_template=pod_template,
    image=flyte.Image.from_debian_base(),
)

@env.task
async def process_data() -> str:
    # Access mounted volume
    with open("/mnt/data/input.txt", "r") as f:
        data = f.read()
    return f"Processed {len(data)} bytes"
```

### GCS/S3 volume mounts

Mount cloud storage directly into your pod for efficient data access:

```python
from kubernetes.client import V1Container, V1PodSpec, V1Volume, V1VolumeMount, V1CSIVolumeSource
import flyte

# GCS example with CSI driver
pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    annotations={
        "gke-gcsfuse/volumes": "true",
        "gke-gcsfuse/cpu-limit": "2",
        "gke-gcsfuse/memory-limit": "1Gi",
    },
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="primary",
                volume_mounts=[V1VolumeMount(name="gcs", mount_path="/mnt/gcs")],
            )
        ],
        volumes=[
            V1Volume(
                name="gcs",
                csi=V1CSIVolumeSource(
                    driver="gcsfuse.csi.storage.gke.io",
                    volume_attributes={"bucketName": "my-bucket"},
                ),
            )
        ],
    ),
)
```

## Sidecar containers

Add sidecar containers to run alongside your task. Common use cases include:

- **Metrics exporters**: Prometheus, Datadog agents
- **Service proxies**: Istio, Linkerd sidecars
- **Data services**: Databases, caches, or specialized services like Nvidia NIMs

```python
from kubernetes.client import V1Container, V1PodSpec
import flyte

pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    pod_spec=V1PodSpec(
        containers=[
            # Primary container (where your task code runs)
            V1Container(name="primary"),

            # Sidecar container
            V1Container(
                name="metrics-sidecar",
                image="prom/pushgateway:latest",
                ports=[{"containerPort": 9091}],
            ),
        ],
    ),
)

env = flyte.TaskEnvironment(
    name="sidecar-example",
    pod_template=pod_template,
    image=flyte.Image.from_debian_base().with_pip_packages("requests"),
)

@env.task
async def task_with_metrics() -> str:
    import requests

    # Send metrics to sidecar
    requests.post("http://localhost:9091/metrics", data="my_metric 42")

    # Your task logic
    return "Task completed with metrics"
```

## Image pull secrets

Configure private registry access:

```python
from kubernetes.client import V1Container, V1PodSpec, V1LocalObjectReference
import flyte

pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    pod_spec=V1PodSpec(
        containers=[V1Container(name="primary")],
        image_pull_secrets=[V1LocalObjectReference(name="my-registry-secret")],
    ),
)
```

## Cluster-specific configuration

Pod templates are often used to configure Kubernetes-specific settings required by your cluster, even when not using multiple containers:

```python
import flyte

pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    annotations={
        "iam.amazonaws.com/role": "my-task-role",  # AWS IAM role
        "cluster-autoscaler.kubernetes.io/safe-to-evict": "false",
    },
    labels={
        "cost-center": "ml-team",
        "project": "recommendations",
    },
)
```

## Important notes

1. **Local execution**: Pod templates only apply to remote execution. When running locally, only your task code executes.

2. **Image building**: Flyte automatically builds and manages the image for your task environment. Images for sidecar containers must be pre-built and available in a registry.

3. **Primary container**: Your task code is automatically injected into the container matching `primary_container_name`. This container must be defined in the `pod_spec.containers` list.

4. **Lifecycle management**: Flyte monitors the primary container and terminates the entire pod when it exits, ensuring sidecar containers don't run indefinitely.

## Best practices

1. **Start simple**: Begin with basic labels and annotations before adding complex sidecars
2. **Test locally first**: Verify your task logic works locally before adding pod customizations
3. **Use environment-specific templates**: Different environments (dev, staging, prod) may need different pod configurations
4. **Set resource limits**: Always set resource requests and limits for sidecars to prevent cluster issues
5. **Security**: Use image pull secrets and least-privilege service accounts

## Learn more

- [Kubernetes Pods Documentation](https://kubernetes.io/docs/concepts/workloads/pods/)
- [Kubernetes Python Client](https://github.com/kubernetes-client/python)
- [V1PodSpec Reference](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#podspec-v1-core)


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/multiple-environments ===

# Multiple environments

In many applications, different tasks within your workflow may require different configurations.
Flyte enables you to manage this complexity by allowing multiple environments within a single workflow.

Multiple environments are useful when:

- Different tasks in your workflow need different dependencies.
- Some tasks require specific CPU/GPU or memory configurations.
- A task requires a secret that other tasks do not (and you want to limit exposure of the secret value).
- You're integrating specialized tools that have conflicting requirements.

## Constraints on multiple environments

To use multiple environments in your workflow you define multiple `TaskEnvironment` instances, each with its own configuration, and then assign tasks to their respective environments.

There are, however, two additional constraints that you must take into account.
If `task_1` in environment `env_1` calls a `task_2` in environment `env_2`, then:

1. `env_1` must declare a deployment-time dependency on `env_2` in the `depends_on` parameter of `TaskEnvironment` that defines `env_1`.
2. The image used in the `TaskEnvironment` of `env_1` must include all dependencies of the module containing the `task_2` (unless `task_2` is invoked as a remote task).

<!-- TODO: Link to remote tasks when that page is live
2. The image used in the `TaskEnvironment` of `env_1` must include all dependencies of the module containing the `task_2` (unless [`task_2` is invoked as a remote task](../task-programming/remote-tasks)).
-->

### Task `depends_on` constraints

The `depends_on` parameter in `TaskEnvironment` is used to provide deployment-time dependencies by establishing a relationship between one `TaskEnvironment` and another.
The system uses this information to determine which environments (and, specifically which images) need to be built in order to be able to run the code.

On `flyte run` (or `flyte deploy`), the system walks the tree defined by the `depends_on` relationships, starting with the environment of the task being invoked (or the environment being deployed, in the case of `flyte deploy`), and prepares each required environment.
Most importantly, it ensures that the container images need for all required environments are available (and if not, it builds them).

This deploy-time determination of what to build is important because it means that for any given `run` or `deploy`, only those environments that are actually required are built.
The alternative strategy of building all environments defined in the set of deployed code can lead to unnecessary and expensive builds, especially when iterating on code.

### Dependency inclusion constraints

When a parent task invokes a child task in a different environment, the container image of the parent task environment must include all dependencies used by the child task.
This is necessary because of the way task invocation works in Flyte:

- When a child task is invoked by function name, that function, necessarily, has to be imported into the parent tasks's Python environment.
- This results in all the dependencies of the child task function also being imported.
- But, nonetheless, the actual execution of the child task occurs in its own environment.

To avoid this requirement, you can invoke a task in another environment _remotely_.

<!-- TODO: Link to remote tasks when that page is live
To avoid this requirement, you can [invoke a task in another environment _remotely_](../task-programming/remote-tasks).
-->

## Example

The following example is a (very) simple mock of an AlphaFold2 pipeline.
It demonstrates a workflow with three tasks, each in its own environment.

The example project looks like this:

```bash
├── msa/
│   ├── __init__.py
│   └── run.py
├── fold/
│   ├── __init__.py
│   └── run.py
├── __init__.py
└── main.py
```

(The source code for this example can be found here:[AlphaFold2 mock example](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/task-configuration/multiple-environments/af2))

In file `msa/run.py` we define the task `run_msa`, which mocks the multiple sequence alignment step of the process:

```python
import flyte
from flyte.io import File

MSA_PACKAGES = ["pytest"]

msa_image = flyte.Image.from_debian_base().with_pip_packages(*MSA_PACKAGES)

msa_env = flyte.TaskEnvironment(name="msa_env", image=msa_image)

@msa_env.task
def run_msa(x: str) -> File:
    f = File.new_remote()
    with f.open_sync("w") as fp:
        fp.write(x)
    return f
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/msa/run.py*

* A dedicated image (`msa_image`) is built using the `MSA_PACKAGES` dependency list, on top of the standard base image.
* A dedicated environment (`msa_env`) is defined for the task, using `msa_image`.
* The task is defined within the context of the `msa_env` environment.

In file `fold/run.py` we define the task `run_fold`, which mocks the fold step of the process:

```python
import flyte
from flyte.io import File

FOLD_PACKAGES = ["ruff"]

fold_image = flyte.Image.from_debian_base().with_pip_packages(*FOLD_PACKAGES)

fold_env = flyte.TaskEnvironment(name="fold_env", image=fold_image)

@fold_env.task
def run_fold(sequence: str, msa: File) -> list[str]:
    with msa.open_sync("r") as f:
        msa_content = f.read()
    return [msa_content, sequence]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/fold/run.py*

* A dedicated image (`fold_image`) is built using the `FOLD_PACKAGES` dependency list, on top of the standard base image.
* A dedicated environment (`fold_env`) is defined for the task, using `fold_image`.
* The task is defined within the context of the `fold_env` environment.

Finally, in file `main.py` we define the task `main` that ties everything together into a workflow.

We import the required modules and functions:

```
import logging
import pathlib

from fold.run import fold_env, fold_image, run_fold
from msa.run import msa_env, MSA_PACKAGES, run_msa

import flyte
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/main.py*

Notice that we import

* The task functions that we will be calling: `run_fold` and `run_msa`.
* The environments of those tasks: `fold_env` and `msa_env`.
* The dependency list of the `run_msa` task: `MSA_PACKAGES`
* The image of the `run_fold` task: `fold_image`

We then assemble the image and the environment:

```
main_image = fold_image.with_pip_packages(*MSA_PACKAGES)

env = flyte.TaskEnvironment(
    name="multi_env",
    depends_on=[fold_env, msa_env],
    image=main_image,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/main.py*

The image for the `main` task (`main_image`) is built by starting with `fold_image` (the image for the `run_fold` task) and adding `MSA_PACKAGES` (the dependency list for the `run_msa` task).
This ensures that `main_image` includes all dependencies needed by both the `run_fold` and `run_msa` tasks.

The environment for the `main` task is defined with:

* The image `main_image`. This ensures that the `main` task has all the dependencies it needs.
* A depends_on list that includes both `fold_env` and `msa_env`. This establishes the deploy-time dependencies on those environments.

Finally, we define the `main` task itself:

```
@env.task
def main(sequence: str) -> list[str]:
    """Given a sequence, outputs files containing the protein structure
    This requires model weights + gpus + large database on aws fsx lustre
    """
    print(f"Running AlphaFold2 for sequence: {sequence}")
    msa = run_msa(sequence)
    print(f"MSA result: {msa}, passing to fold task")
    results = run_fold(sequence, msa)
    print(f"Fold results: {results}")
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/main.py*

Here we call, in turn, the `run_msa` and `run_fold` tasks.
Since we call them directly rather than as remote tasks, we had to ensure that `main_image` includes all dependencies needed by both tasks.

<!-- TODO: Link to remote tasks when that page is live
Note that we call them directly, not as [remote tasks](../task-programming/remote-tasks), which is why we had to ensure that `main_image` includes all dependencies needed by both tasks.
-->

The final piece of the puzzle is the `if __name__ == "__main__":` block that allows us to run the `main` task on the configured Flyte backend:

```
if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, "AAGGTTCCAA")
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/multiple-environments/af2/main.py*

Now you can run the workflow with:

```bash
python main.py
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/retries-and-timeouts ===

# Retries and timeouts

Retries and timeouts are the two primary controls for handling failure on a Flyte task.
Long-running tasks fail and stall in many ways: a transient network blip, a flaky third-party
API, a slow scheduler, a runaway loop, a hung container. Retries decide *whether to try again*;
timeouts decide *how long to wait before giving up*. Used together they keep your workflows
reliable in the face of transient failures while preventing a single stuck attempt from
burning resources indefinitely.

Both can be configured in the `@env.task` decorator or supplied per-call with `override`.
Neither can be set on the `TaskEnvironment` definition itself.

## The action lifecycle

Every attempt of a task moves through a sequence of phases. Knowing the phases makes the
timeout controls obvious, because each timeout bounds a specific stretch of this timeline:

| Phase | What's happening |
| --- | --- |
| **Queued** | The action has been accepted and is waiting to be scheduled onto the cluster. |
| **Waiting for resources** | Scheduled, but waiting for compute (pods, GPUs, quota) to become available. |
| **Initializing** | Resources are in hand; the pod is starting (image pull, init containers, sidecars). |
| **Running** | Your code is actively executing. |
| **Succeeded / Failed / Timed out / Aborted** | Terminal phases. |

The diagram below shows one attempt, plus a second attempt after a retry, and which control
governs each part of the timeline:

```mermaid
gantt
    title Which control covers which part of the timeline
    dateFormat X
    axisFormat %Ss

    section Attempt 1
    Queued                  :q1, 0, 2
    Waiting for resources   :w1, 2, 5
    Initializing            :i1, 5, 6
    Running (your code)     :r1, 6, 11

    section Per-attempt bounds
    max_queued_time         :crit, 0, 5
    max_runtime             :active, 6, 11

    section Attempt 2 (after a retry + backoff)
    Queued -> Running       :q2, 14, 25

    section Across all attempts
    deadline                :done, 0, 27
```

In short:

- **`max_queued_time`** bounds the time spent *waiting to run* (Queued + Waiting for resources). It resets on every attempt.
- **`max_runtime`** bounds the time spent *running your code*. It resets on every attempt.
- **`deadline`** bounds the *total* wall-clock from the first time the action was queued until it reaches a terminal phase: across every attempt, user retries and platform retries alike.

(The brief *Initializing* phase is charged to neither per-attempt bound.)

## Retries

A retry is a fresh attempt at executing a failed action. Each retry runs in a brand-new pod,
so nothing from the failed attempt (local files, in-memory state) carries over.

The code for the retry examples below can be found on
[GitHub](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py).

First we import the required modules and set up a task environment:

```
from datetime import timedelta

import flyte
import flyte.errors

env = flyte.TaskEnvironment(name="retries", resources=flyte.Resources(cpu=1, memory="250Mi"))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py*

### Retry count

The simplest form passes an integer. A "retry" is any attempt after the first, so `retries=3`
means up to **4 attempts** in total (1 original + 3 retries):

```
@env.task(retries=3)
async def call_service() -> str:
    # retries=3 -> up to 4 attempts (1 original + 3 retries).
    # Each retry runs in a fresh pod, so nothing from the failed attempt carries over.
    return await fetch_from_flaky_upstream()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py*

This is the right default for genuinely transient failures (a dropped connection, a brief
`503` from a dependency) where simply trying again is likely to succeed.

### Retries with exponential backoff

Retrying immediately is exactly the wrong thing to do against a downstream that is *already
struggling*: back-to-back retries pile load onto a service that needs room to recover.
A `flyte.RetryStrategy` with a `flyte.Backoff` policy inserts a growing delay between attempts
so a recovering dependency gets breathing room:

```
@env.task(
    retries=flyte.RetryStrategy(
        count=5,
        backoff=flyte.Backoff(
            base=timedelta(seconds=10),  # first retry waits 10s
            factor=2.0,                  # then 20s, 40s, 80s, ...
            cap=timedelta(minutes=5),    # never wait longer than 5m between retries
        ),
    ),
)
async def call_flaky_api() -> str:
    # The delay before the n-th retry (0-indexed) is min(base * factor**n, cap).
    # Backoff gives a recovering downstream room to breathe instead of hammering it.
    return await fetch_from_flaky_upstream()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py*

The delay before the n-th retry (0-indexed) is `min(base * factor**n, cap)`. With the values
above the delays are 10s, 20s, 40s, 80s, then capped at 5m. The `cap` is what keeps an
aggressive `factor` from growing into hours; it is required whenever `factor > 1`.

### Skip retries for failures that can't be fixed

Some failures will never succeed no matter how many times you try: an invalid input, a
malformed config, a permission that was never granted. Retrying them just wastes the budget
(and the wall-clock) before the inevitable failure. Raise
`flyte.errors.NonRecoverableError` to signal that a failure is terminal: the action fails on
the spot, on attempt #1, with no retries consumed even when `retries` is set:

```
@env.task(retries=3)
async def validate_and_process(x: int) -> str:
    if x < 0:
        # A negative input will never succeed, so don't waste the retry budget on it.
        # NonRecoverableError fails the action on attempt #1 — no retries are consumed.
        raise flyte.errors.NonRecoverableError(
            f"Input x={x} is negative — retrying will not help."
        )
    return f"processed({x})"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py*

Finally, configure Flyte and run:

```
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(validate_and_process, x=-5)
    print(run.name)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/retries.py*

### System retries

`retries=N` covers failures of *your task*: exceptions, non-zero exits, timeouts you've configured.
Failures caused by the underlying infrastructure (a node disappearing, ephemeral storage running out, a spot instance getting preempted, and so on) are handled separately. The platform retries these on its own and they do not consume your `retries=N` budget. They **do**, however, count against the `deadline` (see below), because the deadline is an absolute bound on total wall-clock regardless of who triggered the retry.

> [!NOTE]
> Retries run the task from the beginning, so your task logic should be idempotent.
> Avoid relying on local state from a previous attempt, and make any external side effects safe to repeat.

The platform does eventually give up, but only after many retries: a persistent infrastructure problem can churn for a while before the run is terminated. If you spot a task repeatedly failing on the same infrastructure error, abort it manually rather than waiting for the system to give up on its own. You don't configure the system retry budget from Python; it's a platform-level concern.

For workloads on spot/preemptible compute, see also [Spot to on-demand fallback](./interruptible-tasks-and-queues#spot-to-on-demand-fallback), which describes how interruptible tasks transition to on-demand on their final attempt.

## Timeouts

A timeout is a wall-clock bound that, when exceeded, terminates the action with phase
**Timed out**. Without one, a stuck attempt has no natural end: a task waiting on a hung socket
or starved for a GPU that never frees up can sit there for hours. The `timeout` parameter takes
a `flyte.Timeout` value carrying any combination of three independent bounds, each optional: an
unspecified bound is unlimited.

The code for the timeout examples below can be found on
[GitHub](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py).

First, the imports and environment:

```
import asyncio
from datetime import timedelta

import flyte
from flyte import Timeout

env = flyte.TaskEnvironment(name="timeouts", resources=flyte.Resources(cpu=1, memory="250Mi"))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

### `max_runtime`: bound a single attempt's execution

`max_runtime` caps the time an attempt spends in the **Running** phase. It's the answer to
"how long should one run of this code take?" Use it to reap a hung container so a retry can
take over instead of letting a wedged process run forever. It is per-attempt and resets on
each retry.

```
@env.task(timeout=Timeout(max_runtime=timedelta(minutes=30)))
async def train_model() -> str:
    # max_runtime bounds the RUNNING phase of a single attempt. If the task is
    # still running after 30 minutes, this attempt is reaped as TIMED_OUT.
    # The budget is per-attempt: it resets fresh on every retry.
    ...
    return "model trained"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

### `max_queued_time`: fail fast when capacity isn't available

`max_queued_time` caps the time an attempt spends *waiting to run* (the **Queued** and
**Waiting for resources** phases) before execution begins. It answers "how long am I willing
to wait for this to even start?" When a task asks for a scarce resource (a specific GPU, a
large node) that the cluster can't currently supply, this bound makes it fail fast instead of
stalling indefinitely. It is per-attempt and resets on each retry.

```
@env.task(timeout=Timeout(max_queued_time=timedelta(minutes=15)))
async def needs_scarce_gpu() -> str:
    # max_queued_time bounds the time spent waiting to run (QUEUED +
    # WAITING_FOR_RESOURCES). If the cluster can't find capacity within 15
    # minutes, fail fast instead of stalling indefinitely. Per-attempt.
    ...
    return "done"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

### `deadline`: bound the total wall-clock

`deadline` is the strongest of the three: an absolute budget on total wall-clock, measured from
the first time the action was queued to the moment it reaches a terminal phase: across **all**
attempts, including platform-driven system retries. It answers "what is the total time budget
for this work, no matter what?" Use it when a downstream consumer needs a definite outcome by a
certain time: with `retries=5` and `max_runtime=1h` alone, an action could legally consume six
hours plus queue time before giving up. A `deadline` puts a hard ceiling on that.

```
@env.task(timeout=Timeout(deadline=timedelta(hours=2)))
async def must_finish_by() -> str:
    # deadline is an absolute wall-clock budget across ALL attempts, measured
    # from the first time the action was enqueued. Once 2 hours elapse, the
    # action is reaped no matter which phase it is in or how many retries remain.
    ...
    return "done"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

When the `deadline` fires mid-attempt, the action terminates immediately as **Timed out**,
regardless of remaining retry budget or per-attempt timer state.

### Combining the bounds

The three bounds are orthogonal and can be set together. A common shape is a per-attempt
runtime cap, a queue-wait cap, and an absolute ceiling on the whole thing:

```
@env.task(
    timeout=Timeout(
        max_runtime=timedelta(minutes=30),      # per attempt, RUNNING only
        max_queued_time=timedelta(minutes=15),  # per attempt, waiting to run
        deadline=timedelta(hours=2),            # absolute, across all attempts
    ),
)
async def fully_bounded() -> str:
    ...
    return "done"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

For backward compatibility, a bare `int` (seconds) or `timedelta` passed to `timeout` is
interpreted as `max_runtime`:

```
# A bare int (seconds) or timedelta is shorthand for Timeout(max_runtime=...).
@env.task(timeout=timedelta(minutes=30))
async def runtime_only() -> str:
    ...
    return "done"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

## Combining retries and timeouts

Retries and timeouts compose, and the per-attempt versus absolute distinction is what makes the
combination expressive. Because `max_runtime` and `max_queued_time` are per-attempt, they retry
normally: each timed-out attempt counts as a failure and the next attempt gets a fresh budget.
Because `deadline` is absolute, it overrides retries entirely: once it fires, no further
attempts run.

| Timeout | With `retries` set | Without `retries` |
| --- | --- | --- |
| `max_runtime` | Each attempt is reaped at the budget and retried until retries are exhausted. | First timeout is final. |
| `max_queued_time` | Each attempt is reaped pre-Running and retried until retries are exhausted. | First timeout is final. |
| `deadline` | Retries continue until the budget is exhausted **or** the deadline fires, whichever comes first. | Action terminates at the deadline. |

The most useful pattern combines backoff-paced retries with an absolute deadline: keep retrying
a flaky dependency, but never spend more than a fixed total budget on it.

```
@env.task(
    retries=flyte.RetryStrategy(
        count=5,
        backoff=flyte.Backoff(base=timedelta(seconds=30), factor=2.0, cap=timedelta(minutes=5)),
    ),
    timeout=Timeout(
        max_runtime=timedelta(minutes=10),  # cap any single attempt
        deadline=timedelta(hours=1),        # but never spend more than 1h total
    ),
)
async def resilient_work() -> str:
    # Retries continue until either the retry budget is exhausted OR the 1h
    # deadline fires — whichever comes first. The deadline wins ties.
    ...
    return "done"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

Finally, configure Flyte and run:

```
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(fully_bounded)
    print(run.name)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/retries-and-timeouts/timeouts.py*

Together, these controls let your workflows absorb transient failures gracefully while
guaranteeing that broken or starved work is reaped on a schedule you choose rather than left to
run unbounded.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/triggers ===

# Triggers

Triggers allow you to automate and parameterize an execution by scheduling its start time and providing overrides for its task inputs.

> **📝 Note**
>
> In Flyte 1 these were configured with a `LaunchPlan` (the `flytekit.LaunchPlan` API) and `CronSchedule`. Flyte 2 replaces them with `flyte.Trigger` and `flyte.Cron`, described below.

Currently, only **schedule triggers** are supported.
This type of trigger runs a task based on a Cron expression or a fixed-rate schedule.

Support is coming for other trigger types, such as:

* Webhook triggers: Hit an API endpoint to run your task.
* Artifact triggers: Run a task when a specific artifact is produced.

## Triggers are set in the task decorator

A trigger is created by setting the `triggers` parameter in the task decorator to a `flyte.Trigger` object or a list of such objects (triggers are not settable at the `TaskEnvironment` definition or `task.override` levels).

Here is a simple example:

```
import flyte
from datetime import datetime, timezone

env = flyte.TaskEnvironment(name="trigger_env")

@env.task(triggers=flyte.Trigger.hourly())  # Every hour
def hourly_task(trigger_time: datetime, x: int = 1) -> str:
    return f"Hourly example executed at {trigger_time.isoformat()} with x={x}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

Here we use a predefined schedule trigger to run the `hourly_task` every hour.
Other predefined triggers can be used similarly (see **Tasks > Configure tasks > Triggers > Predefined schedule triggers** below).

If you want full control over the trigger behavior, you can define a trigger using the `flyte.Trigger` class directly.

## `flyte.Trigger`

For complete parameter documentation, see the [`Trigger`](../../../api-reference/flyte-sdk/flyte/trigger), [`Cron`](../../../api-reference/flyte-sdk/flyte/cron), and [`FixedRate`](../../../api-reference/flyte-sdk/flyte/fixedrate) API references.

The `Trigger` class allows you to define custom triggers with full control over scheduling and execution behavior. It has the following signature:

```
flyte.Trigger(
    name,
    automation,
    description="",
    auto_activate=True,
    inputs=None,
    env_vars=None,
    interruptible=None,
    overwrite_cache=False,
    queue=None,
    labels=None,
    annotations=None
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

Here's a comprehensive example showing all parameters:

```
comprehensive_trigger = flyte.Trigger(
    name="monthly_financial_report",
    automation=flyte.Cron("0 6 1 * *", timezone="America/New_York"),
    description="Monthly financial report generation for executive team",
    auto_activate=True,
    inputs={
        "report_date": flyte.TriggerTime,
        "report_type": "executive_summary",
        "include_forecasts": True
    },
    env_vars={
        "REPORT_OUTPUT_FORMAT": "PDF",
        "EMAIL_NOTIFICATIONS": "true"
    },
    interruptible=False,  # Critical report, use dedicated resources
    overwrite_cache=True,  # Always fresh data
    queue="financial-reports",
    labels={
        "team": "finance",
        "criticality": "high",
        "automation": "scheduled"
    },
    annotations={
        "compliance.company.com/sox-required": "true",
        "backup.company.com/retain-days": "2555"  # 7 years
    }
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

## The `automation` parameter with `flyte.FixedRate`

You can define a fixed-rate schedule trigger by setting the `automation` parameter of the `flyte.Trigger` to an instance of `flyte.FixedRate`.

The `flyte.FixedRate` has the following signature:

```
flyte.FixedRate(
    interval_minutes,
    start_time=None
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Examples

```
# Every 90 minutes, starting when deployed
every_90_min = flyte.Trigger(
    "data_processing",
    flyte.FixedRate(interval_minutes=90)
)

# Every 6 hours (360 minutes), starting at a specific time
specific_start = flyte.Trigger(
    "batch_job",
    flyte.FixedRate(
        interval_minutes=360,  # 6 hours
        start_time=datetime(2025, 12, 1, 9, 0, 0)  # Start Dec 1st at 9 AM
    )
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

## The `automation` parameter with `flyte.Cron`

You can define a Cron-based schedule trigger by setting the `automation` parameter to an instance of `flyte.Cron`.

The `flyte.Cron` has the following signature:

```
flyte.Cron(
    cron_expression,
    timezone=None
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Examples

```
# Every day at 6 AM UTC
daily_trigger = flyte.Trigger(
    "daily_report",
    flyte.Cron("0 6 * * *")
)

# Every weekday at 9:30 AM Eastern Time
weekday_trigger = flyte.Trigger(
    "business_hours_task",
    flyte.Cron("30 9 * * 1-5", timezone="America/New_York")
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

#### Cron expressions

Here are some common cron expressions you can use:

| Expression     | Description                          |
|----------------|--------------------------------------|
| `0 0 * * *`    | Every day at midnight                |
| `0 9 * * 1-5`  | Every weekday at 9 AM                |
| `30 14 * * 6`  | Every Saturday at 2:30 PM            |
| `0 0 1 * *`    | First day of every month at midnight |
| `0 0 25 * *`   | 25th day of every month at midnight  |
| `0 0 * * 0`    | Every Sunday at midnight             |
| `*/10 * * * *` | Every 10 minutes                     |
| `0 */2 * * *`  | Every 2 hours                        |

For a full guide on Cron syntax, refer to [Crontab Guru](https://crontab.guru/).

## The `inputs` parameter

The `inputs` parameter allows you to provide default values for your task's parameters when the trigger fires.
This is essential for parameterizing your automated executions and passing trigger-specific data to your tasks.

### Basic usage

```
trigger_with_inputs = flyte.Trigger(
    "data_processing",
    flyte.Cron("0 6 * * *"),  # Daily at 6 AM
    inputs={
        "batch_size": 1000,
        "environment": "production",
        "debug_mode": False
    }
)

@env.task(triggers=trigger_with_inputs)
def process_data(batch_size: int, environment: str, debug_mode: bool = True) -> str:
    return f"Processing {batch_size} items in {environment} mode"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Using `flyte.TriggerTime`

The special `flyte.TriggerTime` value is used in the `inputs` to indicate the task parameter into which Flyte will inject the trigger execution timestamp:

```
timestamp_trigger = flyte.Trigger(
    "daily_report",
    flyte.Cron("0 0 * * *"),  # Daily at midnight
    inputs={
        "report_date": flyte.TriggerTime,  # Receives trigger execution time
        "report_type": "daily_summary"
    }
)

@env.task(triggers=timestamp_trigger)
def generate_report(report_date: datetime, report_type: str) -> str:
    return f"Generated {report_type} for {report_date.strftime('%Y-%m-%d')}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Required vs optional parameters

> [!IMPORTANT]
> If your task has parameters without default values, you **must** provide values for them in the trigger inputs, otherwise the trigger will fail to execute.

```python
# ❌ This will fail - missing required parameter 'data_source'
bad_trigger = flyte.Trigger(
    "bad_trigger",
    flyte.Cron("0 0 * * *")
    # Missing inputs for required parameter 'data_source'
)

@env.task(triggers=bad_trigger)
def bad_trigger_taska(data_source: str, batch_size: int = 100) -> str:
    return f"Processing from {data_source} with batch size {batch_size}"

# ✅ This works - all required parameters provided
good_trigger = flyte.Trigger(
    "good_trigger",
    flyte.Cron("0 0 * * *"),
    inputs={
        "data_source": "prod_database",  # Required parameter
        "batch_size": 500  # Override default
    }
)

@env.task(triggers=good_trigger)
def good_trigger_task(data_source: str, batch_size: int = 100) -> str:
    return f"Processing from {data_source} with batch size {batch_size}"
```

### Complex input types

You can pass various data types through trigger inputs:

```
complex_trigger = flyte.Trigger(
    "ml_training",
    flyte.Cron("0 2 * * 1"),  # Weekly on Monday at 2 AM
    inputs={
        "model_config": {
            "learning_rate": 0.01,
            "batch_size": 32,
            "epochs": 100
        },
        "feature_columns": ["age", "income", "location"],
        "validation_split": 0.2,
        "training_date": flyte.TriggerTime
    }
)

@env.task(triggers=complex_trigger)
def train_model(
    model_config: dict,
    feature_columns: list[str],
    validation_split: float,
    training_date: datetime
) -> str:
    return f"Training model with {len(feature_columns)} features on {training_date}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

## Predefined schedule triggers

For common scheduling needs, Flyte provides predefined trigger methods that create Cron-based schedules without requiring you to specify cron expressions manually.
These are convenient shortcuts for frequently used scheduling patterns.

### Available predefined triggers

```
minutely_trigger = flyte.Trigger.minutely()    # Every minute
hourly_trigger = flyte.Trigger.hourly()        # Every hour
daily_trigger = flyte.Trigger.daily()          # Every day at midnight
weekly_trigger = flyte.Trigger.weekly()        # Every week (Sundays at midnight)
monthly_trigger = flyte.Trigger.monthly()      # Every month (1st day at midnight)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

For reference, here's what each predefined trigger is equivalent to:

```python
# These are functionally identical:
flyte.Trigger.minutely() == flyte.Trigger("minutely", flyte.Cron("* * * * *"))
flyte.Trigger.hourly() == flyte.Trigger("hourly", flyte.Cron("0 * * * *"))
flyte.Trigger.daily() == flyte.Trigger("daily", flyte.Cron("0 0 * * *"))
flyte.Trigger.weekly() == flyte.Trigger("weekly", flyte.Cron("0 0 * * 0"))
flyte.Trigger.monthly() == flyte.Trigger("monthly", flyte.Cron("0 0 1 * *"))
```

All predefined trigger methods accept the same parameters as `flyte.Trigger`, plus a `trigger_time_input_key`. For the full parameter list, see the [`Trigger` API reference](../../../api-reference/flyte-sdk/flyte/trigger).

### Trigger time in predefined triggers

By default, predefined triggers will pass the execution time to the parameter `trigger_time` of type `datetime`,if that parameter exists on the task.
If no such parameter exists, the task will still be executed without error.

Optionally, you can customize the parameter name that receives the trigger execution timestamp by setting the `trigger_time_input_key` parameter (in this case the absence of this custom parameter on the task will raise an error at trigger deployment time):

```
@env.task(triggers=flyte.Trigger.daily(trigger_time_input_key="scheduled_at"))
def task_with_custom_trigger_time_input(scheduled_at: datetime) -> str:
    return f"Executed at {scheduled_at}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

## Multiple triggers per task

You can attach multiple triggers to a single task by providing a list of triggers. This allows you to run the same task on different schedules or with different configurations:

```
@env.task(triggers=[
    flyte.Trigger.hourly(),  # Predefined trigger
    flyte.Trigger.daily(),   # Another predefined trigger
    flyte.Trigger("custom", flyte.Cron("0 */6 * * *"))  # Custom trigger every 6 hours
])
def multi_trigger_task(trigger_time: datetime = flyte.TriggerTime) -> str:
    # Different logic based on execution timing
    if trigger_time.hour == 0:  # Daily run at midnight
        return f"Daily comprehensive processing at {trigger_time}"
    else:  # Hourly or custom runs
        return f"Regular processing at {trigger_time.strftime('%H:%M')}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

You can mix and match trigger types, combining predefined triggers with those that use `flyte.Cron`, and `flyte.FixedRate` automations (see below for explanations of these concepts).

## Notifications

You can attach notifications to a trigger using the `notifications` parameter of `flyte.Trigger`.
Notifications fire when a triggered run reaches a terminal execution phase.

```
import flyte
from flyte import notify
from flyte.models import ActionPhase

env = flyte.TaskEnvironment(name="my_task_env")

trigger_with_notifications = flyte.Trigger(
    name="daily_report",
    automation=flyte.Cron("0 9 * * 1-5"),
    notifications=(
        notify.Slack(
            on_phase=ActionPhase.FAILED,
            webhook_url="https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
            message="Run {{.Run.Name}} failed with: {{.Error}}",
        ),
        notify.Email(
            on_phase=ActionPhase.SUCCEEDED,
            recipients=["oncall@example.com"],
            subject="Run {{.Run.Name}} succeeded",
            body="Run: {{.Run.Name}}",
        ),
    ),
)

@env.task(triggers=trigger_with_notifications)
def process_data(date: str) -> str:
    return f"Processed {date}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Execution phases

The `on_phase` parameter accepts a single phase or a tuple of terminal phases from `flyte.models.ActionPhase`:

| Phase | Description                |
|-------|----------------------------|
| `ActionPhase.SUCCEEDED` | Run completed successfully |
| `ActionPhase.FAILED` | Run failed with an error  |
| `ActionPhase.TIMED_OUT` | Run exceeded its timeout  |
| `ActionPhase.ABORTED` | Run was manually aborted  |

To notify on multiple phases with the same notification:

```
notify.Email(
    on_phase=(ActionPhase.FAILED, ActionPhase.ABORTED),
    recipients=["oncall@example.com"],
    subject="Alert: Run completed with phase {{.Phase}}",
    body="Run: {{.Run.Name}}\nError: {{.Error}}",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

### Template variables

All message fields support template variables that are substituted at delivery time:

| Variable           | Description                                            |
|--------------------|--------------------------------------------------------|
| `{{.Run.Project}}` | Project name                                           |
| `{{.Run.Domain}}`  | Domain name                                            |
| `{{.Run.Name}}`    | Run ID                                                 |
| `{{.Phase}}`       | Execution phase                                        |
| `{{.Error}}`       | Error message when failed or abort reason when aborted |

### Slack notifications

`notify.Slack` sends a message to a Slack channel via an [incoming webhook](https://api.slack.com/messaging/webhooks).

**Simple message:**

```
notify.Slack(
    on_phase=ActionPhase.FAILED,
    webhook_url="https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
    message="Run {{.Run.Name}} failed in {{.Run.Project}}/{{.Run.Domain}}: {{.Error}}",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

**Rich formatting with [Block Kit](https://api.slack.com/block-kit):**

Use `blocks` instead of `message` for structured layouts. When `blocks` is provided, `message` is ignored.

```
notify.Slack(
    on_phase=ActionPhase.SUCCEEDED,
    webhook_url="https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
    blocks=[
        {
            "type": "header",
            "text": {"type": "plain_text", "text": "Task Succeeded"},
        },
        {
            "type": "section",
            "fields": [
                {"type": "mrkdwn", "text": "*Run:*\n{{.Run.Name}}"},
                {"type": "mrkdwn", "text": "*Phase:*\n{{.Phase}}"},
            ],
        },
        {"type": "divider"},
        {
            "type": "context",
            "elements": [
                {"type": "mrkdwn", "text": "{{.Run.Project}}/{{.Run.Domain}}"},
            ],
        },
    ],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

### Email notifications

`notify.Email` sends an email notification. You can provide a plain-text `body`, an `html_body`, or both (the email is sent as multipart when both are present).

```
notify.Email(
    on_phase=ActionPhase.FAILED,
    recipients=["oncall@example.com"],
    cc=["team-lead@example.com"],
    subject="ALERT: Run {{.Run.Name}} failed",
    body="Run: {{.Run.Name}}\nError: {{.Error}}",
    html_body="<b>Error:</b> {{.Error}}<br>",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

### Microsoft Teams notifications

`notify.Teams` sends a message to a Teams channel via an incoming webhook. Use `card` for [Adaptive Card](https://adaptivecards.io/designer/) formatting; when `card` is set, `title` and `message` are ignored.

```
notify.Teams(
    on_phase=ActionPhase.FAILED,
    webhook_url="https://outlook.office.com/webhook/YOUR_WEBHOOK_URL",
    title="Task Failed",
    message="Run {{.Run.Name}} failed: {{.Error}}\n",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

### Custom webhook notifications

`notify.Webhook` sends an HTTP request to any endpoint. All string values in `headers` and `body` support template variables.

```
notify.Webhook(
    on_phase=ActionPhase.SUCCEEDED,
    url="https://api.example.com/events",
    method="POST",
    headers={"Authorization": "Bearer my-token"},
    body={
        "event": "task_succeeded",
        "run": "{{.Run.Name}}",
    },
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

## Deploying a task with triggers

We recommend that you define your triggers in code together with your tasks and deploy them together.

The Union UI displays:

* `Owner` - who last deployed the trigger.

* `Last updated` - who last activated or deactivated the trigger and when. Note: If you deploy a trigger with `auto_activate=True`(default), this will match the `Owner`.

* `Last Run` - when was the last run created by this trigger.

For development and debugging purposes, you can adjust and deploy individual triggers from the UI.

To deploy a task with its triggers, you can either use Flyte CLI:

```bash
flyte deploy -p <project> -d <domain> <file_with_tasks_and_triggers.py> env
```

Or in Python:

```python
flyte.deploy(env)
```

Upon deploy, all triggers that are associated with a given task `T` will be automatically switched to apply to the latest version of that task. Triggers on task `T` which are defined elsewhere (i.e. in the UI) will be deleted unless they have been referenced in the task definition of `T`

<!-- TODO
Add link to workflow deployment docs when available.
-->

## Activating and deactivating triggers

By default, triggers are automatically activated upon deployment (`auto_activate=True`).
Alternatively, you can set `auto_activate=False` to deploy inactive triggers.
An inactive trigger will not create runs until activated.

```
env = flyte.TaskEnvironment(name="my_task_env")

custom_cron_trigger = flyte.Trigger(
    "custom_cron",
    flyte.Cron("0 0 * * *"),
    auto_activate=False # Dont create runs yet
)

@env.task(triggers=custom_cron_trigger)
def custom_task() -> str:
    return "Hello, world!"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

This trigger won't create runs until it is explicitly activated.
You can activate a trigger via the Flyte CLI:

```bash
flyte update trigger custom_cron my_task_env.custom_task --activate --project <project> --domain <domain>
```

If you want to stop your trigger from creating new runs, you can deactivate it:

```bash
flyte update trigger custom_cron my_task_env.custom_task --deactivate --project <project> --domain <domain>
```

You can also view and manage your deployed triggers in the Union UI.

## Trigger run timing

The timing of the first run created by a trigger depends on the type of trigger used (Cron-based or Fixed-rate) and whether the trigger is active upon deployment.

### Cron-based triggers

For Cron-based triggers, the first run will be created at the next scheduled time according to the cron expression after trigger activation and similarly thereafter.

* `0 0 * * *` If deployed at 17:00 today, the trigger will first fire 7 hours later (0:00 of the following day) and then every day at 0:00 thereafter.

* `*/15 14 * * 1-5` if today is Tuesday at 17:00, the trigger will fire the next day (Wednesday) at 14:00, 14:15, 14:30, and 14:45 and then the same for every subsequent weekday thereafter.

### Fixed-rate triggers without `start_time`

If no `start_time` is specified, then the first run will be created after the specified interval from the time of activation. No run will be created immediately upon activation, but the activation time will be used as the reference point for future runs.

#### No `start_time`, auto_activate: True

Let's say you define a fixed rate trigger with automatic activation like this:

```
my_trigger = flyte.Trigger("my_trigger", flyte.FixedRate(60))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

In this case, the first run will occur 60 minutes after the successful deployment of the trigger.
So, if you deployed this trigger at 13:15, the first run will occur at 14:15 and so on thereafter.

#### No `start_time`, auto_activate: False

On the other hand, let's say you define a fixed rate trigger without automatic activation like this:

```
my_trigger = flyte.Trigger("my_trigger", flyte.FixedRate(60), auto_activate=False)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

Then you activate it after about 3 hours. In this case the first run will kick off 60 minutes after trigger activation.
If you deployed the trigger at 13:15 and activated it at 16:07, the first run will occur at 17:07.

### Fixed-rate triggers with `start_time`

If a `start_time` is specified, the timing of the first run depends on whether the trigger is active at `start_time` or not.

#### Fixed-rate with `start_time` while active

If a `start_time` is specified, and the trigger is active at `start_time` then the first run will occur at `start_time` and then at the specified interval thereafter.
For example:

```
my_trigger = flyte.Trigger(
    "my_trigger",
    # Runs every 60 minutes starting from October 26th, 2025, 10:00am
    flyte.FixedRate(60, start_time=datetime(2025, 10, 26, 10, 0, 0)),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

If you deploy this trigger on October 24th, 2025, the trigger will wait until October 26th 10:00am and will create the first run at exactly 10:00am.

#### Fixed-rate with `start_time` while inactive

If a start time is specified, but the trigger is activated after `start_time`, then the first run will be created when the next time point occurs that aligns with the recurring trigger interval using `start_time` as the initial reference point.
For example:

```
custom_rate_trigger = flyte.Trigger(
    "custom_rate",
    # Runs every 60 minutes starting from October 26th, 2025, 10:00am
    flyte.FixedRate(60, start_time=datetime(2025, 10, 26, 10, 0, 0)),
    auto_activate=False
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

If activated later than the `start_time`, say on October 28th 12:35pm for example, the first run will be created at October 28th at 1:00pm.

## Deleting triggers

If you decide that you don't need a trigger anymore, you can remove the trigger from the task definition and deploy the task again.

Alternatively, you can use Flyte CLI:

```bash
flyte delete trigger custom_cron my_task_env.custom_task --project <project> --domain <domain>
```

## Schedule time zones

### Setting time zone for a cron schedule

Cron expressions are by default in UTC, but it's possible to specify custom time zones like so:

```
sf_trigger = flyte.Trigger(
    "sf_tz",
    flyte.Cron(
        "0 9 * * *", timezone="America/Los_Angeles"
    ), # Every day at 9 AM PT
    inputs={"start_time": flyte.TriggerTime, "x": 1},
)

nyc_trigger = flyte.Trigger(
    "nyc_tz",
    flyte.Cron(
        "1 12 * * *", timezone="America/New_York"
    ), # Every day at 12:01 PM ET
    inputs={"start_time": flyte.TriggerTime, "x": 1},
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

The above two schedules will fire 1 minute apart, at 9 AM PT and 12:01 PM ET respectively.

### `flyte.TriggerTime` is always in UTC

The `flyte.TriggerTime` value is always in UTC. For timezone-aware logic, convert as needed:

```
@env.task(triggers=flyte.Trigger.minutely(trigger_time_input_key="utc_trigger_time", name="timezone_trigger"))
def timezone_task(utc_trigger_time: datetime) -> str:
    local_time = utc_trigger_time.replace(tzinfo=timezone.utc).astimezone()
    return f"Task fired at {utc_trigger_time} UTC ({local_time} local)"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/triggers/triggers.py*

### Daylight savings time behavior

When Daylight Savings Time (DST) begins and ends, it can impact when the scheduled execution begins.

On the day DST begins, time jumps from 2:00AM to 3:00AM, which means the time of 2:30AM won't exist. In this case, the trigger will not fire until the next 2:30AM, which is the next day.

On the day DST ends, the hour from 1:00AM to 2:00AM repeats, which means the time of 1:30AM will exist twice. If the schedule above was instead set for 1:30AM, it would only run once, on the first occurrence of 1:30AM.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/interruptible-tasks-and-queues ===

# Interruptible tasks

Cloud providers offer discounted compute instances (AWS Spot Instances, GCP Preemptible VMs)
that can be reclaimed at any time. These instances are significantly cheaper than on-demand
instances but come with the risk of preemption.

Setting `interruptible=True` allows Flyte to schedule the task on these spot/preemptible instances
for cost savings:

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    interruptible=True,
)

@env.task
def train_model(data: list) -> dict:
    return {"accuracy": 0.95}
```

## Setting at different levels

`interruptible` can be set at the `TaskEnvironment` level, the `@env.task` decorator level,
and at the `task.override()` invocation level. The more specific level always takes precedence.

This lets you set a default at the environment level and override per-task:

```python
import flyte

# All tasks in this environment are interruptible by default
env = flyte.TaskEnvironment(
    name="my_env",
    interruptible=True,
)

# This task uses the environment default (interruptible)
@env.task
def preprocess(data: list) -> list:
    return [x * 2 for x in data]

# This task overrides to non-interruptible (critical, should not be preempted)
@env.task(interruptible=False)
def save_results(results: dict) -> str:
    return "saved"
```

You can also override at invocation time:

```python
@env.task
async def main(data: list) -> str:
    processed = preprocess(data=data)
    # Run this specific invocation as non-interruptible
    return save_results.override(interruptible=False)(results={"data": processed})
```

## Behavior on preemption

When a spot instance is reclaimed, the task is terminated and rescheduled.
Combine `interruptible=True` with [retries](./retries-and-timeouts) to handle preemptions gracefully:

```python
@env.task(interruptible=True, retries=3)
def train_model(data: list) -> dict:
    return {"accuracy": 0.95}
```

> [!NOTE]
> Retries due to spot preemption do not count against the user-configured retry budget.
> System retries (for preemptions and other system-level failures) are tracked separately.
> This is independent from how the medium (spot vs. on-demand) is chosen for each user-budget attempt. See **Tasks > Configure tasks > Interruptible tasks > Spot to on-demand fallback** below.

## Spot to on-demand fallback

By default, the **final attempt** of an interruptible task runs on an on-demand instance rather than on spot.
This shields long-running or expensive workloads from getting stuck in a preemption loop: once you're on your last attempt, the platform stops gambling on spot capacity.

Two consequences worth knowing:

- A task with `interruptible=True` and no retries is effectively on-demand. Its single attempt is also its final attempt, so it never runs on spot.
- A task with `interruptible=True, retries=2` makes up to 3 attempts: attempts 1 and 2 run on spot, attempt 3 (if reached) runs on on-demand.

If you want spot pricing to apply to *every* attempt, you need to size `retries` so the workload realistically completes before the final attempt. There's no way to opt out of the on-demand fallback on the last attempt.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/task-plugins ===

# Task plugins

Flyte tasks are pluggable by design, allowing you to extend task execution beyond simple containers to support specialized compute frameworks and integrations.

## Default execution: containers

By default, Flyte tasks execute as single containers in Kubernetes. When you decorate a function with `@env.task`, Flyte packages your code into a container and runs it on the cluster. For more advanced scenarios requiring multiple containers in a single pod (such as sidecars for logging or data mounting), you can use [pod templates](./pod-templates), which allow you to customize the entire Kubernetes pod specification.

## Compute plugins

Beyond native container execution, Flyte provides **compute plugins** that enable you to run distributed computing frameworks directly on Kubernetes. These plugins create ephemeral clusters specifically for your task execution, spinning them up on-demand and tearing them down when complete.

### Available compute plugins

Flyte supports several popular distributed computing frameworks through compute plugins:

- **Spark**: Run Apache Spark jobs using the Spark operator
- **Ray**: Execute Ray workloads for distributed Python applications and ML training
- **Dask**: Scale Python workflows with Dask distributed
- **PyTorch**: Run distributed training jobs using PyTorch and Kubeflow's training operator

### How compute plugins work

Compute plugins create temporary, isolated clusters within the same Kubernetes environment as Flyte:

1. **Ephemeral clusters**: Each task execution gets its own cluster, spun up on-demand
2. **Kubernetes operators**: Flyte uses specialized Kubernetes operators (Spark operator, Ray operator, etc.) to manage cluster lifecycle
3. **Native containerization**: The same container image system used for regular tasks works with compute plugins
4. **Per-environment configuration**: You can define the cluster shape (number of workers, resources, etc.) using `plugin_config` in your `TaskEnvironment`

### Using compute plugins

To use a compute plugin, you need to:

1. **Install the plugin package**: Each plugin has a corresponding Python package (e.g., `flyteplugins-ray` for Ray)
2. **Configure the TaskEnvironment**: Set the `plugin_config` parameter with the plugin-specific configuration
3. **Write your task**: Use the framework's native APIs within your task function

#### Example: Ray plugin

Here's how to run a distributed Ray task:

```python
import ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
import flyte

# Define your Ray computation
@ray.remote
def compute_square(x):
    return x * x

# Configure the Ray cluster
ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-workers", replicas=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

# Create a task environment with Ray plugin configuration
image = (
    flyte.Image.from_debian_base(name="ray")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray")
)

ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")),
)

# Use the Ray cluster in your task
@ray_env.task
async def distributed_compute(n: int = 10) -> list[int]:
    futures = [compute_square.remote(i) for i in range(n)]
    return ray.get(futures)
```

When this task runs, Flyte will:

1. Spin up a Ray cluster with 1 head node and 2 worker nodes
2. Execute your task code in the Ray cluster
3. Tear down the cluster after completion

### Using plugins on Union

Most compute plugins are enabled by default on Union or can be enabled upon request. Contact your Account Manager to confirm plugin availability or request specific plugins for your deployment.

## Backend integrations

Beyond compute plugins, Flyte also supports **integrations** with external SaaS services and internal systems through **connectors**. These allow you to interact with:

- **Data warehouses**: Snowflake, BigQuery, Redshift
- **Data platforms**: Databricks
- **Custom services**: Your internal APIs and services

Connectors enable Flyte to delegate task execution to these external systems while maintaining Flyte's orchestration, observability, and data lineage capabilities. See the [connectors documentation](../../../integrations/_index) for more details on available integrations.

## Next steps

For detailed guides on each compute plugin, including configuration options, best practices, and advanced examples, see the [Plugins section](../../../integrations/_index) of the documentation. Each plugin guide covers:

- Installation and setup
- Configuration options
- Resource management
- Advanced use cases
- Troubleshooting tips


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/additional-task-settings ===

# Additional task settings

This page covers task configuration parameters that do not have their own dedicated page:
naming and metadata, default inputs, environment variables, and inline I/O thresholds.

For the full list of all task configuration parameters, see [Configure tasks](./_index).

## Naming and metadata

### `name`

The `name` parameter on `TaskEnvironment` is required.
It is combined with each task function name to form the fully-qualified task name.
For example, if you define a `TaskEnvironment` with `name="my_env"` and a task function `my_task`,
the fully-qualified task name is `my_env.my_task`.

The `name` must use `snake_case` or `kebab-case` and is immutable once set.

### `short_name`

The `short_name` parameter on `@env.task` (and `override()`) overrides the display name of a task in the UI graph view.
By default, the display name is the Python function name.
Overriding `short_name` does not change the fully-qualified task name.

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task(short_name="Train Model")
def train(data: list) -> dict:
    return {"accuracy": 0.95}
```

### `description`

The `description` parameter on `TaskEnvironment` provides a description of the task environment (max 255 characters).
It is used for organizational purposes and can be viewed in the UI.

### `docs`

The `docs` parameter on `@env.task` accepts a `Documentation` object.
If not set explicitly, the documentation is auto-extracted from the task function's docstring.

```python
import flyte
from flyte import Documentation

env = flyte.TaskEnvironment(name="my_env")

@env.task(docs=Documentation(description="Trains a model on the given dataset."))
def train(data: list) -> dict:
    """This docstring is used if docs is not set explicitly."""
    return {"accuracy": 0.95}
```

### `report`

The `report` parameter on `@env.task` controls whether an HTML report is generated for the task.
See [Reports](../task-programming/reports) for details.

### Source-code link (automatic)

In addition to the description and docs above, Flyte automatically attaches a link from each deployed task back to its source code on GitHub or GitLab, when you deploy from inside a checked-out git repository.
The link is rendered next to the task description in the UI.
There is no parameter to set. See [Source-code link discovery](../task-deployment/how-task-deployment-works#6-source-code-link-discovery) for the conditions and caveats.

### `links`

The `links` parameter on `@env.task` (and `override()`) attaches clickable URLs to tasks in the UI.
Use links to connect tasks to external tools like experiment trackers, monitoring dashboards, or logging systems.

Links are defined by implementing the [`Link`](../../../api-reference/flyte-sdk/flyte/link) protocol.
See [Links](../task-programming/links) for full details on creating and using links.

## Default inputs

Task functions support Python default parameter values. When a task parameter has a default, callers can omit it and the default is used.

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
async def process(data: list, batch_size: int = 32, verbose: bool = False) -> dict:
    # batch_size defaults to 32, verbose defaults to False
    ...
```

When running via `flyte run`, parameters with defaults are optional:

```bash
# Uses defaults for batch_size and verbose
flyte run my_file.py process --data '[1, 2, 3]'

# Override a default
flyte run my_file.py process --data '[1, 2, 3]' --batch-size 64
```

When invoking programmatically, Python's normal default argument rules apply:

```python
result = flyte.run(process, data=[1, 2, 3])          # batch_size=32, verbose=False
result = flyte.run(process, data=[1, 2, 3], batch_size=64)  # override
```

Defaults are part of the task's input schema and are visible in the UI when viewing the task.

## Environment variables

The `env_vars` parameter on `TaskEnvironment` injects plain-text environment variables into the task container.
It accepts a `Dict[str, str]`.

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    env_vars={
        "LOG_LEVEL": "DEBUG",
        "API_ENDPOINT": "https://api.example.com",
    },
)

@env.task
def my_task() -> str:
    import os
    return os.environ["API_ENDPOINT"]
```

Environment variables can be overridden at the `task.override()` invocation level
(unless `reusable` is in effect).

Use `env_vars` for non-sensitive configuration values.
For sensitive values like API keys and credentials, use [`secrets`](./secrets) instead.

## Inline I/O threshold

The `max_inline_io_bytes` parameter on `@env.task` (and `override()`) controls the maximum
size for data passed directly in the task request and response
(e.g., primitives, strings, dictionaries).

Data exceeding this threshold raises an `InlineIOMaxBytesBreached` error.

The default value is 10 MiB (`10 * 1024 * 1024` bytes).

This setting does **not** affect [`flyte.io.File`, `flyte.io.Dir`](../task-programming/files-and-directories),
or [`flyte.DataFrame`](../task-programming/dataclasses-and-structures),
which are always offloaded to object storage regardless of size.

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

# Allow up to 50 MiB of inline data
@env.task(max_inline_io_bytes=50 * 1024 * 1024)
def process_large_dict(data: dict) -> dict:
    return {k: v * 2 for k, v in data.items()}
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/logging ===

# Logging

Flyte uses two separate loggers, each with its own level:

- The **framework logger** (`flyte`): Flyte's own internal messages. Lines are prefixed with `[flyte]`. Defaults to `WARNING`.
- The **user logger** (`flyte.user`): the logger you get from `flyte.logger`, for your own application messages. Defaults to `INFO`.

Splitting the two lets you turn up Flyte's internal logging for debugging without flooding your output with framework noise during normal use, and vice versa.

## Set the logging level with environment variables

The quickest way to set the level is via environment variables. They apply to both local and remote runs.

| Variable | Controls | Default |
|----------|----------|---------|
| `LOG_LEVEL` | The framework logger (`flyte`) | `WARNING` |
| `USER_LOG_LEVEL` | The user logger (`flyte.user`) | `INFO` |
| `LOG_FORMAT` | Output format: `console` or `json` | `console` |
| `DISABLE_RICH_LOGGING` | If set (to any value), disables the Rich-formatted console handler | *unset* |
| `FLYTE_RESET_ROOT_LOGGER` | If set to `1`, resets the root logger so third-party library logs are captured as JSON too (see **Tasks > Configure tasks > Logging > JSON logging and third-party libraries**) | *unset* |

For example, to see Flyte's internal debug messages:

```bash
LOG_LEVEL=debug flyte run --local my_workflow.py main
```

Each level variable accepts either a **named level** (`critical`, `error`, `warning` or `warn`, `info`, `debug`, case-insensitive) or a **numeric** Python logging level, such as `10` (`DEBUG`) or `20` (`INFO`):

```bash
LOG_LEVEL=10 USER_LOG_LEVEL=debug flyte run --local my_workflow.py main
```

An unrecognized value falls back to the default for that variable.

## Set the framework verbosity from the CLI

The `flyte` CLI's `-v` flag is a shorthand for the **framework** log level (the `flyte` logger). It's a global option, so it goes *before* the subcommand, and repeating it raises the verbosity:

| Flag | Framework level |
|------|-----------------|
| *(none)* or `-v` | `WARNING` |
| `-vv` | `INFO` |
| `-vvv` | `DEBUG` |

```bash
flyte -vvv run --local my_workflow.py main
```

`-v` controls only the framework logger; it does **not** change your task (user) log level. Set that separately with `--user-log-level` (also a global option, so likewise before the subcommand), or the `USER_LOG_LEVEL` environment variable:

```bash
flyte --user-log-level debug run --local my_workflow.py main
```

## Set the logging level with `flyte.init()`

You can also set logging when you initialize the SDK. The `log_level` and `user_log_level` parameters take numeric Python logging levels:

```python
import logging
import flyte

flyte.init(
    log_level=logging.DEBUG,        # framework logger
    user_log_level=logging.INFO,    # user logger
    log_format="console",           # or "json"
)
```

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `log_level` | `int` | `LOG_LEVEL` env var, else `WARNING` | Level for the framework logger (`flyte`). |
| `user_log_level` | `int` | `USER_LOG_LEVEL` env var, else `INFO` | Level for the user logger (`flyte.user`). |
| `log_format` | `"console"` \| `"json"` | `LOG_FORMAT` env var, else `"console"` | Output format. |
| `reset_root_logger` | `bool` | `False` | When `True`, Flyte clears the root logger's existing handlers and installs its own. When `False` (the default), Flyte leaves your existing root handlers in place and wraps their formatters so third-party log lines also carry the run/action context. |

> [!NOTE]
> Explicit `flyte.init()` arguments take precedence over the environment variables. An argument left unset falls back to the corresponding environment variable, then to the built-in default.

## Write to the user logger

Use `flyte.logger` for your own application messages so they render with the same run/action context as Flyte's output:

```python
import flyte

flyte.logger.info("Processing %d records", len(records))
flyte.logger.debug("Intermediate value: %r", value)
```

## JSON logging and third-party libraries

Setting `LOG_FORMAT=json` (or `log_format="json"`) switches Flyte's output to JSON — but only for the two Flyte loggers, `flyte` and `flyte.user`. Both have propagation disabled, so neither reaches the Python **root logger**. Third-party libraries (`urllib3`, `boto3`, your own dependencies) typically emit through the root logger, so their lines are not converted to JSON: your Flyte output is JSON, but library log lines stay plain text.

This matters when you ship logs to a cloud logging backend — Google Cloud Logging (Stackdriver), AWS CloudWatch, Datadog — that expects one consistent JSON format across every line.

To make third-party lines JSON as well, also set `reset_root_logger=True`. Flyte then clears the root logger's existing handlers and attaches a JSON handler at the root, so every library that propagates to the root logger is captured as JSON:

```python
import flyte

flyte.with_runcontext(
    log_format="json",
    reset_root_logger=True,
).run(main)
```

By default (`reset_root_logger=False`), Flyte does **not** clear the root logger; it only wraps your existing root handlers so their lines carry the run/action context. The root reset is what makes third-party output uniform JSON.

`reset_root_logger` is also available on `flyte.init()`.

### Set it at registration with `FLYTE_RESET_ROOT_LOGGER`

Passing `reset_root_logger=True` to `flyte.with_runcontext()` automatically injects `FLYTE_RESET_ROOT_LOGGER=1` into the task container, so the reset happens in the remote task as well as locally.

To configure it on the task environment instead — resetting the root logger for every run, at registration time rather than per-run — set the environment variable directly. The value must be the string `1`:

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    env_vars={
        "LOG_FORMAT": "json",
        "FLYTE_RESET_ROOT_LOGGER": "1",
    },
)
```

## Related

The same logging parameters can be set per-run via `flyte.with_runcontext()`. See the [Logging parameters](../task-deployment/run-context#logging) in the run context reference.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/overrides ===

# Overrides

Most task configuration is set when you define a task: on the `TaskEnvironment` or in the `@env.task` decorator.
But you often need to change some of that configuration for a **single invocation** of a task: give one call more memory, point it at a different secret, or bump its retries.

The `task.override()` method does exactly this.
It returns a **new task** with the specified parameters changed, leaving the original task definition untouched, so you can invoke the overridden task in place of the original:

```python
import flyte

env = flyte.TaskEnvironment(
    name="training",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

@env.task
async def train(data: str) -> str:
    return f"trained on {data}"

@env.task
async def main() -> str:
    # Invoke train with its environment-level resources.
    baseline = await train("small.csv")

    # Invoke train with overridden resources for this call only.
    heavy = await train.override(
        resources=flyte.Resources(cpu="4", memory="24Gi"),
    )("large.csv")

    return heavy
```

The key idiom is `task.override(...)(args)`: `override()` returns a callable task, which you then invoke with the task's arguments. Note the two sets of parentheses.

## What you can override

`override()` accepts the parameters that are settable at the task-invocation level:

| Parameter | Details |
|-----------|---------|
| **short_name** | [Additional task settings](./additional-task-settings) |
| **resources** | [Resources](./resources) &bull; [`Resources` API ref](../../../api-reference/flyte-sdk/flyte/resources) |
| **cache** | [Caching](./caching) &bull; [`Cache` API ref](../../../api-reference/flyte-sdk/flyte/cache) |
| **retries** | [Retries and timeouts](./retries-and-timeouts) &bull; [`RetryStrategy` API ref](../../../api-reference/flyte-sdk/flyte/retrystrategy) |
| **timeout** | [Retries and timeouts](./retries-and-timeouts) &bull; [`Timeout` API ref](../../../api-reference/flyte-sdk/flyte/timeout) |
| **reusable** | [Reusable containers](./reusable-containers) &bull; [`ReusePolicy` API ref](../../../api-reference/flyte-sdk/flyte/reusepolicy) |
| **env_vars** | [Additional task settings](./additional-task-settings#environment-variables) |
| **secrets** | **Tasks > Configure tasks > Overrides > Overriding secrets** &bull; [Secrets](./secrets) &bull; [`Secret` API ref](../../../api-reference/flyte-sdk/flyte/secret) |
| **max_inline_io_bytes** | [Additional task settings](./additional-task-settings#inline-io-threshold) |
| **pod_template** | [Pod templates](./pod-templates) &bull; [`PodTemplate` API ref](../../../api-reference/flyte-sdk/flyte/podtemplate) |
| **interruptible** | [Interruptible tasks](./interruptible-tasks-and-queues) |
| **links** | [Additional task settings](./additional-task-settings#links) |

For the full parameter interaction matrix showing which parameters can be set at which level, see [Task configuration levels](./_index#task-configuration-levels).

> [!NOTE] Overrides replace, they don't merge
> When you override a collection-valued parameter such as `resources`, `env_vars`, or `secrets`, the value you pass **replaces** the environment's value for that invocation. It is not merged with it.
> To keep some of the original entries, include them in the override.

## What you cannot override

`name`, `image`, `docs`, and the task's interface (its input and output types) **cannot** be overridden. Attempting to override them raises an error.
To run a task with a different image, define it in a separate `TaskEnvironment` (see [Container images](./container-images) and [Multiple environments](./multiple-environments)).

## Overriding when a task is reusable

When a task uses a [reusable container](./reusable-containers) (`reusable` is set), its `resources`, `env_vars`, and `secrets` come from the parent environment and **cannot** be overridden while reuse is active. The container is already running.

To override any of these on a reusable task, turn reuse off in the same `override()` call by passing `reusable="off"`:

```python
result = await my_task.override(
    reusable="off",
    resources=flyte.Resources(cpu="4", memory="8Gi"),
)(data)
```

## Overriding secrets

Just as you can override resources per invocation, you can override the [secrets](./secrets) injected into a task for a single call.
This is useful when the same task needs different credentials depending on how it's invoked: for example, calling an external API with a different key per tenant, or supplying a secret that the task's environment doesn't declare.

Pass `secrets` to `override()` exactly as you would to the `TaskEnvironment`: a secret key, a `Secret` object, or a list of either.

```python
import flyte
from flyte import Secret

env = flyte.TaskEnvironment(
    name="model_calls",
    secrets=Secret("openai-key", as_env_var="LLM_API_KEY"),
)

@env.task
async def call_model(prompt: str) -> str:
    import os
    api_key = os.environ["LLM_API_KEY"]
    ...  # call the model using api_key
    return "response"

@env.task
async def main() -> str:
    # Use the environment's default secret.
    default = await call_model("hello")

    # Override the secret for this invocation, mounting a different
    # store key into the same LLM_API_KEY environment variable.
    alternate = await call_model.override(
        secrets=Secret("anthropic-key", as_env_var="LLM_API_KEY"),
    )("hello")

    return alternate
```

> [!NOTE]
> If the task's environment uses **reusable containers** (`reusable` is set), overriding `secrets`, like `resources` and `env_vars`, requires passing `reusable="off"` in the **same** `override()` call. Otherwise the override is rejected.

As with the environment-level `secrets` parameter, the secret is injected at runtime and accessed inside the task: typically as an environment variable via `os.environ` (or mounted as a file).
See [Secrets](./secrets) for how to create secrets and how they are injected, and remember that the overriding `secrets` value **replaces** the environment's secrets for that invocation rather than adding to them.

> [!NOTE]
> A task can only access a secret if the secret's scope includes the project and domain where the task's `TaskEnvironment` is deployed. Overriding the secret at invocation time does not change this.

> [!WARNING]
> Do not return secret values from tasks. Returned values are stored in plaintext in your data plane's object store and shown in the UI and to downstream tasks, defeating the secret store's protections.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming ===

# Build tasks

This section covers the essential programming patterns and techniques for developing robust Flyte workflows. Once you understand the basics of task configuration, these guides will help you build sophisticated, production-ready data pipelines and machine learning workflows.

## What you'll learn

The task programming section covers key patterns for building effective Flyte workflows:

**Data handling and types**

- **Tasks > Build tasks > Files and directories**: Work with large datasets using Flyte's efficient file and directory types that automatically handle data upload, storage, and transfer between tasks.

- **Tasks > Build tasks > DataFrames**: Pass DataFrames between tasks without downloading data into memory, with support for Pandas, Polars, PyArrow, Dask, and other DataFrame backends.
- **Tasks > Build tasks > Data classes and structures**: Use Python data classes and Pydantic models as task inputs and outputs to create well-structured, type-safe workflows.
- **Tasks > Build tasks > Custom context**: Use custom context to pass metadata through your task execution hierarchy without adding parameters to every task.

**Execution patterns**

- **Tasks > Build tasks > Fanout**: Scale your workflows by running many tasks in parallel, perfect for processing large datasets or running hyperparameter sweeps.
- **Tasks > Build tasks > Mapping over inputs**: Apply the same task to every item of a list with `flyte.map`: in-order results, error handling, concurrency limits, and partials.
- **Tasks > Build tasks > Consuming a message queue**: Pull messages from an external queue such as AWS SQS and fan out processing across a pool of reusable containers.
- **Tasks > Build tasks > Controlling parallel execution**: Limit concurrent task executions using semaphores or `flyte.map` concurrency for rate-limited APIs, GPU quotas, and resource-constrained workflows.
- **Tasks > Build tasks > Streaming map-reduce**: Process fanout results as they complete with `asyncio.as_completed`, reducing in batches incrementally instead of waiting for every task to finish.
- **Tasks > Build tasks > Task dependencies and ordering**: Replicate DAG-like behavior (sequencing, fan-out, fan-in, and fine-grained dependency-driven scheduling) using `asyncio` in Flyte 2's implicit dependency model.
- **Tasks > Build tasks > Structured concurrency with anyio**: Use `anyio` task groups as a top-level structured-concurrency alternative to raw `asyncio`, with automatic sibling cancellation when one task fails.
- **Tasks > Build tasks > External conditions**: Pause a task until an external signal arrives: a human approval, a callback from an external service, or a value supplied at runtime.
- **Tasks > Build tasks > Grouping actions**: Organize related task executions into logical groups for better visualization and management in the UI.
- **Tasks > Build tasks > Run a bioinformatics tool**: Run arbitrary containers in any language without the Flyte SDK installed, using Flyte's copilot sidecar for data flow.
- **Tasks > Build tasks > Remote tasks**: Use previously deployed tasks without importing their code or dependencies, enabling team collaboration and task reuse.
- **Tasks > Configure tasks > Pod templates**: Extend tasks with Kubernetes pod templates to add sidecars, volume mounts, and advanced Kubernetes configurations.
- **Tasks > Build tasks > Abort and cancel actions**: Stop in-progress actions automatically, programmatically, or manually via the CLI and UI.
- **Tasks > Build tasks > Regular async function (not a task)**: Advanced patterns like task forwarding and other specialized task execution techniques.
- **Tasks > Build tasks > Higher-order functions**: Write reusable functions that take tasks as arguments — fallback, memory-scaling retry, circuit breaker, and batch map-reduce wrappers built on Flyte's dynamic execution.

**Development and debugging**

- **Tasks > Build tasks > Notebooks**: Write and iterate on workflows directly in Jupyter notebooks for interactive development and experimentation.
- **Tasks > Build tasks > Test business logic directly**: Test your Flyte tasks using direct invocation for business logic or `flyte.run()` for Flyte-specific features.
- **Tasks > Build tasks > Links**: Add clickable URLs to tasks in the Flyte UI, connecting them to external tools like experiment trackers and monitoring dashboards.
- **Tasks > Build tasks > Reports**: Generate custom HTML reports during task execution to display progress, results, and visualizations in the UI.
- **Tasks > Build tasks > Traces**: Add fine-grained observability to helper functions within your tasks for better debugging and resumption capabilities.
- **Tasks > Build tasks > Intra-task checkpoints**: Save in-progress state within a task (such as a training loop) so retries resume from the last checkpoint instead of starting over.
- **Tasks > Build tasks > Error handling**: Implement robust error recovery strategies, including automatic resource scaling and graceful failure handling.

## When to use these patterns

These programming patterns become essential as your workflows grow in complexity:

- Use **fanout** when you need to process multiple items concurrently or run parameter sweeps.
- Use **mapping over inputs** to apply the same task to every item of a list, and **controlling parallel execution** when you need to limit how many run at the same time.
- Apply **streaming map-reduce** when map tasks have uneven durations or you want to reduce results in batches as they complete, rather than waiting for the entire fanout to finish.
- Implement **error handling** for production workflows that need to recover from infrastructure failures.
- Apply **grouping** to organize complex workflows with many task executions.
- Use **files and directories** when working with large datasets that don't fit in memory.
- Use **DataFrames** to efficiently pass tabular data between tasks across different processing engines.
- Choose **container tasks** when you need to run code in non-Python languages, use legacy containers, or execute AI-generated code in sandboxes.
- Use **remote tasks** to reuse tasks deployed by other teams without managing their dependencies.
- Apply **pod templates** when you need advanced Kubernetes features like sidecars or specialized storage configurations.
- Use **traces** to debug non-deterministic operations like API calls or ML inference.
- Use **intra-task checkpoints** to make long-running training loops resumable across retries, preemptions, and interruptions.
- Use **links** to connect tasks to external tools like Weights & Biases, Grafana, or custom dashboards directly from the Flyte UI.
- Create **reports** to monitor long-running workflows and share results with stakeholders.
- Use **custom context** when you need lightweight, cross-cutting metadata to flow through your task hierarchy without becoming part of the task's logical inputs.
- Write **unit tests** to validate your task logic and ensure type transformations work correctly before deployment.
- Use **abort and cancel** to stop unnecessary actions when conditions change, such as early convergence in HPO or manual intervention.
- Use **external conditions** to insert approval gates or data collection checkpoints into automated workflows.
- Apply **higher-order functions** to factor recurring orchestration logic — retry-on-OOM, fallback, circuit breaking, batching — into reusable wrappers that work with any task.

Each guide includes practical examples and best practices to help you implement these patterns effectively in your own workflows.

## Subpages
- **Tasks > Build tasks > Files and directories**
- **Tasks > Build tasks > Data classes and structures**
- **Tasks > Build tasks > DataFrames**
- **Tasks > Build tasks > Custom types**
- **Tasks > Build tasks > Custom context**
- **Tasks > Build tasks > Abort and cancel actions**
- **Tasks > Build tasks > Run a bioinformatics tool**
- **Tasks > Build tasks > Links**
- **Tasks > Build tasks > Reports**
- **Tasks > Build tasks > Notebooks**
- **Tasks > Build tasks > Remote tasks**
- **Tasks > Build tasks > Error handling**
- **Tasks > Build tasks > Traces**
- **Tasks > Build tasks > Intra-task checkpoints**
- **Tasks > Build tasks > Grouping actions**
- **Tasks > Build tasks > Fanout**
- **Tasks > Build tasks > Mapping over inputs**
- **Tasks > Build tasks > Consuming a message queue**
- **Tasks > Build tasks > Controlling parallel execution**
- **Tasks > Build tasks > Streaming map-reduce**
- **Tasks > Build tasks > Structured concurrency with anyio**
- **Tasks > Build tasks > Task dependencies and ordering**
- **Tasks > Build tasks > External conditions**
- **Tasks > Build tasks > Test business logic directly**
- **Tasks > Build tasks > Higher-order functions**
- **Tasks > Build tasks > Regular async function (not a task)**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/files-and-directories ===

# Files and directories

Flyte provides the `flyte.io.File` and
`flyte.io.Dir` types to represent files and directories, respectively.
Together with [`flyte.io.DataFrame`](./dataframes) they constitute the *offloaded data types* - unlike [materialized types](./dataclasses-and-structures) like data classes, these pass references rather than full data content.

A variable of an offloaded type does not contain its actual data, but rather a reference to the data.
The actual data is stored in the internal blob store of your Union/Flyte instance.
When a variable of an offloaded type is first created, its data is uploaded to the blob store.
It can then be passed from task to task as a reference.
The actual data is only downloaded from the blob stored when the task needs to access it, for example, when the task calls `open()` on a `File` or `Dir` object.

This allows Flyte to efficiently handle large files and directories without needing to transfer the data unnecessarily.
Even very large data objects like video files and DNA datasets can be passed efficiently between tasks.

For the full picture of what gets stored in the bucket versus what stays in the control plane database, see [Where your data lives](../../get-started/core-concepts/where-data-lives).

The `File` and `Dir` classes provide both synchronous and asynchronous methods to interact with the data, so you can use them from either kind of task. See **Tasks > Build tasks > Files and directories > Synchronous and asynchronous APIs** for the full method pairing.

> [!NOTE]
> Because `File` and `Dir` are passed by reference, a downstream cached task does not get a cache hit on identical content stored at a new path. To cache on content, attach a hash at production time - see [Content-based caching for DataFrames, files, and directories](../task-configuration/caching#content-based-caching-for-dataframes-files-and-directories).

## Example usage

The examples below show the basic use-cases of uploading files and directories created locally, and using them as inputs to a task.

```
import asyncio
import tempfile
from pathlib import Path

import flyte
from flyte.io import Dir, File

env = flyte.TaskEnvironment(name="files-and-folders")

@env.task
async def write_file(name: str) -> File:

    # Create a file and write some content to it
    with open("test.txt", "w") as f:
        f.write(f"hello world {name}")

    # Upload the file using flyte
    uploaded_file_obj = await File.from_local("test.txt")
    return uploaded_file_obj
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

The upload happens when the [`File.from_local`](../../../api-reference/flyte-sdk/flyte.io/file#from_local) command is called.
Because the upload would otherwise block execution, `File.from_local` is implemented as an `async` function.
The Flyte SDK frequently uses this class constructor pattern, so you will see it with other types as well.

This is a slightly more complicated task that calls the task above to produce `File` objects.

These are assembled into a directory and the `Dir` object is returned, also via invoking `from_local`.

```
@env.task
async def write_and_check_files() -> Dir:
    coros = []
    for name in ["Alice", "Bob", "Eve"]:
        coros.append(write_file(name=name))

    vals = await asyncio.gather(*coros)
    temp_dir = tempfile.mkdtemp()
    for file in vals:
        async with file.open("rb") as fh:
            contents = await fh.read()
            # Convert bytes to string
            contents_str = contents.decode('utf-8') if isinstance(contents, bytes) else str(contents)
            print(f"File {file.path} contents: {contents_str}")
            new_file = Path(temp_dir) / file.name
            with open(new_file, "w") as out:  # noqa: ASYNC230
                out.write(contents_str)
    print(f"Files written to {temp_dir}")

    # walk the directory and ls
    for path in Path(temp_dir).iterdir():
        print(f"File: {path.name}")

    my_dir = await Dir.from_local(temp_dir)
    return my_dir
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

Finally, these tasks show how to use an offloaded type as an input.
Helper functions like `walk` and `open` have been added to the objects
and do what you might expect.

```
@env.task
async def check_dir(my_dir: Dir):
    print(f"Dir {my_dir.path} contents:")
    async for file in my_dir.walk():
        print(f"File: {file.name}")
        async with file.open("rb") as fh:
            contents = await fh.read()
            # Convert bytes to string
            contents_str = contents.decode('utf-8') if isinstance(contents, bytes) else str(contents)
            print(f"Contents: {contents_str}")

@env.task
async def create_and_check_dir():
    my_dir = await write_and_check_files()
    await check_dir(my_dir=my_dir)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(create_and_check_dir)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

## Synchronous and asynchronous APIs

Every I/O operation on `File` and `Dir` comes in two forms, so you can use the offloaded types from both asynchronous and synchronous tasks:

- In an **asynchronous task** (`async def`), use the coroutine methods: `await` the upload, download, and existence calls, use `async with file.open(...)` to stream, and `async for` to walk a `Dir`. This is the pattern shown in **Tasks > Build tasks > Files and directories > Example usage** above.
- In a **synchronous task** (plain `def`), use the `_sync` variants: `File.from_local_sync()`, `file.open_sync()`, `dir.walk_sync()`, and so on. These block until the operation completes.

The two forms are otherwise equivalent — pick the one that matches how your task is defined. A few constructors do no I/O and so have a single form that is used unchanged from either kind of task: `File.new_remote()`, `File.from_existing_remote()`, `Dir.new_remote()`, and `Dir.from_existing_remote()`.

### File methods

| Asynchronous (in `async def` tasks) | Synchronous (in `def` tasks) | Purpose |
|---|---|---|
| `await File.from_local(path)` | `File.from_local_sync(path)` | Upload a local file to the blob store |
| `File.new_remote()` | `File.new_remote()` | Allocate a new remote file to stream into |
| `File.from_existing_remote(uri)` | `File.from_existing_remote(uri)` | Reference a file that already exists remotely |
| `async with file.open(mode) as fh` | `with file.open_sync(mode) as fh` | Open the file as a stream for reading or writing |
| `await file.download()` | `file.download_sync()` | Download the file to the local filesystem |
| `await file.exists()` | `file.exists_sync()` | Check whether the file exists |

### Dir methods

| Asynchronous (in `async def` tasks) | Synchronous (in `def` tasks) | Purpose |
|---|---|---|
| `await Dir.from_local(path)` | `Dir.from_local_sync(path)` | Upload a local directory to the blob store |
| `Dir.new_remote()` | `Dir.new_remote()` | Allocate a new remote directory to stream into |
| `Dir.from_existing_remote(uri)` | `Dir.from_existing_remote(uri)` | Reference a directory that already exists remotely |
| `async for f in dir.walk()` | `for f in dir.walk_sync()` | Iterate over the files in the directory |
| `await dir.list_files()` | `dir.list_files_sync()` | List the files in the directory (non-recursive) |
| `await dir.get_file(name)` | `dir.get_file_sync(name)` | Get a single file from the directory by name |
| `await dir.download()` | `dir.download_sync()` | Download the whole directory to the local filesystem |
| `await dir.exists()` | `dir.exists_sync()` | Check whether the directory exists |

> [!NOTE]
> `walk_sync()` additionally accepts a `file_pattern` glob (for example `file_pattern="*.txt"`) to filter the files it yields. Both forms accept `recursive` and `max_depth`.

### Synchronous example

The **Tasks > Build tasks > Files and directories > Example usage** uses `await`, `async with`, and `async for`. The same kind of workflow written with the synchronous API uses plain `def` tasks and the `_sync` method names:

```python
import flyte
from flyte.io import File

env = flyte.TaskEnvironment(name="sync-file")

@env.task
def write_file(content: str) -> File:
    # Allocate a new remote file and stream content into it
    f = File.new_remote()
    with f.open_sync("wb") as fh:
        fh.write(content.encode("utf-8"))
    return f

@env.task
def read_file(f: File) -> str:
    # Open the file for reading without downloading the whole object
    with f.open_sync("rb") as fh:
        return fh.read().decode("utf-8")

@env.task
def main() -> str:
    f = write_file(content="hello world")
    return read_file(f)
```

Directories work the same way — use `Dir.from_local_sync()` to upload and `walk_sync()` to iterate:

```python
import os
import tempfile

import flyte
from flyte.io import Dir

env = flyte.TaskEnvironment(name="sync-dir")

@env.task
def upload_dir() -> Dir:
    with tempfile.TemporaryDirectory() as tmp:
        for i in range(3):
            with open(os.path.join(tmp, f"file{i}.txt"), "w") as fh:
                fh.write(f"content {i}")
        # Upload the directory to the blob store
        return Dir.from_local_sync(tmp)

@env.task
def read_dir(d: Dir) -> int:
    count = 0
    # Walk and read every file, all synchronously
    for file in d.walk_sync(recursive=True):
        with file.open_sync("rb") as fh:
            print(f"{file.name}: {fh.read().decode('utf-8')}")
        count += 1
    return count

@env.task
def main() -> int:
    d = upload_dir()
    return read_dir(d)
```

## JSONL files

Flyte provides typed JSON Lines (JSONL) I/O through the `flyteplugins-jsonl` plugin, which extends `File` and `Dir` with the `JsonlFile` and `JsonlDir` types, adding streaming record-level read/write, optional zstd compression, and automatic shard rotation for large datasets.

See the [JSONL integration](../../../integrations/jsonl/_index) guide for installation and usage.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/dataclasses-and-structures ===

# Data classes and structures

Dataclasses and Pydantic models are fully supported in Flyte as **materialized data types**:
Structured data where the full content is serialized and passed between tasks.
Use these as you would normally, passing them as inputs and outputs of tasks.

Unlike **offloaded types** like [`DataFrame`s](./dataframes), [`File`s and `Dir`s](./files-and-directories), data class and Pydantic model data is fully serialized, stored, and deserialized between tasks.
This makes them ideal for configuration objects, metadata, and smaller structured data where all fields should be serializable.

## Example: Combining dataclasses and Pydantic models

This example demonstrates how data classes and Pydantic models work together as materialized data types, showing nested structures and batch processing patterns:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pydantic",
# ]
# main = "main"
# params = ""
# ///

import asyncio
from dataclasses import dataclass
from typing import List

from pydantic import BaseModel
import flyte

env = flyte.TaskEnvironment(name="ex-mixed-structures")

@dataclass
class InferenceRequest:
    feature_a: float
    feature_b: float

@dataclass
class BatchRequest:
    requests: List[InferenceRequest]
    batch_id: str = "default"

class PredictionSummary(BaseModel):
    predictions: List[float]
    average: float
    count: int
    batch_id: str

@env.task
async def predict_one(request: InferenceRequest) -> float:
    """
    A dummy linear model: prediction = 2 * feature_a + 3 * feature_b + bias(=1.0)
    """
    return 2.0 * request.feature_a + 3.0 * request.feature_b + 1.0

@env.task
async def process_batch(batch: BatchRequest) -> PredictionSummary:
    """
    Processes a batch of inference requests and returns summary statistics.
    """
    # Process all requests concurrently
    tasks = [predict_one(request=req) for req in batch.requests]
    predictions = await asyncio.gather(*tasks)

    # Calculate statistics
    average = sum(predictions) / len(predictions) if predictions else 0.0

    return PredictionSummary(
        predictions=predictions,
        average=average,
        count=len(predictions),
        batch_id=batch.batch_id
    )

@env.task
async def summarize_results(summary: PredictionSummary) -> str:
    """
    Creates a text summary from the prediction results.
    """
    return (
        f"Batch {summary.batch_id}: "
        f"Processed {summary.count} predictions, "
        f"average value: {summary.average:.2f}"
    )

@env.task
async def main() -> str:
    batch = BatchRequest(
        requests=[
            InferenceRequest(feature_a=1.0, feature_b=2.0),
            InferenceRequest(feature_a=3.0, feature_b=4.0),
            InferenceRequest(feature_a=5.0, feature_b=6.0),
        ],
        batch_id="demo_batch_001"
    )
    summary = await process_batch(batch)
    result = await summarize_results(summary)
    return result

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataclasses-and-structures/example.py*


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/dataframes ===

# DataFrames

By default, return values in Python are materialized - meaning the actual data is downloaded and loaded into memory. This applies to simple types like integers, as well as more complex types like DataFrames.

To avoid downloading large datasets into memory, Flyte V2 exposes [`flyte.io.dataframe`](../../../api-reference/flyte-sdk/flyte.io/dataframe): a thin,  uniform wrapper type for DataFrame-style objects that allows you to pass a reference to the data, rather than the fully materialized contents.

The `flyte.io.DataFrame` type provides serialization support for common engines like `pandas`, `polars`, `pyarrow`, `dask`, etc.; enabling you to move data between different DataFrame backends.

DataFrame contents are written to the data plane object store and passed between tasks by reference. For the full map of what goes in the bucket versus what stays in the control plane database, see [Where your data lives](../../get-started/core-concepts/where-data-lives).

> [!NOTE]
> Because a DataFrame is passed by reference, a downstream cached task does not get a cache hit on identical content stored at a new path. To cache on content, attach a hash with `flyte.io.HashFunction` - see [Content-based caching for DataFrames, files, and directories](../task-configuration/caching#content-based-caching-for-dataframes-files-and-directories).

## Setting up the environment and sample data

For our example we will start by setting up our task environment with the required dependencies and create some sample data.

```
from typing import Annotated

import numpy as np
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    "dataframe_usage",
    image= flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow", "numpy"),
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

BASIC_EMPLOYEE_DATA = {
    "employee_id": range(1001, 1009),
    "name": ["Alice", "Bob", "Charlie", "Diana", "Ethan", "Fiona", "George", "Hannah"],
    "department": ["HR", "Engineering", "Engineering", "Marketing", "Finance", "Finance", "HR", "Engineering"],
    "hire_date": pd.to_datetime(
        ["2018-01-15", "2019-03-22", "2020-07-10", "2017-11-01", "2021-06-05", "2018-09-13", "2022-01-07", "2020-12-30"]
    ),
}

ADDL_EMPLOYEE_DATA = {
    "employee_id": range(1001, 1009),
    "salary": [55000, 75000, 72000, 50000, 68000, 70000, np.nan, 80000],
    "bonus_pct": [0.05, 0.10, 0.07, 0.04, np.nan, 0.08, 0.03, 0.09],
    "full_time": [True, True, True, False, True, True, False, True],
    "projects": [
        ["Recruiting", "Onboarding"],
        ["Platform", "API"],
        ["API", "Data Pipeline"],
        ["SEO", "Ads"],
        ["Budget", "Forecasting"],
        ["Auditing"],
        [],
        ["Platform", "Security", "Data Pipeline"],
    ],
}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

## Create a raw DataFrame

Now, let's create a task that returns a native Pandas DataFrame:

```
@env.task
async def create_raw_dataframe() -> pd.DataFrame:
    return pd.DataFrame(BASIC_EMPLOYEE_DATA)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

This is the most basic use-case of how to pass DataFrames (of all kinds, not just Pandas).
We simply create the DataFrame as normal, and return it.

Because the task has been declared to return a supported native DataFrame type (in this case `pandas.DataFrame` Flyte will automatically detect it, serialize it correctly and upload it at task completion enabling it to be passed transparently to the next task.

Flyte supports auto-serialization for the following DataFrame types:

* `pandas.DataFrame`
* `pyarrow.Table`
* `dask.dataframe.DataFrame`
* `polars.DataFrame`
* `flyte.io.DataFrame` (see below)

## Create a flyte.io.DataFrame

Alternatively you can also create a `flyte.io.DataFrame` object directly from a native object with the `from_df` method:

```
@env.task
async def create_flyte_dataframe() -> Annotated[flyte.io.DataFrame, "parquet"]:
    pd_df = pd.DataFrame(ADDL_EMPLOYEE_DATA)
    fdf = flyte.io.DataFrame.from_df(pd_df)
    return fdf
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

The `flyte.io.DataFrame` class creates a thin wrapper around objects of any standard DataFrame type. It serves as a generic "any DataFrame type" (a concept that Python itself does not currently offer).

As with native DataFrame types, Flyte will automatically serialize and upload the data at task completion.

The advantage of the unified `flyte.io.DataFrame` wrapper is that you can be explicit about the storage format that makes sense for your use case, by using an `Annotated` type where the second argument encodes format or other lightweight hints. For example, here we specify that the DataFrame should be stored as Parquet:

## Automatically convert between types

You can use Flyte to automatically download and convert the DataFrame between types when needed:

```
@env.task
async def join_data(raw_dataframe: pd.DataFrame, flyte_dataframe: pd.DataFrame) -> flyte.io.DataFrame:
    joined_df = raw_dataframe.merge(flyte_dataframe, on="employee_id", how="inner")
    return flyte.io.DataFrame.from_df(joined_df)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

This task takes two DataFrames as input. We'll pass one raw Pandas DataFrame, and one `flyte.io.DataFrame`.
Flyte automatically converts the `flyte.io.DataFrame` to a Pandas DataFrame (since we declared that as the input type) before passing it to the task.
The actual download and conversion happens only when we access the data, in this case, when we do the merge.

## Downloading DataFrames

When a task receives a `flyte.io.DataFrame`, you can request a concrete backend representation. For example, to download as a pandas DataFrame:

```
@env.task
async def download_data(joined_df: flyte.io.DataFrame):
    downloaded = await joined_df.open(pd.DataFrame).all()
    print("Downloaded Data:\n", downloaded)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

The `open()` call delegates to the DataFrame handler for the stored format and converts to the requested in-memory type.

## Run the example

Finally, we can define a `main` function to run the tasks defined above and a `__main__` block to execute the workflow:

```
@env.task
async def main():
    raw_df = await create_raw_dataframe ()
    flyte_df = await create_flyte_dataframe ()
    joined_df = await join_data (raw_df, flyte_df)
    await download_data (joined_df)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

## Polars DataFrames

The `flyteplugins-polars` package extends Flyte's DataFrame support to `polars.DataFrame` and `polars.LazyFrame`. Install it alongside the core SDK and it registers automatically. No additional configuration required.

```bash
pip install flyteplugins-polars
```

Both types are serialized as Parquet when passed between tasks, just like other DataFrame backends.

### Setup

```
import polars as pl

import flyte

env = flyte.TaskEnvironment(
    name="polars-dataframes",
    image=flyte.Image.from_debian_base(name="polars").with_pip_packages(
        "flyteplugins-polars>=2.0.0", "polars"
    ),
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

EMPLOYEE_DATA = {
    "employee_id": [1001, 1002, 1003, 1004, 1005, 1006],
    "name": ["Alice", "Bob", "Charlie", "Diana", "Ethan", "Fiona"],
    "department": ["Engineering", "Engineering", "Marketing", "Finance", "Finance", "Engineering"],
    "salary": [75000, 72000, 50000, 68000, 70000, 80000],
    "years_experience": [5, 4, 2, 6, 5, 7],
}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

### Eager DataFrames

Use `pl.DataFrame` when you want immediate evaluation. Flyte serializes it to Parquet on output and deserializes it on input:

```
@env.task
async def create_dataframe() -> pl.DataFrame:
    """Create a Polars DataFrame.

    Polars DataFrames are passed between tasks as serialized Parquet files
    stored in the Flyte blob store — no manual upload required.
    """
    return pl.DataFrame(EMPLOYEE_DATA)

@env.task
async def filter_high_earners(df: pl.DataFrame) -> pl.DataFrame:
    """Filter and enrich a Polars DataFrame."""
    return (
        df.filter(pl.col("salary") > 60000)
        .with_columns(
            (pl.col("salary") / pl.col("years_experience")).alias("salary_per_year")
        )
        .sort("salary", descending=True)
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

### Lazy DataFrames

Use `pl.LazyFrame` when you want to defer computation and let Polars optimize the full query plan before executing. Flyte handles serialization the same way as `pl.DataFrame`:

```
@env.task
async def create_lazyframe() -> pl.LazyFrame:
    """Create a Polars LazyFrame.

    LazyFrames defer computation until collected, allowing Polars to
    optimize the full query plan. They are serialized to Parquet just
    like DataFrames when passed between tasks.
    """
    return pl.LazyFrame(EMPLOYEE_DATA)

@env.task
async def aggregate_by_department(lf: pl.LazyFrame) -> pl.DataFrame:
    """Aggregate salary statistics by department using a LazyFrame.

    The query plan is built lazily and executed only when collect() is called.
    """
    return (
        lf.group_by("department")
        .agg(
            pl.col("salary").mean().alias("avg_salary"),
            pl.col("salary").max().alias("max_salary"),
            pl.len().alias("headcount"),
        )
        .sort("avg_salary", descending=True)
        .collect()
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

The `collect()` call in `aggregate_by_department` is what triggers execution of the lazy plan. The `LazyFrame` passed between tasks is serialized as Parquet at that point.

### Run the example

```
@env.task
async def main():
    df = await create_dataframe()
    filtered = await filter_high_earners(df=df)
    print("High earners:")
    print(filtered)

    lf = await create_lazyframe()
    summary = await aggregate_by_department(lf=lf)
    print("Department summary:")
    print(summary)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

## See also

To display a DataFrame as an HTML table in a task report, define a `flyte.types.Renderable` for it — see [Rendering a custom type](./reports#rendering-a-custom-type) on the Reports page.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/handling-custom-types ===

# Custom types

Flyte has a rich type system that handles most Python types automatically. However, there are cases where you may want to pass custom types into a run or between actions. By default, if Flyte doesn't recognize a type, it uses Python pickle to serialize the data. While this works, pickle has several drawbacks:

- **Inefficiency**: Pickle can be very inefficient for certain data types
- **Language compatibility**: Pickle is Python-specific and doesn't work with other languages
- **Version fragility**: Pickled data can break between Python versions
- **Opacity**: Pickled data appears as bytes or file links in the UI, with no automatic form generation

Consider types like Polars DataFrames or PyTorch Tensors. Using pickle for these is extremely inefficient compared to native serialization formats like Parquet or tensor-specific formats.

Flyte SDK addresses this by allowing you to create and share type extensions.

## Types of extensions

Flyte supports two types of type extensions:

1. **Type transformers**: For scalar types (integers, strings, files, directories, custom objects)
2. **DataFrame extensions**: For tabular data types that benefit from DataFrame-specific handling

DataFrame types are special because they have associated metadata (columns, schemas), can be serialized to efficient formats like Parquet, support parallel uploads from engines like Spark, and can be partitioned.

## Creating a type transformer

Type transformers convert between Python types and Flyte's internal representation. Here's how to create one for a custom `PositiveInt` type.

### Step 1: Define your custom type

```python
# custom_type.py
class PositiveInt:
    """A wrapper type that only accepts positive integers."""

    def __init__(self, value: int):
        if not isinstance(value, int):
            raise TypeError(f"Expected int, got {type(value).__name__}")
        if value <= 0:
            raise ValueError(f"Expected positive integer, got {value}")
        self._value = value

    @property
    def value(self) -> int:
        return self._value

    def __repr__(self) -> str:
        return f"PositiveInt({self._value})"
```

### Step 2: Create the type transformer

```python
# transformer.py
from typing import Type

from flyteidl2.core import literals_pb2, types_pb2

from flyte import logger
from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
from my_transformer.custom_type import PositiveInt

class PositiveIntTransformer(TypeTransformer[PositiveInt]):
    """
    Type transformer for PositiveInt that validates and transforms positive integers.
    """

    def __init__(self):
        super().__init__(name="PositiveInt", t=PositiveInt)

    def get_literal_type(self, t: Type[PositiveInt]) -> types_pb2.LiteralType:
        """Returns the Flyte literal type for PositiveInt."""
        return types_pb2.LiteralType(
            simple=types_pb2.SimpleType.INTEGER,
            structure=types_pb2.TypeStructure(tag="PositiveInt"),
        )

    async def to_literal(
        self,
        python_val: PositiveInt,
        python_type: Type[PositiveInt],
        expected: types_pb2.LiteralType,
    ) -> literals_pb2.Literal:
        """Converts a PositiveInt instance to a Flyte Literal."""
        if not isinstance(python_val, PositiveInt):
            raise TypeTransformerFailedError(
                f"Expected PositiveInt, got {type(python_val).__name__}"
            )

        return literals_pb2.Literal(
            scalar=literals_pb2.Scalar(
                primitive=literals_pb2.Primitive(integer=python_val.value)
            )
        )

    async def to_python_value(
        self,
        lv: literals_pb2.Literal,
        expected_python_type: Type[PositiveInt]
    ) -> PositiveInt:
        """Converts a Flyte Literal back to a PositiveInt instance."""
        if not lv.scalar or not lv.scalar.primitive:
            raise TypeTransformerFailedError(
                f"Cannot convert literal {lv} to PositiveInt: missing scalar primitive"
            )

        value = lv.scalar.primitive.integer
        try:
            return PositiveInt(value)
        except (TypeError, ValueError) as e:
            raise TypeTransformerFailedError(
                f"Cannot convert value {value} to PositiveInt: {e}"
            )

    def guess_python_type(
        self,
        literal_type: types_pb2.LiteralType
    ) -> Type[PositiveInt]:
        """Guesses the Python type from a Flyte literal type."""
        if (
            literal_type.simple == types_pb2.SimpleType.INTEGER
            and literal_type.structure
            and literal_type.structure.tag == "PositiveInt"
        ):
            return PositiveInt
        raise ValueError(f"Cannot guess PositiveInt from literal type {literal_type}")
```

### Step 3: Register the transformer

Create a registration function that can be called to register your transformer:

```python
def register_positive_int_transformer():
    """Register the PositiveIntTransformer in the TypeEngine."""
    TypeEngine.register(PositiveIntTransformer())
    logger.info("Registered PositiveIntTransformer in TypeEngine")
```

## Distributing type plugins

To share your type transformer as an installable package, configure it as a Flyte plugin using entry points.

### Configure pyproject.toml

Add the entry point to your `pyproject.toml`:

```toml
[project]
name = "my_transformer"
version = "0.1.0"
description = "Custom type transformer"
requires-python = ">=3.10"
dependencies = []

[project.entry-points."flyte.plugins.types"]
my_transformer = "my_transformer.transformer:register_positive_int_transformer"
```

The entry point group `flyte.plugins.types` tells Flyte to automatically load this transformer when the package is installed.

### Automatic loading

When your plugin package is installed, Flyte automatically loads the type transformer at runtime. This happens during `flyte.init()` or `flyte.init_from_config()`.

## Controlling plugin loading

Loading many type plugins can add overhead to initialization. You can disable automatic plugin loading:

```python
import flyte

# Disable automatic loading of type transformer plugins
flyte.init(load_plugin_type_transformers=False)
```

By default, `load_plugin_type_transformers` is `True`.

## Using custom types in tasks

Once registered, use your custom type like any built-in type:

```python
import flyte
from my_transformer.custom_type import PositiveInt

env = flyte.TaskEnvironment(name="custom_types")

@env.task
async def process_positive(value: PositiveInt) -> int:
    """Process a positive integer."""
    return value.value * 2

if __name__ == "__main__":
    flyte.init_from_config()

    # The custom type works seamlessly
    run = flyte.run(process_positive, value=PositiveInt(42))
    run.wait()
    print(run.outputs()[0])  # 84
```

## DataFrame extensions

For tabular data types, Flyte provides a specialized extension mechanism through `flyte.io.DataFrame`. DataFrame extensions support:

- Automatic conversion to/from Parquet format
- Column metadata and schema information
- Parallel uploads from distributed engines
- Partitioning support

DataFrame extensions use encoders and decoders from `flyte.io.extend`. Documentation for creating DataFrame extensions is coming soon.

## Best practices

1. **Use specific types over pickle**: Define type transformers for any custom types used frequently in your workflows
2. **Keep transformers lightweight**: Avoid expensive operations in `to_literal` and `to_python_value`
3. **Add validation**: Validate data in your transformer to catch errors early
4. **Use meaningful tags**: The `TypeStructure.tag` helps identify your type in the Flyte UI
5. **Be judicious with plugins**: Only install the plugins you need to minimize initialization overhead

## See also

To render a custom type as HTML in a task report, define a `flyte.types.Renderable` for it — see [Rendering a custom type](./reports#rendering-a-custom-type) on the Reports page.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/custom-context ===

# Custom context

Custom context provides a mechanism for implicitly passing configuration and metadata through your entire task execution hierarchy without adding parameters to every task. It is ideal for cross-cutting concerns such as tracing, environment metadata, or experiment identifiers.

Think of custom context as **execution-scoped metadata** that automatically flows from parent to child tasks.

> **📝 Note**
>
> In Flyte 1 the runtime context was accessed with `current_context()` (the `flytekit.current_context` API). Flyte 2 uses `flyte.ctx()`, shown below.

## Overview

Custom context is an implicit key-value configuration map that is automatically available to tasks during execution. It is stored in the blob store of your Union/Flyte instance together with the task’s inputs, making it available across tasks without needing to pass it explicitly.

You can access it in a Flyte task via:

```python
flyte.ctx().custom_context
```

Custom context is fundamentally different from standard task inputs. Task inputs are explicit, strongly typed parameters that you declare as part of a task’s signature. They directly influence the task’s computation and therefore participate in Flyte’s caching and reproducibility guarantees.

Custom context, on the other hand, is implicit metadata. It consists only of string key/value pairs, is not part of the task signature, and does not affect task caching. Because it is injected by the Flyte runtime rather than passed as a formal input, it should be used only for environmental or contextual information, not for data that changes the logical output of a task.

## When to use it and when not to

Custom context is perfect when you need metadata, not domain data, to flow through your tasks.

Good use cases:

- Tracing IDs, span IDs
- Experiment or run metadata
- Environment region, cluster ID
- Logging correlation keys
- Feature flags
- Session IDs for 3rd-party APIs (e.g., an LLM session)

Avoid using for:

- Business/domain data
- Inputs that change task outputs
- Anything affecting caching or reproducibility
- Large blobs of data (keep it small)

It is the cleanest mechanism when you need something available everywhere, but not logically an input to the computation.

## Setting custom context

There are two ways to set custom context for a Flyte run:

1. Set it once for the entire run when you launch (`with_runcontext`): this establishes the base context for the execution
2. Set or override it inside task code using `flyte.custom_context(...)` context manager: this changes the active context for that task block and any nested tasks called from it

Both are legitimate and complementary. The important behavioral rules to understand are:

- `with_runcontext(...)` sets the run-level base. Values provided here are available everywhere unless overridden later. Use this for metadata that should apply to most or all tasks in the run (experiment name, top-level trace id, run id, etc.).
- `flyte.custom_context(...)` is used inside task code to set or override values for that scope. It does affect nested tasks invoked while that context is active. In practice this means you can override run-level entries, add new keys for downstream tasks, or both.
- Merging & precedence: contexts are merged; when the same key appears in multiple places the most recent/innermost value wins (i.e., values set by `flyte.custom_context(...)` override the run-level values from `with_runcontext(...)` for the duration of that block).

### Run-level context

Set base metadata once when starting the run:

```
import flyte

env = flyte.TaskEnvironment("custom-context-example")

@env.task
async def leaf_task() -> str:
    # Reads run-level context
    print("leaf sees:", flyte.ctx().custom_context)
    return flyte.ctx().custom_context.get("trace_id")

@env.task
async def root() -> str:
    return await leaf_task()

if __name__ == "__main__":
    flyte.init_from_config()
    # Base context for the entire run
    flyte.with_runcontext(custom_context={"trace_id": "root-abc", "experiment": "v1"}).run(root)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/custom-context/run_context.py*

Output (every task sees the base keys unless overridden):

```bash
leaf sees: {"trace_id": "root-abc", "experiment": "v1"}
```

### Overriding inside a task (local override that affects nested tasks)

Use `flyte.custom_context(...)` inside a task to override or add keys for downstream calls:

```
@env.task
async def downstream() -> str:
    print("downstream sees:", flyte.ctx().custom_context)
    return flyte.ctx().custom_context.get("trace_id")

@env.task
async def parent() -> str:
    print("parent initial:", flyte.ctx().custom_context)

    # Override the trace_id for the nested call(s)
    with flyte.custom_context(trace_id="child-override"):
        val = await downstream()     # downstream sees trace_id="child-override"

    # After the context block, run-level values are back
    print("parent after:", flyte.ctx().custom_context)
    return val
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/custom-context/override_context.py*

If the run was started with `{"trace_id": "root-abc"}`, this prints:

```bash
parent initial: {"trace_id": "root-abc"}
downstream sees: {"trace_id": "child-override"}
parent after: {"trace_id": "root-abc"}
```

Note that the override affected the nested downstream task because it was invoked while the `flyte.custom_context` block was active.

### Adding new keys for nested tasks

You can add keys (not just override):

```python
with flyte.custom_context(experiment="exp-blue", run_group="g-7"):
    await some_task()   # some_task sees both base keys + the new keys
```

## Accessing custom context

Always via the Flyte runtime:

```python
ctx = flyte.ctx().custom_context
value = ctx.get("key")
```

You can access the custom context using either `flyte.ctx().custom_context` or the shorthand `flyte.get_custom_context()`, which returns the same dictionary of key/value pairs.

Values are always strings, so parse as needed:

```python
timeout = int(ctx["timeout_seconds"])
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/abort-tasks ===

# Abort and cancel actions

When running complex workflows, you may need to stop actions that are no longer needed.
This can happen when one branch of your workflow makes others redundant, when a task fails and its siblings should not continue, or when you need to manually intervene in a running workflow.

Flyte provides three mechanisms for stopping actions:

- **Automatic cleanup**: When a root action completes, all its in-progress descendant actions are automatically aborted.
- **Programmatic cancellation**: Cancel specific `asyncio` tasks from within your workflow code.
- **External abort**: Stop individual actions via the CLI, the UI, or the API.

For background on runs and actions, see [Runs and actions](../../get-started/core-concepts/runs-and-actions).

## Action lifetime

The lifetime of all actions in a [run](../../get-started/core-concepts/runs-and-actions) is tied to the lifetime of the root action (the first task that was invoked).
When the root action exits (whether it succeeds, fails, or returns early) all in-progress descendant actions are automatically aborted and no new actions can be enqueued.

This means you don't need to manually clean up child actions. Flyte handles it for you.

Consider this example where `main` exits after 10 seconds, but it has spawned a `sleep_for` action that is set to run for 30 seconds:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "seconds = 30"
# ///

import asyncio

import flyte

env = flyte.TaskEnvironment(name="action_lifetime")

@env.task
async def do_something():
    print("Doing something")
    await asyncio.sleep(5)
    print("Finished doing something")

@env.task
async def sleep_for(seconds: int):
    print(f"Sleeping for {seconds} seconds")
    try:
        await asyncio.sleep(seconds)
        await do_something()
    except asyncio.CancelledError:
        print("sleep_for was cancelled")
        return
    print(f"Finished sleeping for {seconds} seconds")

@env.task
async def main(seconds: int):
    print("Starting main")
    asyncio.create_task(sleep_for(seconds))
    await asyncio.sleep(10)
    print("Main finished")

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, seconds=30)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/action_lifetime.py*

When `main` returns after 10 seconds, the `sleep_for` action (which still has 20 seconds remaining) is automatically aborted.
The `sleep_for` task receives an `asyncio.CancelledError`, giving it a chance to handle the cancellation gracefully.

## Canceling actions programmatically

As a workflow author, you can cancel specific in-progress actions by canceling their corresponding `asyncio` tasks.
This is useful in scenarios like hyperparameter optimization (HPO), where one action converges to the desired result and the remaining actions can be stopped to save compute.

To cancel actions programmatically:

1. Launch actions using `asyncio.create_task()` and retain references to the returned task objects.
2. When the desired condition is met, call `.cancel()` on the tasks you want to stop.

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "n = 30, f = 10.0"
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment("cancel")

@env.task
async def sleepers(f: float, n: int):
    await asyncio.sleep(f)

@env.task
async def failing_task(f: float):
    raise ValueError("I will fail!")

@env.task
async def main(n: int, f: float):
    sleeping_tasks = []
    for i in range(n):
        sleeping_tasks.append(asyncio.create_task(sleepers(f, i)))

    await asyncio.sleep(f)
    try:
        await failing_task(f)
        await asyncio.gather(*sleeping_tasks)
    except flyte.errors.RuntimeUserError as e:
        if e.code == "ValueError":
            print(f"Received ValueError, canceling {len(sleeping_tasks)} sleeping tasks")
            for t in sleeping_tasks:
                t.cancel()
        return

if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(main, 30, 10.0))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/cancel_tasks.py*

In this code:

* The `main` task launches 30 `sleepers` actions in parallel using `asyncio.create_task()`.
* It then calls `failing_task`, which raises a `ValueError`.
* The error is caught as a `flyte.errors.RuntimeUserError` (since user-raised exceptions are wrapped by Flyte).
* On catching the error, `main` cancels all sleeping tasks by calling `.cancel()` on each one, freeing their compute resources.

This pattern lets you react to runtime conditions and stop unnecessary work. For more on handling errors within workflows, see [Error handling](./error-handling).

## External abort

Sometimes you need to stop an action manually, outside the workflow code itself. You can abort individual actions using the CLI, the UI, or the API.

When an action is externally aborted, the parent action that awaits it receives a `flyte.errors.ActionAbortedError`. You can catch this error to handle the abort gracefully.

### Aborting via the CLI

To abort a specific action:

```bash
flyte abort <run-name> <action-name>
```

Use `--project` and `--domain` to target a specific [project-domain pair](../../get-started/core-concepts/projects-and-domains).
For all available options, see the [CLI reference](../../../api-reference/flyte-cli#flyte-abort).

### Handling external aborts

When using `asyncio.gather()` with `return_exceptions=True`, externally aborted actions return an `ActionAbortedError` instead of raising it. This lets you inspect results and handle aborts on a per-action basis:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "n = 10, sleep_for = 30.0"
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment("external_abort")

@env.task
async def long_sleeper(sleep_for: float):
    await asyncio.sleep(sleep_for)

@env.task
async def main(n: int, sleep_for: float) -> str:
    coros = [long_sleeper(sleep_for) for _ in range(n)]
    results = await asyncio.gather(*coros, return_exceptions=True)
    for i, r in enumerate(results):
        if isinstance(r, flyte.errors.ActionAbortedError):
            print(f"Action [{i}] was externally aborted")
    return "Hello World!"

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, 10, 30.0)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/external_abort.py*

In this code:

* The `main` task launches 10 `long_sleeper` actions in parallel.
* If any action is externally aborted (via the CLI, the UI, or the API) while running, `asyncio.gather` captures the `ActionAbortedError` as a result instead of propagating it.
* The `main` task iterates over the results and logs which actions were aborted.
* Because the abort is handled, `main` can continue executing and return its result normally.

Without `return_exceptions=True`, an external abort would raise `ActionAbortedError` directly, which you can handle with a standard `try...except` block.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/container-tasks ===

Container tasks are one of Flyte's superpowers. They allow you to execute tasks using any container image without requiring the Flyte SDK to be installed in that container. This means you can run code written in any language, execute shell scripts, or even use pre-built containers pulled directly from the internet while still maintaining Flyte's data orchestration capabilities.

## What are container tasks?

A container task is a special type of Flyte task that executes arbitrary container images. Unlike standard `@task` decorated functions that require the Flyte SDK, container tasks can run:

- Code written in any programming language (Rust, Go, Java, R, etc.)
- Legacy containers with unsupported Python versions
- Pre-built bioinformatics or scientific computing containers
- Shell scripts and command-line tools
- Dynamically generated code in sandboxed environments

## How data flows in and out

The magic of container tasks lies in Flyte's **copilot sidecar system**. When you execute a container task, Flyte:

1. Launches your specified container alongside a copilot sidecar container
2. Uses shared Kubernetes pod volumes to pass data between containers
3. Reads inputs from `input_data_dir` and writes outputs to `output_data_dir`
4. Automatically handles serialization and deserialization of typed data

This means you can construct workflows where some tasks are container tasks while others are Python functions, and data will flow between them.

## Basic usage

Here's a simple example that runs a shell command in an Alpine container:

```python
import flyte
from flyte.extras import ContainerTask

greeting_task = ContainerTask(
    name="echo_and_return_greeting",
    image=flyte.Image.from_base("alpine:3.18"),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"name": str},
    outputs={"greeting": str},
    command=[
        "/bin/sh",
        "-c",
        "echo 'Hello, my name is {{.inputs.name}}.' | tee -a /var/outputs/greeting"
    ],
)
```

### Template syntax for inputs

Container tasks support template-style references to inputs using the syntax `{{.inputs.<input_name>}}`. This gets replaced with the actual input value at runtime:

```python
command=["/bin/sh", "-c", "echo 'Processing {{.inputs.user_id}}' > /var/outputs/result"]
```

### Using container tasks in workflows

Container tasks integrate with Python tasks:

```python
container_env = flyte.TaskEnvironment.from_task("container_env", greeting_task)
env = flyte.TaskEnvironment(name="hello_world", depends_on=[container_env])

@env.task
async def say_hello(name: str = "flyte") -> str:
    print("Hello container task")
    return await greeting_task(name=name)
```

## Advanced: Passing files and directories

Container tasks can accept `File` and `Dir` inputs. For these types, use path-based syntax (not template syntax) in your commands:

```python
from flyte.io import File
import pathlib

code_runner = ContainerTask(
    name="python_code_runner",
    image="ghcr.io/astral-sh/uv:debian-slim",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script.py": File, "a": int, "b": int},
    outputs={"result": int},
    command=[
        "/bin/sh",
        "-c",
        "uv run /var/inputs/script.py {{.inputs.a}} {{.inputs.b}} > /var/outputs/result"
    ],
)

@env.task
async def execute_script() -> int:
    path = pathlib.Path(__file__).parent / "my_script.py"
    script_file = await File.from_local(path)
    return await code_runner(**{"script.py": script_file, "a": 10, "b": 20})
```

Note that when passing files, the input key can include the filename (e.g., `"script.py"`), and you reference it in the command as `/var/inputs/script.py`.

## Use case: Agentic sandbox execution

Container tasks are perfect for running AI-generated code in isolated environments. You can generate a data analysis script dynamically and execute it safely:

```python
import flyte
from flyte.extras import ContainerTask
from flyte.io import File
import pathlib

env = flyte.TaskEnvironment(name="agentic_sandbox")

@env.task
async def run_generated_code(script_content: str, param_a: int, param_b: int) -> int:
    # Define a container task that runs arbitrary Python code
    sandbox = ContainerTask(
        name="code_sandbox",
        image="ghcr.io/astral-sh/uv:debian-slim",
        input_data_dir="/var/inputs",
        output_data_dir="/var/outputs",
        inputs={"script": File, "a": int, "b": int},
        outputs={"result": int},
        command=[
            "/bin/sh",
            "-c",
            "uv run --script /var/inputs/script {{.inputs.a}} {{.inputs.b}} > /var/outputs/result"
        ],
    )

    # Save the generated script to a temporary file
    temp_path = pathlib.Path("/tmp/generated_script.py")
    temp_path.write_text(script_content)

    # Execute it in the sandbox
    script_file = await File.from_local(temp_path)
    return await sandbox(script=script_file, a=param_a, b=param_b)
```

This pattern allows you to:

- Generate code using LLMs or other AI systems
- Execute it in a controlled, isolated environment
- Capture results and integrate them back into your workflow
- Maintain full observability and reproducibility

## Use case: Legacy and specialized containers

Many scientific and bioinformatics tools are distributed as pre-built containers. Container tasks let you integrate them directly:

```python
# Run a bioinformatics tool
blast_task = ContainerTask(
    name="run_blast",
    image="ncbi/blast:latest",
    input_data_dir="/data",
    output_data_dir="/results",
    inputs={"query": File, "database": str},
    outputs={"alignments": File},
    command=[
        "blastn",
        "-query", "/data/query",
        "-db", "{{.inputs.database}}",
        "-out", "/results/alignments",
        "-outfmt", "6"
    ],
)

# Run legacy code with an old Python version
legacy_task = ContainerTask(
    name="legacy_python",
    image="python:2.7",  # Unsupported Python version
    input_data_dir="/app/inputs",
    output_data_dir="/app/outputs",
    inputs={"data_file": File},
    outputs={"processed": File},
    command=[
        "python",
        "/legacy_app/process.py",
        "/app/inputs/data_file",
        "/app/outputs/processed"
    ],
)
```

## Use case: Multi-language workflows

Build workflows that span multiple languages:

```python
# Rust task for high-performance computation
rust_task = ContainerTask(
    name="rust_compute",
    image="rust:1.75",
    inputs={"n": int},
    outputs={"result": int},
    input_data_dir="/inputs",
    output_data_dir="/outputs",
    command=["./compute_binary", "{{.inputs.n}}"],
)

# Python task for orchestration
@env.task
async def multi_lang_workflow(iterations: int) -> dict:
    # Call Rust task for heavy computation
    computed = await rust_task(n=iterations)

    # Process results in Python
    processed = await python_analysis_task(computed)

    return {"rust_result": computed, "analysis": processed}
```

## Configuration options

### ContainerTask parameters

- **name**: Unique identifier for the task
- **image**: Container image to use (string or `Image` object)
- **command**: Command to execute in the container (list of strings)
- **inputs**: Dictionary mapping input names to types
- **outputs**: Dictionary mapping output names to types
- **input_data_dir**: Directory where Flyte writes input data (default: `/var/inputs`)
- **output_data_dir**: Directory where Flyte reads output data (default: `/var/outputs`)
- **arguments**: Additional command arguments (list of strings)
- **metadata_format**: Format for metadata serialization (`"JSON"`, `"YAML"`, or `"PROTO"`)
- **local_logs**: Whether to print container logs during local execution (default: `True`)

### Supported input/output types

Container tasks support all standard Flyte types:

- Primitives: `str`, `int`, `float`, `bool`
- Temporal: `datetime.datetime`, `datetime.timedelta`
- File system: `File`, `Dir`
- Complex types: dataclasses, Pydantic models (serialized as JSON/YAML/PROTO)

## Best practices

1. **Use specific image tags**: Prefer `alpine:3.18` over `alpine:latest` for reproducibility
2. **Keep containers focused**: Each container task should do one thing well
3. **Handle errors gracefully**: Ensure your container commands exit with appropriate status codes
4. **Test locally first**: Container tasks can run locally with Docker, making debugging easier
5. **Consider image size**: Smaller images lead to faster task startup times
6. **Document input/output contracts**: Clearly specify what data flows in and out

## Local execution

Container tasks require Docker to be installed and running on your local machine. When you run them locally, Flyte will:

1. Pull the specified image (if not already available)
2. Mount local directories for inputs and outputs
3. Stream container logs to your console
4. Extract outputs after container completion

This makes it easy to develop and test container tasks before deploying to a remote cluster.

## When to use container tasks

Choose container tasks when you need to:

- Run code in languages other than Python
- Execute pre-built tools or legacy applications
- Isolate potentially unsafe code (AI-generated scripts)
- Use specific runtime environments or dependencies
- Integrate external tools without Python wrappers
- Execute shell scripts or command-line utilities

For Python code that can use the Flyte SDK, standard `@task` decorated functions are usually simpler and more efficient.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/links ===

# Links

Links let you add clickable URLs to tasks that appear in the Flyte UI. Use them to connect tasks to external tools like experiment trackers, monitoring dashboards or custom internal services.

![Links in the Flyte UI](../../../_static/images/integrations/wandb/single_node_auto_flyte.png)

You can attach links to tasks in two ways:

- **Statically** in the task decorator with `links=`
- **Dynamically** at call time with `task.override(links=...)`

`Link` is a Python [Protocol](https://docs.python.org/3/library/typing.html#typing.Protocol) that you subclass to define how URLs are generated. The Weights & Biases plugin provides a [built-in link implementation](../../../api-reference/integrations/wandb/wandb) as an example.

## Creating a link

To create a link, subclass `Link` as a dataclass and implement the `get_link()` method. The method returns the URL string to display in the UI:

```python
from dataclasses import dataclass

import flyte
from flyte import Link

@dataclass
class GrafanaLink(Link):
    dashboard_url: str
    name: str = "Grafana"

    def get_link(
        self,
        run_name: str,
        project: str,
        domain: str,
        context: dict,
        parent_action_name: str,
        action_name: str,
        pod_name: str,
        **kwargs,
    ) -> str:
        return f"{self.dashboard_url}?var-pod={pod_name}"

env = flyte.TaskEnvironment(...)

@env.task(links=(GrafanaLink(dashboard_url="https://grafana.example.com/d/abc123"),))
def my_task() -> str:
    return "done"
```

The link appears as a clickable "Grafana" link in the Flyte UI for every execution of `my_task`.

## Using execution metadata

The `get_link()` method receives execution metadata that you can use to construct dynamic URLs. Here's an example modeled on the [built-in Wandb](../../../integrations/wandb/_index) link that uses the `context` dict to resolve a run ID:

```python
from dataclasses import dataclass
from typing import Optional

from flyte import Link

@dataclass
class Wandb(Link):
    project: str
    entity: str
    id: Optional[str] = None
    name: str = "Weights & Biases"

    def get_link(
        self,
        run_name: str,
        project: str,
        domain: str,
        context: dict[str, str],
        parent_action_name: str,
        action_name: str,
        pod_name: str,
        **kwargs,
    ) -> str:
        run_id = self.id or context.get("wandb_id", run_name)
        return f"https://wandb.ai/{self.entity}/{self.project}/runs/{run_id}"
```

The `name` attribute controls the display label in the UI.

See the [`get_link()` API reference](../../../api-reference/flyte-sdk/flyte/link#get_link) for more details. Note that `action_name` and `pod_name` are template variables (`{{.actionName}}` and `{{.podName}}`) that are populated by the backend at runtime.

## Dynamic links with override

Use `task.override(links=...)` to set links at runtime. This is useful when link parameters depend on runtime values like run IDs or configuration:

```python
import os

import flyte
from flyteplugins.wandb import Wandb

env = flyte.TaskEnvironment(...)

WANDB_PROJECT = "my-ml-project"
WANDB_ENTITY = "my-team"

@env.task
def train_model(config: dict) -> dict:
    # Training logic here
    return {"accuracy": 0.95}

@env.task
async def main(wandb_id: str) -> dict:
    result = train_model.override(
        links=(
            Wandb(
                project=WANDB_PROJECT,
                entity=WANDB_ENTITY,
                id=wandb_id,
            ),
        )
    )(config={"lr": 0.001})

    return result

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, wandb_id="my-run-id")
```

The `override` approach lets you attach links with values that are only known at runtime, such as dynamically generated run IDs.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/reports ===

# Reports

The reports feature allows you to display and update custom output in the UI during task execution.

> **📝 Note**
>
> Reports are the Flyte 2 successor to **Decks** in Flyte 1. Where Flyte 1 used `enable_deck=True` and the `flytekit.Deck` API, Flyte 2 uses `report=True` and the `flyte.report` API described below.

First, you set the `report=True` flag in the task decorator. This enables the reporting feature for that task.
Within a task with reporting enabled, a `flyte.report.Report` object is created automatically.

> [!NOTE] Import `flyte.report` explicitly
> `flyte.report` is a submodule that `import flyte` does **not** import automatically.
> You must import it explicitly:
>
> ```python
> import flyte.report
> ```
>
> Without this, calls like `flyte.report.replace()` or `flyte.report.flush()` raise
> `AttributeError: module 'flyte' has no attribute 'report'`, most commonly hit in local or
> notebook runs. This applies to all `flyte.*` submodules: import the specific submodule you use,
> not just the top-level `flyte` package.

A `Report` object contains one or more tabs, each of which contains HTML.
You can write HTML to an existing tab and create new tabs to organize your content.
Initially, the `Report` object has one tab (the default tab, named `main`) with no content.

To write content:

- `flyte.report.log()` appends HTML content directly to the default tab.
- `flyte.report.replace()` replaces the content of the default tab with new HTML.

To get or create a new tab:

- `flyte.report.get_tab()` allows you to specify a unique name for the tab, and it will return the existing tab if it already exists or create a new one if it doesn't.
  It returns a `flyte.report._report.Tab`

You can `log()` or `replace()` HTML on the `Tab` object just as you can directly on the `Report` object.

To access the current `Report` object directly — for example, to enumerate its tabs or assemble the final HTML — call `flyte.report.current_report()`.

Finally, you send the report to the Flyte backend and make it visible in the UI:

- `flyte.report.flush()` dispatches the report (in its current state) to the backend.

You do **not** have to call `flyte.report.flush()` explicitly at the end of a task: when a task with `report=True` finishes, Flyte automatically performs a final flush for you.
Calling `flyte.report.flush()` yourself is only necessary when you want to *stream* updates to the UI while the task is still running (see **Tasks > Build tasks > Reports > Streaming example** below).

## A simple example

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = ""
# ///

import flyte
import flyte.report

env = flyte.TaskEnvironment(name="reports_example")

@env.task(report=True)
async def task1():
    await flyte.report.replace.aio("<p>The quick, brown fox jumps over a lazy dog.</p>")
    tab2 = flyte.report.get_tab("Tab 2")
    tab2.log("<p>The quick, brown dog jumps over a lazy fox.</p>")
    await flyte.report.flush.aio()

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(task1)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/simple.py*

Here we define a task `task1` that uses `flyte.report.replace()` to set the content of the default tab, then creates a new tab named "Tab 2" with `flyte.report.get_tab()` and logs additional HTML content to it.
Finally, `flyte.report.flush()` is called to send the report to the backend.

## A more complex example

Here is another example.
We import the necessary modules, set up the task environment, define the main task with reporting enabled and define the data generation function:

```
import json
import random

import flyte
import flyte.report

env = flyte.TaskEnvironment(
    name="globe_visualization",
)

@env.task(report=True)
async def generate_globe_visualization():
    await flyte.report.replace.aio(get_html_content())
    await flyte.report.flush.aio()

def generate_globe_data():
    """Generate sample data points for the globe"""
    cities = [
        {"city": "New York", "country": "USA", "lat": 40.7128, "lng": -74.0060},
        {"city": "London", "country": "UK", "lat": 51.5074, "lng": -0.1278},
        {"city": "Tokyo", "country": "Japan", "lat": 35.6762, "lng": 139.6503},
        {"city": "Sydney", "country": "Australia", "lat": -33.8688, "lng": 151.2093},
        {"city": "Paris", "country": "France", "lat": 48.8566, "lng": 2.3522},
        {"city": "São Paulo", "country": "Brazil", "lat": -23.5505, "lng": -46.6333},
        {"city": "Mumbai", "country": "India", "lat": 19.0760, "lng": 72.8777},
        {"city": "Cairo", "country": "Egypt", "lat": 30.0444, "lng": 31.2357},
        {"city": "Moscow", "country": "Russia", "lat": 55.7558, "lng": 37.6176},
        {"city": "Beijing", "country": "China", "lat": 39.9042, "lng": 116.4074},
        {"city": "Lagos", "country": "Nigeria", "lat": 6.5244, "lng": 3.3792},
        {"city": "Mexico City", "country": "Mexico", "lat": 19.4326, "lng": -99.1332},
        {"city": "Bangkok", "country": "Thailand", "lat": 13.7563, "lng": 100.5018},
        {"city": "Istanbul", "country": "Turkey", "lat": 41.0082, "lng": 28.9784},
        {"city": "Buenos Aires", "country": "Argentina", "lat": -34.6118, "lng": -58.3960},
        {"city": "Cape Town", "country": "South Africa", "lat": -33.9249, "lng": 18.4241},
        {"city": "Dubai", "country": "UAE", "lat": 25.2048, "lng": 55.2708},
        {"city": "Singapore", "country": "Singapore", "lat": 1.3521, "lng": 103.8198},
        {"city": "Stockholm", "country": "Sweden", "lat": 59.3293, "lng": 18.0686},
        {"city": "Vancouver", "country": "Canada", "lat": 49.2827, "lng": -123.1207},
    ]

    categories = ["high", "medium", "low", "special"]

    data_points = []
    for city in cities:
        data_point = {**city, "value": random.randint(10, 100), "category": random.choice(categories)}
        data_points.append(data_point)

    return data_points
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py*

We then define the HTML content for the report:

```python
def get_html_content():
    data_points = generate_globe_data()
    html_content = f"""
    <!DOCTYPE html>
    <html lang="en">
    ...
    </html>
    """
    return html_content
```

(We exclude it here due to length. You can find it in the [source file](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py)).

Finally, we run the workflow:

```
if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(generate_globe_visualization)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py*

When the workflow runs, the report will be visible in the UI:

![Globe visualization](../../../_static/images/user-guide/globe_visualization.png)

## Streaming example

Above we demonstrated reports that are sent to the UI once, at the end of the task execution.
But, you can also stream updates to the report during task execution and see the display update in real-time.

You do this by calling `flyte.report.flush()` periodically during task execution, instead of just at the end.
As a shortcut, you can also pass `do_flush=True` to `flyte.report.log()` or `flyte.report.replace()` to flush immediately after writing the content.

> [!NOTE]
> In the earlier examples we explicitly call `flyte.report.flush()` to send the report to the UI.
> As noted above, that final flush is optional: it happens automatically when the task completes.
> For streaming reports, on the other hand, calling `flyte.report.flush()` periodically (or passing `do_flush=True`
> to `flyte.report.log()` / `flyte.report.replace()`) is what makes the intermediate updates appear.

First we import the necessary modules, and set up the task environment:

```
import asyncio
import json
import math
import random
import time
from datetime import datetime
from typing import List

import flyte
import flyte.report

env = flyte.TaskEnvironment(name="streaming_reports")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py*

Next we define the HTML content for the report:

```python
DATA_PROCESSING_DASHBOARD_HTML = """
...
"""
```

(We exclude it here due to length. You can find it in the [source file](
https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py)).

Finally, we define the task that renders the report (`data_processing_dashboard`), the driver task of the workflow (`main`), and the run logic:

```
@env.task(report=True)
async def data_processing_dashboard(total_records: int = 50000) -> str:
    """
    Simulates a data processing pipeline with real-time progress visualization.
    Updates every second for approximately 1 minute.
    """
    await flyte.report.log.aio(DATA_PROCESSING_DASHBOARD_HTML, do_flush=True)

    # Simulate data processing
    processed = 0
    errors = 0
    batch_sizes = [800, 850, 900, 950, 1000, 1050, 1100]  # Variable processing rates

    start_time = time.time()

    while processed < total_records:
        # Simulate variable processing speed
        batch_size = random.choice(batch_sizes)

        # Add some processing delays occasionally
        if random.random() < 0.1:  # 10% chance of slower batch
            batch_size = int(batch_size * 0.6)
            await flyte.report.log.aio("""
            <script>addActivity("⚠️ Detected slow processing batch, optimizing...");</script>
            """, do_flush=True)
        elif random.random() < 0.05:  # 5% chance of error
            errors += random.randint(1, 5)
            await flyte.report.log.aio("""
            <script>addActivity("❌ Processing errors detected, retrying failed records...");</script>
            """, do_flush=True)
        else:
            await flyte.report.log.aio(f"""
            <script>addActivity("✅ Successfully processed batch of {batch_size} records");</script>
            """, do_flush=True)

        processed = min(processed + batch_size, total_records)
        current_time = time.time()
        elapsed = current_time - start_time
        rate = int(batch_size) if elapsed < 1 else int(processed / elapsed)
        success_rate = ((processed - errors) / processed) * 100 if processed > 0 else 100

        # Update dashboard
        await flyte.report.log.aio(f"""
        <script>
            updateDashboard({processed}, {total_records}, {rate}, {success_rate});
        </script>
        """, do_flush=True)

        print(f"Processed {processed:,} records, Errors: {errors}, Rate: {rate:,}"
              f" records/sec, Success Rate: {success_rate:.2f}%", flush=True)
        await asyncio.sleep(1)  # Update every second

        if processed >= total_records:
            break

    # Final completion message
    total_time = time.time() - start_time
    avg_rate = int(total_records / total_time)

    await flyte.report.log.aio(f"""
    <script>addActivity("🎉 Processing completed successfully!");</script>
    <div style="background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; padding: 20px; border-radius: 8px; margin-top: 20px;">
        <h3>🎉 Processing Complete!</h3>
        <ul>
            <li><strong>Total Records:</strong> {total_records:,}</li>
            <li><strong>Processing Time:</strong> {total_time:.1f} seconds</li>
            <li><strong>Average Rate:</strong> {avg_rate:,} records/second</li>
            <li><strong>Success Rate:</strong> {success_rate:.2f}%</li>
            <li><strong>Errors Handled:</strong> {errors}</li>
        </ul>
    </div>
    """, do_flush=True)
    print(f"Data processing completed: {processed:,} records processed with {errors} errors.", flush=True)

    return f"Processed {total_records:,} records successfully"

@env.task
async def main():
    """
    Main task to run both reports.
    """
    await data_processing_dashboard(total_records=50000)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py*

The key to the live update ability is the `while` loop that appends Javascript to the report. The Javascript calls execute on append to the document and update it.

When the workflow runs, you can see the report updating in real-time in the UI:

![Data Processing Dashboard](../../../_static/images/user-guide/data_processing_dashboard.png)

## Rendering a custom type

The examples above build report HTML by hand. When you have a **custom type** — or a DataFrame, or a `StructuredDataset` — that you render the same way in many places, you can define a reusable **renderer** for the type and attach it to the type, instead of repeating the HTML-building logic at every call site.

A renderer is any class that satisfies the `flyte.types.Renderable` protocol: it implements a single `to_html(self, value) -> str` method that returns an HTML fragment for a value of your type.

```python
from flyte.types import Renderable

class Molecule:
    def __init__(self, name: str, smiles: str):
        self.name = name
        self.smiles = smiles

class MoleculeRenderer(Renderable):
    """A Renderable for the Molecule type."""

    def to_html(self, mol: Molecule) -> str:
        return f"<h2>{mol.name}</h2><pre>{mol.smiles}</pre>"
```

You attach the renderer to the type with `typing.Annotated`, then dispatch a value through its attached renderer with `flyte.types.TypeEngine.to_html()`. Log the resulting HTML to the report just like any other content:

```python
from typing import Annotated

import flyte
import flyte.report
from flyte.types import TypeEngine

env = flyte.TaskEnvironment(name="custom_renderer")

# Attaching the renderer to the type is the "registration".
RenderedMolecule = Annotated[Molecule, MoleculeRenderer()]

@env.task(report=True)
async def show_molecule() -> Molecule:
    mol = Molecule("caffeine", "CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

    # Dispatch the value through the renderer attached to RenderedMolecule.
    html = TypeEngine.to_html(mol, RenderedMolecule)
    await flyte.report.log.aio(html)
    await flyte.report.flush.aio()

    return mol

if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(show_molecule).url)
```

`TypeEngine.to_html()` finds the `Renderable` attached to the type via `Annotated`, calls its `to_html()`, and returns the HTML string — which you then send to the report with `flyte.report.log()` (or `replace()`). The same pattern works for a DataFrame or `StructuredDataset`: annotate the type with a renderer that turns the frame into an HTML table.

> [!NOTE]
> The report contains only what you explicitly `log()` or `replace()`. Returning a value whose type has a renderer attached does **not** by itself add it to the report — render the value and log the HTML, as shown above.

Flyte's SDK also implements a few renderers of this kind internally — for pandas and PyArrow DataFrames and for Markdown strings. These aren't exposed as public API (only the `flyte.types.Renderable` protocol is), so treat them as examples of the same pattern rather than importable helpers.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/notebooks ===

# Notebooks

Flyte is designed to work with Jupyter notebooks, allowing you to write and execute workflows directly within a notebook environment.

## Iterating on and running a workflow

Download the following notebook file and open it in your favorite Jupyter environment: [interactive.ipynb](../../../_static/public/interactive.ipynb)

<!-- TODO: add back when working
📥 [interactive.ipynb](/_static/public/interactive.ipynb)
-->

In this example we have a simple workflow defined in our notebook.
You can iterate on the code in the notebook while running each cell in turn.

Note that the `flyte.init()` call at the top of the notebook looks like this:

```python
flyte.init(
    endpoint="https://union.example.com",
    org="example_org",
    project="example_project",
    domain="development",
)
```

You will have to adjust it to match your Union server endpoint, organization, project, and domain.

## Accessing runs and downloading logs

Similarly, you can download the following notebook file and open it in your favorite Jupyter environment: [remote.ipynb](../../../_static/public/remote.ipynb)

<!-- TODO: add back when working
📥 [remote.ipynb](/_static/public/remote.ipynb)
-->

In this example we use the `flyte.remote` package to list existing runs, access them, and download their details and logs.

For a guide on working with runs, actions, inputs, and outputs, see [Interact with runs and actions](../task-deployment/interacting-with-runs).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/remote-tasks ===

# Remote tasks

Remote tasks let you use previously deployed tasks without importing their code or dependencies. This enables teams to share and reuse tasks without managing complex dependency chains or container images.

## Prerequisites

Remote tasks must be deployed before you can use them. See the [task deployment guide](../task-deployment/_index) for details.

## Basic usage

Use `flyte.remote.Task.get()` to reference a deployed task:

```python
import flyte
import flyte.remote

env = flyte.TaskEnvironment(name="my_env")

# Get the latest version of a deployed task
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Use it in your task
@env.task
async def my_task(data_path: str) -> flyte.io.DataFrame:
    # Call the reference task like any other task
    result = await data_processor(input_path=data_path)
    return result
```

You can run this directly without deploying it:

```bash
flyte run my_workflow.py my_task --data_path s3://my-bucket/data.parquet
```

## Understanding lazy loading

Remote tasks use **lazy loading** to keep module imports fast and enable flexible client configuration. When you call `flyte.remote.Task.get()`, it returns a lazy reference that doesn't actually fetch the task from the server until the first invocation.

### When tasks are fetched

The remote task is fetched from the server only when:

- You call `flyte.run()` with the task
- You call `flyte.deploy()` with code that uses the task
- You invoke the task with the `()` operator inside another task
- You explicitly call `.fetch()` on the lazy reference

```python
import flyte.remote

# This does NOT make a network call - returns a lazy reference
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# The task is fetched here when you invoke it
run = flyte.run(data_processor, input_path="s3://my-bucket/data.parquet")
```

### Benefits of lazy loading

**Fast module loading**: Since no network calls are made during import, your Python modules load quickly even when referencing many remote tasks.

**Late binding**: You can call `flyte.init()` after importing remote tasks, and the correct client will be bound when the task is actually invoked:

```python
import flyte
import flyte.remote

# Load remote task reference at module level
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Initialize the client later
flyte.init_from_config()

# The task uses the client configured above
run = flyte.run(data_processor, input_path="s3://data.parquet")
```

### Error handling

Because of lazy loading, if a referenced task doesn't exist, you won't get an error when calling `get()`. Instead, the error occurs during invocation, raising a `flyte.errors.RemoteTaskNotFoundError`:

```python
import flyte
import flyte.remote
import flyte.errors

# This succeeds even if the task doesn't exist
data_processor = flyte.remote.Task.get(
    "nonexistent.task",
    auto_version="latest"
)

try:
    # Error occurs here during invocation
    run = flyte.run(data_processor, input_path="s3://data.parquet")
except flyte.errors.RemoteTaskNotFoundError as e:
    print(f"Task not found or invocation failed: {e}")
    # Handle the error - perhaps use a fallback task
    # or notify the user that the task needs to be deployed
```

You can also catch errors when using remote tasks within other tasks:

```python
import flyte.errors

@env.task
async def pipeline_with_fallback(data_path: str) -> dict:
    try:
        # Try to use the remote task
        result = await data_processor(input_path=data_path)
        return {"status": "success", "result": result}
    except flyte.errors.RemoteTaskNotFoundError as e:
        # Fallback to local processing
        print(f"Remote task failed: {e}, using local fallback")
        return {"status": "fallback", "result": local_process(data_path)}
    except flyte.errors.RemoteTaskUsageError as e:
        raise ValueError(f"Bad Usage of remote task, maybe arguments dont match!")
```

### Eager fetching with `fetch()`

While lazy loading is convenient, you can explicitly fetch a task upfront using the `fetch()` method. This is useful for:

- **Catching errors early**: Validate that the task exists before execution starts
- **Caching**: Avoid the network call on first invocation when running multiple times
- **Service initialization**: Pre-load tasks when your service starts

```python
import flyte
import flyte.remote
import flyte.errors

# Get the lazy reference
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

try:
    # Eagerly fetch the task details
    task_details = data_processor.fetch()

    # Now the task is cached - subsequent calls won't hit the remote service
    # You can pass either the original reference or task_details to flyte.run
    run1 = flyte.run(data_processor, input_path="s3://data1.parquet")
    run2 = flyte.run(task_details, input_path="s3://data2.parquet")

except flyte.errors.RemoteTaskNotFoundError as e:
    print(f"Task not found failed at startup: {e}")
    raise
except flyte.errors.RemoteTaskUsageError as e:
    print(f"Task run validation failed....")
    # Handle the error before any execution attempts
```

For async contexts, use `await fetch.aio()`:

```python
import flyte.remote

async def initialize_service():
    processor_ref = flyte.remote.Task.get(
        "data_team.spark_analyzer",
        auto_version="latest"
    )

    try:
        # Fetch asynchronously
        task_details = await processor_ref.fetch.aio()
        print(f"Task {task_details.name} loaded successfully")
        return processor_ref  # Return the cached reference
    except flyte.errors.RemoteTaskNotFoundError as e:
        print(f"Failed to load task: {e}")
        raise

# Initialize once at service startup
cached_processor = None

async def startup():
    global cached_processor
    cached_processor = await initialize_service()

# Later in your service
async def process_request(data_path: str):
    # The task is already cached from initialization
    # No network call on first invocation
    run = flyte.run(cached_processor, input_path=data_path)
    return run
```

**When to use eager fetching**:

- **Service startup**: Fetch all remote tasks during initialization to validate they exist and cache them
- **Multiple invocations**: If you'll invoke the same task many times, fetch once to cache it
- **Fail-fast validation**: Catch configuration errors before execution begins

**When lazy loading is better**:

- **Single-use tasks**: If you only invoke the task once, lazy loading is simpler
- **Import-time overhead**: Keep imports fast by deferring network calls
- **Conditional usage**: If the task may not be needed, don't fetch it upfront

### Module-level vs dynamic loading

**Module-level loading (recommended)**: Load remote tasks at the module level for cleaner, more maintainable code:

```python
import flyte.remote

# Module-level - clear and maintainable
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

@env.task
async def my_task(data_path: str):
    return await data_processor(input_path=data_path)
```

**Dynamic loading**: You can also load remote tasks dynamically within a task if needed:

```python
@env.task
async def dynamic_pipeline(task_name: str, data_path: str):
    # Load the task based on runtime parameters
    processor = flyte.remote.Task.get(
        f"data_team.{task_name}",
        auto_version="latest"
    )

    try:
        result = await processor(input_path=data_path)
        return result
    except flyte.errors.RemoteTaskNotFoundError as e:
        raise ValueError(f"Task {task_name} not found: {e}")
```

## Complete example

This example shows how different teams can collaborate using remote tasks.

### Team A: Spark environment

Team A maintains Spark-based data processing tasks:

```python
# spark_env.py
from dataclasses import dataclass
import flyte

env = flyte.TaskEnvironment(name="spark_env")

@dataclass
class AnalysisResult:
    mean_value: float
    std_dev: float

@env.task
async def analyze_data(data_path: str) -> AnalysisResult:
    # Spark code here (not shown)
    return AnalysisResult(mean_value=42.5, std_dev=3.2)

@env.task
async def compute_score(result: AnalysisResult) -> float:
    # More Spark processing
    return result.mean_value / result.std_dev
```

Deploy the Spark environment:

```bash
flyte deploy spark_env/
```

### Team B: ML environment

Team B maintains PyTorch-based ML tasks:

```python
# ml_env.py
from pydantic import BaseModel
import flyte

env = flyte.TaskEnvironment(name="ml_env")

class PredictionRequest(BaseModel):
    feature_x: float
    feature_y: float

class Prediction(BaseModel):
    score: float
    confidence: float
    model_version: str

@env.task
async def run_inference(request: PredictionRequest) -> Prediction:
    # PyTorch model inference (not shown)
    return Prediction(
        score=request.feature_x * 2.5,
        confidence=0.95,
        model_version="v2.1"
    )
```

Deploy the ML environment:

```bash
flyte deploy ml_env/
```

### Team C: Orchestration

Team C builds a workflow using remote tasks from both teams without needing Spark or PyTorch dependencies:

```python
# orchestration_env.py
import flyte.remote

env = flyte.TaskEnvironment(name="orchestration")

# Reference tasks from other teams
analyze_data = flyte.remote.Task.get(
    "spark_env.analyze_data",
    auto_version="latest"
)

compute_score = flyte.remote.Task.get(
    "spark_env.compute_score",
    auto_version="latest"
)

run_inference = flyte.remote.Task.get(
    "ml_env.run_inference",
    auto_version="latest"
)

@env.task
async def orchestrate_pipeline(data_path: str) -> float:
    # Use Spark tasks without Spark dependencies
    analysis = await analyze_data(data_path=data_path)

    # Access attributes from the result
    # (Flyte creates a fake type that allows attribute access)
    print(f"Analysis: mean={analysis.mean_value}, std={analysis.std_dev}")

    data_score = await compute_score(result=analysis)

    # Use ML task without PyTorch dependencies
    # Pass Pydantic models as dictionaries
    prediction = await run_inference(
        request={
            "feature_x": analysis.mean_value,
            "feature_y": data_score
        }
    )

    # Access Pydantic model attributes
    print(f"Prediction: {prediction.score} (confidence: {prediction.confidence})")

    return prediction.score
```

Run the orchestration task directly (no deployment needed):

**Using Python API**:

```python
if __name__ == "__main__":
    flyte.init_from_config()

    run = flyte.run(
        orchestrate_pipeline,
        data_path="s3://my-bucket/data.parquet"
    )

    print(f"Execution URL: {run.url}")
    # You can wait for the execution
    run.wait()

    # You can then retrieve the outputs
    print(f"Pipeline result: {run.outputs()}")
```

**Using CLI**:

```bash
flyte run orchestration_env.py orchestrate_pipeline --data_path s3://my-bucket/data.parquet
```

## Invoke remote tasks in a script.

You can also run any remote task directly using a script in a similar way

```python
import flyte
import flyte.models
import flyte.remote

flyte.init_from_config()

# Fetch the task
remote_task = flyte.remote.Task.get("package-example.calculate_average", auto_version="latest")

# Create a run, note keyword arguments are required currently. In the future this will accept positional args based on the declaration order, but, we still recommend to use keyword args.
run = flyte.run(remote_task, numbers=[1.0, 2.0, 3.0])

print(f"Execution URL: {run.url}")
# you can view the phase

print(f"Current Phase: {run.phase}")
# You can wait for the execution
run.wait()

# Only available after flyte >= 2.0.0b39
print(f"Current phase: {run.phase}")

# Phases can be compared to
if run.phase == flyte.models.ActionPhase.SUCCEEDED:
    print(f"Run completed!")

# You can then retrieve the outputs
print(f"Pipeline result: {run.outputs()}")
```

## Why use remote tasks?

Remote tasks solve common collaboration and dependency management challenges:

**Cross-team collaboration**: Team A has deployed a Spark task that analyzes large datasets. Team B needs this analysis for their ML pipeline but doesn't want to learn Spark internals, install Spark dependencies, or build Spark-enabled container images. With remote tasks, Team B simply references Team A's deployed task.

**Platform reusability**: Platform teams can create common, reusable tasks (data validation, feature engineering, model serving) that other teams can use without duplicating code or managing complex dependencies.

**Microservices for data workflows**: Remote tasks work like microservices for long-running tasks or agents, enabling secure sharing while maintaining isolation.

## When to use remote tasks

Use remote tasks when you need to:

- Use functionality from another team without their dependencies
- Share common tasks across your organization
- Build reusable platform components
- Avoid dependency conflicts between different parts of your workflow
- Create modular, maintainable data pipelines

## How remote tasks work

### Security model

Remote tasks run in the **caller's project and domain** using the caller's compute resources, but execute with the **callee's service accounts, IAM roles, and secrets**. This ensures:

- Tasks are secure from misuse
- Resource usage is properly attributed
- Authentication and authorization are maintained
- Collaboration remains safe and controlled

### Type system

Remote tasks use Flyte's default types as inputs and outputs. Flyte's type system translates data between tasks without requiring the original dependencies:

| Remote Task Type | Flyte Type |
|-------------------|------------|
| DataFrames (`pandas`, `polars`, `spark`, etc.) | `flyte.io.DataFrame` |
| Object store files | `flyte.io.File` |
| Object store directories | `flyte.io.Dir` |
| Pydantic models | Dictionary (Flyte creates a representation) |

Any DataFrame type (pandas, polars, spark) automatically becomes `flyte.io.DataFrame`, allowing data exchange between tasks using different DataFrame libraries. You can also write custom integrations or explore Flyte's plugin system for additional types.

For Pydantic models specifically, you don't need the exact model locally. Pass a dictionary as input, and Flyte will handle the translation.

## Versioning options

Reference tasks support flexible versioning:

**Specific version**:

```python
task = flyte.remote.Task.get(
    "team_a.process_data",
    version="v1.2.3"
)
```

**Latest version** (`auto_version="latest"`):

```python
# Always use the most recently deployed version
task = flyte.remote.Task.get(
    "team_a.process_data",
    auto_version="latest"
)
```

**Current version** (`auto_version="current"`):

```python
# Use the same version as the calling task's deployment
# Useful when all environments deploy with the same version
# Can only be used from within a task context
task = flyte.remote.Task.get(
    "team_a.process_data",
    auto_version="current"
)
```

## Customizing remote tasks

Remote tasks can be customized by overriding various properties without modifying the original deployed task. This allows you to adjust resource requirements, retry strategies, caching behavior, and more based on your specific use case.

### Available overrides

The `override()` method on remote tasks accepts the following parameters:

- **short_name** (`str`): A short name for the task instance
- **resources** (`flyte.Resources`): CPU, memory, GPU, and storage limits
- **retries** (`int | flyte.RetryStrategy`): Number of retries or retry strategy
- **timeout** (`flyte.TimeoutType`): Task execution timeout
- **env_vars** (`Dict[str, str]`): Environment variables to set
- **secrets** (`flyte.SecretRequest`): Secrets to inject
- **max_inline_io_bytes** (`int`): Maximum size for inline IO in bytes
- **cache** (`flyte.Cache`): Cache behavior and settings
- **queue** (`str`): Execution queue to use

### Override examples

**Increase resources for a specific use case**:

```python
import flyte.remote

# Get the base task
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Override with more resources for large dataset processing
large_data_processor = data_processor.override(
    resources=flyte.Resources(
        cpu="16",
        memory="64Gi",
        storage="200Gi"
    )
)

@env.task
async def process_large_dataset(data_path: str):
    # Use the customized version
    return await large_data_processor(input_path=data_path)
```

**Add retries and timeout**:

```python
# Override with retries and timeout for unreliable operations
reliable_processor = data_processor.override(
    retries=3,
    timeout="2h"
)

@env.task
async def robust_pipeline(data_path: str):
    return await reliable_processor(input_path=data_path)
```

**Configure caching**:

```python
# Override cache settings
cached_processor = data_processor.override(
    cache=flyte.Cache(
        behavior="override",
        version_override="v2",
        serialize=True
    )
)
```

**Set environment variables and secrets**:

```python
# Override with custom environment and secrets
custom_processor = data_processor.override(
    env_vars={
        "LOG_LEVEL": "DEBUG",
        "REGION": "us-west-2"
    },
    secrets=flyte.SecretRequest(
        secrets={"api_key": "my-secret-key"}
    )
)
```

**Multiple overrides**:

```python
# Combine multiple overrides
production_processor = data_processor.override(
    short_name="prod_spark_analyzer",
    resources=flyte.Resources(cpu="8", memory="32Gi"),
    retries=5,
    timeout="4h",
    env_vars={"ENV": "production"},
    queue="high-priority"
)

@env.task
async def production_pipeline(data_path: str):
    return await production_processor(input_path=data_path)
```

### Chain overrides

You can chain multiple `override()` calls to incrementally adjust settings:

```python
# Start with base task
processor = flyte.remote.Task.get("data_team.analyzer", auto_version="latest")

# Add resources
processor = processor.override(resources=flyte.Resources(cpu="4", memory="16Gi"))

# Add retries for production
if is_production:
    processor = processor.override(retries=5, timeout="2h")

# Use the customized task
result = await processor(input_path="s3://data.parquet")
```

## Best practices

### 1. Use meaningful task names

Remote tasks are accessed by name, so use clear, descriptive naming:

```python
# Good
customer_segmentation = flyte.remote.Task.get("ml_platform.customer_segmentation")

# Avoid
task1 = flyte.remote.Task.get("team_a.task1")
```

### 2. Document task interfaces

Since remote tasks abstract away implementation details, clear documentation of inputs, outputs, and behavior is essential:

```python
@env.task
async def process_customer_data(
    customer_ids: list[str],
    date_range: tuple[str, str]
) -> flyte.io.DataFrame:
    """
    Process customer data for the specified date range.

    Args:
        customer_ids: List of customer IDs to process
        date_range: Tuple of (start_date, end_date) in YYYY-MM-DD format

    Returns:
        DataFrame with processed customer features
    """
    ...
```

### 3. Prefer module-level loading

Load remote tasks at the module level rather than inside functions for cleaner code:

```python
import flyte.remote

# Good - module level
data_processor = flyte.remote.Task.get("team.processor", auto_version="latest")

@env.task
async def my_task(data: str):
    return await data_processor(input=data)
```

This approach:

- Makes dependencies clear and discoverable
- Reduces code duplication
- Works well with lazy loading (no performance penalty)

Dynamic loading within tasks is also supported when you need runtime flexibility.

### 4. Handle versioning thoughtfully

- Use `auto_version="latest"` during development for rapid iteration
- Use specific versions in production for stability and reproducibility
- Use `auto_version="current"` when coordinating multienvironment deployments

### 5. Deploy remote tasks first

Always deploy the remote tasks before using them. Tasks that reference them can be run directly without deployment:

Deploy the remote task environments first:

```bash
flyte deploy spark_env/
flyte deploy ml_env/
```

Then run the orchestration task directly (no deployment needed):

```bash
flyte run orchestration_env.py orchestrate_pipeline
```

If you want to deploy the orchestration task as well (for scheduled runs or to be referenced by other tasks), deploy it after its dependencies:

```bash
flyte deploy orchestration_env/
```

## Limitations

1. **Lazy error detection**: Because of lazy loading, errors about missing or invalid tasks only occur during invocation, not when calling `get()`. You'll receive a `flyte.errors.RemoteTaskNotFoundError` if the task doesn't exist and `flyte.errors.RemoteTaskUsageError` if it can't be invoked in the way you are passing either arguments or overrides.

2. **Type fidelity**: While Flyte translates types, you work with Flyte's representation of Pydantic models, not the exact original types

3. **Deployment order**: Referenced tasks must be deployed before tasks that reference them can be invoked

4. **Context requirement**: Using `auto_version="current"` requires running within a task context

5. **Dictionary inputs**: Pydantic models must be passed as dictionaries, which loses compile-time type checking

6. **No positional arguments**: Remote tasks currently only support keyword arguments (this may change in future versions)

## Next steps

- Learn about [task deployment](../task-deployment/_index)
- Explore [task environments and configuration](../task-configuration/_index)


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/error-handling ===

# Error handling

One of the key features of Flyte 2 is the ability to recover from user-level errors in a workflow execution.
This includes out-of-memory errors, timeouts, oversized inline I/O, and other exceptions.

In a distributed system with heterogeneous compute, certain types of errors are expected and even, in a sense, acceptable.
Flyte 2 recognizes this and allows you to handle them gracefully as part of your workflow logic.

This ability is a direct result of the fact that workflows are now written in regular Python,
giving you all the power and flexibility of Python error handling.
When a task fails, Flyte surfaces the failure to the calling task as a typed exception that you can catch with a
standard `try...except` block and respond to however you like: retry with more resources, fall back to a different
code path, or clean up and re-raise.

## How Flyte represents failures

When a downstream task fails, the failure propagates to the awaiting parent task as an exception from the
`flyte.errors` module. Every native exception derives from a small hierarchy of base classes:

- `flyte.errors.BaseRuntimeError`: the root of all Flyte runtime errors.
- `flyte.errors.RuntimeUserError`: the failure was caused by your code (a bug, an exception you raised, an
  out-of-memory condition, and so on). An exception you raise inside a task, say a `ValueError`, is wrapped and
  surfaces to the parent as a `flyte.errors.RuntimeUserError`.
- `flyte.errors.RuntimeSystemError`: the failure was caused by the platform rather than your code.
- `flyte.errors.RuntimeUnknownError`: the failure could not be classified as a user or system error.
Every concrete error carries a `code` attribute (a short, stable string identifier, often the exception's class name, e.g. `"TaskTimeoutError"`) that you can
inspect when logging or branching. Because the errors form a hierarchy, you can catch broadly
(`except flyte.errors.RuntimeUserError`) or narrowly (`except flyte.errors.OOMError`), depending on how specific
your recovery logic needs to be.

## Catching and recovering from errors

The most common pattern is to catch a specific exception and re-run the failing task with a different
configuration. The following example intentionally triggers an out-of-memory error, catches the
`flyte.errors.OOMError`, and retries the task with more memory:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = ""
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment(name="fail", resources=flyte.Resources(cpu=1, memory="250Mi"))

@env.task
async def oomer(x: int):
    large_list = [0] * 100000000
    print(len(large_list))

@env.task
async def always_succeeds() -> int:
    await asyncio.sleep(1)
    return 42

@env.task
async def main() -> int:
    try:
        await oomer(2)
    except flyte.errors.OOMError as e:
        print(f"Failed with oom trying with more resources: {e}, of type {type(e)}, {e.code}")
        try:
            await oomer.override(resources=flyte.Resources(cpu=1, memory="1Gi"))(5)
        except flyte.errors.OOMError as e:
            print(f"Failed with OOM Again giving up: {e}, of type {type(e)}, {e.code}")
            raise e
    finally:
        await always_succeeds()

    return await always_succeeds()

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/error-handling/error_handling.py*

In this code, we do the following:

* Import the necessary modules, including `flyte.errors`.
* Set up the task environment with a modest resource allocation of 1 CPU and 250 MiB of memory.
* Define two tasks: `oomer`, which allocates a large list and is likely to run out of memory, and
  `always_succeeds`, which always returns cleanly.
* Define the `main` task (the top-level workflow task) that contains the failure-recovery logic.

The `try...except` block in `main` runs `oomer`. If it exhausts memory, `main` catches the
`flyte.errors.OOMError` and retries by calling `oomer.override(resources=...)` with a larger memory allocation.
If the retry also runs out of memory, `main` gives up and re-raises the error. The `finally` block runs
`always_succeeds` regardless of the outcome.

This type of dynamic error handling lets you gracefully recover from user-level errors in your workflows using
patterns you already know from ordinary Python. For a complete, self-tuning version of this pattern that caches
the optimal memory setting across runs, see the
[`resource_tuner` example](https://github.com/flyteorg/flyte-sdk/blob/main/examples/advanced/resource_tuner.py).

> [!NOTE] Programmatic recovery vs. automatic retries
> Catching an exception and re-running a task is *programmatic* recovery: you decide what to do differently on
> the next attempt. This is distinct from Flyte's *automatic* retries (`retries=N` on a task), which simply
> re-run the same attempt unchanged. The two compose: automatic retries handle transient failures, while a
> `try...except` handles failures you want to respond to deliberately. See
> [Retries and timeouts](../task-configuration/retries-and-timeouts).

## Limiting inline I/O

Small task inputs and outputs are passed *inline* (embedded directly in the task's metadata) rather than offloaded
to blob storage. This is fast, but very large inline values are undesirable, so each task has a ceiling on the
size of its inline I/O. You set this ceiling with the `max_inline_io_bytes` parameter on `@env.task`, and Flyte
raises a `flyte.errors.InlineIOMaxBytesBreached` when an input or output exceeds it:

```python
import flyte
import flyte.errors

env = flyte.TaskEnvironment(
    name="large_inline_io",
    resources=flyte.Resources(cpu=1, memory="250Mi"),
)

@env.task(max_inline_io_bytes=100 * 1024)  # Limit inline I/O to 100 KiB
async def printer_task(x: str) -> str:
    print(f"Printer task received: {x}")
    return x

@env.task
async def large_inline_io() -> str:
    small = await printer_task("Hello, world!")
    print(f"Small string result: {small}")

    # A large string that exceeds the 100 KiB inline limit
    large_string = "A" * 10**6  # ~1 MiB
    try:
        return await printer_task(large_string)
    except flyte.errors.InlineIOMaxBytesBreached as e:
        print(f"Inline I/O limit breached: {e}")
        raise
```

The small string passes through, but the ~1 MiB string breaches the 100 KiB limit and raises
`flyte.errors.InlineIOMaxBytesBreached`. When you expect large values, raise `max_inline_io_bytes` or pass the
data as a `flyte.io.File` or `flyte.io.Dir` so it is offloaded to blob storage instead of travelling inline. A
runnable version of this example is available as
[`large_inline_io.py`](https://github.com/flyteorg/flyte-sdk/blob/main/examples/advanced/large_inline_io.py).

## Natively-supported exceptions

Flyte raises typed exceptions for the failure modes it recognizes, so you can catch exactly the condition you
care about. The most commonly caught errors are:

| Exception | Raised when |
|---|---|
| `flyte.errors.OOMError` | A task exceeds its memory allocation. |
| `flyte.errors.TaskTimeoutError` | A task runs longer than its configured timeout. |
| `flyte.errors.InlineIOMaxBytesBreached` | An input or output exceeds the task's `max_inline_io_bytes` limit. |
| `flyte.errors.RetriesExhaustedError` | A task fails after all of its automatic retries are used up. |
| `flyte.errors.TaskInterruptedError` | A task running on interruptible (spot) compute is preempted. |
| `flyte.errors.ActionAbortedError` | An action is aborted externally via the CLI, UI, or API. |
| `flyte.errors.ImagePullBackOffError` | The task's container image cannot be pulled. |
| `flyte.errors.NonRecoverableError` | A failure that should not be retried, regardless of the retry budget. |

These all derive from `flyte.errors.RuntimeUserError`, so a single `except flyte.errors.RuntimeUserError`
catches any of them when you want uniform handling. This is only a selection. For the complete catalog of
catchable exception classes, see the [`flyte.errors` API reference](../../../api-reference/flyte-sdk/flyte.errors/_index).

## Related pages

- [Retries and timeouts](../task-configuration/retries-and-timeouts): configure automatic retries and execution time limits.
- [Abort and cancel actions](./abort-tasks): stop actions programmatically or externally, and handle `flyte.errors.ActionAbortedError`.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/traces ===

# Traces

The `@flyte.trace` decorator provides fine-grained observability and resumption capabilities for functions called within your Flyte workflows.
Traces are used on **helper functions** that tasks call to perform specific operations like API calls, data processing, or computations.
Traces are particularly useful for [managing the challenges of non-deterministic behavior in workflows](../../migration/flyte-2/gotchas-and-caveats#non-deterministic-behavior), allowing you to track execution details and resume from failures.

## What are traced functions for?

At the top level, Flyte workflows are composed of **tasks**. But it is also common practice to break down complex task logic into smaller, reusable functions by defining helper functions that tasks call to perform specific operations.

Any helper functions defined or imported into the same file as a task definition are automatically uploaded to the Flyte environment alongside the task when it is deployed.

At the task level, observability and resumption of failed executions is provided by caching, but what if you want these capabilities at a more granular level, for the individual operations that tasks perform?

This is where **traced functions** come in. By decorating helper functions with `@flyte.trace`, you enable:

- **Detailed observability**: Track execution time, inputs/outputs, and errors for each function call.
- **Fine-grained resumption**: If a workflow fails, resume from the last successful traced function instead of re-running the entire task.
Each traced function is effectively a checkpoint within its task.

Here is an example:

```
import asyncio

import flyte

env = flyte.TaskEnvironment("env")

@flyte.trace
async def call_llm(prompt: str) -> str:
    await asyncio.sleep(0.1)
    return f"LLM response for: {prompt}"

@flyte.trace
async def process_data(data: str) -> dict:
    await asyncio.sleep(0.2)
    return {"processed": data, "status": "completed"}

@env.task
async def research_workflow(topic: str) -> dict:
    llm_result = await call_llm(f"Generate research plan for: {topic}")
    processed_data = await process_data(llm_result)
    return {"topic": topic, "result": processed_data}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/task_vs_trace.py*

## What gets traced

Traces capture detailed execution information:

- **Execution time**: How long each function call takes.
- **Inputs and outputs**: Function parameters and return values.
- **Checkpoints**: State that enables workflow resumption.

### Errors are not recorded

Only successful trace executions are recorded in the checkpoint system. When a traced function fails, the exception propagates up to your task code where you can handle it with standard error handling patterns.

### Supported function types

The trace decorator works with:

- **Asynchronous functions**: Functions defined with `async def`.
- **Generator functions**: Functions that `yield` values.
- **Async generators**: Functions that `async yield` values.

> [!NOTE]
> Currently tracing only works for asynchronous functions. Tracing of synchronous functions is coming soon.

```
@flyte.trace
async def async_api_call(topic: str) -> dict:
    # Asynchronous API call
    await asyncio.sleep(0.1)
    return {"data": ["item1", "item2", "item3"], "status": "success"}

@flyte.trace
async def stream_data(items: list[str]):
    # Async generator function for streaming
    for item in items:
        await asyncio.sleep(0.02)
        yield f"Processing: {item}"

@flyte.trace
async def async_stream_llm(prompt: str):
    # Async generator for streaming LLM responses
    chunks = ["Research shows", " that machine learning", " continues to evolve."]
    for chunk in chunks:
        await asyncio.sleep(0.05)
        yield chunk

@env.task
async def research_workflow(topic: str) -> dict:
    llm_result = await async_api_call(topic)

    # Collect async generator results
    processed_data = []
    async for item in stream_data(llm_result["data"]):
        processed_data.append(item)

    llm_stream = []
    async for chunk in async_stream_llm(f"Summarize research on {topic}"):
        llm_stream.append(chunk)

    return {
        "topic": topic,
        "processed_data": processed_data,
        "llm_summary": "".join(llm_stream)
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/function_types.py*

## Task orchestration pattern

The typical Flyte workflow follows this pattern:

```
@flyte.trace
async def search_web(query: str) -> list[dict]:
    # Search the web and return results
    await asyncio.sleep(0.1)
    return [{"title": f"Article about {query}", "content": f"Content on {query}"}]

@flyte.trace
async def summarize_content(content: str) -> str:
    # Summarize content using LLM
    await asyncio.sleep(0.1)
    return f"Summary of {len(content.split())} words"

@flyte.trace
async def extract_insights(summaries: list[str]) -> dict:
    # Extract insights from summaries
    await asyncio.sleep(0.1)
    return {"insights": ["key theme 1", "key theme 2"], "count": len(summaries)}

@env.task
async def research_pipeline(topic: str) -> dict:
    # Each helper function creates a checkpoint
    search_results = await search_web(f"research on {topic}")

    summaries = []
    for result in search_results:
        summary = await summarize_content(result["content"])
        summaries.append(summary)

    final_insights = await extract_insights(summaries)

    return {
        "topic": topic,
        "insights": final_insights,
        "sources_count": len(search_results)
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/pattern.py*

**Benefits of this pattern:**

- If `search_web` succeeds but `summarize_content` fails, resumption skips the search step
- Each operation is independently observable and debuggable
- Clear separation between workflow coordination (task) and execution (traced functions)

## Relationship to caching and checkpointing

Understanding how traces work with Flyte's other execution features:

| Feature | Scope | Purpose | Default Behavior |
|---------|-------|---------|------------------|
| **Task Caching** | Entire task execution (`@env.task`) | Skip re-running tasks with same inputs | Enabled (`cache="auto"`) |
| **Traces** | Individual helper functions | Observability and fine-grained resumption | Manual (requires `@flyte.trace`) |
| **Checkpointing** | Workflow state | Resume workflows from failure points | Automatic when traces are used |

### How they work together

<!-- TODO
Lets use better typing for all of these examples, we have the opportunity to make this right for our users
-->

```
@flyte.trace
async def traced_data_cleaning(dataset_id: str) -> List[str]:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.2)
    return [f"cleaned_record_{i}_{dataset_id}" for i in range(100)]

@flyte.trace
async def traced_feature_extraction(data: List[str]) -> dict:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.3)
    return {
        "features": [f"feature_{i}" for i in range(10)],
        "feature_count": len(data),
        "processed_samples": len(data)
    }

@flyte.trace
async def traced_model_training(features: dict) -> dict:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.4)
    sample_count = features["processed_samples"]
    # Mock accuracy based on sample count
    accuracy = min(0.95, 0.7 + (sample_count / 1000))
    return {
        "accuracy": accuracy,
        "epochs": 50,
        "model_size": "125MB"
    }

@env.task(cache="auto")  # Task-level caching enabled
async def data_pipeline(dataset_id: str) -> dict:
    # 1. If this exact task with these inputs ran before,
    #    the entire task result is returned from cache

    # 2. If not cached, execution begins and each traced function
    #    creates checkpoints for resumption
    cleaned_data = await traced_data_cleaning(dataset_id)      # Checkpoint 1
    features = await traced_feature_extraction(cleaned_data)   # Checkpoint 2
    model_results = await traced_model_training(features)      # Checkpoint 3

    # 3. If workflow fails at step 3, resumption will:
    #    - Skip traced_data_cleaning (checkpointed)
    #    - Skip traced_feature_extraction (checkpointed)
    #    - Re-run only traced_model_training

    return {"dataset_id": dataset_id, "accuracy": model_results["accuracy"]}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/caching_vs_checkpointing.py*

### Execution flow

1. **Task Submission**: Task is submitted with input parameters
2. **Cache Check**: Flyte checks if identical task execution exists in cache
3. **Cache Hit**: If cached, return cached result immediately (no traces needed)
4. **Cache Miss**: Begin fresh execution
5. **Trace Checkpoints**: Each `@flyte.trace` function creates resumption points
6. **Failure Recovery**: If workflow fails, resume from last successful checkpoint
7. **Task Completion**: Final result is cached for future identical inputs

<!--
Clarify what actually happens on error vs success with traces

## Error handling and observability

Traces capture comprehensive execution information for debugging and monitoring:

```
@flyte.trace
async def risky_api_call(endpoint: str, data: dict) -> dict:
    """API call that might fail - traces capture errors."""
    try:
        response = await api_client.post(endpoint, json=data)
        return response.json()
    except Exception as e:
        # Error is automatically captured in trace
        logger.error(f"API call failed: {e}")
        raise

@env.task
async def error_handling() -> dict:
    try:
        result = await risky_api_call("/process", {"invalid": "data"})
        return {"status": "success", "result": result}
    except Exception as e:
        # The error is recorded in the trace for debugging
        return {"status": "error", "message": str(e)}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/error_handling.py*

**What traces capture:**
- **Execution time**: Duration of each function call
- **Inputs and outputs**: Function parameters and return values
- **Checkpoints**: State that enables workflow resumption from successful executions
- **Action IDs**: Unique identifiers for each execution

**Error handling:**
- Errors from traced functions are not recorded in checkpoints
- Exceptions propagate to your task code for standard error handling
- The error_handling example shows how to catch and handle these exceptions in your task

TODO:
Ketan Umare:
we should show an example where tasks and traces can be used interchangeably

## Examples in practice

### LLM pipeline with traces

```python
import flyte

env = flyte.TaskEnvironment("llm-pipeline")

@flyte.trace
async def call_llm(prompt: str, model: str = "gpt-4") -> str:
    """Call LLM with specified model."""
    response = await llm_client.chat(prompt, model=model)
    return response

@flyte.trace
async def extract_entities(text: str) -> list[str]:
    """Extract named entities from text."""
    entities = await nlp_service.extract_entities(text)
    return entities

@env.task
async def process_documents(documents: list[str]) -> dict:
    """Process multiple documents through LLM pipeline."""
    results = []

    for doc in documents:
        # Each call is traced for monitoring and resumption
        summary = await call_llm(f"Summarize: {doc}")
        entities = await extract_entities(summary)

        results.append({
            "document": doc,
            "summary": summary,
            "entities": entities
        })

    return {"processed_documents": results, "total_count": len(results)}
```

This comprehensive tracing system provides visibility into your workflow execution while enabling robust error recovery and resumption capabilities.
-->


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/intra-task-checkpoints ===

# Intra-task checkpoints

Long-running tasks (model training especially) can fail partway through: a spot
instance is reclaimed, a pod is evicted, an out-of-memory error kills the process.
When the task is retried, it normally starts over from the beginning.

Intra-task checkpoints let a task save in-progress state to object storage as it
runs and load that state at the start of the next attempt, so a retry resumes from
where the previous attempt left off instead of repeating completed work.

## The checkpoint object

Inside a running task, `flyte.ctx().checkpoint` returns a `flyte.Checkpoint` (or `None`
when checkpointing isn't configured) bound to the action's checkpoint location in object storage:

- **Save**: `await checkpoint.save(...)` (async tasks) or `checkpoint.save_sync(...)`
  (sync tasks and synchronous framework callbacks). Accepts raw `bytes`, a file path,
  or a directory path; a directory is stored as a single compressed archive.
- **Load**: `await checkpoint.load()` or `checkpoint.load_sync()`. Returns a local
  `pathlib.Path` to the restored file or directory tree, or `None` when there is no
  previous checkpoint (i.e. on the first attempt).
- `flyte.latest_checkpoint(root, glob_pattern="**/last.ckpt")` finds the newest checkpoint file under a restored
  directory tree, useful for frameworks like PyTorch Lightning that write
  `last.ckpt` files into a directory.

Checkpoints only matter when the task can run more than once, so give the task
retries with `@env.task(retries=...)`. Each retry attempt sees the checkpoint saved
by the attempt before it.

> [!NOTE] Checkpoints vs. caching vs. traces
>
> - **Task caching** skips an entire task when it has already run with the same inputs.
> - **[Traces](./traces)** checkpoint at the boundaries of helper functions called by a task.
> - **Intra-task checkpoints** save state *within* a single task body (mid-loop,
>   mid-epoch) across retry attempts of the same action.

## Basic usage

The simplest checkpoint is a raw byte payload. This task counts up to
`n_iterations`, saving its progress on every iteration. A simulated failure kills
it partway through; the retry loads the saved counter and continues rather than
restarting from zero:

### Async

```
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic")

RETRIES = 3

@env.task(retries=RETRIES)
async def use_checkpoint(n_iterations: int = 10) -> int:
    checkpoint = flyte.ctx().checkpoint

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load() returns None.
    path = await checkpoint.load()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        await checkpoint.save(f"{index + 1}".encode())
    return index
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint.py*

### Sync

```
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic_sync")

RETRIES = 3

@env.task(retries=RETRIES)
def use_checkpoint(n_iterations: int = 10) -> int:
    checkpoint = flyte.ctx().checkpoint

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load_sync() returns None.
    path = checkpoint.load_sync()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        checkpoint.save_sync(f"{index + 1}".encode())
    return index
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint_sync.py*

Running this with `n_iterations=10` produces three failed attempts and one
successful one. Each attempt fails later than the last, because each one starts
from the checkpoint its predecessor saved.

## Checkpointing ML training frameworks

The same pattern applies to real training loops, whatever the framework:

1. **Load** the previous attempt's checkpoint at the start of the task; if one
   exists, restore the model/optimizer state and work out where to resume.
2. **Save** a checkpoint at a regular interval (every epoch or every N steps)
   as training progresses.

Frameworks with their own checkpoint files (PyTorch Lightning, Hugging Face
`Trainer`) already write them to a local directory; there you hook their callback
system and mirror that directory to the Flyte checkpoint, then feed the restored
directory back to the framework's native resume mechanism.

### PyTorch

Save the model state dict, optimizer state, and epoch counter with `torch.save`
after each epoch, and restore all three with `torch.load` on retry:

<br>

```
@env.task(retries=RETRIES)
async def train_linear(epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    model = nn.Linear(4, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    # Resume model, optimizer, and epoch from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        blob = torch.load(prev, map_location="cpu", weights_only=False)
        model.load_state_dict(blob["model"])
        opt.load_state_dict(blob["opt"])
        start = int(blob["epoch"]) + 1
    else:
        start = 0

    wpath = pathlib.Path("pytorch_linear") / "training.pt"
    wpath.parent.mkdir(parents=True, exist_ok=True)

    failure_interval = epochs // RETRIES
    for epoch in range(start, epochs):
        x = torch.randn(8, 4)
        y = torch.randn(8, 1)
        loss = torch.nn.functional.mse_loss(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()

        if epoch > start and epoch % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at epoch {epoch}")

        # Save model, optimizer, and epoch state to object storage.
        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
            wpath,
        )
        await checkpoint.save(wpath)

    with torch.no_grad():
        return float(model(torch.ones(1, 4)).squeeze().item())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_checkpoint.py*

### PyTorch Lightning

Lightning already writes `last.ckpt` through its `ModelCheckpoint` callback.
Subclass it to mirror the checkpoint directory to Flyte after each epoch
(Lightning callbacks are synchronous, so use `flyte.Checkpoint.save_sync`):

<br>

```
class FlyteLightningCheckpointCallback(ModelCheckpoint):
    """A `ModelCheckpoint` that mirrors `dirpath` to the Flyte checkpoint after each epoch."""

    def __init__(self, flyte_checkpoint: flyte.Checkpoint, *, dirpath: str | pathlib.Path, **kwargs) -> None:
        super().__init__(dirpath=str(dirpath), **kwargs)
        self._flyte_checkpoint = flyte_checkpoint

    @override
    def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
        super().on_train_epoch_end(trainer, pl_module)
        if self.dirpath:
            # Lightning callbacks are synchronous, so use save_sync
            self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

In the task, restore the previous tree, pick the newest `last.ckpt` with
`flyte.latest_checkpoint(restored_root, glob_pattern="**/last.ckpt")`, and hand it to `Trainer.fit(ckpt_path=...)`; Lightning
restores the model, optimizer, and epoch from there:

<br>

CODE3

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

### Hugging Face Trainer

`transformers.Trainer` writes `checkpoint-<step>` directories under its
`output_dir` (here, every epoch via `save_strategy="epoch"`). A `TrainerCallback`
mirrors that directory to Flyte after each save:

<br>

```
class FlyteTrainerCheckpointCallback(TrainerCallback):
    """Mirror the Trainer's `output_dir` to the Flyte checkpoint after each epoch."""

    def __init__(self, checkpoint: flyte.Checkpoint, output_dir: pathlib.Path) -> None:
        self._checkpoint = checkpoint
        self._output_dir = output_dir

    def on_epoch_end(self, args, state, control, **kwargs) -> None:
        # Trainer callbacks are synchronous, so use save_sync
        self._checkpoint.save_sync(self._output_dir)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

On retry, restore the tree, locate the last Hugging Face checkpoint with
`get_last_checkpoint`, and pass it to `trainer.train(resume_from_checkpoint=...)`:

<br>

CODE4

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

### scikit-learn

For estimators that support incremental training with `partial_fit`, pickle the
estimator together with a progress counter after each training chunk. A retry
unpickles the bundle and continues from the next chunk:

<br>

CODE5

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/sklearn_partial_checkpoint.py*

### Unsloth

LoRA fine-tuning with [Unsloth](https://unsloth.ai/) and `trl.SFTTrainer` uses the
same callback-and-resume pattern as the Hugging Face `Trainer`, since `SFTTrainer`
is built on it. Unsloth requires an NVIDIA, AMD, or Intel GPU, so the task
environment requests one:

<br>

CODE6

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/unsloth_sft_checkpoint.py*

> [!NOTE] Simulated failures in the runnable examples
> The full example files for the basic, PyTorch, and scikit-learn cases inject a
> failure at a regular interval (`failure_interval`) so you can watch the retries
> resume from the checkpoint. In production code you would drop those lines; real
> failures (preemptions, OOMs, crashes) trigger the same resume path.

## How checkpoints are stored

Each action attempt gets a checkpoint prefix in the object store configured for
your cluster. `flyte.Checkpoint.save` uploads a file as-is, stores a directory as
a gzip-compressed tarball, and accepts raw `bytes` as a single blob.
`flyte.Checkpoint.load` downloads the previous attempt's object into a local
temporary workspace and returns the path: a restored directory tree, or the path
to the single restored file.

Saving repeatedly overwrites the same checkpoint object, so the cost of frequent
checkpointing is upload bandwidth, not unbounded storage growth. Checkpoint how
often you can afford to lose work: every epoch is typical for training loops.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/grouping-actions ===

# Grouping actions

Groups are an organizational feature in Flyte that allow you to logically cluster related task invocations (called "actions") for better visualization and management in the UI.
Groups help you organize task executions into manageable, hierarchical structures regardless of whether you're working with large fanouts or smaller, logically-related sets of operations.

## What are groups?

Groups provide a way to organize task invocations into logical units in the Flyte UI.
When you have multiple task executions (whether from large [fanouts](./fanout), sequential operations, or any combination of tasks), groups help organize them into manageable units.

### The problem groups solve

Without groups, complex workflows can become visually overwhelming in the Flyte UI:

- Multiple task executions appear as separate nodes, making it hard to see the high-level structure
- Related operations are scattered throughout the workflow graph
- Debugging and monitoring becomes difficult when dealing with many individual task executions

Groups solve this by:

- **Organizing actions**: Multiple task executions within a group are presented as a hierarchical "folder" structure
- **Improving UI visualization**: Instead of many individual nodes cluttering the view, you see logical groups that can be collapsed or expanded
- **Aggregating status information**: Groups show aggregated run status (success/failure) of their contained actions when you hover over them in the UI
- **Maintaining execution parallelism**: Tasks still run concurrently as normal, but are organized for display

### How groups work

Groups are declared using the `flyte.group` context manager.
Any task invocations that occur within the `with flyte.group()` block are automatically associated with that group:

```python
with flyte.group("my-group-name"):
    # All task invocations here belong to "my-group-name"
    result1 = await task_a(data)
    result2 = await task_b(data)
    result3 = await task_c(data)
```

The key points about groups:

1. **Context-based**: Use the `with flyte.group("name"):` context manager.
2. **Organizational tool**: Task invocations within the context are grouped together in the UI.
3. **UI folders**: Groups appear as collapsible/expandable folders in the Flyte UI run tree.
4. **Status aggregation**: Hover over a group in the UI to see aggregated success/failure information.
5. **Execution unchanged**: Tasks still execute in parallel as normal; groups only affect organization and visualization.

**Important**: Groups do not aggregate outputs. Each task execution still produces its own individual outputs. Groups are purely for organization and UI presentation.

## Common grouping patterns

### Sequential operations

Group related sequential operations that logically belong together:

```
@env.task
async def data_pipeline(raw_data: str) -> str:
    with flyte.group("data-validation"):
        validated_data = await process_data(raw_data, "validate_schema")
        validated_data = await process_data(validated_data, "check_quality")
        validated_data = await process_data(validated_data, "remove_duplicates")

    with flyte.group("feature-engineering"):
        features = await process_data(validated_data, "extract_features")
        features = await process_data(features, "normalize_features")
        features = await process_data(features, "select_features")

    with flyte.group("model-training"):
        model = await process_data(features, "train_model")
        model = await process_data(model, "validate_model")
        final_model = await process_data(model, "save_model")

    return final_model
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Parallel processing with groups

Groups work well with parallel execution patterns:

```
@env.task
async def parallel_processing_example(n: int) -> str:
    tasks = []

    with flyte.group("parallel-processing"):
        # Collect all task invocations first
        for i in range(n):
            tasks.append(process_item(i, "transform"))

        # Execute all tasks in parallel
        results = await asyncio.gather(*tasks)

    # Convert to string for consistent return type
    return f"parallel_results: {results}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Multi-phase workflows

Use groups to organize different phases of complex workflows:

```
@env.task
async def multi_phase_workflow(data_size: int) -> str:
    # First phase: data preprocessing
    preprocessed = []
    with flyte.group("preprocessing"):
        for i in range(data_size):
            preprocessed.append(process_item(i, "preprocess"))
        phase1_results = await asyncio.gather(*preprocessed)

    # Second phase: main processing
    processed = []
    with flyte.group("main-processing"):
        for result in phase1_results:
            processed.append(process_item(result, "transform"))
        phase2_results = await asyncio.gather(*processed)

    # Third phase: postprocessing
    postprocessed = []
    with flyte.group("postprocessing"):
        for result in phase2_results:
            postprocessed.append(process_item(result, "postprocess"))
        final_results = await asyncio.gather(*postprocessed)

    # Convert to string for consistent return type
    return f"multi_phase_results: {final_results}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Conditional grouping

Groups can be used with conditional logic:

```
@env.task
async def conditional_processing(use_advanced_features: bool, input_data: str) -> str:
    base_result = await process_data(input_data, "basic_processing")

    if use_advanced_features:
        with flyte.group("advanced-features"):
            enhanced_result = await process_data(base_result, "advanced_processing")
            optimized_result = await process_data(enhanced_result, "optimize_result")
            return optimized_result
    else:
        with flyte.group("basic-features"):
            simple_result = await process_data(base_result, "simple_processing")
            return simple_result
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

## Key insights

Groups are primarily an organizational and UI visualization tool; they don't change how your tasks execute or aggregate their outputs, but they help organize related task invocations (actions) into collapsible folder-like structures for better workflow management and display. The aggregated status information (success/failure rates) is visible when hovering over group folders in the UI.

Groups make your Flyte workflows more maintainable and easier to understand, especially when working with complex workflows that involve multiple logical phases or large numbers of task executions. They serve as organizational "folders" in the UI's call stack tree, allowing you to collapse sections to reduce visual distraction while still seeing aggregated status information on hover.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/fanout ===

# Fanout

Flyte is designed to scale, allowing you to run workflows with large fanouts.
When you need to execute many tasks in parallel (such as processing a large dataset or running hyperparameter sweeps), Flyte provides powerful patterns to implement these operations efficiently.

> **📝 Note**
>
> In Flyte 1, mapping a task over many inputs used `map_task()` (the `flytekit.map_task` API). In Flyte 2, fan out with `asyncio.gather()` or `flyte.map()`.

This page covers the general `asyncio.gather` fanout pattern. For applying the *same* task to every
item of a list (the direct successor to Flyte 1's `map_task`), see [Mapping over inputs](./map).

That page also covers concurrency limits and error handling.

## Understanding fanout

A "fanout" pattern occurs when you spawn multiple tasks concurrently.
Each task runs in its own container and contributes an output that you later collect.
The most common way to implement this is using the [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather) function.

In Flyte terminology, each individual task execution is called an "action": this represents a specific invocation of a task with particular inputs. When you call a task multiple times in a loop, you create multiple actions.

## Example

We start by importing our required packages, defining our Flyte environment, and creating a simple task that fetches user data from a mock API.

```
import asyncio
from typing import List, Tuple

import flyte

env = flyte.TaskEnvironment("fanout_env")

@env.task
async def fetch_data(user_id: int) -> dict:
    """Simulate fetching user data from an API - good for parallel execution."""
    # Simulate network I/O delay
    await asyncio.sleep(0.1)
    return {
        "user_id": user_id,
        "name": f"User_{user_id}",
        "score": user_id * 10,
        "data": f"fetched_data_{user_id}"
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

### Parallel execution

Next we implement the most common fanout pattern, which is to collect task invocations and execute them in parallel using `asyncio.gather()`:

```
@env.task
async def parallel_data_fetching(user_ids: List[int]) -> List[dict]:
    """Fetch data for multiple users in parallel - ideal for I/O bound operations."""
    tasks = []

    # Collect all fetch tasks - these can run in parallel since they're independent
    for user_id in user_ids:
        tasks.append(fetch_data(user_id))

    # Execute all fetch operations in parallel
    results = await asyncio.gather(*tasks)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

### Running the example

To actually run our example, we create a main guard that initializes Flyte and runs our main driver task:

```
if __name__ == "__main__":
    flyte.init_from_config()
    user_ids = [1, 2, 3, 4, 5]
    r = flyte.run(parallel_data_fetching, user_ids)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

## How Flyte handles concurrency and parallelism

In the example we use a standard `asyncio.gather()` pattern.
When this pattern is used in a normal Python environment, the tasks would execute **concurrently** (cooperatively sharing a single thread through the event loop), but not in true **parallel** (multiple CPU cores simultaneously).

However, **Flyte transforms this concurrency model into true parallelism**. When you use `asyncio.gather()` in a Flyte task:

1. **Flyte acts as a distributed event loop**: Instead of scheduling coroutines on a single machine, Flyte schedules each task action to run in its own container across the cluster
2. **Concurrent becomes parallel**: What would be cooperative multitasking in regular Python becomes true parallel execution across multiple machines
3. **Native Python patterns**: You use familiar `asyncio` patterns, but Flyte automatically distributes the work

This means that when you write:

```python
results = await asyncio.gather(fetch_data(1), fetch_data(2), fetch_data(3))
```

Instead of three coroutines sharing one CPU, you get three separate containers running simultaneously, each with their own CPU, memory, and resources. Flyte bridges the gap between Python's concurrency model and distributed parallel computing, allowing for massive scalability while maintaining the familiar async/await programming model.

## Iterative fanout: recursive feature elimination

Fanout isn't limited to a single parallel burst; you can fan out **repeatedly**, using the results of one round to shape the next.
A good real-world example is [recursive feature elimination (RFE)](https://github.com/flyteorg/flyte-sdk/blob/main/examples/ml/rfe.py), a feature-selection technique that repeatedly trains a model with one candidate feature held out, drops the feature whose removal least hurts the score, and repeats until a single feature remains.
Every iteration is itself a fanout: for each remaining feature, a `train` action runs in parallel with that feature dropped, scored by cross-validation.

The `train` task evaluates the model with a single feature held out and returns its cross-validated score:

```python
@worker.task
async def train(features: list[str], drop: str) -> float:
    features.remove(drop)

    X, y = fetch_california_housing(as_frame=True, return_X_y=True)
    fold = KFold(n_splits=5, random_state=42, shuffle=True)
    model = LinearRegression()

    scores = cross_val_score(estimator=model, X=X[features], y=y, cv=fold, scoring="r2")
    return float(scores.mean())
```

The `rfe` driver task runs the elimination loop.
Each round wraps its fanout in a `flyte.group` context (see [Grouping actions](./grouping-actions)) so the iterations appear as collapsible folders in the UI, and uses `asyncio.gather()` to evaluate every candidate feature in parallel:

```python
@worker.task
async def rfe():
    x, _y = fetch_california_housing(as_frame=True, return_X_y=True)
    features = list(x.columns)

    for i in range(len(features) - 1):
        with flyte.group(f"iteration-{i}"):
            runs = {feature: train(list(features), drop=feature) for feature in features}
            values = await asyncio.gather(*(runs[feature] for feature in runs))
            scores = dict(zip(runs.keys(), values))
            best = max(scores, key=scores.get)
            features.remove(best)
```

Because each `train` call becomes its own action, every iteration's candidate evaluations run as separate containers in true parallel, while grouping keeps the nested rounds organized in the run tree.

> **📝 Note**
>
> The full runnable example lives in the [Flyte SDK repository](https://github.com/flyteorg/flyte-sdk/blob/main/examples/ml/rfe.py). From a local checkout of the `flyte-sdk` repository, run it with `uv run --prerelease=allow examples/ml/rfe.py` (the command uses a repo-relative path).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/map ===

# Mapping over inputs

`flyte.map` applies a single task to every item of one or more input iterables, running the
invocations in parallel across the cluster and yielding their results **in input order**.
It is the structured way to [fan out](./fanout) uniform work: instead of assembling a list of
coroutines by hand and passing them to `asyncio.gather`, you hand `flyte.map` the task and the
inputs and it produces one action per item.

Use `flyte.map` when every item goes through the *same* task. For fanning out across *different*
tasks, or for full control over how invocations are assembled, use `asyncio.gather`. See
[Fanout](./fanout).

## Minimal example

From a **synchronous** task, iterate the results with a plain `for` loop:

```python
from typing import List

import flyte

env = flyte.TaskEnvironment(name="map-example")

@env.task
def process(x: int) -> str:
    return f"result-{x}"

@env.task
def main(n: int) -> List[str]:
    results: List[str] = []
    for r in flyte.map(process, range(n)):
        if isinstance(r, Exception):
            raise r
        results.append(r)
    return results

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, 10)
    print(run.url)
```

Each item in `range(n)` becomes its own action, running in its own container, and the results come
back in the same order as the inputs.

## Mapping from an async task: `flyte.map.aio`

`flyte.map` returns a synchronous iterator. Inside an **async** task, use `flyte.map.aio`, which
returns an async iterator you consume with `async for`:

```python
@env.task
async def main(n: int) -> List[str]:
    results: List[str] = []
    async for r in flyte.map.aio(process, range(n)):
        if isinstance(r, Exception):
            raise r
        results.append(r)
    return results
```

`flyte.map.aio` works over both async and sync tasks, so you can call an existing synchronous task
in parallel from an async context without rewriting it: useful when migrating a Flyte 1.x
`map_task` or integrating legacy sync code.

## Signature and parameters

```python
flyte.map(
    func,               # the task (or functools.partial) to apply to each item
    *args,              # one or more iterables, zipped item-by-item into func's arguments
    group_name=None,    # optional name for the group of mapped actions (UI grouping)
    concurrency=0,      # max actions in flight at once; 0 means unbounded (all at once)
    return_exceptions=True,
)
```

- **`func`**: the task to map. It receives one item per invocation. To hold some arguments
  constant across the map, wrap it with `functools.partial` (see **Tasks > Build tasks > Mapping over inputs > Binding constant arguments with `functools.partial`**).
- **`*args`**: one or more input iterables. With multiple iterables they are **zipped**: the
  *i*-th invocation receives the *i*-th element of each, matching `func`'s positional parameters in
  order.
- **`group_name`**: groups the resulting actions under a single label in the UI (see
  [Grouping actions](./grouping-actions)).
- **`concurrency`**: the maximum number of actions in flight at any moment. `0` (the default)
  submits everything at once. A positive value bounds the fan-out with a worker pool, so memory
  stays proportional to `concurrency` rather than to the total number of items. See
  [Controlling parallel execution](./controlling-parallelism).
- **`return_exceptions`**: when `True` (the default), a failed invocation yields the raised
  exception as its result instead of aborting the whole map; check each result with
  `isinstance(r, Exception)`. When `False`, the first failure stops iteration and raises.

Results are always yielded **in the order of the inputs**, regardless of the order in which the
individual actions finish.

## Limiting concurrency

For rate-limited APIs, GPU quotas, or connection limits, cap how many actions run at once with the
`concurrency` parameter:

```python
async for r in flyte.map.aio(call_llm_api, prompts, concurrency=3):
    ...
```

Only three actions are in flight at a time; as each completes, the next input is submitted. For a
full comparison of `flyte.map(concurrency=N)` against `asyncio.Semaphore`, see
[Controlling parallel execution](./controlling-parallelism).

## Handling errors

By default (`return_exceptions=True`) the map runs to completion even if some invocations fail, and
each failure surfaces as an exception object in the results stream:

```python
@env.task
def maybe_fail(x: int) -> str:
    if x == 2:
        raise ValueError("bad input")
    return f"ok-{x}"

@env.task
def main(n: int) -> None:
    for r in flyte.map(maybe_fail, range(n)):
        if isinstance(r, Exception):
            print(f"error: {r}")
        else:
            print(r)
```

Set `return_exceptions=False` to fail fast instead: iteration raises on the first failed action.

## Binding constant arguments with `functools.partial`

Often you want to map over one argument while holding others constant. Bind the constants with
`functools.partial`, leaving exactly one parameter free. That's the one `flyte.map` varies:

```python
from functools import partial

import flyte

env = flyte.TaskEnvironment(name="map-partial")

@env.task
def score(compound_id: str, model_name: str, batch_id: str) -> str:
    return f"{compound_id}:{model_name}:{batch_id}"

@env.task
def main() -> None:
    compounds = [str(i) for i in range(3)]
    scorer = partial(score, model_name="v2", batch_id="run-42")
    # compound_id is the only parameter left unbound, so it is what map varies.
    results = list(flyte.map(scorer, compounds))
    print("\n".join(results))
```

`flyte.map` inserts each mapped value **positionally, right after the partial's bound positional
arguments**, and requires **exactly one** parameter to be left unbound. Above, `model_name` and
`batch_id` are bound as keywords, so the mapped value fills the first slot: `compound_id`. To vary a
*later* parameter, bind the ones before it positionally and the ones after it by keyword. For
example, `partial(score, "compound-1", batch_id="run-42")` maps `model_name`. `flyte.map` raises a
`TypeError` if more or fewer than one parameter is left unbound, or if the mapped positional slot is
also bound as a keyword.

## When to use `flyte.map`

Reach for `flyte.map` when:

- Every item goes through the **same** task.
- You want built-in, in-order result collection and per-item error capture.
- You want simple, declarative concurrency control via the `concurrency` parameter.

Use [`asyncio.gather`](./fanout) instead when you are fanning out across **different** tasks in one
batch, or when you need full control over how the coroutines are assembled. Use an
[`asyncio.Semaphore`](./controlling-parallelism) when different task types in the same batch need
different concurrency limits.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/consuming-a-message-queue ===

# Consuming a message queue

A common production pattern is a **queue consumer**: a long-running task that pulls messages
from an external message queue (such as [AWS SQS](https://aws.amazon.com/sqs/)) and processes
each message concurrently.
Flyte 2 expresses this naturally by combining three building blocks you have already seen:

- **Async tasks**: the consumer loop is an `async def` task that awaits I/O against the queue.
- [**Fanout**](./fanout): each received message is dispatched to its own `process_message` action with `asyncio.create_task()`, so processing runs in parallel across the cluster.
- [**Reusable containers**](../task-configuration/reusable-containers): a `ReusePolicy` keeps a warm pool of replicas ready, so messages are processed without per-message container cold-start.

The complete, runnable source for this example, a producer (`generator.py`) and a consumer
(`processor.py`), lives in the Flyte SDK repository under
[`examples/queue-reader`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/queue-reader).

> [!NOTE]
> This example relies on [reusable containers](../task-configuration/reusable-containers) (`flyte.ReusePolicy`), which are only available when running your Flyte code on a Union backend.
> See [Reusable containers](../task-configuration/reusable-containers) for details.

> [!NOTE]
> This example reads from AWS SQS and therefore requires an SQS queue and AWS credentials
> available to the running task (here the queue is passed as an ARN via the `QUEUE_ARN`
> environment variable). The Flyte pattern shown below applies to any external queue: swap
> the SQS client calls for your queue's client.

## The consumer

### Define the task environment

The consumer runs in a [reusable](../task-configuration/reusable-containers) `TaskEnvironment`.
`replicas=3` keeps a warm pool of at least two replicas to avoid starvation while the parent
consumer task occupies one, and `idle_ttl=300` shuts the pool down after five minutes of
inactivity.
The image is built from the script's own inline dependencies with
`flyte.Image.from_uv_script`, plus the `unionai-reuse` runtime library that reusable containers
require:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte",
#    "aioboto3>=11.3.0",
#    "asyncio",
# ]
# ///

import asyncio
import json
import os
from typing import List

import aioboto3

import flyte

env = flyte.TaskEnvironment(
    name="sqs_processor",
    resources=flyte.Resources(memory="500Mi", cpu=1),
    image=flyte.Image.from_uv_script(
        __file__,
        name="flyte",
    ).with_pip_packages("unionai-reuse>=0.1.3"),
    reusable=flyte.ReusePolicy(
        replicas=3,  # 1 for the consumer loop + 2 workers, so processing never starves
        idle_ttl=300,  # Idle time to keep the task environment alive
    ),
)

# The queue is passed as an ARN via the QUEUE_ARN environment variable.
DEFAULT_QUEUE_ARN = os.getenv("QUEUE_ARN")

def get_queue_url_from_arn(queue_arn: str) -> str:
    """Convert an SQS ARN to a queue URL."""
    parts = queue_arn.split(":")
    region = parts[3]
    account = parts[4]
    queue_name = parts[5]

    return f"https://sqs.{region}.amazonaws.com/{account}/{queue_name}"
```

### Process a single message

Each message is handled by its own task. These tasks run in parallel across the reusable pool,
bounded by the number of worker replicas: with `replicas=3` and the default `concurrency=1`, the
parent consumer loop occupies one replica and the other two each process a single message at a
time, so about two messages are handled concurrently. To let a single replica handle more than one
message at once, raise `concurrency` above 1:

```python
@env.task
async def process_message(message: dict) -> str:
    """Process a single message asynchronously and return the extracted word."""
    body = json.loads(message["Body"])
    word = body.get("word", "unknown")
    print(f"Task Processing message {body.get('message_id')}: {word}")
    return word
```

### The consumer loop

The driver task long-polls the queue, and for each message it receives it **dispatches a
`process_message` action with `asyncio.create_task()`** rather than awaiting it inline. This
is what fans the work out in parallel. It deletes each message once processing has started, then
awaits all dispatched tasks with `asyncio.gather()`:

```python
@env.task
async def main(queue_arn: str = DEFAULT_QUEUE_ARN, max_messages: int = 10) -> List[str]:
    queue_url = get_queue_url_from_arn(queue_arn)
    session = aioboto3.Session(region_name="us-east-2")

    results = []
    tasks = []
    messages_received = 0

    async with session.client("sqs") as sqs:
        while messages_received < max_messages:
            response = await sqs.receive_message(
                QueueUrl=queue_url,
                AttributeNames=["All"],
                MaxNumberOfMessages=1,   # one message at a time
                WaitTimeSeconds=20,      # long-polling timeout (max 20 seconds)
            )

            messages = response.get("Messages", [])
            if not messages:
                continue

            message = messages[0]
            messages_received += 1

            # Fan out: dispatch processing as a parallel action.
            process_task = asyncio.create_task(process_message(message))
            tasks.append(process_task)

            # Delete the message once we've started processing it.
            await sqs.delete_message(QueueUrl=queue_url, ReceiptHandle=message["ReceiptHandle"])

    # Wait for all dispatched processing tasks to complete.
    if tasks:
        completed_tasks = await asyncio.gather(*tasks)
        results.extend(completed_tasks)

    return results
```

### Run it

Initialize Flyte from your config and run the consumer remotely:

```python
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, queue_arn=DEFAULT_QUEUE_ARN, max_messages=10)
    print(run.url)
```

## The producer

To exercise the consumer, the example includes a standalone
[`generator.py`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/queue-reader) that
pushes ten JSON messages onto the same SQS queue with `boto3`. It is an ordinary Python script,
not a Flyte task. Any producer that writes to the queue will do.

## Notes and gotchas

- **Delete after receive, not after processing completes.** The example deletes each message as
  soon as it dispatches the processing task. If a `process_message` action can fail and you need
  at-least-once semantics, delete the message only after the task succeeds instead.
- **`max_messages` bounds the run.** The consumer loop here stops after `max_messages`. For a
  continuously running consumer, drive it on a [trigger](../task-configuration/triggers) or remove
  the bound and manage the task lifecycle explicitly.
- **Reusable containers require a Union backend.** See
  [Reusable containers](../task-configuration/reusable-containers) for the `ReusePolicy` parameters
  (`replicas`, `concurrency`, `idle_ttl`, `scaledown_ttl`) and their capacity math.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/controlling-parallelism ===

# Controlling parallel execution

When you [fan out](./fanout) to many tasks, you often need to limit how many run at the same time.
Common reasons include rate-limited APIs, GPU quotas, database connection limits, or simply avoiding overwhelming a downstream service.

Flyte 2 provides two ways to control concurrency:
[`asyncio.Semaphore`](https://docs.python.org/3/library/asyncio-sync.html#asyncio.Semaphore) for fine-grained control,
and `flyte.map` with a built-in `concurrency` parameter for simpler cases.

## The problem: unbounded parallelism

A straightforward `asyncio.gather` launches every task at once.
If you are calling an external API that allows only a few concurrent requests, this can cause throttling or errors:

```
import asyncio

import flyte

env = flyte.TaskEnvironment("controlling_parallelism")

@env.task
async def call_llm_api(prompt: str) -> str:
    """Simulate calling a rate-limited LLM API."""
    # In a real workflow, this would call an external API.
    # The API might allow only a few concurrent requests.
    await asyncio.sleep(0.5)
    return f"Response to: {prompt}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

```
@env.task
async def process_all_at_once(prompts: list[str]) -> list[str]:
    """Send all requests in parallel with no concurrency limit.

    This can overwhelm a rate-limited API, causing errors or throttling.
    """
    results = await asyncio.gather(*[call_llm_api(p) for p in prompts])
    return list(results)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

With eight prompts, this fires eight concurrent API calls.
That works fine when there are no limits, but will fail when the API enforces a concurrency cap.

## Using asyncio.Semaphore

An `asyncio.Semaphore` acts as a gate: only a fixed number of tasks can pass through at a time.
The rest wait until a slot opens up.

```
@env.task
async def process_batch_with_semaphore(
    prompts: list[str],
    max_concurrent: int = 3,
) -> list[str]:
    """Process prompts in parallel, limiting concurrency with a semaphore.

    At most `max_concurrent` calls to the API run at any given time.
    The remaining tasks wait until a slot is available.
    """
    semaphore = asyncio.Semaphore(max_concurrent)

    async def limited_call(prompt: str) -> str:
        async with semaphore:
            return await call_llm_api(prompt)

    results = await asyncio.gather(*[limited_call(p) for p in prompts])
    return list(results)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

The pattern is:

1. Create a semaphore with the desired limit.
2. Wrap each task call in an inner async function that acquires the semaphore before calling and releases it after.
3. Pass all wrapped calls to `asyncio.gather`.

All eight tasks are submitted immediately, but the Flyte orchestrator only allows three to run in parallel.
As each one completes, the next waiting task starts.

> [!NOTE]
> The semaphore controls how many tasks execute concurrently on the Flyte cluster.
> Each task still runs in its own container with its own resources: the semaphore simply limits how many containers are active at a time.

## Using flyte.map with concurrency

For uniform work (applying the same task to a list of inputs), `flyte.map` with the `concurrency` parameter is simpler:

CODE2

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

This achieves the same concurrency limit with less boilerplate.

For the full `flyte.map` treatment (signature, return order, error handling, and partials), see
[Mapping over inputs](./map).

## Running the example

CODE3

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

## When to use each approach

Use **`flyte.map(concurrency=N)`** when:

- Every item goes through the same task.
- You want the simplest possible code.

Use **`asyncio.Semaphore`** when:

- You need different concurrency limits for different task types within the same workflow.
- You want to combine concurrency control with error handling (e.g., `asyncio.gather(*tasks, return_exceptions=True)`).
- You are calling multiple different tasks in one parallel batch.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/streaming-map-reduce ===

# Streaming map-reduce

When you [fan out](./fanout) with [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather), you wait for **every** task to finish before doing anything with the results.
For a map-reduce workload that is wasteful: the reduce step sits idle until the slowest mapper returns.

A better pattern is to process results **as they complete** — accumulating them into batches and kicking off reduce operations incrementally, while the remaining map tasks are still running.
This is a *gradual* (or *streaming*) map-reduce, and it is built on the standard-library [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) function.

## When to use it

Reach for streaming map-reduce when:

- Map tasks have **uneven durations**, so waiting for the slowest one wastes time the faster ones could spend reducing.
- You are processing a **large number of items** and want to reduce in batches rather than holding every intermediate result in memory at once.
- The reduce step is **associative** — batch results can themselves be reduced into a final result (counts, sums, aggregations, embeddings, inference outputs).

If you simply need all results before a single reduce, plain `asyncio.gather` (see [Fanout](./fanout)) is simpler.
If your goal is to *cap* how many map tasks run at once, see [Controlling parallel execution](./controlling-parallelism); the two patterns compose.

## Example

We define an environment and two tasks: one that maps over a single item, and one that reduces a batch of results.

```python
import asyncio
import random

import flyte

env = flyte.TaskEnvironment(
    name="streaming_map_reduce",
    resources=flyte.Resources(cpu="1"),
)

@env.task
async def process_item(item: str) -> str:
    print(f"Processing {item}", flush=True)
    # Simulate varying processing times so results finish out of order.
    await asyncio.sleep(random.uniform(1, 5))
    return f"processed_{item}"

@env.task
async def reduce_batch(items: list[str]) -> str:
    print(f"Reducing batch of {len(items)} items")
    return f"reduced_batch_of_{len(items)}_items"
```

### The driver task

The driver fans out all the map tasks up front, then walks the results in completion order with `asyncio.as_completed`.
Each time a batch fills up, it launches a `reduce_batch` action **without blocking** — the loop keeps consuming newly completed map results while the reduce runs.

```python
@env.task
async def streaming_reduce_processing() -> str:
    input_items = [f"item_{i}" for i in range(100)]

    # Fan out: start every item task immediately.
    tasks = [asyncio.create_task(process_item(item)) for item in input_items]

    batch_size = 10
    accumulated_values: list[str] = []
    reducers: list[asyncio.Task] = []

    print(f"Started {len(tasks)} tasks, will reduce in batches of {batch_size}")

    # Consume results as each task finishes, rather than waiting for all of them.
    for task in asyncio.as_completed(tasks):
        result = await task
        accumulated_values.append(result)

        # Once a batch has accumulated, kick off a reduce without blocking the loop.
        if len(accumulated_values) >= batch_size:
            print(f"Triggering reduce for batch of {len(accumulated_values)}")
            reducer_task = asyncio.create_task(reduce_batch(accumulated_values.copy()))
            reducers.append(reducer_task)
            accumulated_values.clear()

    # Reduce any stragglers that did not fill a full batch.
    if accumulated_values:
        print(f"Handling final batch of {len(accumulated_values)} stragglers")
        reducers.append(asyncio.create_task(reduce_batch(accumulated_values)))

    # Wait for every batch reduce to finish.
    reduced_results = await asyncio.gather(*reducers)

    # Combine the batch results into a single final result.
    final_result = await reduce_batch(reduced_results)

    print(f"Completed {len(reducers)} reduce operations, final result: {final_result}")
    return final_result
```

### Running the example

```python
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(streaming_reduce_processing)
    print(run.url)
```

## How it works

The key building blocks are all standard `asyncio`:

1. **`asyncio.create_task(process_item(item))`** schedules each map action. Because `process_item` is a Flyte task, each of these runs in its own container on the cluster — the fanout is real distributed parallelism, not single-machine concurrency (see [Fanout](./fanout) for how Flyte turns `asyncio` into distributed execution).
2. **`asyncio.as_completed(tasks)`** yields the task handles in the order they *finish*, not the order they were submitted. This is what lets the driver react to the fastest map results first.
3. **`asyncio.create_task(reduce_batch(...))`** launches each reduce as its own Flyte action and appends it to `reducers` without awaiting it, so map consumption and reduction overlap.
4. **`asyncio.gather(*reducers)`** joins all the in-flight batch reduces before the final combine step.

The result is a pipeline where reduce work begins as soon as the first batch of map results is ready, instead of after the last map task returns.

> [!NOTE]
> `as_completed` returns awaitables in completion order but gives you no control over *how many* map tasks run at once — it schedules all of them.
> To bound the map fanout as well, combine this pattern with an `asyncio.Semaphore` or `flyte.map(concurrency=...)` from [Controlling parallel execution](./controlling-parallelism).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/structured-concurrency-anyio ===

# Structured concurrency with anyio

Flyte builds a task's dependency graph from what you `await` — not from any particular async library.
`asyncio` is the default and the one used throughout the [Fanout](./fanout), [Controlling parallel execution](./controlling-parallelism), and [Task dependencies and ordering](./task-dependencies) guides, but it is not the only option.
Any structured-concurrency runtime that drives coroutines works, and [`anyio`](https://anyio.readthedocs.io/) is a popular one.
Its **task groups** give you a top-level alternative to raw `asyncio.gather` / `asyncio.create_task`, with clearer lifetime and error-propagation semantics.

Use `anyio` when you want structured concurrency — a scope that owns the tasks it spawns, waits for all of them on exit, and cancels the siblings automatically if one fails — instead of tracking `asyncio.create_task` handles by hand.

## The task-group pattern

An `anyio` task group is an `async with` block. You spawn work into it with `start_soon`, and the block does not exit until every spawned task has finished. Because task groups don't return the spawned tasks' values directly, this example uses [`aioresult`](https://aioresult.readthedocs.io/)'s `ResultCapture` to collect each result.

We define a reusable environment and a simple per-item task. `anyio` and `aioresult` are ordinary pip dependencies, so we add them to the image:

```python
from dataclasses import dataclass

import aioresult
import anyio

import flyte

env = flyte.TaskEnvironment(
    name="anyio_batch",
    resources=flyte.Resources(cpu="1"),
    image=flyte.Image.from_debian_base(name="anyio").with_pip_packages("anyio", "aioresult"),
)

@dataclass
class InferenceRequest:
    feature_a: float
    feature_b: float

@env.task
async def predict_one(request: InferenceRequest) -> float:
    # A dummy linear model: 2 * feature_a + 3 * feature_b + bias(=1.0)
    return 2.0 * request.feature_a + 3.0 * request.feature_b + 1.0
```

The driver task fans out over the batch inside a task group:

```python
@env.task
async def predict_batch(requests: list[InferenceRequest]) -> list[float]:
    captured = []
    async with anyio.create_task_group() as tg:
        # Start each prediction; they run at the same time.
        for req in requests:
            captured.append(aioresult.ResultCapture.start_soon(tg, predict_one, req))
    # The `async with` block has exited, so every task has completed.
    return [c.result() for c in captured]
```

What happens here mirrors an `asyncio.gather` fanout, but with structured-concurrency guarantees:

1. **`start_soon` schedules each `predict_one`** into the group. As with `asyncio`, Flyte runs each action in its own container, so the batch executes in true parallel across the cluster — the runtime you use to express concurrency does not change how Flyte distributes the work.
2. **Leaving the `async with` block is the fan-in edge.** The group blocks until all spawned tasks finish, exactly as `await asyncio.gather(...)` would. `predict_batch` cannot return until every prediction is in.
3. **`ResultCapture` collects the return values**, which you read with `.result()` after the group closes.

> [!NOTE]
> Task groups give you cancellation for free: if any task in the group raises, `anyio` cancels the remaining siblings and propagates the error out of the `async with` block. You get the "cancel the rest on failure" behavior that requires manual `.cancel()` bookkeeping with `asyncio` (see [Abort and cancel actions](./abort-tasks#canceling-actions-programmatically)).

## When to use anyio

Reach for `anyio` when:

- You want **structured concurrency** — spawned work is scoped to a block, awaited on exit, and cancelled together on error — rather than manually pairing `asyncio.create_task` handles with `asyncio.gather`.
- Your code (or a library you depend on) already uses `anyio` or `trio`, and you want one consistent concurrency model.

Stay with `asyncio` when:

- You just need to fan out and collect results — `await asyncio.gather(...)` is simpler (see [Fanout](./fanout)).
- You need fine-grained, dependency-driven scheduling where different consumers await different producers (see [Task dependencies and ordering](./task-dependencies)).

Either way, the underlying model is the same: Flyte reads the dependency graph from your `await`s and turns concurrent coroutines into distributed parallel actions.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/task-dependencies ===

# Task dependencies and ordering

Flyte 1 built a workflow's DAG (directed acyclic graph) explicitly: you declared nodes and wired their edges with the `>>` operator or `create_node`.
Flyte 2 has no such API. Instead, **the dependency graph is inferred from the data you `await`**.
When you await one task's result and pass it into another, Flyte records the edge; tasks that share no data run independently.
This page shows how to express the ordering patterns you used to build by hand — sequencing, fan-out, fan-in, and fine-grained dependency-driven scheduling — using ordinary Python `asyncio`.

If you are coming from Flyte 1, read [Parallelism and fan-out](../../migration/flyte-2/parallelism) first for the migration mapping.

## The dependency graph is implicit

There is nothing special to learn: a data dependency *is* the edge.

```python
import asyncio

import flyte

env = flyte.TaskEnvironment(name="pipeline")

@env.task
async def extract() -> str:
    return "raw"

@env.task
async def transform(data: str) -> str:
    return f"transformed({data})"

@env.task
async def load(data: str) -> str:
    return f"loaded({data})"

@env.task
async def main() -> str:
    raw = await extract()          # runs first
    clean = await transform(raw)   # waits for extract — it consumes `raw`
    return await load(clean)       # waits for transform — it consumes `clean`
```

Each `await` means "wait for this to finish before continuing," so a chain of `await`s that pass results downstream runs sequentially — exactly like a linear Flyte 1 workflow. You never declare the edges; passing `raw` into `transform` and `clean` into `load` *is* the DAG.

## Ordering without a data dependency

Sometimes you need task `B` to run after task `A` even though `B` does not consume `A`'s output — for example, `A` writes to a store that `B` reads out-of-band, or `A` must finish before you send a notification.

Because ordering comes from `await`, you force it simply by awaiting `A` before invoking `B`:

CODE0

You do not need a special "run after" construct — a preceding `await` is the ordering primitive.

## Fan-out and fan-in

**Fan-out** launches independent tasks concurrently; **fan-in** collects their results into a single downstream task.
Use [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather) to await several tasks at once — Flyte runs each in its own container in parallel (see [Fanout](./fanout)):

CODE1

The `await asyncio.gather(...)` establishes the fan-in edge: `combine` cannot start until all three upstream tasks have produced their results.

## Dependency-driven scheduling

The pattern that most often motivates "replicating DAG behavior" is **fine-grained scheduling**: a diamond or fork where each downstream task should start the moment *its own* upstreams finish, without waiting for unrelated slow tasks.

Consider three producers of different durations and four consumers with different dependencies:

- `needs_short` depends on `short` only
- `needs_medium` depends on `medium` only
- `needs_long` depends on `long` only
- `needs_all` depends on all three

A single `await asyncio.gather(short, medium, long)` before starting any consumer would make every consumer wait for the slowest producer. To let each consumer start as early as possible, start the producers as [`asyncio.create_task`](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task) handles, then wrap each consumer in a small helper coroutine that awaits only the handles it needs. Launch all the helpers together with `asyncio.gather`:

CODE2

`needs_short` starts about a second in, as soon as `short_producer` returns — it does not wait for the 10-second `long_producer`. Awaiting an `asyncio` task handle more than once is safe: the handle caches its result, so `long_task` can feed both `run_needs_long` and `run_needs_all` without re-running the producer.

> [!NOTE]
> Reach for helper coroutines that each `await` their specific handles, rather than a manual completion loop that inspects [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) and dispatches downstream tasks by hand. Hand-rolled dispatch loops are easy to get wrong — a mis-tracked "has this fired yet?" check can launch the same downstream task twice. Let the dependency edges fall out of `await` instead.

## When to reach for `as_completed`

Use [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) when you want to process results **in completion order** — for example, streaming each result into a running reduction as it lands — rather than to encode a fixed dependency graph:

CODE3

For a worked streaming/reduce example, see [Fanout](./fanout) and [Controlling parallel execution](./controlling-parallelism).

## Summary

- Flyte 2 has no explicit DAG-construction API; dependencies come from the data you `await`.
- Sequence tasks by awaiting them in order — a preceding `await` orders even tasks that share no data.
- Fan out with `asyncio.gather`; fan in by awaiting several results into one downstream task.
- For fine-grained scheduling, keep producer handles from `asyncio.create_task` and have each consumer await only the handles it depends on, so it starts as early as possible.
- Prefer letting `await` express the graph over hand-rolled completion-tracking loops.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/conditions ===

# External conditions

An **external condition** is a first-class action that pauses a running task until an external
signal arrives: a human approval, a callback from an external service, or a value supplied at
runtime. The paused action stays observable, resumable, and governable like any other action, so
you no longer need polling loops or side processes to wait on something the workflow can't produce
itself.

Inside a task, `await flyte.new_condition.aio(...)` registers a condition action and returns a
handle; `await handle.wait.aio()` blocks the task until the condition is signaled and returns the
typed payload. (`new_condition` and `wait` are sync-by-default; in an `async def` task use their
`.aio()` form.)

## Supported types

A condition declares a `data_type`, which determines what a signal must supply and what `wait()`
returns:

| `data_type` | `wait()` returns | A signal value of |
|---|---|---|
| `bool` (default) | `True` / `False` | `true` / `false` |
| `int` | Python `int` | an integer literal |
| `float` | Python `float` | a decimal literal |
| `str` | Python `str` | any string |

## Example: human approval

A typed approval gate with a timeout (the most common use case):

```python
from datetime import timedelta
import flyte

env = flyte.TaskEnvironment("approvals")

@env.task
async def etl_pipeline():
    staged = await transform()

    approval = await flyte.new_condition.aio(
        "prod_write_approval",
        prompt="Approve writing staged data to production?",
        data_type=bool,
        timeout=timedelta(hours=24),
    )
    if not await approval.wait.aio():
        raise RuntimeError("Pipeline rejected by reviewer")

    await write_to_prod(staged)
```

The task pauses at `await approval.wait.aio()` until someone signals the condition (see
**Tasks > Build tasks > External conditions > Signaling a condition**). If the timeout elapses with no signal, `wait()`
raises `flyte.errors.ConditionTimedoutError`.

## Example: string input at runtime

A condition can collect a typed value (not just a yes/no) and feed it back into the workflow.
Here the task waits for a free-form string before continuing:

```python
import flyte

env = flyte.TaskEnvironment("conditions")

@env.task
async def deploy_with_reason():
    reason = await flyte.new_condition.aio(
        "deploy_reason",
        prompt="Enter a deployment reason to continue:",
        data_type=str,
    )
    note: str = await reason.wait.aio()
    # `note` now holds the string a human supplied — use it downstream.
    await record_audit(note)
CODE1python
flyte.new_condition(
    name,
    prompt="Approve?",
    prompt_type="text",
    data_type=bool,
    description="",
    timeout=None,
    webhook=None,
)
```

| Parameter | Type | Default | Description |
|---|---|---|---|
| `name` | `str` | required | Identifier for the condition within the parent action. Signal it with this name (`flyte signal condition <run> <name>`) or look it up with `flyte.remote.Condition.get("<name>", ...)`. |
| `prompt` | `str` | `"Approve?"` | Human-readable text shown in the UI signal form. |
| `prompt_type` | `"text"` \| `"markdown"` | `"text"` | How the prompt is rendered. |
| `data_type` | `type` | `bool` | Payload type: one of `bool`, `int`, `float`, `str`. Determines what `wait()` returns and what a signal must supply. |
| `description` | `str` | `""` | Longer explanation rendered alongside the prompt. |
| `timeout` | `timedelta` \| `int` \| `float` \| `None` | `None` | Maximum wait. If it elapses with no signal, `wait()` raises `flyte.errors.ConditionTimedoutError`. |

An optional advanced `webhook` parameter accepts a `flyte.ConditionWebhook` so the backend POSTs a
callback URL when the condition is created; see the API reference for details.

## Signaling a condition

A condition is satisfied by delivering exactly one typed signal of its declared `data_type`.

### From the CLI

CODE2

The value is coerced to the condition's declared `data_type` (`true`/`false` for `bool`, integer
literals for `int`, decimal literals for `float`, any string for `str`).

### From Python (remote)

CODE3

`flyte.remote.Condition.listall(run_name=...)` enumerates the conditions on a run.

## Timeout with a fallback

CODE4

## Errors

| Situation | Raised |
|---|---|
| Timeout elapses before a signal | `flyte.errors.ConditionTimedoutError` |
| Creating a condition whose `name` already exists in the action | `flyte.errors.ConditionAlreadyExistsError` |
| Condition fails during execution | `flyte.errors.ConditionFailedError` |
| Signal value doesn't match `data_type` | `TypeError` (client-side, before any call) |
| `wait()` called outside a task context | `RuntimeError` |


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/unit-testing ===

Unit testing is essential for ensuring your Flyte tasks work correctly. Flyte 2.0 provides flexible testing approaches that allow you to test both your business logic and Flyte-specific features like type transformations and caching.

## Understanding task invocation

When working with functions decorated with `@env.task`, there are two ways to invoke them, each with different behavior:

### Direct function invocation

When you call a task directly like a regular Python function:

```python
result = my_task(x=10, y=20)
```

**Flyte features are NOT invoked**, including:

- Type transformations and serialization
- Caching
- Data validation

This behaves exactly like calling a regular Python function, making it ideal for testing your business logic.

### Using `flyte.run()`

When you invoke a task using `flyte.run()`:

```python
run = flyte.run(my_task, x=10, y=20)
result = run.outputs()
```

**Flyte features ARE invoked**, including:

- Type transformations and serialization
- Data validation
- Type checking (raises `flyte.errors` if types are not supported or restricted)

This allows you to test Flyte-specific behavior like serialization and caching.

## Testing business logic

For most unit tests, you want to verify your business logic works correctly. Use **direct function invocation** for this:

```python
import flyte

env = flyte.TaskEnvironment("my_env")

@env.task
def add(a: int, b: int) -> int:
    return a + b

def test_add():
    result = add(a=3, b=5)
    assert result == 8
```

### Testing async tasks

Async tasks work the same way with direct invocation:

```python
import pytest

@env.task
async def subtract(a: int, b: int) -> int:
    return a - b

@pytest.mark.asyncio
async def test_subtract():
    result = await subtract(a=10, b=4)
    assert result == 6
```

### Testing nested tasks

When tasks call other tasks, direct invocation continues to work without any Flyte overhead:

```python
@env.task
def nested(a: int, b: int) -> int:
    return add(a, b)  # Calls the add task directly

def test_nested():
    result = nested(3, 5)
    assert result == 8
```

## Testing type transformations and serialization

When you need to test how Flyte handles data types, serialization, or caching, use `flyte.run()`:

```python
@pytest.mark.asyncio
async def test_add_with_flyte_run():
    run = flyte.run(add, 3, 5)
    assert run.outputs() == 8
```

### Testing type restrictions

Some types may not be supported or may be restricted. Use `flyte.run()` to test that these restrictions are enforced:

```python
from typing import Tuple
import flyte.errors

@env.task
def not_supported_types(x: Tuple[str, str]) -> str:
    return x[0]

@pytest.mark.asyncio
async def test_not_supported_types():
    # Direct invocation works fine
    result = not_supported_types(x=("a", "b"))
    assert result == "a"

    # flyte.run enforces type restrictions
    with pytest.raises(flyte.errors.RestrictedTypeError):
        flyte.run(not_supported_types, x=("a", "b"))
```

### Testing nested tasks with serialization

You can also test nested task execution with Flyte's full machinery:

```python
@pytest.mark.asyncio
async def test_nested_with_run():
    run = flyte.run(nested, 3, 5)
    assert run.outputs() == 8
```

## Testing traced functions

Functions decorated with `@flyte.trace` can be tested similarly to tasks:

```python
@flyte.trace
async def traced_multiply(a: int, b: int) -> int:
    return a * b

@pytest.mark.asyncio
async def test_traced_multiply():
    result = await traced_multiply(a=6, b=7)
    assert result == 42
```

## Best practices

1. **Test logic with direct invocation**: For most unit tests, call tasks directly to test your business logic without Flyte overhead.

2. **Test serialization with `flyte.run()`**: Use `flyte.run()` when you need to verify:
   - Type transformations work correctly
   - Data serialization/deserialization
   - Caching behavior
   - Type restrictions are enforced

3. **Use standard testing frameworks**: Flyte tasks work with pytest, unittest, and other Python testing frameworks.

4. **Test async tasks properly**: Use `@pytest.mark.asyncio` for async tasks and await their results.

5. **Mock external dependencies**: Use standard Python mocking techniques for external services, databases, etc.

## Quick reference

| Test Scenario | Method | Example |
|--------------|--------|---------|
| Business logic (sync) | Direct call | `result = task(x=10)` |
| Business logic (async) | Direct await | `result = await task(x=10)` |
| Type transformations | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Data serialization | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Caching behavior | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Type restrictions | `flyte.run()` + pytest.raises | `pytest.raises(flyte.errors.RestrictedTypeError)` |

## Example test suite

Here's a complete example showing different testing approaches:

```python
import pytest
import flyte
import flyte.errors

env = flyte.TaskEnvironment("test_env")

@env.task
def add(a: int, b: int) -> int:
    return a + b

@env.task
async def subtract(a: int, b: int) -> int:
    return a - b

# Test business logic directly
def test_add_logic():
    result = add(a=3, b=5)
    assert result == 8

@pytest.mark.asyncio
async def test_subtract_logic():
    result = await subtract(a=10, b=4)
    assert result == 6

# Test with Flyte serialization
@pytest.mark.asyncio
async def test_add_serialization():
    run = flyte.run(add, 3, 5)
    assert run.outputs() == 8

@pytest.mark.asyncio
async def test_subtract_serialization():
    run = flyte.run(subtract, a=10, b=4)
    assert run.outputs() == 6
```

## Future improvements

The Flyte SDK team is actively working on improvements for advanced unit testing scenarios, particularly around initialization and setup for complex test cases. Additional utilities and patterns may be introduced in future releases to further simplify unit testing.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/higher-order-functions ===

# Higher-order functions

A *higher-order function* is a function that takes other functions as arguments or returns them. Because Flyte 2 tasks execute as native Python and can be [passed as arguments](./other-features#passing-tasks-and-functions-as-arguments) like any other callable, you can write higher-order functions that operate on **tasks themselves**: reusable orchestration components that wrap a task with retry, fallback, batching, or fault-tolerance logic, without changing the task's business logic.

This is possible because Flyte 2 workflows run as ordinary Python:

- **Tasks are callables.** You can accept a task as a parameter and `await` it, `.override(...)` its resources, or hand it to `asyncio`.
- **Arbitrary nesting.** A task can invoke other tasks at any depth, so an orchestration wrapper can drive a task from inside another task.
- **Native control flow.** Loops, conditionals, and `try`/`except` work directly on task results (task outputs are plain Python objects, not promises), so a wrapper can inspect a result or catch an exception and react.

> [!NOTE] Higher-order functions are plain functions, not tasks
> The wrappers below are **not** decorated with `@env.task`. They are regular `async` Python functions that orchestrate tasks. You call them from inside a driver task (an `@env.task`), which is where the actual task invocations happen. Keep the reusable orchestration logic in a plain function so it can be applied to any task.

The patterns on this page are drawn from the runnable [`higher_order_patterns`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/higher_order_patterns) examples in the Flyte SDK repository.

## Fallback runner

Run a primary task and, if it fails with a matching exception, automatically fall back to an alternative task. Useful for degrading to a cheaper model, a different region, or a simpler algorithm when the preferred path fails.

```python
from typing import Callable, List, Optional, Type, TypeVar

R = TypeVar("R")

async def run_with_fallback(
    primary_task: Callable[..., R],
    fallback_task: Callable[..., R],
    *args,
    fallback_exceptions: Optional[List[Type[Exception]]] = None,
    **kwargs,
) -> R:
    try:
        return await primary_task(*args, **kwargs)
    except Exception as e:
        # Fall back only on the exceptions we opted into (None means any).
        should_fallback = fallback_exceptions is None or any(
            isinstance(e, exc) for exc in fallback_exceptions
        )
        if not should_fallback:
            raise
        return await fallback_task(*args, **kwargs)
```

Call it from a driver task, passing the two tasks as arguments:

```python
import flyte
import flyte.errors

env = flyte.TaskEnvironment("fallback")

@env.task
async def primary(x: int) -> int:
    # Business logic that may fail, e.g. raise ValueError(...) on bad input.
    ...

@env.task
async def backup(x: int) -> int:
    ...

@env.task
async def main(x: int) -> int:
    return await run_with_fallback(primary, backup, x, fallback_exceptions=[flyte.errors.RuntimeUserError])
```

Note the `fallback_exceptions` list holds `flyte.errors` types, not bare Python exceptions. An exception raised inside a task does not reach the parent as its original Python type: Flyte wraps it as a `flyte.errors` type (a `ValueError` raised in a task surfaces to the caller as a `flyte.errors.RuntimeUserError` whose `code` is `"ValueError"`). So `isinstance`/type-matching in a wrapper must target the `flyte.errors.*` hierarchy; matching on `ValueError` here would never fire and the fallback would never run. See [Error handling](./error-handling) for how failures propagate.

## Retry with increasing memory (OOM retrier)

Retry a task with progressively larger memory allocations when it hits an out-of-memory error, so you don't have to hard-code a worst-case memory request. The wrapper uses `.override()` to raise the task's `flyte.Resources` on each attempt and catches `flyte.errors.OOMError`.

```python
import flyte
import flyte.errors

async def retry_with_memory(
    task_fn,
    *args,
    initial_memory_mi: int = 250,
    increment_mi: int = 200,
    max_memory_mi: int = 4096,
    cpu: int = 1,
    **kwargs,
):
    current = initial_memory_mi
    while current <= max_memory_mi:
        try:
            return await task_fn.override(
                resources=flyte.Resources(cpu=cpu, memory=f"{current}Mi")
            )(*args, **kwargs)
        except flyte.errors.OOMError:
            if current >= max_memory_mi:
                break
            current = min(current + increment_mi, max_memory_mi)
    raise RuntimeError(f"Task still OOMing at {max_memory_mi}Mi")
```

Because the wrapper only takes the task and its arguments, it works with any task:

```python
@env.task
async def process(data: list[int]) -> int:
    # Business logic that may run out of memory on large inputs.
    return sum(data)

@env.task
async def main(data: list[int]) -> int:
    return await retry_with_memory(process, data, initial_memory_mi=500, max_memory_mi=8192)
```

See [Error handling](./error-handling) for more on `flyte.errors.OOMError` and resource-based recovery.

## Circuit breaker

Run a task over many items in parallel, but stop early ("open the circuit") once failures exceed a threshold, so a systemic problem doesn't burn resources on every remaining item. It launches all invocations with `asyncio.create_task`, processes them as they complete, and cancels the rest when the limit is crossed.

```python
import asyncio
from typing import Callable, List, Optional, TypeVar

T = TypeVar("T")
R = TypeVar("R")

class CircuitBreakerError(Exception):
    """Raised when too many failures occur."""

async def circuit_breaker_execute(
    task_fn: Callable[[T], R], items: List[T], max_failures: int = 3
) -> List[Optional[R]]:
    tasks = [asyncio.create_task(task_fn(item)) for item in items]
    results: List[Optional[R]] = [None] * len(items)
    failures = 0
    pending = set(tasks)

    while pending:
        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
        for task in done:
            idx = tasks.index(task)
            if task.exception():
                failures += 1
                if failures > max_failures:
                    for remaining in pending:
                        remaining.cancel()
                    raise CircuitBreakerError(
                        f"{failures} failures exceed limit of {max_failures}"
                    )
            else:
                results[idx] = task.result()
    return results
```

Failed items come back as `None`; if the failure threshold is crossed, the remaining tasks are cancelled and `CircuitBreakerError` is raised. See [Fanout](./fanout) for the basics of running tasks in parallel and [Controlling parallel execution](./controlling-parallelism) for bounding concurrency.

## Auto batcher

Split a large input into batches, run a map task over each batch in parallel, then combine the results with a reduce step. This bounds how many invocations are in flight at once while still processing everything.

```python
import asyncio
from typing import Any, Callable, List, TypeVar

T = TypeVar("T")
R = TypeVar("R")

def create_batches(data: List[T], batch_size: int) -> List[List[T]]:
    return [data[i : i + batch_size] for i in range(0, len(data), batch_size)]

async def batch_map_reduce(
    map_fn: Callable[[T], R],
    reduce_fn: Callable[[List[R]], Any],
    data: List[T],
    batch_size: int = 10,
) -> Any:
    all_results: List[R] = []
    for batch in create_batches(data, batch_size):
        coros = [asyncio.create_task(map_fn(item)) for item in batch]
        all_results.extend(await asyncio.gather(*coros))
    return reduce_fn(all_results)
```

The map step is a task; the reduce step can be a task or a plain function:

```python
@env.task
async def square(x: int) -> int:
    return x * x

@env.task
async def main(data: list[int]) -> int:
    return await batch_map_reduce(square, sum, data, batch_size=25)
```

For a first-class parallel-map primitive, see `flyte.map` in [Fanout](./fanout).

## Composing the patterns

Because each wrapper is just a function that takes a task, you can layer them (for example, wrap a task in the OOM retrier and then hand *that* to the fallback runner) to build orchestration behavior out of small, reusable pieces without touching the underlying task code.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/other-features ===

This section covers advanced programming patterns and techniques for working with Flyte tasks.

## Task forwarding

When one task calls another task using the normal invocation syntax (e.g., `await inner_task(x)`), Flyte creates a durable action that's recorded in the UI with data passed through the metadata store. However, if you want to execute a task in the same Python VM without this overhead, use the `.forward()` method.

**When to use**: You want to avoid durability overhead and execute task logic directly in the current VM.

```python
import flyte

env = flyte.TaskEnvironment("my-env")

@env.task
async def inner_task(x: int) -> int:
    return x + 1

@env.task
async def outer_task(x: int) -> int:
    # Executes in same VM, no durable action created
    v = await inner_task.forward(x=10)

    # Creates a durable action, recorded in UI
    return await inner_task(v)
```

The `.forward()` method works with both sync and async tasks:

```python
@env.task
def sync_inner_task(x: int) -> int:
    return x + 1

@env.task
def sync_outer_task(x: int) -> int:
    # Direct execution, no remote call
    v = sync_inner_task.forward(x=10)
    return sync_inner_task(v)
```

## Passing tasks and functions as arguments

You can pass both Flyte tasks and regular Python functions as arguments to other tasks. Flyte handles this through pickling, so the code appears as pickled data in the UI.

```python
import typing
import flyte

env = flyte.TaskEnvironment("udfs")

@env.task
async def add_one_udf(x: int) -> int:
    return x + 1

# Regular async function (not a task)
async def fn_add_two_udf(x: int) -> int:
    return x + 2

@env.task
async def run_udf(x: int, udf: typing.Callable[[int], typing.Awaitable[int]]) -> int:
    return await udf(x)

@env.task
async def main() -> list[int]:
    # Pass a Flyte task as an argument
    result_one = await run_udf(5, add_one_udf)

    # Pass a regular function as an argument
    result_two = await run_udf(5, fn_add_two_udf)

    return [result_one, result_two]
```

**Note**: Both tasks and regular functions are serialized via pickling when passed as arguments.

## Custom action names

By default, actions in the UI use the task's function name. You can provide custom, user-friendly names using the `short_name` parameter.

### Set at task definition

```python
import flyte

env = flyte.TaskEnvironment("friendly_names")

@env.task(short_name="my_task")
async def some_task() -> str:
    return "Hello, Flyte!"
```

### Override at call time

```python
@env.task(short_name="entrypoint")
async def main() -> str:
    # Uses the default short_name "my_task"
    s = await some_task()

    # Overrides to use "my_name" for this specific action
    return s + await some_task.override(short_name="my_name")()
```

This is useful when the same task is called multiple times with different contexts, making the UI more readable.

## Invoking async functions from sync tasks

When migrating from Flyte 1.x to 2.0, you may have legacy sync code that needs to call async functions. Use `nest_asyncio.apply()` to enable `asyncio.run()` within sync tasks.

```python
import asyncio
import nest_asyncio
import flyte

env = flyte.TaskEnvironment(
    "async_in_sync",
    image=flyte.Image.from_debian_base().with_pip_packages("nest_asyncio"),
)

# Apply at module level
nest_asyncio.apply()

async def async_function() -> str:
    await asyncio.sleep(1)
    return "done"

@env.task
def sync_task() -> str:
    # Now you can use asyncio.run() in a sync task
    return asyncio.run(async_function())
```

**Important**:

- Call `nest_asyncio.apply()` at the module level before defining tasks
- Add `nest_asyncio` to your image dependencies
- This is particularly useful during migration when you have mixed sync/async code

## Async and sync task interoperability

When migrating from older sync-based code to async tasks, or when working with mixed codebases, you need to call sync tasks from async parent tasks. Flyte provides the `.aio` method on every task (even sync ones) to enable this.

### Calling sync tasks from async tasks

Every sync task automatically has an `.aio` property that returns an async-compatible version:

```python
import flyte

env = flyte.TaskEnvironment("mixed-tasks")

@env.task
def sync_task(x: int) -> str:
    """Legacy sync task"""
    return f"Processed {x}"

@env.task
async def async_task(x: int) -> str:
    """New async task that calls legacy sync task"""
    # Use .aio to call sync task from async context
    result = await sync_task.aio(x)
    return result
```

### Using with `flyte.map.aio()`

When you need to call sync tasks in parallel from an async context, use `flyte.map.aio()`:

```python
from typing import List
import flyte

env = flyte.TaskEnvironment("map-example")

@env.task
def sync_process(x: int) -> str:
    """Synchronous processing task"""
    return f"Task {x}"

@env.task
async def async_main(n: int) -> List[str]:
    """Async task that maps over sync task"""
    results = []

    # Map over sync task from async context
    async for result in flyte.map.aio(sync_process, range(n)):
        if isinstance(result, Exception):
            raise result
        results.append(result)

    return results
```

**Why this matters**: This pattern is powerful when migrating from Flyte 1.x or integrating legacy sync tasks with new async code. You don't need to rewrite all sync tasks at once; they can be called from async contexts.

## Using AnyIO in async tasks

Flyte async tasks support `anyio` for structured concurrency as an alternative to `asyncio.gather()`.

```python
import anyio
import aioresult
import flyte

env = flyte.TaskEnvironment(
    "anyio_example",
    image=flyte.Image.from_debian_base().with_pip_packages("anyio", "aioresult"),
)

@env.task
async def process_item(x: int) -> int:
    return x * 2

@env.task
async def batch_process(items: list[int]) -> list[int]:
    captured_results = []

    async with anyio.create_task_group() as tg:
        # Start multiple tasks concurrently
        for item in items:
            captured_results.append(
                aioresult.ResultCapture.start_soon(tg, process_item, item)
            )

    # Extract results
    return [r.result() for r in captured_results]
```

**Note**: You can use anyio's task groups, timeouts, and other structured concurrency primitives within Flyte async tasks.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment ===

# Run and deploy tasks

You have seen how to configure and build the tasks that compose your project.
Now you need to decide how to execute them on your Flyte backend.

Flyte offers two distinct approaches for getting your tasks onto the backend:

**Use `flyte run` when you're iterating and experimenting:**

- Quickly test changes during development
- Try different parameters or code modifications
- Debug issues without creating permanent artifacts
- Prototype new ideas rapidly

**Use `flyte deploy` when your project is ready to be formalized:**

- Freeze a stable version of your tasks for repeated use
- Share tasks with team members or across environments
- Move from experimentation to a more structured workflow
- Create a permanent reference point (not necessarily production-ready)

This section explains both approaches and when to use each one.

## Ephemeral deployment and immediate execution

The `flyte run` CLI command and the `flyte.run()` SDK function are used to **ephemerally deploy** and **immediately execute** a task on the backend in a single step.
The task can be re-run and its execution and outputs can be observed in the **Runs list** UI, but it is not permanently added to the **Tasks list** on the backend.

Let's say you have the following file called `greeting.py`:

```python
# greeting.py

import flyte

env = flyte.TaskEnvironment(name="greeting_env")

@env.task
async def greet(message: str) -> str:
    return f"{message}!"
```

### Programmatic

You can run the task programmatically using the `flyte.run()` function:

```python
# greeting.py

import flyte

env = flyte.TaskEnvironment(name="greeting_env")

@env.task
async def greet(message: str) -> str:
    return f"{message}!"

if __name__ == "__main__":
    flyte.init_from_config()
    result = flyte.run(greet, message="Good morning!")
    print(f"Result: {result}")
```

Here we add a `__main__` block to the `greeting.py` file that initializes the Flyte SDK from the configuration file and then calls `flyte.run()` with the `greet` task and its argument.
Now you can run the `greet` task on the backend just by executing the `greeting.py` file locally as a script:

```bash
python greeting.py
```

### CLI

The general form of the command for running a task from a local file is:

```bash
flyte run <file_path> <task_name> <args>
```

So, to run the `greet` task defined in the `greeting.py` file, you would run:

```bash
flyte run greeting.py greet --message "Good morning!"
```

This command:

1. **Temporarily deploys** the task environment named `greeting_env` (held by the variable `env`) that contains the `greet` task.
2. **Executes** the `greet` function with argument `message` set to `"Good morning!"`. Note that `message` is the actual parameter name defined in the function signature.
3. **Returns** the execution results and displays them in the terminal.

For how to pass inputs of other types (datetimes, durations, enums, files, dataclasses, and more) on the command line, see **Tasks > Run and deploy tasks > Run command options > Task argument passing > Passing inputs by type**.

For more details on how `flyte run` and `flyte.run()` work under the hood, see **Tasks > Run and deploy tasks > How task run works**.

## Persistent deployment

The `flyte deploy` CLI command and the `flyte.deploy()` SDK function are used to **persistently deploy** a task environment (and all its contained tasks) to the backend.
The tasks within the deployed environment will appear in the **Tasks list** UI on the backend and can then be executed multiple times without needing to redeploy them.

### Programmatic

You can deploy programmatically using the `flyte.deploy()` function:

```python
# greeting.py

import flyte

env = flyte.TaskEnvironment(name="greeting_env")

@env.task
async def greet(message: str) -> str:
    return f"{message}!"

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(env)
    print(deployments[0].summary_repr())
```

Now you can deploy the `greeting_env` task environment (and therefore the `greet()` task) just by executing the `greeting.py` file locally as a script.

```bash
python greeting.py
```

### CLI

The general form of the command for deploying a task environment from a local file is:

```bash
flyte deploy <file_path> <task_environment_variable>
```

So, using the same `greeting.py` file as before, you can deploy the `greeting_env` task environment like this:

```bash
flyte deploy greeting.py env
```

This command deploys the task environment *assigned to the variable `env`* in the `greeting.py` file, which is the `TaskEnvironment` named `greeting_env`.

Notice that you must specify the *variable* to which the `TaskEnvironment` is assigned (`env` in this case), not the name of the environment itself (`greeting_env`).

Deploying a task environment deploys all tasks defined within it. Here, that means all functions decorated with `@env.task`.
In this case there is just one: `greet()`.

For more details on how `flyte deploy` and `flyte.deploy()` work under the hood, see **Tasks > Run and deploy tasks > How task deployment works**.

## Running already deployed tasks

If you have already deployed your task environment, you can run its tasks without redeploying by using the `flyte run` CLI command or the `flyte.run()` SDK function in a slightly different way. Alternatively, you can always initiate execution of a deployed task from the UI.

### Programmatic

You can run already-deployed tasks programmatically using the `flyte.run()` function.
For example, to run the previously deployed `greet` task from the `greeting_env` environment:

```python
# greeting.py

import flyte

env = flyte.TaskEnvironment(name="greeting_env")

@env.task
async def greet(message: str) -> str:
    return f"{message}!"

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env)
    task = flyte.remote.Task.get("greeting_env.greet", auto_version="latest")
    result = flyte.run(task, message="Good morning!")
    print(f"Result: {result}")
```

When you execute this script locally, it will:

- Deploy the `greeting_env` task environment as before.
- Retrieve the already-deployed `greet` task using `flyte.remote.Task.get()`, specifying its full task reference as a string: `"greeting_env.greet"`.
- Call `flyte.run()` with the retrieved task and its argument.

For more details on how running already-deployed tasks works, see **Tasks > Run and deploy tasks > How task run works > Running deployed tasks**.

### CLI

To run a permanently deployed task using the `flyte run` CLI command, use the special `deployed-task` keyword followed by the task reference in the format `{environment_name}.{task_name}`. For example, to run the previously deployed `greet` task from the `greeting_env` environment:

```bash
flyte run deployed-task greeting_env.greet --message "World"
```

Notice that now that the task environment is deployed, you use its name (`greeting_env`), not by the variable name to which it was assigned in source code (`env`).
The task environment name plus the task name (`greet`) are combined with a dot (`.`) to form the full task reference: `greeting_env.greet`.
The special `deployed-task` keyword tells the CLI that you are referring to a task that has already been deployed. In effect, it replaces the file path argument used for ephemeral runs.

When executed, this command will run the already-deployed `greet` task with argument `message` set to `"World"`. You will see the result printed in the terminal. You can also, of course, observe the execution in the **Runs list** UI.

To execute a deployed task in a different project or domain than your configured defaults, use `--run-project` and `--run-domain`:

```bash
flyte run --run-project prod-project --run-domain production deployed-task greeting_env.greet --message "World"
```

For all `flyte run` options, see **Tasks > Run and deploy tasks > Run command options**.

## Configuring runs with `flyte.with_runcontext()`

Both `flyte run` and `flyte.run()` accept a range of invocation-time parameters that control where the run executes, where outputs are stored, caching behavior, and more.
Programmatically, these are set with `flyte.with_runcontext()` before calling `.run()`.
Inside a running task, `flyte.ctx()` provides read access to the same context.

For the full parameter reference, see **Tasks > Run and deploy tasks > Run context**.

<!--
TODO: Add link to Flyte remote documentation when available
For details on Flyte remote functionality, see the [Flyte remote]().
-->

## Subpages
- **Tasks > Run and deploy tasks > How task run works**
- **Tasks > Run and deploy tasks > Interact with runs and actions**
- **Tasks > Run and deploy tasks > View logs**
- **Tasks > Run and deploy tasks > Work with local data**
- **Tasks > Run and deploy tasks > Run command options**
- **Tasks > Run and deploy tasks > Build a custom CLI**
- **Tasks > Run and deploy tasks > How task deployment works**
- **Tasks > Run and deploy tasks > Deploy command options**
- **Tasks > Run and deploy tasks > Code packaging for remote execution**
- **Tasks > Run and deploy tasks > Deployment patterns**
- **Tasks > Run and deploy tasks > Run context**
- **Tasks > Run and deploy tasks > Entrypoint tasks**
- **Tasks > Run and deploy tasks > Run a Python script**
- **Tasks > Run and deploy tasks > Run with notifications**
- **Tasks > Run and deploy tasks > Rerun a run**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/how-task-run-works ===

# How task run works

The `flyte run` command and `flyte.run()` SDK function support three primary execution modes:

1. **Ephemeral deployment + run**: Automatically prepare task environments ephemerally and execute tasks (development shortcut)
2. **Run deployed task**: Execute permanently deployed tasks without redeployment
3. **Local execution**: Run tasks on your local machine for development and testing

Additionally, you can run deployed tasks through the Flyte/Union UI for interactive execution and monitoring.

## Ephemeral deployment + run: The development shortcut

The most common development pattern combines ephemeral task preparation and execution in a single command, automatically handling the temporary deployment process when needed.

### Programmatic

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
async def my_task(name: str) -> str:
    return f"Hello, {name}!"

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy and run in one step
    result = flyte.run(my_task, name="World")
    print(f"Result: {result}")
    print(f"Execution URL: {result.url}")
```

### CLI

```bash
flyte run my_example.py my_task --name "World"
```

With explicit project and domain:

```bash
flyte run --project my-project --domain development my_example.py my_task --name "World"
```

With deployment options:

```bash
flyte run --version v1.0.0 --copy-style all my_example.py my_task --name "World"
```

**How it works:**

1. **Environment discovery**: Flyte loads the specified Python file and identifies task environments
2. **Ephemeral preparation**: Temporarily prepares the task environment for execution (similar to deployment but not persistent)
3. **Task execution**: Immediately runs the specified task with provided arguments in the ephemeral environment
4. **Result return**: Returns execution results and monitoring URL
5. **Cleanup**: The ephemeral environment is not stored permanently in the backend

**Benefits of ephemeral deployment + run:**

- **Development efficiency**: No separate permanent deployment step required
- **Always current**: Uses your latest code changes without polluting the backend
- **Clean development**: Ephemeral environments don't clutter your task registry
- **Integrated workflow**: Single command for complete development cycle

## Running deployed tasks

For production workflows or when you want to use stable deployed versions, you can run tasks that have been **permanently deployed** with `flyte deploy` without triggering any deployment process.

### Programmatic

```python
import flyte

flyte.init_from_config()

# Method 1: Using remote task reference
deployed_task = flyte.remote.Task.get("my_env.my_task", version="v1.0.0")
result = flyte.run(deployed_task, name="World")

# Method 2: Get latest version
deployed_task = flyte.remote.Task.get("my_env.my_task", auto_version="latest")
result = flyte.run(deployed_task, name="World")
```

### CLI

```bash
flyte run deployed-task my_env.my_task --name "World"
```

With a specific project and domain:

```bash
flyte run --project prod --domain production deployed-task my_env.my_task --batch_size 1000
```

**Task reference format:** `{environment_name}.{task_name}`

- `environment_name`: The `name` property of your `TaskEnvironment`
- `task_name`: The function name of your task

>[!NOTE]
> When you deploy a task environment with `flyte deploy`, you specify the `TaskEnvironment` by the variable to which it is assigned.
> Once deployed, you refer to it by its `name` property.

**Benefits of running deployed tasks:**

- **Performance**: No deployment overhead, faster execution startup
- **Stability**: Uses tested, stable versions of your code
- **Production safety**: Isolated from local development changes
- **Version control**: Explicit control over which code version runs

## Local execution

For development, debugging, and testing, you can run tasks locally on your machine without any backend interaction.

### Programmatic

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
async def my_task(name: str) -> str:
    return f"Hello, {name}!"

# Method 1: No client configured (defaults to local)
result = flyte.run(my_task, name="World")

# Method 2: Explicit local mode
flyte.init_from_config()  # Client configured
result = flyte.with_runcontext(mode="local").run(my_task, name="World")
```

### CLI

```bash
flyte run --local my_example.py my_task --name "World"
```

With development data:

```bash
flyte run --local data_pipeline.py process_data --input_path "/local/data" --debug true
```

**Benefits of local execution:**

- **Rapid development**: Instant feedback without network latency
- **Debugging**: Full access to local debugging tools
- **Offline development**: Works without backend connectivity
- **Resource efficiency**: Uses local compute resources

## Running tasks through the Union UI

If you are running your Flyte code on a Union backend, the UI provides an interactive way to run deployed tasks with form-based input and real-time monitoring.

### Accessing task execution in the Union UI

1. **Navigate to tasks**: Go to your project → domain → Tasks section
2. **Select task**: Choose the task environment and specific task
3. **Launch execution**: Click "Launch" to open the execution form
4. **Provide inputs**: Fill in task parameters through the web interface
5. **Monitor progress**: Watch real-time execution progress and logs

**UI execution benefits:**

- **User-friendly**: No command-line expertise required
- **Visual monitoring**: Real-time progress visualization
- **Input validation**: Built-in parameter validation and type checking
- **Execution history**: Easy access to previous runs and results
- **Sharing**: Shareable execution URLs for collaboration

Here is a short video demonstrating task execution through the Union UI:

📺 [Watch on YouTube](https://www.youtube.com/watch?v=id="8jbau9yGoDg)

## Execution flow and architecture

### Fast registration architecture

Flyte v2 uses "fast registration" to enable rapid development cycles:

#### How it works

1. **Container images** contain the runtime environment and dependencies
2. **Code bundles** contain your Python source code (stored separately)
3. **At runtime**: Code bundles are downloaded and injected into running containers

#### Benefits

- **Rapid iteration**: Update code without rebuilding images
- **Resource efficiency**: Share images across multiple deployments
- **Version flexibility**: Run different code versions with same base image
- **Caching optimization**: Separate caching for images vs. code

#### When code gets injected

At task execution time, the fast registration process follows these steps:

1. **Container starts** with the base image containing runtime environment and dependencies
2. **Code bundle download**: The Flyte agent downloads your Python code bundle from storage
3. **Code extraction**: The code bundle is extracted and mounted into the running container
4. **Task execution**: Your task function executes with the injected code

### Ephemeral preparation logic

When using ephemeral deploy + run mode, Flyte determines whether temporary preparation is needed:

```mermaid
graph TD
    A[flyte run command] --> B{Need preparation?}
    B -->|Yes| C[Ephemeral preparation]
    B -->|No| D[Use cached preparation]
    C --> E[Execute task]
    D --> E
    E --> F[Cleanup ephemeral environment]
```

### Execution modes comparison

| Mode | Deployment | Performance | Use Case | Code Version |
|------|------------|-------------|-----------|--------------|
| Ephemeral Deploy + Run | Ephemeral (temporary) | Medium | Development, testing | Latest local |
| Run Deployed | None (uses permanent deployment) | Fast | Production, stable runs | Deployed version |
| Local | None | Variable | Development, debugging | Local |
| UI | None | Fast | Interactive, collaboration | Deployed version |


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/interacting-with-runs ===

# Interact with runs and actions

When a task is launched, the resulting execution is called a **run**.
Because tasks typically call other tasks, a run will almost always involve multiple sub-task executions. Each such execution is called an **action**.

Through the Flyte SDK and CLI, you can interact with the run and its actions to monitor progress, retrieve results, and access data. This section explains how to work with runs and actions programmatically and through the CLI.

## Understanding runs and actions

Runs are not declared explicitly in the code of the entry point task.
Instead, they are simply a result of the task being invoked in a specific way:

* User with `flyte run`
* User via the UI
* Other code calling `flyte.run()`
* [Trigger](../task-configuration/triggers)

When a task is invoked in one of these ways, it creates a run to represent the execution of that task and all its nested tasks, considered together.
Each task execution within that run is represented by an **action**.
The entry point task execution is represented by the main action (usually called `a0`), and then every nested call of one task from another creates an additional action.

```mermaid
graph TD
    A[Run] --> B[Action a0 - Main task]
    B --> C[Action a1 - Nested task]
    B --> D[Action a2 - Nested task]
    D --> E[Action a3 - Deeply nested task]
```

Because what constitutes a run depends only on how a task is invoked, the same task can execute as a deeply nested action in one run and the main action in another run.
Unlike Flyte 1, there is no explicit `@workflow` construct in Flyte 2; instead, "workflows" are defined implicitly by the structure of task composition and the entry point chosen at runtime.

> [!NOTE]
> Despite there being no explicit `@workflow` decorator, you'll often see the assemblage of tasks referred to as a "workflow" in documentation and discussions. The top-most task in a run is sometimes referred to as the "parent", "driver", or "entry point" task of the "workflow".
> In these docs we will sometime use "workflow" informally to refer to the collection of tasks (considered statically) involved in a run.

### Key concepts

- **Attempts**: Each action can have multiple attempts due to retries. Retries occur for two reasons:
  - User-configured retries for handling transient failures
  - Automatic system retries for infrastructure issues

- **Phases**: Both runs and actions progress through phases (e.g., QUEUED, RUNNING, SUCCEEDED, FAILED) until reaching a terminal state

- **Durability**: Flyte is a durable execution engine, so every input, output, failure, and attempt is recorded for each action. All data is persisted, allowing you to retrieve information about runs and actions even after completion

## Working with runs

Runs are created when you execute tasks using `flyte run` or `flyte.run()`. For details on running tasks, see [how task run works](./how-task-run-works). To learn about running previously deployed remote tasks, see [remote tasks](../task-programming/remote-tasks).

### Retrieving a run

### Programmatic

Use `flyte.remote.Run.get()` to retrieve information about a run:

```python
import flyte

flyte.init_from_config()

# Get a run by name
run = flyte.remote.Run.get("my_run_name")

# Access basic information
print(run.url)        # UI URL for the run
print(run.action.phase)  # Phase of the main action
```

### CLI

Get a specific run:

```bash
flyte get run my_run_name
```

List all runs:

```bash
flyte get run
```

Use `--project` and `--domain` to scope results to a specific [project-domain pair](../../get-started/core-concepts/projects-and-domains).
For all available options, see the [CLI reference](../../../api-reference/flyte-cli#flyte-get-run).

### Filtering runs by label

If runs were launched with [labels](./run-context#identity-and-resources) — arbitrary `key=value` metadata attached at launch — you can filter the run list by those labels. Multiple label filters combine with **AND** semantics: a run must match every label you specify to be returned.

### CLI

Filter by one or more `key=value` labels:

```bash
flyte get run --with-label team=ml --with-label env=prod
```

Filter by the presence of a label key, regardless of its value:

```bash
flyte get run --with-label-key team
```

### Programmatic

`flyte.remote.Run.listall()` accepts the same filters and returns an async iterator of runs:

```python
import flyte

flyte.init_from_config()

# Runs that carry BOTH labels (AND semantics)
async for run in flyte.remote.Run.listall(with_labels={"team": "ml", "env": "prod"}):
    print(run.name)

# Runs that have the "team" label key set to any value
async for run in flyte.remote.Run.listall(with_label_keys=["team"]):
    print(run.name)
```

### UI

In the run list, each run shows its labels as chips in the **Labels** column. Click a label chip to filter the list down to the runs that share that label.

To attach labels to a run when you launch it, see [Run context](./run-context#identity-and-resources).

### Watching run progress

Monitor a run as it progresses through phases:

```python
# Wait for run to complete
run = flyte.run(my_task, input_data="test")
run.wait()  # Blocks until terminal state

# Check if done
if run.action.done():
    print("Run completed!")
```

### Getting detailed run information

Use `flyte.remote.RunDetails` for comprehensive information including nested actions and metadata:

```python
run_details = flyte.remote.RunDetails.get(name="my_run_name")

# Access detailed information
print(run_details.pb2)  # Full protobuf representation
```

## Working with actions

Actions represent individual task executions within a run. Each action has a unique identifier within its parent run.

### Retrieving an action

### Programmatic

```python
# Get a specific action by run name and action name
action = flyte.remote.Action.get(
    run_name="my_run_name",
    name="a0"  # Main action
)

# Access action information
print(action.phase)       # Current phase
print(action.task_name)   # Task being executed
print(action.start_time)  # Execution start time
```

### CLI

Get a specific action:

```bash
flyte get action my_run_name a0
```

List all actions for a run:

```bash
flyte get action my_run_name
```

For all available options, see the [CLI reference](../../../api-reference/flyte-cli#flyte-get-action).

### Nested actions

Deeply nested actions are uniquely identified by their path under the run:

```python
# Get a nested action
nested_action = flyte.remote.Action.get(
    run_name="my_run_name",
    name="a1"  # Nested action identifier
)
```

### Getting detailed action information

Use `flyte.remote.ActionDetails` for comprehensive action information:

```python
action_details = flyte.remote.ActionDetails.get(
    run_name="my_run_name",
    name="a0"
)

# Access detailed information
print(action_details.pb2)  # Full protobuf representation
```

## Accessing logs

Each action captures the logs emitted by its task, per attempt. Stream them from the CLI with `flyte get logs`, or view them on the action in the console. For details and all available options, see [View logs](./view-logs).

```bash
# Logs for the run's main action
flyte get logs my_run_name

# Logs for a specific action
flyte get logs my_run_name a1
```

## Retrieving inputs and outputs

### Programmatic

Both `Run` and `Action` objects provide methods to retrieve inputs and outputs:

```python
run = flyte.remote.Run.get("my_run_name")

# Get inputs - returns ActionInputs (dict-like)
inputs = run.inputs()
print(inputs)  # {"param_name": value, ...}

# Get outputs - returns tuple
outputs = run.outputs()
print(outputs)  # (result1, result2, ...)

# Single output
single_output = outputs[0]

# No outputs are represented as (None,)
```

**Important notes:**

- **Inputs**: Returned as `flyte.remote.ActionInputs`, a dictionary with parameter names as keys and values as the actual data passed in
- **Outputs**: Always returned as `flyte.remote.ActionOutputs` tuple, even for single outputs or no outputs
- **No outputs**: Represented as `(None,)`
- **Availability**: Outputs are only available if the action completed successfully
- **Type safety**: Flyte's rich type system converts data to an intermediate representation, allowing retrieval even without the original dependencies installed

### CLI

Get inputs and outputs for a run:

```bash
flyte get io my_run_name
```

Get inputs and outputs for a specific action:

```bash
flyte get io my_run_name a1
```

For all available options, see the [CLI reference](../../../api-reference/flyte-cli#flyte-get-io).

### Handling failures

If an action fails, outputs are not available, but you can retrieve error information:

```python
action = flyte.remote.Action.get(run_name="my_run_name", name="a0")

if action.phase == flyte.models.ActionPhase.FAILED:
    # Outputs will raise an error
    try:
        outputs = action.outputs()
    except RuntimeError as e:
        print("Action failed, outputs not available")

    # Get error details instead
    action_details = flyte.remote.ActionDetails.get(
        run_name="my_run_name",
        name="a0"
    )
    print(action_details.pb2.error_info)
```

## Understanding data storage

Flyte handles different types of data differently, as explained in [data flow](../../run-scaling/data-flow):

- **Parameterized data** (primitives, small objects): Returned directly in inputs/outputs
- **Large data** (files, directories, DataFrames, models): Stored in cloud storage (S3, GCS, Azure Blob Storage)

When you retrieve outputs containing large data, Flyte returns references rather than the actual data. To access the actual raw data, you need proper cloud storage permissions and configuration.

## Accessing large data from cloud storage

To download and work with files, directories, and DataFrames stored in cloud object storage, you must configure storage access with appropriate credentials.

### S3 storage access

To access data stored in Amazon S3:

**1. Set environment variables:**

```bash
export FLYTE_AWS_ACCESS_KEY_ID="your-access-key-id"
export FLYTE_AWS_SECRET_ACCESS_KEY="your-secret-access-key"
```

These are standard AWS credential environment variables that Flyte recognizes. They are your IAM user's access keys. If you don't already have them, follow the AWS guide on [managing access keys for IAM users](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_access-keys.html) to create them.

**2. Initialize Flyte with S3 storage configuration:**

```python
import flyte
import flyte.storage

# Auto-configure from environment variables
flyte.init_from_config(
    storage=flyte.storage.S3.auto(region="us-east-2")
)

# Or provide credentials explicitly
flyte.init_from_config(
    storage=flyte.storage.S3(
        access_key_id="your-access-key-id",
        secret_access_key="your-secret-access-key",
        region="us-east-2"
    )
)
```

**3. Access data from outputs:**

```python
run = flyte.remote.Run.get("my_run_name")
outputs = run.outputs()

# Outputs containing files, dataframes, etc. can now be downloaded
dataframe = outputs[0]
df = await dataframe.open(pd.DataFrame).all()
```

### GCS storage access

To access data stored in Google Cloud Storage:

**1. Set environment variables:**

```bash
export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service-account-key.json"
```

This is the standard Google Cloud authentication method using service account credentials.

**2. Initialize Flyte with GCS storage configuration:**

```python
import flyte
import flyte.storage

# Auto-configure from environment
flyte.init_from_config(
    storage=flyte.storage.GCS.auto()
)

# Or configure explicitly
flyte.init_from_config(
    storage=flyte.storage.GCS()
)
```

**3. Access data from outputs:**

```python
run = flyte.remote.Run.get("my_run_name")
outputs = run.outputs()

# Download data as needed
file_output = outputs[0]
# Work with file output
```

### Azure Blob Storage access

To access data stored in Azure Blob Storage (ABFS):

**1. Set environment variables:**

For storage account key authentication:

```bash
export AZURE_STORAGE_ACCOUNT_NAME="your-storage-account"
export AZURE_STORAGE_ACCOUNT_KEY="your-account-key"
```

For service principal authentication:

```bash
export AZURE_TENANT_ID="your-tenant-id"
export AZURE_CLIENT_ID="your-client-id"
export AZURE_CLIENT_SECRET="your-client-secret"
export AZURE_STORAGE_ACCOUNT_NAME="your-storage-account"
```

**2. Initialize Flyte with Azure storage configuration:**

```python
import flyte
import flyte.storage

# Auto-configure from environment variables
flyte.init_from_config(
    storage=flyte.storage.ABFS.auto()
)

# Or provide credentials explicitly
flyte.init_from_config(
    storage=flyte.storage.ABFS(
        account_name="your-storage-account",
        account_key="your-account-key"
    )
)

# Or use service principal
flyte.init_from_config(
    storage=flyte.storage.ABFS(
        account_name="your-storage-account",
        tenant_id="your-tenant-id",
        client_id="your-client-id",
        client_secret="your-client-secret"
    )
)
```

**3. Access data from outputs:**

```python
run = flyte.remote.Run.get("my_run_name")
outputs = run.outputs()

# Download data as needed
directory_output = outputs[0]
# Work with directory output
```

## Complete example

Here's a complete example showing how to launch a run and interact with it:

```python
import flyte
import flyte.storage

# Initialize with storage access
flyte.init_from_config(
    storage=flyte.storage.S3.auto(region="us-east-2")
)

# Define and run a task
env = flyte.TaskEnvironment(name="data_processing")

@env.task
async def process_data(input_value: str) -> str:
    return f"Processed: {input_value}"

# Launch the run
run = flyte.run(process_data, input_value="test_data")

# Monitor progress
print(f"Run URL: {run.url}")
run.wait()

# Check status
if run.action.done():
    print(f"Run completed with phase: {run.action.phase}")

    # Get inputs and outputs
    inputs = run.inputs()
    print(f"Inputs: {inputs}")

    outputs = run.outputs()
    print(f"Outputs: {outputs}")

    # Access the result
    result = outputs[0]
    print(f"Result: {result}")
```

## API reference

### Key classes

- `flyte.remote.Run` - Represents a run with basic information
- `flyte.remote.RunDetails` - Detailed run information including all actions
- `flyte.remote.Action` - Represents an action with basic information
- `flyte.remote.ActionDetails` - Detailed action information including error details
- `flyte.remote.ActionInputs` - Dictionary-like object containing action inputs
- `flyte.remote.ActionOutputs` - Tuple containing action outputs

### CLI commands

For complete CLI documentation and all available options, see the [Flyte CLI reference](../../../api-reference/flyte-cli):

- [`flyte get run`](../../../api-reference/flyte-cli#flyte-get-run) - Get run information
- [`flyte get action`](../../../api-reference/flyte-cli#flyte-get-action) - Get action information
- [`flyte get io`](../../../api-reference/flyte-cli#flyte-get-io) - Get inputs and outputs
- [`flyte get logs`](../../../api-reference/flyte-cli#flyte-get-logs) - Get action logs

### Storage configuration

- `flyte.storage.S3` - Amazon S3 configuration
- `flyte.storage.GCS` - Google Cloud Storage configuration
- `flyte.storage.ABFS` - Azure Blob Storage configuration

For more details on data flow and storage, see [data flow](../../run-scaling/data-flow).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/view-logs ===

# View logs

Every action in a run captures the logs its task emits while executing.
Because Flyte is a durable execution engine, these logs are persisted per action and per attempt, so you can retrieve them while a run is in progress or after it has reached a terminal state.

There are two ways to view logs:

- **The CLI**: stream logs for a run or a specific action with `flyte get logs`.
- **The console**: open the run in the UI and inspect logs on any of its actions.

## Stream logs with the CLI

Use `flyte get logs` to stream the logs for a run or action:

```bash
flyte get logs <run_name> [<action_name>]
```

If you provide only the run name, Flyte streams the logs for the run's parent (main) action:

```bash
flyte get logs my_run
```

To see the logs for a specific action within the run, provide the action name as the second argument. Action names such as `a0` (the main action) and `a1`, `a2`, … (nested actions) identify each task execution in the run. See [Interact with runs and actions](./interacting-with-runs#understanding-runs-and-actions):

```bash
flyte get logs my_run a0
```

### Raw vs. pretty output

By default, logs are shown in raw format and scroll the terminal as they arrive.

To instead tail the logs in an auto-scrolling box that shows only the most recent lines, pass `--pretty`. Use `--lines`/`-l` to set how many lines the box keeps in view (default `30`); this limit only applies in pretty mode:

```bash
flyte get logs my_run a0 --pretty --lines 50
```

To prepend a timestamp to each log line, add `--show-ts`:

```bash
flyte get logs my_run a0 --show-ts
```

## View the logs for a specific attempt

An action can run more than once: Flyte records a separate attempt for each user-configured retry and for each automatic system retry (see [Key concepts](./interacting-with-runs#key-concepts)).

By default, `flyte get logs` shows the logs for the **latest** attempt. To inspect an earlier attempt, pass its number with `--attempt`/`-a`:

```bash
flyte get logs my_run a0 --attempt 1
```

## Filter out system logs

Alongside your task's own output, the logs include system messages emitted by the Flyte runtime. To hide those and show only the logs produced by your task, pass `--filter-system`:

```bash
flyte get logs my_run a0 --filter-system
```

## Scope to a project and domain

Like other `flyte get` commands, `flyte get logs` resolves the run within your configured [project and domain](../../get-started/core-concepts/projects-and-domains). Override them for a single invocation with `--project`/`-p` and `--domain`/`-d`:

```bash
flyte get logs my_run a0 --project my-project --domain development
```

## Command options

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--lines`, `-l` | integer | `30` | Number of lines to show; only applies with `--pretty`. |
| `--show-ts` | boolean | `False` | Show timestamps. |
| `--pretty` | boolean | `False` | Show logs in an auto-scrolling box, limited to `--lines` lines. |
| `--attempt`, `-a` | integer | latest | Attempt number to show logs for; defaults to the latest attempt. |
| `--filter-system` | boolean | `False` | Filter all system logs from the output. |
| `--project`, `-p` | text | configured | Project to which this command applies. |
| `--domain`, `-d` | text | configured | Domain to which this command applies. |

For the full command reference, see [`flyte get logs`](../../../api-reference/flyte-cli#flyte-get-logs).

## View logs in the console

Logs are also available in the Flyte console. Open the run. The `url` attribute of a `flyte.remote.Run` gives its console link:

```python
import flyte

flyte.init_from_config()

run = flyte.remote.Run.get("my_run")
print(run.url)  # Console URL for the run
```

From the run view, select an action to see its logs, phases, inputs, and outputs. Each attempt of an action has its own logs, as with the CLI.

> [!NOTE]
> The exact log-viewing controls in the console depend on your deployment and console version. If the layout differs from what is described here, verify the current behavior in your live console.

## Related

- [Interact with runs and actions](./interacting-with-runs): retrieve runs, actions, inputs, and outputs.
- [Flyte CLI reference](../../../api-reference/flyte-cli#flyte-get-logs): the complete `flyte get logs` reference.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/work-with-local-data ===

# Work with local data

When running Flyte tasks that take inputs like DataFrames, files, or directories, data is passed between actions through the configured blob store. For details on how data flows through your workflows, see [data flow](../../run-scaling/data-flow).

Flyte provides several built-in types for handling data:

- `flyte.io.DataFrame` for tabular data
- `flyte.io.File` for individual files
- `flyte.io.Dir` for directories

You can also create custom type extensions for specialized data types. See [custom types](../task-programming/handling-custom-types) for details.

## Local execution

One of the most powerful features of Flyte is the ability to work with data entirely locally, without creating a remote run. When you run tasks in local mode, all inputs, outputs, and intermediate data stay on your local machine.

```python
import flyte

env = flyte.TaskEnvironment(name="local_data")

@env.task
async def process_data(data: str) -> str:
    return f"Processed: {data}"

# Run locally - no remote storage needed
run = flyte.with_runcontext(mode="local").run(process_data, data="test")
run.wait()
print(run.outputs()[0])
```

For more details on local execution, see [how task run works](./how-task-run-works#local-execution).

## Uploading local data to remote runs

When you want to send local data to a remote task, you need to upload it first. Flyte provides a secure data uploading system that handles this automatically. The same system used for [code bundling](./packaging) can upload files, DataFrames, and directories.

To upload local data, use the Flyte core representation for that type with the `from_local_sync()` method.

### Uploading DataFrames

Use `flyte.io.DataFrame.from_local_sync()` to upload a local DataFrame:

```python
from typing import Annotated

import pandas as pd

import flyte
import flyte.io

img = flyte.Image.from_debian_base()
img = img.with_pip_packages("pandas", "pyarrow")

env = flyte.TaskEnvironment(
    "dataframe_usage",
    image=img,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

@env.task
async def process_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Process a DataFrame and return the result."""
    df["processed"] = True
    return df

if __name__ == "__main__":
    flyte.init_from_config()

    # Create a local pandas DataFrame
    local_df = pd.DataFrame({
        "name": ["Alice", "Bob", "Charlie"],
        "value": [10, 20, 30]
    })

    # Upload the local DataFrame for remote execution
    uploaded_df = flyte.io.DataFrame.from_local_sync(local_df)

    # Pass to a remote task
    run = flyte.run(process_dataframe, df=uploaded_df)
    print(f"Run URL: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

### Uploading files

Use `flyte.io.File.from_local_sync()` to upload a local file:

```python
import tempfile

import flyte
from flyte.io import File

env = flyte.TaskEnvironment(name="file-local")

@env.task
async def process_file(file: File) -> str:
    """Read and process a file."""
    async with file.open("rb") as f:
        content = bytes(await f.read())
        return content.decode("utf-8")

if __name__ == "__main__":
    flyte.init_from_config()

    # Create a temporary local file
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as temp:
        temp.write("Hello, Flyte!")
        temp_path = temp.name

    # Upload the local file for remote execution
    file = File.from_local_sync(temp_path)

    # Pass to a remote task
    run = flyte.run(process_file, file=file)
    print(f"Run URL: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

### Uploading directories

Use `flyte.io.Dir.from_local_sync()` to upload a local directory:

```python
import os
import tempfile

import flyte
from flyte.io import Dir

env = flyte.TaskEnvironment(name="dir-local")

@env.task
async def process_dir(dir: Dir) -> dict[str, str]:
    """Process a directory and return file contents."""
    file_contents = {}
    async for file in dir.walk(recursive=False):
        if file.name.endswith(".py"):
            async with file.open("rb") as f:
                content = bytes(await f.read())
                file_contents[file.name] = content.decode("utf-8")[:100]
    return file_contents

if __name__ == "__main__":
    flyte.init_from_config()

    # Create a temporary directory with test files
    with tempfile.TemporaryDirectory() as temp_dir:
        for i in range(3):
            with open(os.path.join(temp_dir, f"file{i}.py"), "w") as f:
                f.write(f"print('Hello from file {i}!')")

        # Upload the local directory for remote execution
        dir = Dir.from_local_sync(temp_dir)

        # Pass to a remote task
        run = flyte.run(process_dir, dir=dir)
        print(f"Run URL: {run.url}")
        run.wait()
        print(run.outputs()[0])
```

## Passing outputs between runs

If you're passing outputs from a previous run to a new run, no upload is needed. Flyte's data is represented using native references that point to storage locations, so passing them between runs works automatically:

```python
import flyte

flyte.init_from_config()

# Get outputs from a previous run
previous_run = flyte.remote.Run.get("my_previous_run")
previous_output = previous_run.outputs()[0]  # Already a Flyte reference

# Pass directly to a new run - no upload needed
new_run = flyte.run(my_task, data=previous_output)
```

## Performance considerations

The `from_local_sync()` method uses HTTP to upload data. This is convenient but not the most performant option for large datasets.

**Best suited for:**

- Small to medium test datasets
- Development and debugging
- Quick prototyping

**For larger data uploads**, configure cloud storage access and use `flyte.storage` directly:

```python
import flyte
import flyte.storage

# Configure storage access
flyte.init_from_config(
    storage=flyte.storage.S3.auto(region="us-east-2")
)
```

For details on configuring storage access, see [interact with runs and actions](./interacting-with-runs#accessing-large-data-from-cloud-storage).

## Summary

| Scenario | Approach |
|----------|----------|
| Local development and testing | Use local execution mode |
| Small test data to remote tasks | Use `from_local_sync()` |
| Passing data between runs | Pass outputs directly (automatic) |
| Large datasets | Configure `flyte.storage` for direct cloud access |


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/run-command-options ===

# Run command options

The `flyte run` command provides the following options:

**`flyte run [OPTIONS] <PATH>|deployed-task <TASK_NAME>`**

| Option                      | Short | Type   | Default                   | Description                                            |
|-----------------------------|-------|--------|---------------------------|--------------------------------------------------------|
| `--project`                 | `-p`  | text   | *from config*             | Project to run tasks in                                |
| `--domain`                  | `-d`  | text   | *from config*             | Domain to run tasks in                                 |
| `--local`                   |       | flag   | `false`                   | Run the task locally                                   |
| `--copy-style`              |       | choice | `loaded_modules\|all\|none` | Code bundling strategy                                 |
| `--root-dir`                |       | path   | *current dir*             | Override source root directory                         |
| `--raw-data-path`           |       | text   |                           | Override the output location for offloaded data types. |
| `--service-account`         |       | text   |                           | Kubernetes service account.                            |
| `--name`                    |       | text   |                           | Name of the run.                                       |
| `--label`                   |       | text   |                           | User-defined `key=value` label on the run. Repeatable. |
| `--follow`                  | `-f`  | flag   | `false`                   | Wait and watch logs for the parent action.             |
| `--image`                   |       | text   |                           | Image to be used in the run (format: `name=uri`).      |
| `--no-sync-local-sys-paths` |       | flag   | `false`                   | Disable synchronization of local sys.path entries.      |
| `--run-project`             |       | text   | *from config*             | Execute deployed task in this project (`deployed-task` only). |
| `--run-domain`              |       | text   | *from config*             | Execute deployed task in this domain (`deployed-task` only).  |

## `--project`, `--domain`

**`flyte run --domain <DOMAIN> --project <PROJECT> <PATH>|deployed-task <TASK_NAME>`**

You can specify `--project` and `--domain` which will override any defaults defined in your `config.yaml`:

```bash
flyte run my_example.py my_task
```

Specify a target project and domain:

```bash
flyte run --project my-project --domain development my_example.py my_task
```

## `--run-project`, `--run-domain`

**`flyte run --run-project <PROJECT> --run-domain <DOMAIN> deployed-task <TASK_REF>`**

When using the `deployed-task` subcommand, `--run-project` and `--run-domain` specify the [project-domain pair](../../get-started/core-concepts/projects-and-domains) in which to *execute* the task. This lets you run a deployed task in a different project or domain than the one configured in your `config.yaml`:

```bash
flyte run --run-project prod-project --run-domain production deployed-task my_env.my_task
```

If not provided, these default to the `task.project` and `task.domain` values in your configuration file. These options only apply to the `deployed-task` subcommand and are ignored for file-based runs.

## `--local`

**`flyte run --local <PATH> <TASK_NAME>`**

The `--local` option runs tasks locally instead of submitting them to the remote Flyte backend:

```bash
flyte run --local my_example.py my_task --input "test_data"
```

Compare with remote execution:

```bash
flyte run my_example.py my_task --input "test_data"
```

### When to use local execution

- **Development and testing**: Quick iteration without deployment overhead
- **Debugging**: Full access to local debugging tools and environment
- **Resource constraints**: When remote resources are unavailable or expensive
- **Data locality**: When working with large local datasets

## `--copy-style`

**`flyte run --copy-style [loaded_modules|all|none] <PATH> <TASK_NAME>`**

The `--copy-style` option controls code bundling for remote execution.
This applies to the ephemeral preparation step of the `flyte run` command and works similarly to `flyte deploy`:

Smart bundling (default) includes only imported project modules:

```bash
flyte run --copy-style loaded_modules my_example.py my_task
```

Include all project files:

```bash
flyte run --copy-style all my_example.py my_task
```

No code bundling (task must be pre-deployed):

```bash
flyte run --copy-style none deployed-task my_deployed_task
```

### Copy style options

- **`loaded_modules` (default)**: Bundles only imported Python modules from your project
- **`all`**: Includes all files in the project directory
- **`none`**: No bundling; requires permanently deployed tasks

## `--root-dir`

**`flyte run --root-dir <DIRECTORY> <PATH> <TASK_NAME>`**

Override the source directory for code bundling and import resolution:

Run from a monorepo root with a specific root directory:

```bash
flyte run --root-dir ./services/ml ./services/ml/my_example.py my_task
```

Handle cross-directory imports:

```bash
flyte run --root-dir .. my_example.py my_workflow
```

This applies to the ephemeral preparation step of the `flyte run` command.
It works identically to the `flyte deploy` command's `--root-dir` option.

## `--raw-data-path`

**`flyte run --raw-data-path <PATH> <SOURCE> <TASK_NAME>`**

Override the default output location for offloaded data types (large objects, DataFrames, etc.):

Use a custom S3 location for large outputs:

```bash
flyte run --raw-data-path s3://my-bucket/custom-path/ my_example.py process_large_data
```

Use a local directory for development:

```bash
flyte run --local --raw-data-path ./output/ my_example.py my_task
```

### Use cases

- **Custom storage locations**: Direct outputs to specific S3 buckets or paths
- **Cost optimization**: Use cheaper storage tiers for temporary data
- **Access control**: Ensure outputs go to locations with appropriate permissions
- **Local development**: Store large outputs locally when testing

## `--service-account`

**`flyte run --service-account <ACCOUNT_NAME> <PATH> <TASK_NAME>`**

Specify a Kubernetes service account for task execution:

```bash
flyte run --service-account ml-service-account my_example.py train_model
flyte run --service-account data-reader-sa my_example.py load_data
```

### Use cases

- **Cloud resource access**: Service accounts with permissions for S3, GCS, etc.
- **Security isolation**: Different service accounts for different workload types
- **Compliance requirements**: Enforcing specific identity and access policies

## `--name`

**`flyte run --name <EXECUTION_NAME> <PATH> <TASK_NAME>`**

Provide a custom name for the execution run:

```bash
flyte run --name "daily-training-run-2024-12-02" my_example.py train_model
flyte run --name "experiment-lr-0.01-batch-32" my_example.py hyperparameter_sweep
```

### Benefits of custom names

- **Easy identification**: Find specific runs in the Flyte console
- **Experiment tracking**: Include key parameters or dates in names
- **Automation**: Programmatically generate meaningful names for scheduled runs

## `--label`

**`flyte run --label <KEY>=<VALUE> <PATH> <TASK_NAME>`**

Attach one or more user-defined `key=value` labels to the run for filtering and organizing runs. The flag is repeatable:

```bash
flyte run --label team=ml --label env=prod my_example.py train_model
```

Later, list or filter runs by these labels with `flyte get run --with-label team=ml` (see [Filtering runs by label](./interacting-with-runs#filtering-runs-by-label)).

## `--follow`

**`flyte run --follow <PATH> <TASK_NAME>`**

Wait and watch logs for the execution in real-time:

```bash
flyte run --follow my_example.py long_running_task
```

Combine with other options:

```bash
flyte run --follow --name "training-session" my_example.py train_model
```

### Behavior

- **Log streaming**: Real-time output from task execution
- **Blocking execution**: Command waits until task completes
- **Exit codes**: Returns appropriate exit code based on task success/failure

## `--image`

**`flyte run --image <IMAGE_MAPPING> <PATH> <TASK_NAME>`**

Override container images during ephemeral preparation, same as the equivalent `flyte deploy` option:

Override a specific named image:

```bash
flyte run --image gpu=ghcr.io/org/gpu:v2.1 my_example.py gpu_task
```

Override the default image:

```bash
flyte run --image ghcr.io/org/custom:latest my_example.py my_task
```

Multiple image overrides:

```bash
flyte run \
  --image base=ghcr.io/org/base:v1.0 \
  --image gpu=ghcr.io/org/gpu:v2.0 \
  my_example.py multi_env_workflow
```

### Image mapping formats

- **Named mapping**: `name=uri` overrides images created with `Image.from_ref_name("name")`
- **Default mapping**: `uri` overrides the default "auto" image
- **Multiple mappings**: Use multiple `--image` flags for different image references

## `--no-sync-local-sys-paths`

**`flyte run --no-sync-local-sys-paths <PATH> <TASK_NAME>`**

Disable synchronization of local `sys.path` entries to the remote execution environment during ephemeral preparation.
Identical to the `flyte deploy` command's `--no-sync-local-sys-paths` option:

```bash
flyte run --no-sync-local-sys-paths my_example.py my_task
```

This advanced option works identically to the deploy command equivalent, useful for:

- **Container isolation**: Prevent local development paths from affecting remote execution
- **Custom environments**: When containers have pre-configured Python paths
- **Security**: Avoiding exposure of local directory structures

## Task argument passing

A task's inputs are passed **after** the task name. On the CLI every input is a **named option**, `--<input_name>`, followed by a value when the type requires one, where `<input_name>` is the exact parameter name in the task's function signature:

```bash
flyte run my_file.py my_task --name "World" --count 5 --verbose
```

The equivalent SDK call passes the same inputs as keyword arguments:

```python
result = flyte.run(my_task, name="World", count=5, verbose=True)
```

A few rules apply to every input:

- **The option name matches the parameter name exactly, including underscores.** A parameter `event_time` is passed as `--event_time` (not `--event-time`).
- **Quote any value that contains spaces**, for example `--name "Ada Lovelace"`.
- **An input with a default is optional**; an input with no default is required.

### Passing inputs by type

Flyte parses each value according to the parameter's Python type. The type-specific syntax is summarized below and detailed in the following sections.

| Python type | CLI syntax | Example |
|---|---|---|
| `str` | Plain text (quote if it contains spaces) | `--name "Ada"` |
| `int` | Integer literal | `--count 5` |
| `float` | Decimal literal | `--rate 0.01` |
| `bool` | A flag (see **Tasks > Run and deploy tasks > Run command options > Task argument passing > Boolean inputs**) | `--verbose` |
| `datetime.datetime` / `datetime.date` | ISO 8601, `now`, `today`, or a relative expression | `--start 2024-01-15` |
| `datetime.timedelta` | ISO 8601 duration or a human-readable duration | `--timeout PT2H30M` |
| `enum.Enum` | The **member name** (case-sensitive) | `--color RED` |
| `list` / `dict` | A JSON literal (or a path to a `.json`/`.yaml` file) | `--nums '[1, 2, 3]'` |
| dataclass / Pydantic model / `TypedDict` / `NamedTuple` | A JSON literal (or a path to a `.json`/`.yaml` file) | `--config '{"lr": 0.01}'` |
| `flyte.io.File` / `flyte.io.Dir` | A local path or a remote URI | `--data ./input.csv` |
| `flyte.io.DataFrame` | A path to a `.parquet` or `.csv` file (local or remote) | `--df ./data.parquet` |

The `now`/`today`, relative-datetime, human-duration, and JSON-string forms are **CLI conveniences**. When you call `flyte.run()` programmatically, pass the corresponding native Python objects instead (a `datetime.datetime`, a `datetime.timedelta`, a dataclass instance, and so on).

### Boolean inputs

A `bool` input is a **flag**, not a value-taking option. Do not write `--debug true`.

- If the parameter defaults to `False` (or has no default), pass the bare flag to set it `True`, and omit it to leave it `False`:

  ```bash
  flyte run my_file.py my_task --verbose      # verbose=True
  flyte run my_file.py my_task                # verbose=False
  ```

- If the parameter defaults to `True`, Flyte also registers a `--no-` form so you can turn it off:

  ```bash
  flyte run my_file.py my_task --no-cache     # cache=False
  ```

### Datetime and duration inputs

A `datetime.datetime` (or `datetime.date`) input accepts an ISO 8601 timestamp, the keywords `now` or `today`, or a **relative expression** of the form `<base> <+|-> <ISO 8601 duration>` (the spaces around the sign are required, so quote the value):

```bash
flyte run my_file.py my_task --start 2024-01-15
flyte run my_file.py my_task --start "2024-01-15T13:30:00"
flyte run my_file.py my_task --start now
flyte run my_file.py my_task --start "now - P1D"        # 24 hours ago
flyte run my_file.py my_task --start "today + P1DT2H"   # tomorrow, 02:00
```

A `datetime.timedelta` input accepts an ISO 8601 duration or a human-readable duration (`<minutes>:<seconds>`, or a value with a unit such as `days`, `hours`, `minutes`, `seconds`):

```bash
flyte run my_file.py my_task --timeout P1DT2H30M    # 1 day, 2 hours, 30 minutes
flyte run my_file.py my_task --timeout "10 days"
flyte run my_file.py my_task --timeout "1 minute"
flyte run my_file.py my_task --timeout 1:24         # 1 minute, 24 seconds
```

### Enum inputs

For an `enum.Enum` input, pass the **name** of the member (case-sensitive), not its value:

```python
# my_file.py
import enum

class Color(enum.Enum):
    RED = "red"
    GREEN = "green"
```

```bash
flyte run my_file.py my_task --color RED
```

### List, dict, dataclass, and other structured inputs

Collections (`list`, `dict`) and structured types (dataclasses, Pydantic models, `TypedDict`, `NamedTuple`) are passed as a **JSON literal**:

```bash
flyte run my_file.py my_task --nums '[1, 2, 3]'
flyte run my_file.py my_task --mapping '{"a": 1, "b": 2}'
flyte run my_file.py my_task --config '{"lr": 0.01, "epochs": 5}'
```

For anything larger than a one-liner, pass a **path to a `.json` or `.yaml` file** instead. Flyte reads and parses the file:

```bash
flyte run my_file.py my_task --config ./config.yaml
```

### File, directory, and dataframe inputs

A `flyte.io.File`, `flyte.io.Dir`, or `flyte.io.DataFrame` input accepts either a **local path** or a **remote URI**. A local path is uploaded to the run's data store before execution; a remote URI (for example `s3://…` or `gs://…`) is used as-is. Under `--local`, the local path is used directly without uploading:

```bash
flyte run my_file.py my_task --data ./input.csv
flyte run my_file.py my_task --data s3://my-bucket/input.csv
flyte run my_file.py my_task --df ./data.parquet
```

### Optional inputs

An `Optional[...]` input (or any input with a default) is not required. Omit the option to leave it at its default:

```bash
flyte run my_file.py my_task --count 5      # `name` omitted, uses its default
```

## Positional arguments

**The CLI has no positional form for task inputs.** `flyte run` always passes inputs as named options (`--<input_name>`), as shown above. There is no `flyte run my_file.py my_task "World" 5` form.

Positional arguments apply only when you invoke a task **programmatically**. `flyte.run()`, and a direct (native) call to one task from inside another, accept positional arguments and map them to the task's parameters in signature order:

CODE32

Two caveats apply:

- **Deployed tasks are keyword-only.** A task retrieved with `flyte.remote.Task.get()` (or run via the `deployed-task` CLI subcommand) does **not** accept positional arguments. Pass its inputs by keyword (or, on the CLI, as `--<input_name>` options).
- **Don't provide the same input both positionally and by keyword.**

## SDK options

The core `flyte run` functionality is also available programmatically through the `flyte.run()` function.
For SDK-level configuration of all run parameters (storage, caching, identity, logging, and more),
see [Run context](./run-context).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/custom-cli ===

# Build a custom CLI

The built-in `flyte run` command (see [Run command options](./run-command-options))
turns your task's parameters into `--<input_name>` options automatically. That's
the fastest way to run a task from the command line, but it gives you the CLI
that Flyte generates.

When you want your **own** command-line interface (custom argument names, grouped
options, subcommands, config files, `--help` text you control), build it with an
argument-parsing library of your choice and hand the parsed values to
`flyte.run()`. Because `flyte.run()` is an ordinary Python function, any parser
works: [`tyro`](https://brentyi.github.io/tyro/), `argparse`, `click`, or `hydra`.

## The pattern

A custom CLI wrapper is three steps:

1. **Parse the command line** into a config object (a dataclass, a Pydantic model,
   or plain arguments) with the parser of your choice.
2. **Initialize Flyte** with `flyte.init_from_config()`.
3. **Run the task** with `flyte.run()`, passing the parsed config as the task's
   input.

## Example: a typed CLI with `tyro`

[`tyro`](https://brentyi.github.io/tyro/) generates a fully-typed CLI directly
from a dataclass, so you describe your parameters once and get parsing,
validation, and `--help` for free.

```python
# /// script
# dependencies = [
#    "tyro",
#    "flyte",
# ]
# ///

from dataclasses import dataclass

import tyro

import flyte

env = flyte.TaskEnvironment(
    name="custom_cli",
    image=flyte.Image.from_uv_script(__file__, name="flyte"),
)

@dataclass
class Config:
    foo: int
    bar: str = "default"

@env.task
async def main(config: Config):
    print(f"foo: {config.foo}, bar: {config.bar}")

if __name__ == "__main__":
    # Generate a CLI and instantiate `Config` with its two arguments: `foo` and `bar`.
    config = tyro.cli(Config)

    flyte.init_from_config()
    run = flyte.run(main, config)
    print(run.url)
```

Run it like any script, and `tyro` exposes `foo` and `bar` as CLI options and prints
`--help` for you:

CODE0

Here `tyro.cli(Config)` does the parsing, `flyte.init_from_config()` loads your
`config.yaml` (endpoint, project, domain, and so on), and
`flyte.run(main, config)` deploys and runs the task with the parsed config as its
input. The `config` object is passed **positionally** to `main`, mapping to its
`config` parameter. These are ordinary Python positional arguments.

## Using a different parser

The pattern is identical whichever library you reach for: only step 1 changes.
With the standard library's `argparse`:

CODE1

Swap `argparse` for `click` or `hydra` the same way: parse however you like, then
call `flyte.run()` with the resulting values (positionally or by keyword). This is
also the entry point for richer CLIs: subcommands that each run a different task,
`hydra` config composition, environment-driven defaults, and so on.

## When to use which

- **Reach for the built-in `flyte run`** when you just need to run a task from the
  command line and Flyte's generated `--<input_name>` options are enough. See
  [Run command options](./run-command-options).
- **Build a custom CLI** when you need control over the interface itself (your own
  option names, subcommands, config-file loading, or help text), or when the CLI is
  a first-class part of a tool you're shipping.

For configuring the run itself (storage, caching, identity, logging) rather than the
task's inputs, see [Run context](./run-context).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/how-task-deployment-works ===

# How task deployment works

This section explains how the `flyte deploy` command and the `flyte.deploy()` SDK function work under the hood to deploy tasks to your Flyte backend.

When you perform a deployment, here's what happens:

## 1. Module loading and task environment discovery

In the first step, Flyte determines which files to load in order to search for task environments, based on the command-line options provided:

### Single file (default)

```bash
flyte deploy my_example.py env
```

- The file `my_example.py` is executed,
- All declared `TaskEnvironment` objects in the file are instantiated,
  but only the one assigned to the variable `env` is selected for deployment.

### `--all` option

```bash
flyte deploy --all my_example.py
```

- The file `my_example.py` is executed,
- All declared `TaskEnvironment` objects in the file are instantiated and selected for deployment.
- No specific variable name is required.

### `--recursive` option

```bash
flyte deploy --recursive ./directory
```

- The directory is recursively traversed and all Python files are executed and all `TaskEnvironment` objects are instantiated.
- All `TaskEnvironment` objects across all files are selected for deployment.

## 2. Task analysis and serialization

- For every task environment selected for deployment, all of its tasks are identified.
- Task metadata is extracted: parameter types, return types, and resource requirements.
- Each task is serialized into a Flyte `TaskTemplate`.
- Dependency graphs between environments are built (see below).

## 3. Task environment dependency resolution

In many cases, a task in one environment may invoke a task in another environment, establishing a dependency between the two environments.
For example, if `env_a` has a task that calls a task in `env_b`, then `env_a` depends on `env_b`.
This means that when deploying `env_a`, `env_b` must also be deployed to ensure that all tasks can be executed correctly.

To handle this, `TaskEnvironment`s can declare dependencies on other `TaskEnvironment`s using the `depends_on` parameter.
During deployment, the system performs the following steps to resolve these dependencies:

1. Starting with specified environment(s)
2. Recursively discovering all transitive dependencies
3. Including all dependencies in the deployment plan
4. Processing dependencies depth-first to ensure correct order

```python
# Define environments with dependencies
prep_env = flyte.TaskEnvironment(name="preprocessing")
ml_env = flyte.TaskEnvironment(name="ml_training", depends_on=[prep_env])
viz_env = flyte.TaskEnvironment(name="visualization", depends_on=[ml_env])

# Deploy only viz_env - automatically includes ml_env and prep_env
deployment = flyte.deploy(viz_env, version="v2.0.0")

# Or deploy multiple environments explicitly
deployment = flyte.deploy(data_env, ml_env, viz_env, version="v2.0.0")
```

For detailed information about working with multiple environments, see [Multiple Environments](../task-configuration/multiple-environments).

## 4. Code bundle creation and upload

Once the task environments and their dependencies are resolved, Flyte proceeds to package your code into a bundle based on the `copy_style` option:

### `--copy_style loaded_modules` (default)

This is the smart bundling approach that analyzes which Python modules were actually imported during the task environment discovery phase.
It examines the runtime module registry (`sys.modules`) and includes only those modules that meet specific criteria:
they must have source files located within your project directory (not in system locations like `site-packages`), and they must not be part of the Flyte SDK itself.
This selective approach results in smaller, faster-to-upload bundles that contain exactly the code needed to run your tasks, making it ideal for most development and production scenarios.

### `--copy_style all`

This bundling strategy takes a directory-walking approach, recursively traversing your entire project directory and including every file it encounters.
Unlike the smart bundling that only includes imported Python modules, this method captures all project files regardless of whether they were imported during discovery.
This is particularly useful for projects that use dynamic imports, load configuration files or data assets at runtime, or have dependencies that aren't captured through normal Python import mechanisms.

### `--copy_style none`

This option completely skips code bundle creation, meaning no source code is packaged or uploaded to cloud storage.
When using this approach, you must provide an explicit version parameter since there's no code bundle to generate a version from.
This strategy is designed for scenarios where your code is already baked into custom container images, eliminating the need for separate code injection during task execution.
It results in the fastest deployment times but requires more complex image management workflows.

### `--root-dir` option

By default, Flyte uses your current working directory as the root for code bundling.
You can override this with `--root-dir` to specify a different base directory - particularly useful for monorepos or when deploying from subdirectories. This affects all copy styles: `loaded_modules` will look for imported modules relative to the root directory, `all` will walk the directory tree starting from the root, and the root directory setting works with any copy style. See the [Deploy command options](./deploy-command-options#--root-dir) for detailed usage examples.

After the code bundle is created (if applicable), it is uploaded to a cloud storage location (like S3 or GCS) accessible by your Flyte backend. It is now ready to be run.

## 5. Image building

If your `TaskEnvironment` specifies [custom images](../task-configuration/container-images), Flyte builds and pushes container images before deploying tasks.
The build process varies based on your configuration and backend type:

### Local image building

When `image.builder` is set to `local` in [your `config.yaml`](../../get-started/run-modes/running-devbox#configure), images are built on your local machine using Docker. This approach:

- Requires Docker to be installed and running on your development machine
- Uses Docker BuildKit to build images from generated Dockerfiles or your custom Dockerfile
- Pushes built images to the container registry specified in your `Image` configuration
- Is the only option available for Flyte OSS instances

### Remote image building

When `image.builder` is set to `remote` in your `config.yaml`, images are built on cloud infrastructure. This approach:

- Builds images using Union's ImageBuilder service (currently only available for Union backends, not OSS Flyte)
- Requires no local Docker installation or configuration
- Can push to Union's internal registry or external registries you specify
- Provides faster, more consistent builds by leveraging cloud resources

> [!NOTE]
> Remote building is currently exclusive to Union backends. OSS Flyte installations must use `local`

## 6. Source-code link discovery

While each task is being serialized in step 2, Flyte attempts to attach a link from the task back to the source line of its Python function.
The link is rendered next to the task description in the UI, so anyone viewing a deployed task can jump directly to the code that defines it.

This is fully automatic. There is no decorator argument, no config flag, and no opt-in step. If the conditions below are met, the link appears.

### How the link is built

Flyte inspects the local repository at deploy time using the standard `git` CLI:

- The repository root, current commit SHA, working-tree-clean status, and remote push URL are read via `git rev-parse` and `git remote get-url --push origin`.
- The file path is taken from the task function's `__code__.co_filename`, made relative to the repo root.
- The line number is the line of the function definition (the line just after the `@env.task` decorator).
- For a GitHub remote, the URL takes the form `https://github.com/<owner>/<repo>/blob/<sha>/<file>#L<line>`. GitLab uses the equivalent `/-/blob/` form.

The `#L<line>` anchor is only included when the working tree is clean. A dirty tree still produces a valid blob URL, but without the line jump, because the local file may no longer match the committed file.

### Conditions

For the link to appear, all of the following must hold:

- The `git` CLI is installed and the directory you deploy from is inside a git repository.
- The repository has a remote. The push URL of `origin` is preferred; otherwise the first remote, alphabetically, is used.
- The remote host is `github.com` or `gitlab.com`. Other hosts (Bitbucket, self-hosted Gitea, GitHub Enterprise on a custom domain, etc.) currently produce no link.
- The task's source file lives under the repository root.

If any condition fails, the deploy still succeeds: only the source-code link is omitted.

> [!NOTE]
> Flyte does not check whether the current commit has been pushed.
> If you deploy from a clean local commit that is not yet on the remote, the URL will resolve to a missing SHA on GitHub or GitLab.
> Push your commit before deploying if you want the link to be followable.

For private repositories, the link is still generated; viewers need to be authenticated to the host to follow it.

See `flyte.git.GitStatus` for the underlying API.

## Understanding option relationships

It's important to understand how the various deployment options work together.
The **discovery options** (`--recursive` and `--all`) operate independently of the **bundling options** (`--copy-style`),
giving you flexibility in how you structure your deployments.

Environment discovery determines which files Flyte will examine to find `TaskEnvironment` objects,
while code bundling controls what gets packaged and uploaded for execution.
You can freely combine these approaches.
For example, discovering environments recursively across your entire project while using smart bundling to include only the necessary code modules.

When multiple environments are discovered, they all share the same code bundle, which is efficient for related services or components that use common dependencies:

```bash
flyte deploy --recursive --copy-style loaded_modules ./project
```

> [!NOTE]
> All discovered environments share the same code bundle.

For a full overview of all deployment options, see **Flyte CLI > flyte > flyte deploy**.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/deploy-command-options ===

# Deploy command options

The `flyte deploy` command provides extensive configuration options:

**`flyte deploy [OPTIONS] <PATH> [TASK_ENV_VARIABLE]`**

| Option                      | Short | Type   | Default                   | Description                                       |
|-----------------------------|-------|--------|---------------------------|---------------------------------------------------|
| `--project`                 | `-p`  | text   | *from config*             | Project to deploy to                              |
| `--domain`                  | `-d`  | text   | *from config*             | Domain to deploy to                               |
| `--version`                 |       | text   | *auto-generated*          | Explicit version tag for deployment               |
| `--dry-run`/`--dryrun`      |       | flag   | `false`                   | Preview deployment without executing              |
| `--all`                     |       | flag   | `false`                   | Deploy all environments in specified path         |
| `--recursive`               | `-r`  | flag   | `false`                   | Deploy environments recursively in subdirectories |
| `--copy-style`              |       | choice | `loaded_modules\|all\|none` | Code bundling strategy                            |
| `--root-dir`                |       | path   | *current dir*             | Override source root directory                    |
| `--image`                   |       | text   |                           | Image URI mappings (format: `name=uri`)           |
| `--ignore-load-errors`      | `-i`  | flag   | `false`                   | Continue deployment despite module load failures  |
| `--no-sync-local-sys-paths` |       | flag   | `false`                   | Disable local `sys.path` synchronization          |

## `--project`, `--domain`

**`flyte deploy --domain <DOMAIN> --project <PROJECT> <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

You can specify `--project` and `--domain` which will override any defaults defined in your `config.yaml`:

```bash
flyte deploy my_example.py env
```

Specify a target project and domain:

```bash
flyte deploy --project my-project --domain development my_example.py env
```

## `--version`

**`flyte deploy --version <VERSION> <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--version` option controls how deployed tasks are tagged and identified in the Flyte backend:

Auto-generated version (default):

```bash
flyte deploy my_example.py env
```

Explicit version:

```bash
flyte deploy --version v1.0.0 my_example.py env
```

> [!NOTE]
> An explicit version is required when using `--copy-style none`, since there is no code bundle to generate a hash from.

```bash
flyte deploy --copy-style none --version v1.0.0 my_example.py env
```

### When versions are used

- **Explicit versioning**: Provides human-readable task identification (e.g., `v1.0.0`, `prod-2024-12-01`)
- **Auto-generated versions**: When no version is specified, Flyte creates an MD5 hash from the code bundle, environment configuration, and image cache
- **Version requirement**: `copy-style none` mandates explicit versions since there's no code bundle to hash
- **Task referencing**: Versions enable precise task references in `flyte run deployed-task` and workflow invocations

## `--dry-run`

**`flyte deploy --dry-run <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--dry-run` option allows you to preview what would be deployed without actually performing the deployment:

```bash
flyte deploy --dry-run my_example.py env
```

## `--all` and `--recursive`

**`flyte deploy --all <SOURCE_FILE>`**

**`flyte deploy --recursive <DIRECTORY_PATH>`**

Control which environments get discovered and deployed:

**Single environment (default):**

```bash
flyte deploy my_example.py env
```

**All environments in file:**

```bash
flyte deploy --all my_example.py
```

**Recursive directory deployment:**

```bash
flyte deploy --recursive ./src
```

Combine with comprehensive bundling:

```bash
flyte deploy --recursive --copy-style all ./project
```

## `--copy-style`

**`flyte deploy --copy_style [loaded_modules|all|none] <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--copy-style` option controls what gets packaged:

### `--copy-style loaded_modules` (default)

```bash
flyte deploy --copy-style loaded_modules my_example.py env
```

- **Includes**: Only imported Python modules from your project
- **Excludes**: Site-packages, system modules, Flyte SDK
- **Best for**: Most projects (optimal size and speed)

### `--copy-style all`

```bash
flyte deploy --copy-style all my_example.py env
```

- **Includes**: All files in project directory
- **Best for**: Projects with dynamic imports or data files

### `--copy-style none`

```bash
flyte deploy --copy-style none --version v1.0.0 my_example.py env
```

- **Requires**: Explicit version parameter
- **Best for**: Pre-built container images with baked-in code

## `--root-dir`

**`flyte deploy --root-dir <DIRECTORY> <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--root-dir` option overrides the default source directory that Flyte uses as the base for code bundling and import resolution.
This is particularly useful for monorepos and projects with complex directory structures.

### Default behavior (without `--root-dir`)

- Flyte uses the current working directory as the root
- Code bundling starts from this directory
- Import paths are resolved relative to this location

### Common use cases

**Monorepos:**
Deploy a service from the monorepo root:

```bash
flyte deploy --root-dir ./services/ml ./services/ml/my_example.py env
```

Deploy from anywhere in the monorepo:

```bash
cd ./docs/
flyte deploy --root-dir ../services/ml ../services/ml/my_example.py env
```

**Cross-directory imports:**
When a workflow imports modules from sibling directories (e.g., `project/workflows/my_example.py` imports `project/src/utils.py`):

```bash
cd project/workflows/
flyte deploy --root-dir .. my_example.py env
```

**Working directory independence:**

```bash
flyte deploy --root-dir /path/to/project /path/to/project/my_example.py env
```

### How it works

1. **Code bundling**: Files are collected starting from `--root-dir` instead of the current working directory
2. **Import resolution**: Python imports are resolved relative to the specified root directory
3. **Path consistency**: Ensures the same directory structure in local and remote execution environments
4. **Dependency packaging**: Captures all necessary modules that may be located outside the workflow file's immediate directory

### Example with complex project structure

```
my-project/
├── services/
│   ├── ml/
│   │   └── my_example.py     # imports shared.utils
│   └── api/
└── shared/
    └── utils.py
```

```bash
flyte deploy --root-dir ./my-project ./my-project/services/ml/my_example.py env
```

This ensures that both `services/ml/` and `shared/` directories are included in the code bundle, allowing the workflow to successfully import `shared.utils` during remote execution.

## `--image`

**`flyte deploy --image <IMAGE_MAPPING> <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--image` option allows you to override image URIs at deployment time without modifying your code. Format: `imagename=imageuri`

### Named image mappings

```bash
flyte deploy --image base=ghcr.io/org/base:v1.0 my_example.py env
```

Multiple named image mappings:

```bash
flyte deploy \
  --image base=ghcr.io/org/base:v1.0 \
  --image gpu=ghcr.io/org/gpu:v2.0 \
  my_example.py env
```

### Default image mapping

```bash
flyte deploy --image ghcr.io/org/default:latest my_example.py env
```

### How it works

- Named mappings (e.g., `base=URI`) override images created with `Image.from_ref_name("base")`.
- Unnamed mappings (e.g., just `URI`) override the default "auto" image.
- Multiple `--image` flags can be specified.
- Mappings are resolved during the image building phase of deployment.

## `--ignore-load-errors`

**`flyte deploy --ignore-load-errors <SOURCE_PATH> <TASK_ENV_VARIABLE>`**

The `--ignore-load-errors` option allows the deployment process to continue even if some modules fail to load during the environment discovery phase. This is particularly useful for large projects or monorepos where certain modules may have missing dependencies or other issues that prevent them from being imported successfully.

```bash
flyte deploy --recursive --ignore-load-errors ./large-project
```

## `--no-sync-local-sys-paths`

**`flyte deploy --no-sync-local-sys-paths <SOURCE_FILE> <TASK_ENV_VARIABLE>`**

The `--no-sync-local-sys-paths` option disables the automatic synchronization of local `sys.path` entries to the remote container environment. This is an advanced option for specific deployment scenarios.

### Default behavior (path synchronization enabled)

- Flyte captures local `sys.path` entries that are under the root directory
- These paths are passed to the remote container via the `_F_SYS_PATH` environment variable
- At runtime, the remote container adds these paths to its `sys.path`, maintaining the same import environment

### When to disable path synchronization

```bash
flyte deploy --no-sync-local-sys-paths my_example.py env
```

### Use cases for disabling

- **Custom container images**: When your container already has the correct `sys.path` configuration
- **Conflicting path structures**: When local development paths would interfere with container paths
- **Security concerns**: When you don't want to expose local development directory structures
- **Minimal environments**: When you want precise control over what gets added to the container's Python path

### How it works

- **Enabled (default)**: Local paths like `./my_project/utils` get synchronized and added to remote `sys.path`
- **Disabled**: Only the container's native `sys.path` is used, along with the deployed code bundle

Most users should leave path synchronization enabled unless they have specific requirements for container path isolation or are using pre-configured container environments.

## SDK deployment options

The core deployment functionality is available programmatically through the `flyte.deploy()` function, though some CLI-specific options are not applicable:

```python
import flyte

env = flyte.TaskEnvironment(name="my_env")

@env.task
async def process_data(data: str) -> str:
    return f"Processed: {data}"

if __name__ == "__main__":
    flyte.init_from_config()

    # Comprehensive deployment configuration
    deployment = flyte.deploy(
        env,                          # Environment to deploy
        dryrun=False,                 # Set to True for dry run
        version="v1.2.0",             # Explicit version tag
        copy_style="loaded_modules"   # Code bundling strategy
    )
    print(f"Deployment successful: {deployment[0].summary_repr()}")
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/packaging ===

# Code packaging for remote execution

When you run Flyte tasks remotely, your code needs to be available in the execution environment. Flyte SDK provides two main approaches for packaging your code:

1. **Code bundling** - Bundle code dynamically at runtime
2. **Container-based deployment** - Embed code directly in container images

## Quick comparison

| Aspect | Code bundling | Container-based |
|--------|---------------|-----------------|
| **Speed** | Fast (no image rebuild) | Slower (requires image build) |
| **Best for** | Rapid development, iteration | Production, immutable deployments |
| **Code changes** | Immediate effect | Requires image rebuild |
| **Setup** | Automatic by default | Manual configuration needed |
| **Reproducibility** | Excellent (hash-based versioning) | Excellent (immutable images) |
| **Rollback** | Requires version control | Tag-based, straightforward |

---

## Code bundling

**Default approach** - Automatically bundles and uploads your code to remote storage at runtime.

### How it works

When you run `flyte run` or call `flyte.run()`, Flyte automatically:

1. **Scans loaded modules** from your codebase
2. **Creates a tarball** (gzipped, without timestamps for consistent hashing)
3. **Uploads to blob storage** (S3, GCS, Azure Blob)
4. **Deduplicates** based on content hashes
5. **Downloads in containers** at runtime

This process happens transparently - every container downloads and extracts the code bundle before execution.

> [!NOTE]
> Code bundling is optimized for speed:
>
> - Bundles are created without timestamps for consistent hashing
> - Identical code produces identical hashes, enabling deduplication
> - Only modified code triggers new uploads
> - Containers cache downloaded bundles
>
> **Reproducibility:** Flyte automatically versions code bundles based on content hash. The same code always produces the same hash, guaranteeing reproducibility without manual versioning. However, version control is still recommended for rollback capabilities.

### Automatic code bundling

**Default behavior** - Bundles all loaded modules automatically.

#### What gets bundled

Flyte includes modules that are:

- ✅ **Loaded when environment is parsed** (imported at module level)
- ✅ **Part of your codebase** (not system packages)
- ✅ **Within your project directory**
- ❌ **NOT lazily loaded** (imported inside functions)
- ❌ **NOT system-installed packages** (e.g., from site-packages)

#### Example: Basic automatic bundling

```python
# app.py
import flyte
from my_module import helper  # ✅ Bundled automatically

env = flyte.TaskEnvironment(
    name="default",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "numpy")
)

@env.task
def process_data(x: int) -> int:
    # This import won't be bundled (lazy load)
    from another_module import util  # ❌ Not bundled automatically
    return helper.transform(x)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(process_data, x=42)
    print(run.url)
```

When you run this:

```bash
flyte run app.py process_data --x 42
```

Flyte automatically:

1. Bundles `app.py` and `my_module.py`
2. Preserves the directory structure
3. Uploads to blob storage
4. Makes it available in the remote container

#### Project structure example

```
my_project/
├── app.py              # Main entry point
├── tasks/
│   ├── __init__.py
│   ├── data_tasks.py   # Flyte tasks
│   └── ml_tasks.py
└── utils/
    ├── __init__.py
    ├── preprocessing.py # Business logic
    └── models.py
```

```python
# app.py
import flyte
from tasks.data_tasks import load_data    # ✅ Bundled
from tasks.ml_tasks import train_model    # ✅ Bundled
# utils modules imported in tasks are also bundled

@flyte.task
def pipeline(dataset: str) -> float:
    data = load_data(dataset)
    accuracy = train_model(data)
    return accuracy

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline, dataset="train.csv")
```

**All modules are bundled with their directory structure preserved.**

### Manual code bundling

Control exactly what gets bundled by configuring the copy style.

#### Copy styles

Three options available:

1. **`"auto"`** (default) - Bundle loaded modules only
2. **`"all"`** - Bundle everything in the working directory
3. **`"none"`** - Skip bundling entirely (requires code in container)

#### Using `copy_style="all"`

Bundle all files under your project directory:

```python
import flyte

flyte.init_from_config()

# Bundle everything in current directory
run = flyte.with_runcontext(copy_style="all").run(
    my_task,
    input_data="sample.csv"
)
```

Or via CLI:

```bash
flyte run --copy-style=all app.py my_task --input-data sample.csv
```

**Use when:**

- You have data files or configuration that tasks need
- You use dynamic imports or lazy loading
- You want to ensure all project files are available

#### Using `copy_style="none"`

Skip code bundling (see **Tasks > Run and deploy tasks > Code packaging for remote execution > Container-based deployment**):

```python
run = flyte.with_runcontext(copy_style="none").run(my_task, x=10)
```

### Including additional files with `include`

Code bundling discovers Python modules by following imports. That's the right behavior for source code, but it won't pick up non-Python assets like HTML templates, SQL files, small reference datasets, prompt files, or configuration files. Your task imports Python, not an `.html` file, so those assets never show up in the bundle.

The `include` parameter on `TaskEnvironment` lets you attach these extra files to the environment's bundle explicitly:

```python
import flyte

env = flyte.TaskEnvironment(
    name="html_template_report",
    image=flyte.Image.from_debian_base(python_version=(3, 12)),
    include=("report_template.html",),
)
```

At bundling time, Flyte resolves each entry, unions it with whatever the `copy_style` discovered, and ships everything in the same tarball. The files land in the container at the same path they occupy in your project, so the task can read them with a normal relative path:

```python
from datetime import datetime, timezone
from pathlib import Path

import flyte
import flyte.report

env = flyte.TaskEnvironment(
    name="html_template_report",
    image=flyte.Image.from_debian_base(python_version=(3, 12)),
    include=("report_template.html",),
)

@env.task(report=True)
async def generate_template_report() -> str:
    template_path = Path(__file__).parent / "report_template.html"
    template = template_path.read_text()
    body = template.format(
        title="Hello",
        generated_at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
    )
    flyte.report.get_tab("Main").log(body)
    await flyte.report.flush.aio()
    return "ok"
```

#### How paths are resolved

- **Relative paths** are anchored at the directory of the file where the `TaskEnvironment` is instantiated, not at `root_dir` or the current working directory. `include=("report_template.html",)` looks for `report_template.html` next to the Python file that declared the env.
- **Absolute paths** are used as-is.
- **Directories** are included recursively.
- **Glob patterns** are expanded against the declaring file's directory.

`include` supplements the `copy_style` discovery; it does not replace it. Files listed here are bundled *in addition to* the Python modules that `copy_style` picks up. It also works alongside `copy_style="none"`. In that case, only the include entries are bundled.

#### When to use `include`

Good fits:

- **HTML templates** rendered into reports (see the [reports examples](https://github.com/flyteorg/flyte-sdk/tree/main/examples/reports) in the SDK).
- **Small configuration files** (YAML, JSON, TOML) that the task reads at runtime.
- **SQL files, prompt templates, or other text assets** versioned alongside the task.
- **Small reference data files** (a few MB of lookup tables, fixtures, etc.).
- **Small model files** that you want versioned with the task code.

> [!WARNING]
> **Don't bundle large files with `include`.** Every container that runs the task downloads the full bundle before it starts, so large includes directly inflate cold-start latency and eat network bandwidth on every execution. As a rule of thumb, if a file is more than a few MB, or if it changes independently of your code, it doesn't belong in `include`.
>
> For larger assets, prefer one of these instead:
>
> 1. **Store in object storage (most recommended)**: keep the file in S3, GCS, or Azure Blob and read it with `flyte.io.File("s3://...")` or the cloud SDK of your choice. The file is fetched only when needed, cached independently of the code bundle, and can be updated without redeploying.
> 2. **Bake into the container image**: use **Tasks > Run and deploy tasks > Code packaging for remote execution > Container-based deployment > Image source copying methods** when the asset is stable and tied to a specific image version. The file is downloaded once when the image is pulled, not on every run.

### Controlling the root directory

The `root_dir` parameter controls which directory serves as the bundling root.

#### Why root directory matters

1. **Determines what gets bundled** - All code paths are relative to root_dir
2. **Preserves import structure** - Python imports must match the bundle structure
3. **Affects path resolution** - Files and modules are located relative to root_dir

#### Setting root directory

##### Via CLI

```bash
flyte run --root-dir /path/to/project app.py my_task
```

##### Programmatically

```python
import pathlib
import flyte

flyte.init_from_config(
    root_dir=pathlib.Path(__file__).parent
)
```

#### Root directory use cases

##### Use case 1: Multi-module project

```
project/
├── src/
│   ├── workflows/
│   │   └── pipeline.py
│   └── utils/
│       └── helpers.py
└── config.yaml
```

```python
# src/workflows/pipeline.py
import pathlib
import flyte
from utils.helpers import process  # Relative import from project root

# Set root to project root (not src/)
flyte.init_from_config(
    root_dir=pathlib.Path(__file__).parent.parent.parent
)

@flyte.task
def my_task():
    return process()
```

**Root set to `project/` so imports like `from utils.helpers` work correctly.**

##### Use case 2: Shared utilities

```
workspace/
├── shared/
│   └── common.py
└── project/
    └── app.py
```

```python
# project/app.py
import flyte
import pathlib
from shared.common import shared_function  # Import from parent directory

# Set root to workspace/ to include shared/
flyte.init_from_config(
    root_dir=pathlib.Path(__file__).parent.parent
)
```

##### Use case 3: Monorepo

```
monorepo/
├── libs/
│   ├── data/
│   └── models/
└── services/
    └── ml_service/
        └── workflows.py
```

```python
# services/ml_service/workflows.py
import flyte
import pathlib
from libs.data import loader  # Import from monorepo root
from libs.models import predictor

# Set root to monorepo/ to include libs/
flyte.init_from_config(
    root_dir=pathlib.Path(__file__).parent.parent.parent
)
```

#### Root directory best practices

1. **Set root_dir at project initialization** before importing any task modules
2. **Use absolute paths** with `pathlib.Path(__file__).parent` navigation
3. **Match your import structure** - if imports are relative to project root, set root_dir to project root
4. **Keep consistent** - use the same root_dir for both `flyte run` and `flyte.init()`

### Code bundling examples

#### Example: Standard Python package

```
my_package/
├── pyproject.toml
├── src/
│   └── my_package/
│       ├── __init__.py
│       ├── main.py
│       ├── data/
│       │   ├── loader.py
│       │   └── processor.py
│       └── models/
│           └── analyzer.py
```

```python
# src/my_package/main.py
import flyte
import pathlib
from my_package.data.loader import fetch_data
from my_package.data.processor import clean_data
from my_package.models.analyzer import analyze

env = flyte.TaskEnvironment(
    name="pipeline",
    image=flyte.Image.from_debian_base().with_uv_project(
        pyproject_file=pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"
    )
)

@env.task
async def fetch_task(url: str) -> dict:
    return await fetch_data(url)

@env.task
def process_task(raw_data: dict) -> list[dict]:
    return clean_data(raw_data)

@env.task
def analyze_task(data: list[dict]) -> str:
    return analyze(data)

if __name__ == "__main__":
    import flyte.git

    # Set root to project root for proper imports
    flyte.init_from_config(
        flyte.git.config_from_root(),
        root_dir=pathlib.Path(__file__).parent.parent.parent
    )

    # All modules bundled automatically
    run = flyte.run(analyze_task, data=[{"value": 1}, {"value": 2}])
    print(f"Run URL: {run.url}")
```

**Run with:**

```bash
cd my_package
flyte run src/my_package/main.py analyze_task --data '[{"value": 1}]'
```

#### Example: Dynamic environment based on domain

```python
# environment_picker.py
import flyte

def create_env():
    """Create different environments based on domain."""
    if flyte.current_domain() == "development":
        return flyte.TaskEnvironment(
            name="dev",
            image=flyte.Image.from_debian_base(),
            env_vars={"ENV": "dev", "DEBUG": "true"}
        )
    elif flyte.current_domain() == "staging":
        return flyte.TaskEnvironment(
            name="staging",
            image=flyte.Image.from_debian_base(),
            env_vars={"ENV": "staging", "DEBUG": "false"}
        )
    else:  # production
        return flyte.TaskEnvironment(
            name="prod",
            image=flyte.Image.from_debian_base(),
            env_vars={"ENV": "production", "DEBUG": "false"},
            resources=flyte.Resources(cpu="2", memory="4Gi")
        )

env = create_env()

@env.task
async def process(n: int) -> int:
    import os
    print(f"Running in {os.getenv('ENV')} environment")
    return n * 2

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(process, n=5)
    print(run.url)
```

**Why this works:**

- `flyte.current_domain()` is set correctly when Flyte re-instantiates modules remotely
- Environment configuration is deterministic and reproducible
- Code automatically bundled with domain-specific settings

> [!NOTE]
> `flyte.current_domain()` only works after `flyte.init()` is called:
>
> - ✅ Works with `flyte run` and `flyte deploy` (auto-initialize)
> - ✅ Works in `if __name__ == "__main__"` after explicit `flyte.init()`
> - ❌ Does NOT work at module level without initialization

### When to use code bundling

✅ **Use code bundling when:**

- Rapid development and iteration
- Frequently changing code
- Multiple developers testing changes
- Jupyter notebook workflows
- Quick prototyping and experimentation

❌ **Consider container-based instead when:**

- Need easy rollback to previous versions (container tags are simpler than finding git commits)
- Working with air-gapped environments (no blob storage access)
- Code changes require coordinated dependency updates

---

## Container-based deployment

**Advanced approach** - Embed code directly in container images for immutable deployments.

### How it works

Instead of bundling code at runtime:

1. **Build container image** with code copied inside
2. **Disable code bundling** with `copy_style="none"`
3. **Container has everything** needed at runtime

**Trade-off:** Every code change requires a new image build (slower), but provides complete reproducibility.

### Configuration

Three key steps:

#### 1. Set `copy_style="none"`

Disable runtime code bundling:

```python
flyte.with_runcontext(copy_style="none").run(my_task, n=10)
```

Or via CLI:

```bash
flyte run --copy-style=none app.py my_task --n 10
```

#### 2. Copy code into image

Use `Image.with_source_file()` or `Image.with_source_folder()`:

```python
import pathlib
import flyte

env = flyte.TaskEnvironment(
    name="embedded",
    image=flyte.Image.from_debian_base().with_source_folder(
        src=pathlib.Path(__file__).parent,
        copy_contents_only=True
    )
)
```

#### 3. Set correct `root_dir`

Match your image copy configuration:

```python
flyte.init_from_config(
    root_dir=pathlib.Path(__file__).parent
)
```

### Image source copying methods

#### `with_source_file()` - copy individual files

Copy a single file into the container:

```python
image = flyte.Image.from_debian_base().with_source_file(
    src=pathlib.Path(__file__),
    dst="/app/main.py"
)
```

**Use for:**

- Single-file workflows
- Copying configuration files
- Adding scripts to existing images

#### `with_source_folder()` - copy directories

Copy entire directories into the container:

```python
image = flyte.Image.from_debian_base().with_source_folder(
    src=pathlib.Path(__file__).parent,
    dst="/app",
    copy_contents_only=False  # Copy folder itself
)
```

**Parameters:**

- `src`: Source directory path
- `dst`: Destination path in container (optional, defaults to workdir)
- `copy_contents_only`: If `True`, copies folder contents; if `False`, copies folder itself

##### `copy_contents_only=True` (Recommended)

Copies only the contents of the source folder:

```python
# Project structure:
# my_project/
#   ├── app.py
#   └── utils.py

image = flyte.Image.from_debian_base().with_source_folder(
    src=pathlib.Path(__file__).parent,
    copy_contents_only=True
)

# Container will have:
# /app/app.py
# /app/utils.py

# Set root_dir to match:
flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
```

##### `copy_contents_only=False`

Copies the folder itself with its name:

```python
# Project structure:
# workspace/
#   └── my_project/
#       ├── app.py
#       └── utils.py

image = flyte.Image.from_debian_base().with_source_folder(
    src=pathlib.Path(__file__).parent,  # Points to my_project/
    copy_contents_only=False
)

# Container will have:
# /app/my_project/app.py
# /app/my_project/utils.py

# Set root_dir to parent to match:
flyte.init_from_config(root_dir=pathlib.Path(__file__).parent.parent)
```

### Complete container-based example

```python
# full_build.py
import pathlib
import flyte
from dep import helper  # Local module

# Configure environment with source copying
env = flyte.TaskEnvironment(
    name="full_build",
    image=flyte.Image.from_debian_base()
        .with_pip_packages("numpy", "pandas")
        .with_source_folder(
            src=pathlib.Path(__file__).parent,
            copy_contents_only=True
        )
)

@env.task
def square(x: int) -> int:
    return x ** helper.get_exponent()

@env.task
def main(n: int) -> list[int]:
    return list(flyte.map(square, range(n)))

if __name__ == "__main__":
    import flyte.git

    # Initialize with matching root_dir
    flyte.init_from_config(
        flyte.git.config_from_root(),
        root_dir=pathlib.Path(__file__).parent
    )

    # Run with copy_style="none" and explicit version
    run = flyte.with_runcontext(
        copy_style="none",
        version="v1.0.0"  # Explicit version for image tagging
    ).run(main, n=10)

    print(f"Run URL: {run.url}")
    run.wait()
```

**Project structure:**

```
project/
├── full_build.py
├── dep.py          # Local dependency
└── .flyte/
    └── config.yaml
```

**Run with:**

```bash
python full_build.py
```

This will:

1. Build a container image with `full_build.py` and `dep.py` embedded
2. Tag it as `v1.0.0`
3. Push to registry
4. Execute remotely without code bundling

### Using externally built images

When containers are built outside of Flyte (e.g., in CI/CD), use `Image.from_ref_name()`:

#### Step 1: Build your image externally

```dockerfile
# Dockerfile
FROM python:3.11-slim

WORKDIR /app

# Copy your code
COPY src/ /app/

# Install dependencies
RUN pip install flyte pandas numpy

# Ensure flyte executable is available
RUN flyte --help
```

Build and push the image:

```bash
docker build -t myregistry.com/my-app:v1.2.3 .
docker push myregistry.com/my-app:v1.2.3
```

#### Step 2: Reference image by name

```python
# app.py
import flyte

env = flyte.TaskEnvironment(
    name="external",
    image=flyte.Image.from_ref_name("my-app-image")  # Reference name
)

@env.task
def process(x: int) -> int:
    return x * 2

if __name__ == "__main__":
    flyte.init_from_config()

    # Pass actual image URI at deploy/run time
    run = flyte.with_runcontext(
        copy_style="none",
        images={"my-app-image": "myregistry.com/my-app:v1.2.3"}
    ).run(process, x=10)
```

Or via CLI:

```bash
flyte run \
  --copy-style=none \
  --image my-app-image=myregistry.com/my-app:v1.2.3 \
  app.py process --x 10
```

**For deployment:**

```bash
flyte deploy \
  --image my-app-image=myregistry.com/my-app:v1.2.3 \
  app.py
```

#### Why use reference names?

1. **Decouples code from image URIs** - Change images without modifying code
2. **Supports multiple environments** - Different images for dev/staging/prod
3. **Integrates with CI/CD** - Build images in pipelines, reference in code
4. **Enables image reuse** - Multiple tasks can reference the same image

#### Example: Multi-environment deployment

```python
import flyte
import os

# Code references image by name
env = flyte.TaskEnvironment(
    name="api",
    image=flyte.Image.from_ref_name("api-service")
)

@env.task
def api_call(endpoint: str) -> dict:
    # Implementation
    return {"status": "success"}

if __name__ == "__main__":
    flyte.init_from_config()

    # Determine image based on environment
    environment = os.getenv("ENV", "dev")
    image_uri = {
        "dev": "myregistry.com/api-service:dev",
        "staging": "myregistry.com/api-service:staging",
        "prod": "myregistry.com/api-service:v1.2.3"
    }[environment]

    run = flyte.with_runcontext(
        copy_style="none",
        images={"api-service": image_uri}
    ).run(api_call, endpoint="/health")
```

### Container-based best practices

1. **Always set explicit versions** when using `copy_style="none"`:

   ```python
   flyte.with_runcontext(copy_style="none", version="v1.0.0")
   ```

2. **Match `root_dir` to `copy_contents_only`**:
   - `copy_contents_only=True` → `root_dir=Path(__file__).parent`
   - `copy_contents_only=False` → `root_dir=Path(__file__).parent.parent`

3. **Ensure `flyte` executable is in container** - Add to PATH or install flyte package

4. **Use `.dockerignore`** to exclude unnecessary files:

   ```
   # .dockerignore
   __pycache__/
   *.pyc
   .git/
   .venv/
   *.egg-info/
   ```

5. **Test containers locally** before deploying:

   ```bash
   docker run -it myimage:latest /bin/bash
   python -c "import mymodule"  # Verify imports work
   ```

### When to use container-based deployment

✅ **Use container-based when:**

- Deploying to production
- Need immutable, reproducible environments
- Working with complex system dependencies
- Deploying to air-gapped or restricted environments
- CI/CD pipelines with automated builds
- Code changes are infrequent

❌ **Don't use container-based when:**

- Rapid development and frequent code changes
- Quick prototyping
- Interactive development (Jupyter notebooks)
- Learning and experimentation

---

## Choosing the right approach

### Decision tree

```
Are you iterating quickly on code?
├─ Yes → Use Code Bundling (Default)
│         (Development, prototyping, notebooks)
│         Both approaches are fully reproducible via hash/tag
└─ No  → Do you need easy version rollback?
          ├─ Yes → Use Container-based
          │         (Production, CI/CD, straightforward tag-based rollback)
          └─ No  → Either works
                    (Code bundling is simpler, container-based for air-gapped)
```

### Hybrid approach

You can use different approaches for different tasks:

```python
import flyte
import pathlib

# Fast iteration for development tasks
dev_env = flyte.TaskEnvironment(
    name="dev",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas")
    # Code bundling (default)
)

# Immutable containers for production tasks
prod_env = flyte.TaskEnvironment(
    name="prod",
    image=flyte.Image.from_debian_base()
        .with_pip_packages("pandas")
        .with_source_folder(pathlib.Path(__file__).parent, copy_contents_only=True)
    # Requires copy_style="none"
)

@dev_env.task
def experimental_task(x: int) -> int:
    # Rapid development with code bundling
    return x * 2

@prod_env.task
def stable_task(x: int) -> int:
    # Production with embedded code
    return x ** 2

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)

    # Use code bundling for dev task
    dev_run = flyte.run(experimental_task, x=5)

    # Use container-based for prod task
    prod_run = flyte.with_runcontext(
        copy_style="none",
        version="v1.0.0"
    ).run(stable_task, x=5)
```

---

## Troubleshooting

### Import errors

**Problem:** `ModuleNotFoundError` when task executes remotely

**Solutions:**

1. **Check loaded modules** - Ensure modules are imported at module level:

   ```python
   # ✅ Good - bundled automatically
   from mymodule import helper

   @flyte.task
   def my_task():
       return helper.process()
   ```

   ```python
   # ❌ Bad - not bundled (lazy load)
   @flyte.task
   def my_task():
       from mymodule import helper
       return helper.process()
   ```

2. **Verify `root_dir`** matches your import structure:

   ```python
   # If imports are: from mypackage.utils import foo
   # Then root_dir should be parent of mypackage/
   flyte.init_from_config(root_dir=pathlib.Path(__file__).parent.parent)
   ```

3. **Use `copy_style="all"`** to bundle everything:

   ```bash
   flyte run --copy-style=all app.py my_task
   ```

### Code changes not reflected

**Problem:** Remote execution uses old code despite local changes

> [!NOTE]
> This is rare with code bundling - Flyte automatically versions based on content hash, so code changes should be detected automatically. This issue typically occurs with caching problems or when using `copy_style="none"`.

**Solutions:**

1. **Use explicit version bump** (mainly for container-based deployments):

   ```python
   run = flyte.with_runcontext(version="v2").run(my_task)
   ```

2. **Check if `copy_style="none"`** is set - this requires image rebuild:

   ```python
   # If using copy_style="none", rebuild image
   run = flyte.with_runcontext(
       copy_style="none",
       version="v2"  # Bump version to force rebuild
   ).run(my_task)
   ```

### Files missing in container

**Problem:** Task can't find data files or configs

**Solutions:**

1. **Use `copy_style="all"`** to bundle all files:

   ```bash
   flyte run --copy-style=all app.py my_task
   ```

2. **Copy files explicitly in image**:

   ```python
   image = flyte.Image.from_debian_base().with_source_file(
       src=pathlib.Path("config.yaml"),
       dst="/app/config.yaml"
   )
   ```

3. **Store data in remote storage** instead of bundling:

   ```python
   @flyte.task
   def my_task():
       # Read from S3/GCS instead of local files
       import flyte.io
       data = flyte.io.File("s3://bucket/data.csv").open().read()
   ```

### Container build failures

**Problem:** Image build fails with `copy_style="none"`

**Solutions:**

1. **Check `root_dir` matches `copy_contents_only`**:

   ```python
   # copy_contents_only=True
   image = Image.from_debian_base().with_source_folder(
       src=Path(__file__).parent,
       copy_contents_only=True
   )
   flyte.init(root_dir=Path(__file__).parent)  # Match!
   ```

2. **Ensure `flyte` executable available**:

   ```python
   image = Image.from_debian_base()  # Has flyte pre-installed
   ```

3. **Check file permissions** in source directory:

   ```bash
   chmod -R +r project/
   ```

### Version conflicts

**Problem:** Multiple versions of same image causing confusion

**Solutions:**

1. **Use explicit versions**:

   ```python
   run = flyte.with_runcontext(
       copy_style="none",
       version="v1.2.3"  # Explicit, not auto-generated
   ).run(my_task)
   ```

2. **Clean old images**:

   ```bash
   docker image prune -a
   ```

3. **Use semantic versioning** for clarity:

   ```python
   version = "v1.0.0"  # Major.Minor.Patch
   ```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/deployment-patterns ===

# Deployment patterns

Once you understand the basics of task deployment, you can use various deployment patterns to handle different project structures, dependency management approaches, and deployment requirements. This section covers the most common patterns with practical examples.

## Overview of deployment patterns

Flyte supports multiple deployment patterns to accommodate different project structures and requirements:

1. ****Tasks > Run and deploy tasks > Deployment patterns > Simple file deployment**** - Single file with tasks and environments
2. ****Tasks > Run and deploy tasks > Deployment patterns > Custom Dockerfile deployment**** - Full control over container environment
3. ****Tasks > Run and deploy tasks > Deployment patterns > PyProject package deployment**** - Structured Python packages with dependencies and async tasks
4. ****Tasks > Run and deploy tasks > Deployment patterns > Package structure deployment**** - Organized packages with shared environments
5. ****Tasks > Run and deploy tasks > Deployment patterns > Full build deployment**** - Complete code embedding in containers
6. ****Tasks > Run and deploy tasks > Deployment patterns > Python path deployment**** - Multi-directory project structures
7. ****Tasks > Run and deploy tasks > Deployment patterns > Dynamic environment deployment**** - Environment selection based on domain context

Each pattern serves specific use cases and can be combined as needed for complex projects.

## Simple file deployment

The simplest deployment pattern involves defining both your tasks and task environment in a single Python file. This pattern works well for:

- Prototyping and experimentation
- Simple tasks with minimal dependencies
- Educational examples and tutorials

### Example structure

```python
import flyte

env = flyte.TaskEnvironment(name="simple_env")

@env.task
async def my_task(name: str) -> str:
    return f"Hello, {name}!"

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/simple_file.py*

### Deployment commands

Deploy the environment:

```bash
flyte deploy my_example.py env
```

Run the task ephemerally:

```bash
flyte run my_example.py my_task --name "World"
```

### When to use

- Quick prototypes and experiments
- Single-purpose scripts
- Learning Flyte basics
- Tasks with no external dependencies

## Custom Dockerfile deployment

When you need full control over the container environment, you can specify a custom Dockerfile. This pattern is ideal for:

- Complex system dependencies
- Specific OS or runtime requirements
- Custom base images
- Multi-stage builds

### Example structure

```dockerfile
# syntax=docker/dockerfile:1.5
FROM ghcr.io/astral-sh/uv:0.8 as uv
FROM python:3.12-slim-bookworm

USER root

# Copy in uv so that later commands don't have to mount it in
COPY --from=uv /uv /usr/bin/uv

# Configure default envs
ENV UV_COMPILE_BYTECODE=1 \
    UV_LINK_MODE=copy \
    VIRTUALENV=/opt/venv \
    UV_PYTHON=/opt/venv/bin/python \
    PATH="/opt/venv/bin:$PATH"

# Create a virtualenv with the user specified python version
RUN uv venv /opt/venv --python=3.12

WORKDIR /root

# Install dependencies
COPY requirements.txt .
RUN uv pip install --pre -r /root/requirements.txt
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dockerfile/Dockerfile*

```python
from pathlib import Path

import flyte

env = flyte.TaskEnvironment(
    name="docker_env",
    image=flyte.Image.from_dockerfile(
        # relative paths in python change based on where you call, so set it relative to this file
        Path(__file__).parent / "Dockerfile",
        registry="ghcr.io/flyteorg",
        name="docker_env_image",
    ),
)

@env.task
def main(x: int) -> int:
    return x * 2

if __name__ == "__main__":
    import flyte.git

    flyte.init_from_config(flyte.git.config_from_root())

    run = flyte.run(main, x=10)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dockerfile/dockerfile_env.py*

### Alternative: Dockerfile in different directory

You can also reference Dockerfiles from subdirectories:

```python
from pathlib import Path

import flyte

env = flyte.TaskEnvironment(
    name="docker_env_in_dir",
    image=flyte.Image.from_dockerfile(
        # relative paths in python change based on where you call, so set it relative to this file
        Path(__file__).parent.parent / "Dockerfile.workdir",
        registry="ghcr.io/flyteorg",
        name="docker_env_image",
    ),
)

@env.task
def main(x: int) -> int:
    return x * 2

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, x=10)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dockerfile/src/docker_env_in_dir.py*

```dockerfile
# syntax=docker/dockerfile:1.5
FROM ghcr.io/astral-sh/uv:0.8 as uv
FROM python:3.12-slim-bookworm

USER root

# Copy in uv so that later commands don't have to mount it in
COPY --from=uv /uv /usr/bin/uv

# Configure default envs
ENV UV_COMPILE_BYTECODE=1 \
    UV_LINK_MODE=copy \
    VIRTUALENV=/opt/venv \
    UV_PYTHON=/opt/venv/bin/python \
    PATH="/opt/venv/bin:$PATH"

# Create a virtualenv with the user specified python version
RUN uv venv /opt/venv --python=3.12

WORKDIR /app

# Install dependencies
COPY requirements.txt .
RUN uv pip install --pre -r /app/requirements.txt
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dockerfile/Dockerfile.workdir*

### Key considerations

- **Path handling**: Use `Path(__file__).parent` for relative Dockerfile paths

  ```python
  # relative paths in python change based on where you call, so set it relative to this file
  Path(__file__).parent / "Dockerfile"
  ```

- **Registry configuration**: Specify a registry for image storage
- **Build context**: The directory containing the Dockerfile becomes the build context
- **Flyte installation**: Ensure Flyte is installed in the container and available on `$PATH`

  ```dockerfile
  # Install Flyte in your Dockerfile
  RUN pip install flyte
  ```

- **Dependencies**: Include all application requirements in the Dockerfile or requirements.txt

### When to use

- Need specific system packages or tools
- Custom base image requirements
- Complex installation procedures
- Multi-stage build optimization

## PyProject package deployment

For structured Python projects with proper package management, use the PyProject pattern. This approach demonstrates a **realistic Python project structure** that provides:

- Proper dependency management with `pyproject.toml` and external packages like `httpx`
- Clean separation of business logic and Flyte tasks across multiple modules
- Professional project structure with `src/` layout
- Async task execution with API calls and data processing
- Entrypoint patterns for both command-line and programmatic execution

### Example structure

```
pyproject_package/
├── pyproject.toml          # Project metadata and dependencies
├── README.md              # Documentation
└── src/
    └── pyproject_package/
        ├── __init__.py     # Package initialization
        ├── main.py         # Entrypoint script
        ├── data/
        │   ├── __init__.py
        │   ├── loader.py   # Data loading utilities (no Flyte)
        │   └── processor.py # Data processing utilities (no Flyte)
        ├── models/
        │   ├── __init__.py
        │   └── analyzer.py # Analysis utilities (no Flyte)
        └── tasks/
            ├── __init__.py
            └── tasks.py    # Flyte task definitions
```

### Business logic modules

The business logic is completely separate from Flyte and can be used independently:

#### Data loading (`data/loader.py`)

```python
import json
from pathlib import Path
from typing import Any

import httpx

async def fetch_data_from_api(url: str) -> list[dict[str, Any]]:
    async with httpx.AsyncClient() as client:
        response = await client.get(url, timeout=10.0)
        response.raise_for_status()
        return response.json()

def load_local_data(file_path: str | Path) -> dict[str, Any]:
    path = Path(file_path)

    if not path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")

    with path.open("r") as f:
        return json.load(f)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/src/pyproject_package/data/loader.py*

#### Data processing (`data/processor.py`)

```python
import asyncio
from typing import Any

from pydantic import BaseModel, Field, field_validator

class DataItem(BaseModel):
    id: int = Field(gt=0, description="Item ID must be positive")
    value: float = Field(description="Item value")
    category: str = Field(min_length=1, description="Item category")

    @field_validator("category")
    @classmethod
    def category_must_be_lowercase(cls, v: str) -> str:
        return v.lower()

def clean_data(raw_data: dict[str, Any]) -> dict[str, Any]:
    # Remove None values
    cleaned = {k: v for k, v in raw_data.items() if v is not None}

    # Validate items if present
    if "items" in cleaned:
        validated_items = []
        for item in cleaned["items"]:
            try:
                validated = DataItem(**item)
                validated_items.append(validated.model_dump())
            except Exception as e:
                print(f"Skipping invalid item {item}: {e}")
                continue
        cleaned["items"] = validated_items

    return cleaned

def transform_data(data: dict[str, Any]) -> list[dict[str, Any]]:
    items = data.get("items", [])

    # Add computed fields
    transformed = []
    for item in items:
        transformed_item = {
            **item,
            "value_squared": item["value"] ** 2,
            "category_upper": item["category"].upper(),
        }
        transformed.append(transformed_item)

    return transformed

async def aggregate_data(items: list[dict[str, Any]]) -> dict[str, Any]:
    # Simulate async processing
    await asyncio.sleep(0.1)

    aggregated: dict[str, dict[str, Any]] = {}

    for item in items:
        category = item["category"]

        if category not in aggregated:
            aggregated[category] = {
                "count": 0,
                "total_value": 0.0,
                "values": [],
            }

        aggregated[category]["count"] += 1
        aggregated[category]["total_value"] += item["value"]
        aggregated[category]["values"].append(item["value"])

    # Calculate averages
    for category, v in aggregated.items():
        total = v["total_value"]
        count = v["count"]
        v["average_value"] = total / count if count > 0 else 0.0

    return {"categories": aggregated, "total_items": len(items)}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/src/pyproject_package/data/processor.py*

#### Analysis (`models/analyzer.py`)

```python
from typing import Any

import numpy as np

def calculate_statistics(data: list[dict[str, Any]]) -> dict[str, Any]:
    if not data:
        return {
            "count": 0,
            "mean": 0.0,
            "median": 0.0,
            "std_dev": 0.0,
            "min": 0.0,
            "max": 0.0,
        }

    values = np.array([item["value"] for item in data])

    stats = {
        "count": len(values),
        "mean": float(np.mean(values)),
        "median": float(np.median(values)),
        "std_dev": float(np.std(values)),
        "min": float(np.min(values)),
        "max": float(np.max(values)),
        "percentile_25": float(np.percentile(values, 25)),
        "percentile_75": float(np.percentile(values, 75)),
    }

    return stats

def generate_report(stats: dict[str, Any]) -> str:
    report_lines = [
        "=" * 60,
        "DATA ANALYSIS REPORT",
        "=" * 60,
    ]

    # Basic statistics section
    if "basic" in stats:
        basic = stats["basic"]
        report_lines.extend(
            [
                "",
                "BASIC STATISTICS:",
                f"  Count:       {basic.get('count', 0)}",
                f"  Mean:        {basic.get('mean', 0.0):.2f}",
                f"  Median:      {basic.get('median', 0.0):.2f}",
                f"  Std Dev:     {basic.get('std_dev', 0.0):.2f}",
                f"  Min:         {basic.get('min', 0.0):.2f}",
                f"  Max:         {basic.get('max', 0.0):.2f}",
                f"  25th %ile:   {basic.get('percentile_25', 0.0):.2f}",
                f"  75th %ile:   {basic.get('percentile_75', 0.0):.2f}",
            ]
        )

    # Category aggregations section
    if "aggregated" in stats and "categories" in stats["aggregated"]:
        categories = stats["aggregated"]["categories"]
        total_items = stats["aggregated"].get("total_items", 0)

        report_lines.extend(
            [
                "",
                "CATEGORY BREAKDOWN:",
                f"  Total Items: {total_items}",
                "",
            ]
        )

        for category, cat_stats in sorted(categories.items()):
            report_lines.extend(
                [
                    f"  Category: {category.upper()}",
                    f"    Count:         {cat_stats.get('count', 0)}",
                    f"    Total Value:   {cat_stats.get('total_value', 0.0):.2f}",
                    f"    Average Value: {cat_stats.get('average_value', 0.0):.2f}",
                    "",
                ]
            )

    report_lines.append("=" * 60)

    return "\n".join(report_lines)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/src/pyproject_package/models/analyzer.py*

These modules demonstrate:

- **No Flyte dependencies** - can be tested and used independently
- **Pydantic models** for data validation with custom validators
- **Async patterns** with proper context managers and error handling
- **NumPy integration** for statistical calculations
- **Professional error handling** with timeouts and validation

### Flyte orchestration layer

The Flyte tasks orchestrate the business logic with proper async execution:

```python
import pathlib
from typing import Any

import flyte
from pyproject_package.data import loader, processor
from pyproject_package.models import analyzer

UV_PROJECT_ROOT = pathlib.Path(__file__).parent.parent.parent.parent

env = flyte.TaskEnvironment(
    name="data_pipeline",
    image=flyte.Image.from_debian_base().with_uv_project(pyproject_file=UV_PROJECT_ROOT / "pyproject.toml"),
    resources=flyte.Resources(memory="512Mi", cpu="500m"),
)

@env.task
async def fetch_task(url: str) -> list[dict[str, Any]]:
    print(f"Fetching data from: {url}")
    data = await loader.fetch_data_from_api(url)
    print(f"Fetched {len(data)} top-level keys")
    return data

@env.task
async def process_task(raw_data: dict[str, Any]) -> list[dict[str, Any]]:
    print("Cleaning data...")
    cleaned = processor.clean_data(raw_data)

    print("Transforming data...")
    transformed = processor.transform_data(cleaned)

    print(f"Processed {len(transformed)} items")
    return transformed

@env.task
async def analyze_task(processed_data: list[dict[str, Any]]) -> str:
    print("Aggregating data...")
    aggregated = await processor.aggregate_data(processed_data)

    print("Calculating statistics...")
    stats = analyzer.calculate_statistics(processed_data)

    print("Generating report...")
    report = analyzer.generate_report({"basic": stats, "aggregated": aggregated})

    print("\n" + report)
    return report

@env.task
async def pipeline(api_url: str) -> str:
    # Chain tasks together
    raw_data = await fetch_task(url=api_url)
    processed_data = await process_task(raw_data=raw_data[0])
    report = await analyze_task(processed_data=processed_data)

    return report
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/src/pyproject_package/tasks/tasks.py*

### Entrypoint configuration

The main entrypoint demonstrates proper initialization and execution patterns:

```python
import pathlib

import flyte
from pyproject_package.tasks.tasks import pipeline

def main():
    # Initialize Flyte connection
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent.parent)

    # Example API URL with mock data
    # In a real scenario, this would be a real API endpoint
    example_url = "https://jsonplaceholder.typicode.com/posts"

    # For demonstration, we'll use mock data instead of the actual API
    # to ensure the example works reliably
    print("Starting data pipeline...")
    print(f"Target API: {example_url}")

    # To run remotely, uncomment the following:
    run = flyte.run(pipeline, api_url=example_url)
    print(f"\nRun Name: {run.name}")
    print(f"Run URL: {run.url}")
    run.wait()

if __name__ == "__main__":
    main()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/src/pyproject_package/main.py*

### Dependencies and configuration

```toml
[project]
name = "pyproject-package"
version = "0.1.0"
description = "Example Python package with Flyte tasks and modular business logic"
readme = "README.md"
authors = [
    { name = "Ketan Umare", email = "kumare3@users.noreply.github.com" }
]
requires-python = ">=3.10"
dependencies = [
    "flyte>=2.0.0b52",
    "httpx>=0.27.0",
    "numpy>=1.26.0",
    "pydantic>=2.0.0",
]

[project.scripts]
run-pipeline = "pyproject_package.main:main"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pyproject_package/pyproject.toml*

### Key features

- **Async task chains**: Tasks can be chained together with proper async/await patterns
- **External dependencies**: Demonstrates integration with external libraries (`httpx`, `pyyaml`)
- **uv integration**: Uses `.with_uv_project()` for dependency management
- **Resource specification**: Shows how to set memory and CPU requirements
- **Proper error handling**: Includes timeout and error handling in API calls

### Key learning points

1. **Separation of concerns**: Business logic (`data/`, `models/`) separate from orchestration (`main.py`)
2. **Reusable code**: Non-Flyte modules can be tested independently and reused
3. **Async support**: Demonstrates async Flyte tasks for I/O-bound operations
4. **Dependency management**: Shows how external packages integrate with Flyte
5. **Realistic structure**: Mirrors real-world Python project organization
6. **Entrypoint script**: Shows how to create runnable entry points

### Usage patterns

**Run locally:**

```bash
python -m pyproject_package.main
```

**Deploy to Flyte:**

```bash
flyte deploy .
```

**Run remotely:**

```bash
python -m pyproject_package.main  # Uses remote execution
```

### What this example demonstrates

- Multiple files and modules in a package
- Async Flyte tasks with external API calls
- Separation of business logic from orchestration
- External dependencies (`httpx`, `numpy`, `pydantic`)
- **Data validation with Pydantic models** for data processing
- **Professional error handling** with try/catch for data validation
- **Timeout configuration** for external API calls (`timeout=10.0`)
- **Async context managers** for proper resource management (`async with httpx.AsyncClient()`)
- Entrypoint script pattern with `project.scripts`
- Realistic project structure with `src/` layout
- Task chaining and data flow
- How non-Flyte code integrates with Flyte tasks

### When to use

- Production-ready, maintainable projects
- Projects requiring external API integration
- Complex data processing pipelines
- Team development with proper separation of concerns
- Applications needing async execution patterns

## Package structure deployment

For organizing Flyte workflows in a package structure with shared task environments and utilities, use this pattern. It's particularly useful for:

- Multiple workflows that share common environments and utilities
- Organized code structure with clear module boundaries
- Projects where you want to reuse task environments across workflows

### Example structure

```
lib/
├── __init__.py
└── workflows/
    ├── __init__.py
    ├── workflow1.py    # First workflow
    ├── workflow2.py    # Second workflow
    ├── env.py          # Shared task environment
    └── utils.py        # Shared utilities
```

### Key concepts

- **Shared environments**: Define task environments in `env.py` and import across workflows
- **Utility modules**: Common functions and utilities shared between workflows
- **Root directory handling**: Use `--root-dir` flag for proper Python path configuration

### Running with root directory

When running workflows with a package structure, specify the root directory:

```bash
flyte run --root-dir . lib/workflows/workflow1.py process_workflow
flyte run --root-dir . lib/workflows/workflow2.py math_workflow --n 6
```

### How `--root-dir` works

The `--root-dir` flag automatically configures the Python path (`sys.path`) to ensure:

1. **Local execution**: Package imports work correctly when running locally
2. **Consistent behavior**: Same Python path configuration locally and at runtime
3. **No manual PYTHONPATH**: Eliminates need to manually export environment variables
4. **Runtime packaging**: Flyte packages and copies code correctly to execution environment
5. **Runtime consistency**: The same package structure is preserved in the runtime container

### Alternative: Using a Python project

For larger projects, create a proper Python project with `pyproject.toml`:

```toml
# pyproject.toml
[project]
name = "lib"
version = "0.1.0"

[build-system]
requires = ["setuptools>=45", "wheel"]
build-backend = "setuptools.build_meta"
```

Then install in editable mode:

```bash
pip install -e .
```

After installation, you can run workflows without `--root-dir`:

```bash
flyte run lib/workflows/workflow1.py process_workflow
```

However, for deployment and remote execution, still use `--root-dir` for consistency:

```bash
flyte run --root-dir . lib/workflows/workflow1.py process_workflow
flyte deploy --root-dir . lib/workflows/workflow1.py
```

### When to use

- Multiple related workflows in one project
- Shared task environments and utilities
- Team projects with multiple contributors
- Applications requiring organized code structure
- Projects that benefit from proper Python packaging

## Full build deployment

When you need complete reproducibility and want to embed all code directly in the container image, use the full build pattern. This disables Flyte's fast deployment system in favor of traditional container builds.

### Overview

By default, Flyte uses a fast deployment system that:

- Creates a tar archive of your files
- Skips the full image build and push process
- Provides faster iteration during development

However, sometimes you need to **completely embed your code into the container image** for:

- Full reproducibility with immutable container images
- Environments where fast deployment isn't available
- Production deployments with all dependencies baked in
- Air-gapped or restricted deployment environments

### Key configuration

```python
import pathlib

from dep import foo

import flyte

env = flyte.TaskEnvironment(
    name="full_build",
    image=flyte.Image.from_debian_base().with_source_folder(
        pathlib.Path(__file__).parent,
        copy_contents_only=True  # Avoid nested folders
    ),
)

@env.task
def square(x) -> int:
    return x ** foo()

@env.task
def main(n: int) -> list[int]:
    return list(flyte.map(square, range(n)))

if __name__ == "__main__":
    # copy_contents_only=True requires root_dir=parent, False requires root_dir=parent.parent
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext(copy_style="none", version="x").run(main, n=10)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/full_build/main.py*

### Local dependency example

The main.py file imports from a local dependency that gets included in the build:

```python
def foo() -> int:
    return 1
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/full_build/dep.py*

### Critical configuration components

1. **Set `copy_style` to `"none"`**:

   ```python
   flyte.with_runcontext(copy_style="none", version="x").run(main, n=10)
   ```

   This disables Flyte's fast deployment system and forces a full container build.

2. **Set a custom version**:

   ```python
   flyte.with_runcontext(copy_style="none", version="x").run(main, n=10)
   ```

   The `version` parameter should be set to a desired value (not auto-generated) for consistent image tagging.

3. **Configure image source copying**:

   ```python
   image=flyte.Image.from_debian_base().with_source_folder(
       pathlib.Path(__file__).parent,
       copy_contents_only=True
   )
   ```

   Use `.with_source_folder()` to specify what code to copy into the container.

4. **Set `root_dir` correctly**:

   ```python
   flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
   ```

   - If `copy_contents_only=True`: Set `root_dir` to the source folder (contents are copied)
   - If `copy_contents_only=False`: Set `root_dir` to parent directory (folder is copied)

### Configuration options

#### Option A: Copy folder structure

```python
# Copies the entire folder structure into the container
image=flyte.Image.from_debian_base().with_source_folder(
    pathlib.Path(__file__).parent,
    copy_contents_only=False  # Default
)

# When copy_contents_only=False, set root_dir to parent.parent
flyte.init_from_config(root_dir=pathlib.Path(__file__).parent.parent)
```

#### Option B: Copy contents only (recommended)

```python
# Copies only the contents of the folder (flattens structure)
# This is useful when you want to avoid nested folders - for example all your code is in the root of the repo
image=flyte.Image.from_debian_base().with_source_folder(
    pathlib.Path(__file__).parent,
    copy_contents_only=True
)

# When copy_contents_only=True, set root_dir to parent
flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
```

### Version management best practices

When using `copy_style="none"`, always specify an explicit version:

- Use semantic versioning: `"v1.0.0"`, `"v1.1.0"`
- Use build numbers: `"build-123"`
- Use git commits: `"abc123"`

Avoid auto-generated versions to ensure reproducible deployments.

### Performance considerations

- **Full builds take longer** than fast deployment
- **Container images will be larger** as they include all source code
- **Better for production** where immutability is important
- **Use during development** when testing the full deployment pipeline

### When to use

✅ **Use full build when:**

- Deploying to production environments
- Need immutable, reproducible container images
- Working with complex dependency structures
- Deploying to air-gapped or restricted environments
- Building CI/CD pipelines

❌ **Don't use full build when:**

- Rapid development and iteration
- Working with frequently changing code
- Development environments where speed matters
- Simple workflows without complex dependencies

### Troubleshooting

**Common issues:**

1. **Import errors**: Check your `root_dir` configuration matches `copy_contents_only`
2. **Missing files**: Ensure all dependencies are in the source folder
3. **Version conflicts**: Use explicit, unique version strings
4. **Build failures**: Check that the base image has all required system dependencies

**Debug tips:**

- Add print statements to verify file paths in containers
- Use `docker run -it <image> /bin/bash` to inspect built images
- Check Flyte logs for build errors and warnings
- Verify that relative imports work correctly in the container context

## Python path deployment

For projects where workflows are separated from business logic across multiple directories, use the Python path pattern with proper `root_dir` configuration.

### Example structure

```
pythonpath/
├── workflows/
│   └── workflow.py      # Flyte workflow definitions
├── src/
│   └── my_module.py     # Business logic modules
├── run.sh               # Execute from project root
└── run_inside_folder.sh # Execute from workflows/ directory
```

### Implementation

```python
import pathlib

from src.my_module import env, say_hello

import flyte

env = flyte.TaskEnvironment(
    name="workflow_env",
    depends_on=[env],
)

@env.task
async def greet(name: str) -> str:
    return await say_hello(name)

if __name__ == "__main__":
    current_dir = pathlib.Path(__file__).parent
    flyte.init_from_config(root_dir=current_dir.parent)
    r = flyte.run(greet, name="World")
    print(r.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pythonpath/workflows/workflow.py*

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_module",
)

@env.task
async def say_hello(name: str) -> str:
    return f"Hello, {name}!"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/pythonpath/src/my_module.py*

### Task environment dependencies

Note how the workflow imports both the task environment and the task function:

```python
from src.my_module import env, say_hello

env = flyte.TaskEnvironment(
    name="workflow_env",
    depends_on=[env],  # Depends on the imported environment
)
```

This pattern allows sharing task environments across modules while maintaining proper dependency relationships.

### Key considerations

- **Import resolution**: `root_dir` enables proper module imports across directories
- **File packaging**: Flyte packages all files starting from `root_dir`
- **Execution flexibility**: Works regardless of where you execute the script
- **PYTHONPATH handling**: Different behavior for CLI vs direct Python execution

### CLI vs direct Python execution

#### Using Flyte CLI with `--root-dir` (Recommended)

When using `flyte run` with `--root-dir`, you don't need to export PYTHONPATH:

```bash
flyte run --root-dir . workflows/workflow.py greet --name "World"
```

The CLI automatically:

- Adds the `--root-dir` location to `sys.path`
- Resolves all imports correctly
- Packages files from the root directory for remote execution

#### Using Python directly

When running Python scripts directly, you must set PYTHONPATH manually:

```bash
PYTHONPATH=.:$PYTHONPATH python workflows/workflow.py
```

This is because:

- Python doesn't automatically know about your project structure
- You need to explicitly tell Python where to find your modules
- The `root_dir` parameter handles remote packaging, not local path resolution

### Best practices

1. **Always set `root_dir`** when workflows import from multiple directories
2. **Use pathlib** for cross-platform path handling
3. **Set `root_dir` to your project root** to ensure all dependencies are captured
4. **Test both execution patterns** to ensure deployment works from any directory

### Common pitfalls

- **Forgetting `root_dir`**: Results in import errors during remote execution
- **Wrong `root_dir` path**: May package too many or too few files
- **Not setting PYTHONPATH when using Python directly**: Use `flyte run --root-dir .` instead
- **Mixing execution methods**: If you use `flyte run --root-dir .`, you don't need PYTHONPATH

### When to use

- Legacy projects with established directory structures
- Separation of concerns between workflows and business logic
- Multiple workflow definitions sharing common modules
- Projects with complex import hierarchies

**Note:** This pattern is an escape hatch for larger projects where code organization requires separating workflows from business logic. Ideally, structure projects with `pyproject.toml` for cleaner dependency management.

## Dynamic environment deployment

For environments that need to change based on deployment context (development vs production), use dynamic environment selection based on Flyte domains.

### Domain-based environment selection

Use `flyte.current_domain()` to deterministically create different task environments based on the deployment domain:

```python
# NOTE: flyte.init() invocation at the module level is strictly discouraged.
# At runtime, Flyte controls initialization and configuration files are not present.

import os

import flyte

def create_env():
    if flyte.current_domain() == "development":
        return flyte.TaskEnvironment(name="dev", image=flyte.Image.from_debian_base(), env_vars={"MY_ENV": "dev"})
    return flyte.TaskEnvironment(name="prod", image=flyte.Image.from_debian_base(), env_vars={"MY_ENV": "prod"})

env = create_env()

@env.task
async def my_task(n: int) -> int:
    print(f"Environment Variable MY_ENV = {os.environ['MY_ENV']}", flush=True)
    return n + 1

@env.task
async def entrypoint(n: int) -> int:
    print(f"Environment Variable MY_ENV = {os.environ['MY_ENV']}", flush=True)
    return await my_task(n)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dynamic_environments/environment_picker.py*

### Why this pattern works

**Environment reproducibility in local and remote clusters is critical.** Flyte re-instantiates modules in remote clusters, so `current_domain()` will be set correctly based on where the code executes.

✅ **Do use `flyte.current_domain()`** - Flyte automatically sets this based on the execution context

❌ **Don't use environment variables directly** - They won't yield correct results unless manually passed to the downstream system

### How it works

1. Flyte sets the domain context when initializing
2. `current_domain()` returns the domain string (e.g., "development", "staging", "production")
3. Your code deterministically configures resources based on this domain
4. When Flyte executes remotely, it re-instantiates modules with the correct domain context
5. The same environment configuration logic runs consistently everywhere

### Important constraints

`flyte.current_domain()` only works **after** `flyte.init()` is called:

- ✅ Works with `flyte run` and `flyte deploy` CLI commands (they init automatically)
- ✅ Works when called from `if __name__ == "__main__"` after explicit `flyte.init()`
- ❌ Does NOT work at module level without initialization

**Critical:** `flyte.init()` invocation at the module level is **strictly discouraged**. The reason is that at runtime, Flyte controls the initialization and configuration files are not present at runtime.

### Alternative: environment variable approach

For cases where you need to pass domain information as environment variables to the container runtime, use this approach:

```python
import os

import flyte

def create_env(domain: str):
    # Pass domain as environment variable so tasks can see which domain they're running in
    if domain == "development":
        return flyte.TaskEnvironment(name="dev", image=flyte.Image.from_debian_base(), env_vars={"DOMAIN_NAME": domain})
    return flyte.TaskEnvironment(name="prod", image=flyte.Image.from_debian_base(), env_vars={"DOMAIN_NAME": domain})

env = create_env(os.getenv("DOMAIN_NAME", "development"))

@env.task
async def my_task(n: int) -> int:
    print(f"Environment Variable MY_ENV = {os.environ['DOMAIN_NAME']}", flush=True)
    return n + 1

@env.task
async def entrypoint(n: int) -> int:
    print(f"Environment Variable MY_ENV = {os.environ['DOMAIN_NAME']}", flush=True)
    return await my_task(n)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(entrypoint, n=5)
    print(r.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dynamic_environments_with_envvars/environment_picker.py*

#### Key differences from domain-based approach

- **Environment variable access**: The domain name is available inside tasks via `os.environ['DOMAIN_NAME']`
- **External control**: Can be controlled via system environment variables before execution
- **Runtime visibility**: Tasks can inspect which environment they're running in during execution
- **Default fallback**: Uses `"development"` as default when `DOMAIN_NAME` is not set

#### Usage with environment variables

Set the environment variable and run:

```bash
export DOMAIN_NAME=production
flyte run environment_picker.py entrypoint --n 5
```

Or set it inline:

```bash
DOMAIN_NAME=development flyte run environment_picker.py entrypoint --n 5
```

#### When to use environment variables vs domain-based

**Use environment variables when:**

- Tasks need runtime access to environment information
- External systems set environment configuration
- You need flexibility to override environment externally
- Debugging requires visibility into environment selection

**Use domain-based approach when:**

- Environment selection should be automatic based on Flyte domain
- You want tighter integration with Flyte's domain system
- No need for runtime environment inspection within tasks

You can vary multiple aspects based on context:

- **Base images**: Different images for dev vs prod
- **Environment variables**: Configuration per environment
- **Resource requirements**: Different CPU/memory per domain
- **Dependencies**: Different package versions
- **Registry settings**: Different container registries

### Usage patterns

```bash
flyte run environment_picker.py entrypoint --n 5
flyte deploy environment_picker.py
```

For programmatic usage, ensure proper initialization:

```python
import flyte

flyte.init_from_config()
from environment_picker import entrypoint

if __name__ == "__main__":
    r = flyte.run(entrypoint, n=5)
    print(r.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/deployment-patterns/dynamic_environments/main.py*

### When to use dynamic environments

**General use cases:**

- Multi-environment deployments (dev/staging/prod)
- Different resource requirements per environment
- Environment-specific dependencies or settings
- Context-sensitive configuration needs

**Domain-based approach for:**

- Automatic environment selection tied to Flyte domains
- Simpler configuration without external environment variables
- Integration with Flyte's built-in domain system

**Environment variable approach for:**

- Runtime visibility into environment selection within tasks
- External control over environment configuration
- Debugging and logging environment-specific behavior
- Integration with external deployment systems that set environment variables

## Dependency and version management

Flyte declares its dependencies with loose version ranges so that installing `flyte` doesn't over-constrain the rest of your project. The downside is that a new release of a transitive dependency can sometimes introduce a regression in an environment that isn't otherwise pinned. To get reproducible installs, pin your dependency versions using one of the approaches below. Both are optional. Neither changes how a plain `pip install flyte` behaves.

### Recommended: use a lockfile

If you manage your project with `uv` (or any tool that produces a lockfile), commit the lockfile and install from it:

```bash
uv add flyte          # resolves and records exact versions in uv.lock
uv sync --frozen      # installs exactly what uv.lock specifies, no re-resolution
```

Commit `uv.lock` to version control. A lockfile records the exact version of every direct and transitive dependency, so every install resolves to the same versions locally, in CI, and in your Flyte task images. When you build task images with `.with_uv_project()` (see **Tasks > Run and deploy tasks > Deployment patterns > PyProject package deployment**), the image is built from the locked environment, so what you test is what runs remotely.

### No lockfile? Install against Flyte's published constraints

If you don't use a lockfile tool, Flyte publishes a constraints file with every release. It lists the exact transitive dependency versions that release was tested against. Pass it to your installer with `-c`:

```bash
# Pin to the versions tested for a specific release
pip install "flyte==<version>" \
  -c "https://github.com/flyteorg/flyte-sdk/releases/download/v<version>/constraints-<version>.txt"

# Or always track the latest release's constraints
pip install flyte \
  -c "https://github.com/flyteorg/flyte-sdk/releases/latest/download/constraints.txt"
```

`uv pip install` accepts the same `-c` flag. A constraints file only pins the versions of packages that get installed. It never pulls a package into your environment on its own, and using it is optional. It gives you Flyte's tested dependency set without adopting a lockfile.

### Choosing an approach

| Approach | Reproducibility | Best for |
|----------|-----------------|----------|
| Lockfile (`uv.lock`) | Full. Every dependency pinned in your own repo | Projects already using `uv` or another lockfile tool |
| Published constraints (`-c`) | Pinned to Flyte's per-release tested versions | Projects that can't or don't want to maintain a lockfile |
| Default ranges (no pinning) | None | Quick prototypes and experiments |

## Best practices

### Project organization

1. **Separate concerns**: Keep business logic separate from Flyte task definitions
2. **Use proper imports**: Structure projects for clean import patterns
3. **Version control**: Include all necessary files in version control
4. **Documentation**: Document deployment requirements and patterns

### Image management

1. **Registry configuration**: Use consistent registry settings across environments
2. **Image tagging**: Use meaningful tags for production deployments
3. **Base image selection**: Choose appropriate base images for your needs
4. **Dependency management**: Keep container images lightweight but complete

### Configuration management

1. **Root directory**: Set `root_dir` appropriately for your project structure
2. **Path handling**: Use `pathlib.Path` for cross-platform compatibility
3. **Environment variables**: Use environment-specific configurations
4. **Secrets management**: Handle sensitive data appropriately

### Development workflow

1. **Local testing**: Test tasks locally before deployment
2. **Incremental development**: Use `flyte run` for quick iterations
3. **Production deployment**: Use `flyte deploy` for permanent deployments
4. **Monitoring**: Monitor deployed tasks and environments

## Choosing the right pattern

| Pattern | Use Case | Complexity | Best For |
|---------|----------|------------|----------|
| Simple file | Quick prototypes, learning | Low | Single tasks, experiments |
| Custom Dockerfile | System dependencies, custom environments | Medium | Complex dependencies |
| PyProject package | Professional projects, async pipelines | Medium-High | Production applications |
| Package structure | Multiple workflows, shared utilities | Medium | Organized team projects |
| Full build | Production, reproducibility | High | Immutable deployments |
| Python path | Legacy structures, separated concerns | Medium | Existing codebases |
| Dynamic environment | Multi-environment, domain-aware deployments | Medium | Context-aware deployments |

Start with simpler patterns and evolve to more complex ones as your requirements grow. Many projects will combine multiple patterns as they scale and mature.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/run-context ===

# Run context

Every Flyte run has a **run context**: a set of invocation-time parameters that control where the run executes, where its outputs are stored, how caching behaves, and more.

There are two sides to run context:

- **Write side**: `flyte.with_runcontext()` sets run parameters before the run starts (programmatic) or via CLI flags.
- **Read side**: `flyte.ctx()` accesses run parameters inside a running task.

## Configuring a run with `flyte.with_runcontext()`

`flyte.with_runcontext()` returns a runner object. Call `.run(task, ...)` on it to start the run with the specified context:

```
import flyte

env = flyte.TaskEnvironment("run-context-example")

@env.task
async def process(n: int) -> int:
    return n * 2

@env.task
async def root() -> int:
    return await process(21)

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.with_runcontext(
        name="my-run",
        project="my-project",
        domain="development",
    ).run(root)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/run-context/run_context.py*

All parameters are optional. Unset parameters inherit from the configuration file (`config.yaml`) or system defaults.

### Execution target

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `mode` | `"local"` \| `"remote"` \| `"hybrid"` | *from config* | Where the run executes. `"remote"` runs on the Flyte backend; `"local"` runs in-process. |
| `project` | `str` | *from config* | Project to run in. |
| `domain` | `str` | *from config* | Domain to run in (e.g. `"development"`, `"production"`). |
| `name` | `str` | *auto-generated* | Custom name for the run, visible in the UI. |
| `version` | `str` | *from code bundle* | Version string for the ephemeral task deployment. |
| `queue` | `str` | *from config* | Cluster queue to schedule tasks on. |
| `interruptible` | `bool` | *per-task setting* | Override the interruptible setting for all tasks in the run. `True` allows spot/preemptible instances; `False` forces non-interruptible instances. |

### Storage

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `raw_data_path` | `str` | *from config* | Storage prefix for offloaded data types ([Files](../task-programming/files-and-directories), [Dirs](../task-programming/files-and-directories), [DataFrames](../task-programming/dataframes), checkpoints). Accepts `s3://`, `gs://`, or local paths. |
| `run_base_dir` | `str` | *auto-generated* | Base directory for the run's inputs, outputs, and intermediate per-action artifacts in the data plane object store. Distinct from `raw_data_path`. |

For the difference between what `raw_data_path` controls (offloaded values) and what stays at the deployment-configured location (`inputs.pb`, `outputs.pb`, Decks) or in the control plane database, see [Where your data lives](../../get-started/core-concepts/where-data-lives).

To direct all task outputs to a specific bucket for a run:

```
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.with_runcontext(
        # Store all task outputs in a dedicated S3 prefix for this run
        raw_data_path="s3://my-bucket/runs/experiment-42/",
    ).run(root)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/run-context/run_context.py*

The equivalent CLI flag is `--raw-data-path`. See [Run command options](./run-command-options#--raw-data-path) for CLI usage.

### Caching

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `overwrite_cache` | `bool` | `False` | Re-execute all tasks even if a cached result exists, and overwrite the cache with new results. |
| `disable_run_cache` | `bool` | `False` | Skip cache lookups and writes entirely for this run. |
| `cache_lookup_scope` | `"global"` \| ... | `"global"` | Scope for cache lookups. |

### Identity and resources

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `service_account` | `str` | *from config* | Kubernetes service account for task pods. |
| `env_vars` | `Dict[str, str]` | `None` | Additional environment variables to inject into task containers. |
| `labels` | `Dict[str, str]` | `None` | User-defined `key=value` labels attached to the run — used to filter/organize runs (`flyte get run --with-label`) and propagated to the task pods as Kubernetes labels. |
| `annotations` | `Dict[str, str]` | `None` | User-defined `key=value` annotations attached to the run and propagated to the task pods as Kubernetes annotations (not filterable; no CLI flag). |

Labels tag a run with arbitrary `key=value` metadata so you can find and group related runs later, and are also propagated to the run's task pods as Kubernetes labels (available for cluster-level monitoring, routing, or policies). Set them programmatically with `with_runcontext(labels={...})`, or from the CLI with the repeatable `--label` flag:

```bash
flyte run --label team=ml --label env=prod my_example.py main
```

To list and filter runs by their labels, see [Filtering runs by label](./interacting-with-runs#filtering-runs-by-label).

### Logging

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `log_level` | `int` | *from config* | Python log level for the framework logger (`flyte`), e.g. `logging.DEBUG`. |
| `user_log_level` | `int` | *from config* | Python log level for the user logger (`flyte.user`, i.e. `flyte.logger`). |
| `log_format` | `"console"` \| `"json"` | `"console"` | Log output format. |
| `reset_root_logger` | `bool` | `False` | If `True`, clear the root logger's existing handlers and install Flyte's own. If `False` (the default), leave existing root handlers in place and wrap their formatters with the run/action context. |

For setting the logging level (including via environment variables), the framework-vs-user logger split, and capturing third-party library logs as JSON with `reset_root_logger`, see [Logging](../task-configuration/logging).

### Code bundling

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `copy_style` | `"loaded_modules"` \| `"all"` \| `"none"` | `"loaded_modules"` | Code bundling strategy. See [Run command options](./run-command-options#--copy-style). |
| `dry_run` | `bool` | `False` | Build and upload the code bundle without executing the run. |
| `copy_bundle_to` | `Path` | `None` | When `dry_run=True`, copy the bundle to this local path. |
| `interactive_mode` | `bool` | *auto-detected* | Override interactive mode detection (set automatically for Jupyter notebooks). |
| `preserve_original_types` | `bool` | `False` | Keep native DataFrame types (e.g. `pd.DataFrame`) rather than converting to `flyte.io.DataFrame` when deserializing outputs. |

### Context propagation

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `custom_context` | `Dict[str, str]` | `None` | Metadata propagated through the entire task hierarchy. Readable inside any task via `flyte.ctx().custom_context`. See [Custom context](../task-programming/custom-context). |

---

## Reading context inside a task with `flyte.ctx()`

Inside a running task, `flyte.ctx()` returns a `TaskContext` object with information about the current execution. Outside of a task, it returns `None`.

```
@env.task
async def inspect_context() -> str:
    ctx = flyte.ctx()
    action = ctx.action
    return (
        f"run={action.run_name}, "
        f"action={action.name}, "
        f"mode={ctx.mode}, "
        f"in_cluster={ctx.is_in_cluster()}"
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/run-context/run_context.py*

### `TaskContext` fields

| Field | Type | Description |
|-------|------|-------------|
| `action` | `ActionID` | Identity of this specific action (task invocation) within the run. |
| `mode` | `"local"` \| `"remote"` \| `"hybrid"` | Execution mode of the current run. |
| `version` | `str` | Version of the deployed task code bundle. |
| `raw_data_path` | `str` | Storage prefix where offloaded outputs are written. |
| `run_base_dir` | `str` | Base directory for the run's inputs, outputs, and intermediate per-action artifacts in the data plane object store. |
| `custom_context` | `Dict[str, str]` | Propagated context metadata from `with_runcontext()`. |
| `disable_run_cache` | `bool` | Whether run caching is disabled for this run. |
| `is_in_cluster()` | method | Returns `True` when `mode == "remote"`. Useful for branching local/remote behavior. |

### `ActionID` fields

The `ctx.action` object identifies this specific task invocation:

| Field | Type | Description |
|-------|------|-------------|
| `name` | `str` | Unique identifier for this action. |
| `run_name` | `str` | Name of the parent run (defaults to `name` if not set). |
| `project` | `str \| None` | Project the action runs in. |
| `domain` | `str \| None` | Domain the action runs in. |
| `org` | `str \| None` | Organization. |

### Naming external resources

`ctx.action.run_name` is useful for tying external tool runs (experiment trackers, dashboards) to the corresponding Flyte run:

```
import wandb  # type: ignore[import]

@env.task
async def train_model(epochs: int) -> float:
    ctx = flyte.ctx()
    # Use run_name to tie the W&B run to this Flyte run
    run = wandb.init(
        project="my-project",
        name=ctx.action.run_name,
        config={"epochs": epochs},
    )
    # ... training logic ...
    loss = 0.42
    run.log({"loss": loss})
    run.finish()
    return loss
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/run-context/run_context.py*

This ensures that when you look up a run in Weights & Biases (or any other tool), its name matches what you see in the Flyte UI.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/entrypoints ===

# Entrypoint tasks

As a project grows, the number of deployed tasks in a project/domain grows with it. Every `@env.task` you deploy shows up as a flat entry in the **Tasks list** alongside every other task, including internal helpers, leaf tasks, and intermediate steps that users aren't meant to invoke directly. Finding the handful of tasks that are actually meant to be *run* becomes a navigation problem.

**Entrypoint tasks** solve this by letting you mark which tasks in an environment are the ones intended to be invoked by humans or other services. It's a purely declarative hint (it doesn't change how the task executes), but it makes the task easy to surface in the CLI, UI, and remote API.

## When to use entrypoints

Mark a task as an entrypoint when:

- **It is meant to be run directly**, by a user, a teammate, or a service, rather than only being called from another task.
- **You want it to be discoverable** in the UI's entrypoint view or via filtered CLI queries.
- **You are sharing tasks with team members** and want the "start here" tasks to stand out from internal helpers.

Conversely, don't mark a task as an entrypoint if it only exists as a subtask called by another task, or if it's a utility helper not intended to be invoked on its own.

> [!NOTE]
> Entrypoints and [triggers](../task-configuration/triggers) solve related but distinct problems. A **trigger** automates *when* a task runs (on a schedule, on an event). An **entrypoint** declares *which* tasks are meant to be run at all. A task can be both: a scheduled entrypoint is simply a named starting point that also runs on a cron.

## Mark a task as an entrypoint

Pass `entrypoint=True` to the `@env.task` decorator:

```python
import flyte

env = flyte.TaskEnvironment(
    name="hello_world",
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

@env.task
async def square(i: int) -> int:
    return i * i

@env.task(entrypoint=True)
async def say_hello_nested(data: str = "default string", n: int = 3) -> str:
    vals = [await square(i=i) for i in range(n)]
    return f"Hello {data} {vals}"
```

Here `say_hello_nested` is the intended starting point; `square` is an internal helper that `say_hello_nested` calls. Both get deployed, but only `say_hello_nested` is marked as an entrypoint.

The flag is also available on `task.override()`, so you can promote or demote a task as an entrypoint at invocation time without editing the source:

```python
promoted = square.override(entrypoint=True)
```

The default is `entrypoint=False`. If you don't mark anything, nothing is treated as an entrypoint.

## Discover entrypoint tasks

Once deployed, entrypoint tasks can be surfaced in three ways.

### CLI

Filter `flyte get task` to only entrypoint tasks:

```bash
flyte get task --entrypoint
```

This lists every task in the configured project/domain that was deployed with `entrypoint=True`.

### Programmatic

Use `flyte.remote.Task.listall()` with the `entrypoint` filter:

```python
import flyte

flyte.init_from_config()

for task in flyte.remote.Task.listall(entrypoint=True):
    print(task.name, task.version)
```

This is useful for building your own catalogs, dashboards, or tooling on top of the Flyte API.

### UI

The **Tasks list** in the UI has an **entrypoints** toggle that filters the view to only tasks marked as entrypoints. Use this when you're browsing a project and want to see only the tasks meant to be run.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/run-python-script ===

# Run a Python script

`flyte run python-script` executes an arbitrary `.py` file on a remote
Flyte cluster without wrapping it in a workflow or `@env.task`. Point it
at a script, request resources, and the SDK builds an image, bundles
the code, and runs it.

## When to use it

Reach for `python-script` when:

- You have a working script and want to run it on real cluster hardware
  (GPUs, big memory, long jobs) without rewriting it as a `@env.task`.
- You're prototyping, benchmarking, or reproducing a one-off job and
  don't need typed I/O, caching, or DAG semantics.
- You want to give a collaborator something they can run themselves
  with a single CLI command.

Reach for a normal `@env.task` workflow instead when you need typed
inputs and outputs, cross-task dependencies, caching, retries, or
parameterized re-runs from the UI.

## Quickstart

```bash
flyte run python-script hello.py
```

This uses the default Debian base image, 1 CPU, and 2 GiB of memory,
enough for "hello world". Add flags to request more.

```bash
flyte run python-script train.py \
    --gpu 1 --gpu-type A100 --memory 64Gi \
    --packages torch,transformers \
    --output-dir output
```

Run-level options (`--follow`, `--name`, `--project`, `--domain`) go
on the `flyte run` command, before `python-script`:

```bash
flyte run --project my-proj --domain development python-script ...
```

## What you can specify

| Flag | Default | Purpose |
|---|---|---|
| `--cpu` | `1` | CPU cores to request. |
| `--memory` | `2Gi` | Memory to request. |
| `--gpu` | `0` | Number of GPUs (omit for CPU-only). |
| `--gpu-type` | `T4` | Accelerator class: `T4`, `A100`, `H100`, `L4`, `A10G`, etc. Only used when `--gpu > 0`. |
| `--image` | *(auto)* | Container image URI or a named image from your config. Mutually exclusive with `--packages`. |
| `--packages` | *(none)* | Comma-separated pip packages layered onto the default base image, e.g. `torch,transformers`. Mutually exclusive with `--image`. |
| `--timeout` | `3600` | Task timeout in seconds. The subprocess is killed 60 s before this. |
| `--extra-args` | *(none)* | Comma-separated values passed to the script as `sys.argv`. |
| `--output-dir` | *(none)* | Directory path *inside the container* to upload as the task's output after the script finishes. |
| `--include-files` | *(none)* | Path or glob (relative to the script's directory) to bundle alongside the script. Repeat the flag to pass multiple entries. |
| `--queue` | *(config)* | Flyte queue / cluster override. |

## Handling dependencies

Three modes, in order of increasing control:

### 1. Standard library only

No flag needed. The default image is a plain Debian with Python;
anything in the stdlib (`json`, `os`, `subprocess`, `csv`, `datetime`,
…) works without further setup.

```bash
flyte run python-script hello.py
```

### 2. Pip packages on top of the default image

```bash
flyte run python-script analyze.py --packages numpy,pandas,matplotlib
```

The SDK layers a `pip install` step onto the default image and caches
the result. Re-runs with the same package set skip the build. Changing
the set invalidates the cache and rebuilds.

### 3. Bring your own image

When you need a specific Python version, system packages, or a
pre-baked environment, pass a full image URI:

```bash
flyte run python-script job.py \
    --image 1234.dkr.ecr.us-east-2.amazonaws.com/my-org/job:v1.2.3
```

Your image **must have `flyte` installed** so the Flyte runtime inside
the pod can launch the script. A Dockerfile starting from
`python:3.12-slim-bookworm` with `pip install flyte` works.

If your image is registered in your Flyte config under a short name,
pass the short name instead of a URI and the SDK looks it up in the
image map.

## Handling inputs and outputs

`python-script` does **not** use Flyte's typed I/O. Your script is a
regular Python program: it reads `sys.argv` and writes to a directory
on disk.

### Passing inputs

Use `--extra-args` for positional arguments. The comma-separated values
become `sys.argv[1:]` inside the container:

```bash
flyte run python-script job.py --extra-args "alice,42,--verbose"
```

```python
# job.py
import sys
name, count, *flags = sys.argv[1:]
```

For structured inputs (JSON payloads, files), either inline the data
as a single `--extra-args` value that your script parses, or have the
script download from an object storage URI it knows about (for example,
an S3 URI).

### Capturing outputs

Use `--output-dir` to point at a directory the script writes to. After
the script exits with code 0, the SDK uploads the entire directory
recursively (including subdirectories and binary files) as a
`flyte.io.Dir` on the task's output:

```bash
flyte run python-script job.py --output-dir output
```

```python
# job.py
import os
os.makedirs("output", exist_ok=True)
with open("output/result.csv", "w") as f:
    f.write("col1,col2\n1,2\n")
```

The uploaded directory appears in the run's outputs panel as a blob
store URI (for example, an S3 URI). Browse or download it from the UI.

`stdout` (with `stderr` merged) is also available on the task output
as a ~80-line tail, so you can peek at the end-of-run output without
reopening the log stream.

## Bundling: what gets uploaded

By default the CLI only bundles the script you pointed at. If your
script imports sibling modules or reads local config/data files, list
them with `--include-files` (repeat the flag for multiple entries):

```bash
flyte run python-script train.py \
    --include-files "*.py" \
    --include-files "configs/settings.yaml"
```

Entries are paths or globs resolved relative to the script's directory;
absolute paths are passed through unchanged. Given:

```
my_job/
  train.py       # flyte run python-script my_job/train.py --include-files "*.py"
  data_utils.py  # importable as: from data_utils import ...
  model.py       # importable as: from model import ...
```

the glob `"*.py"` picks up both sibling modules and makes them
importable inside the container.

For anything more elaborate (transitive package trees, large data
assets, system dependencies), bake the files into a custom image with
`--image`, or switch to a regular `@env.task` workflow.

## Python API

The same feature is available as `flyte.run_python_script()` for
programmatic use:

```python
from pathlib import Path
import flyte

flyte.init_from_config("config.yaml")

run = flyte.run_python_script(
    Path("train.py"),
    gpu=1,
    gpu_type="A100",
    memory="64Gi",
    image=["torch", "transformers"],  # list = pip packages on base image
    output_dir="output",
    extra_args=["--epochs", "10"],
    include_files=["*.py", "configs/settings.yaml"],
)
print(run.url)
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/run-with-notifications ===

# Run with notifications

You can attach notifications to a single run by passing them to `flyte.with_runcontext()`.
Notifications fire when the run reaches the terminal execution phase; no trigger or persistent deployment is required.

```
import os
import flyte
from flyte import notify
from flyte.models import ActionPhase

env = flyte.TaskEnvironment(name="notify_example")

SLACK_WEBHOOK_URL = os.environ["SLACK_WEBHOOK_URL"]
NOTIFICATION_EMAIL = os.environ["NOTIFICATION_EMAIL"]

@env.task
def compute(x: int, y: int) -> int:
    return x + y

if __name__ == "__main__":
    result = flyte.with_runcontext(
        notifications=(
            notify.Slack(
                on_phase=ActionPhase.SUCCEEDED,
                webhook_url=SLACK_WEBHOOK_URL,
                message="Run {{.Run.Name}} succeeded.",
            ),
            notify.Email(
                on_phase=ActionPhase.FAILED,
                recipients=[NOTIFICATION_EMAIL],
                subject="ALERT: Run {{.Run.Name}} failed",
                body="Run: {{.Run.Name}}\nError: {{.Error}}",
            ),
        ),
    ).run(compute, x=3, y=7)
    print(f"Result: {result}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-deployment/run-with-notifications/run_with_notifications.py*

Pass a single notification or a tuple of notifications. The notification types in `flyte.notify` are `Slack`, `Email`, `Teams`, `Webhook`, and `NamedDelivery`.

> [!NOTE]
> To attach notifications to every run created by a scheduled trigger, set `notifications` on the `flyte.Trigger` object instead. See [Notifications](../task-configuration/triggers#notifications).

## Execution phases

The `on_phase` parameter accepts a single phase or a tuple of phases from `flyte.models.ActionPhase`:

| Phase | Description |
|-------|-------------|
| `ActionPhase.SUCCEEDED` | Run completed successfully |
| `ActionPhase.FAILED` | Run failed with an error |
| `ActionPhase.TIMED_OUT` | Run exceeded its timeout |
| `ActionPhase.ABORTED` | Run was manually aborted |

## Template variables

All message fields support template variables substituted at delivery time:

| Variable | Description |
|----------|-------------|
| `{{.Run.Project}}` | Project name |
| `{{.Run.Domain}}` | Domain name |
| `{{.Run.Name}}` | Run ID |
| `{{.Phase}}` | Execution phase |
| `{{.Error}}` | Error message (failed) or abort reason (aborted) |

## Slack notifications

`notify.Slack` sends a message to a Slack channel via an [incoming webhook](https://api.slack.com/messaging/webhooks).

**Simple message:**

```
notify.Slack(
    on_phase=ActionPhase.FAILED,
    webhook_url="https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
    message="Run {{.Run.Name}} failed in {{.Run.Project}}/{{.Run.Domain}}: {{.Error}}",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

**Rich formatting with [Block Kit](https://api.slack.com/block-kit):**

Use `blocks` instead of `message` for structured layouts. When `blocks` is provided, `message` is ignored.

```
notify.Slack(
    on_phase=ActionPhase.SUCCEEDED,
    webhook_url="https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
    blocks=[
        {
            "type": "header",
            "text": {"type": "plain_text", "text": "Task Succeeded"},
        },
        {
            "type": "section",
            "fields": [
                {"type": "mrkdwn", "text": "*Run:*\n{{.Run.Name}}"},
                {"type": "mrkdwn", "text": "*Phase:*\n{{.Phase}}"},
            ],
        },
        {"type": "divider"},
        {
            "type": "context",
            "elements": [
                {"type": "mrkdwn", "text": "{{.Run.Project}}/{{.Run.Domain}}"},
            ],
        },
    ],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

## Email notifications

`notify.Email` sends an email notification. Provide `body` for plain text, `html_body` for HTML, or both (sent as multipart).

```
notify.Email(
    on_phase=ActionPhase.FAILED,
    recipients=["oncall@example.com"],
    cc=["team-lead@example.com"],
    subject="ALERT: Run {{.Run.Name}} failed",
    body="Run: {{.Run.Name}}\nError: {{.Error}}",
    html_body="<b>Error:</b> {{.Error}}<br>",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/notifications/notifications.py*

To receive emails locally while developing, start a debug SMTP server before running your script:

```bash
# Python >= 3.12
pip install aiosmtpd
python -m aiosmtpd -n -l localhost:1025
```

The server prints received emails to stdout. Port 1025 needs no special privileges; to listen on the standard SMTP port 25 instead, run the command with `sudo`. Either way, set the SMTP port in your configuration to match.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-deployment/rerun-runs ===

# Rerun a run

Every run in Flyte is durable: its task, code, inputs, and configuration are all
recorded on the backend. That means you can launch a brand-new run from any previous one without
having the original code checked out locally. This is useful for retrying a transient failure,
reproducing a result, or tweaking a few inputs and running again.

You can rerun at two levels of granularity:

- **The whole run**: re-execute the entry point task (the `a0` action) and everything it calls.
- **A single action**: re-execute one nested action on its own.

And you can trigger a rerun from three places: the **UI**, the **CLI**, or **programmatically**
with the Python SDK.

## Rerun from the UI

### Rerun an entire run

Open the run in the UI. In the top-right corner of the run view, click **Rerun**.

![The Rerun button in the top-right of the run view](../../../_static/images/user-guide/task-deployment/rerun-runs/rerun-button.png)

This opens the launch form pre-filled with the original run's inputs, environment variables, and
code. You can either:

- Trigger the run as-is to reproduce the original execution exactly, or
- Edit the inputs, environment variables, and other launch settings in the form before triggering
  to run a variation.

![The launch form, pre-filled with the run's recorded inputs](../../../_static/images/user-guide/task-deployment/rerun-runs/rerun-launch-form.png)

### Rerun a single action

You can also rerun an individual action without rerunning the whole workflow. Navigate to the
action in the run's action list, open its details view, and use the **Rerun action** option (in the
action menu in the top-right of the action details panel).

![The action menu with the Rerun action option](../../../_static/images/user-guide/task-deployment/rerun-runs/rerun-action-menu.png)

This launches a new run starting from that action, using the action's recorded inputs. As with a
full rerun, you can adjust the inputs in the launch form first.

## Rerun from the CLI

The CLI offers two complementary commands depending on whether you want the **original code** or
**new local code**.

### `flyte rerun`: rerun with the original code and inputs

`flyte rerun` fetches the prior run's task **and** inputs from the backend and launches a new run.
You don't need the original code checked out locally. Everything is pulled from the platform:

```bash
# Rerun with the prior run's exact code and inputs
flyte rerun <run-name>

# Give the new run a name and stream its parent-action logs
flyte rerun <run-name> --name retry-1 --follow
```

Common options:

| Option | Description |
|---|---|
| `-p`, `--project` | Project for the new run (defaults to your config). |
| `-d`, `--domain` | Domain for the new run (defaults to your config). |
| `--name` | Name for the new run (a random name is generated if unset). |
| `-e`, `--env KEY=VALUE` | Override an environment variable for the new run. Repeatable. |
| `--label KEY=VALUE` | Set a label on the new run. Repeatable. |
| `-f`, `--follow` | Stream the parent action's logs after launch. |

> [!NOTE]
> `flyte rerun` reuses the prior run's inputs as-is. To change input values from the command line,
> use the programmatic `flyte.rerun(<run-name>, key=value)` form shown below.

### `flyte run --rerun-from`: rerun with new local code

When you've changed your code locally but want to reuse a prior run's inputs, use `flyte run` with
the `--rerun-from` flag. This deploys **your local code** and feeds it the inputs from the prior
run, so you don't need to re-specify any per-task input flags:

```bash
flyte run --rerun-from <run-name> main.py main
```

`--rerun-from` is remote-only; it cannot be combined with `--local`.

### Choosing between the two

| Command | Code | Inputs |
|---|---|---|
| `flyte run <file> <task>` | local | from CLI |
| `flyte run --rerun-from <run> <file> <task>` | local | prior run's |
| `flyte rerun <run>` | fetched from backend | prior run's |

## Rerun programmatically

Use `flyte.rerun()` to rerun from Python. Like the CLI, it fetches the prior run's task and inputs
from the backend, so no local code is required:

```python
import flyte

flyte.init_from_config()

# Rerun a prior run with its exact inputs (task + inputs fetched from the platform):
flyte.rerun("ul56wcvgqrb9vzhzz5l2")

# Change input parameters — they are converted against the prior run's interface:
flyte.rerun("ul56wcvgqrb9vzhzz5l2", x_list=[1, 2, 3])

# Substitute new code while reusing the original run's inputs:
flyte.rerun("ul56wcvgqrb9vzhzz5l2", task_template=fixed_task)
```

`flyte.rerun()` returns a `flyte.remote.Run`, just like `flyte.run()`, so you can monitor it, wait
on it, and retrieve its outputs in the same way. See
[Interact with runs and actions](./interacting-with-runs) for details.

To control launch settings (name, project, domain, environment variables, labels) use
`flyte.with_runcontext(...).rerun(...)`:

```python
flyte.with_runcontext(
    name="retry-1",
    env_vars={"LOG_LEVEL": "20"},
).rerun("ul56wcvgqrb9vzhzz5l2")
```

## Related

- [Interact with runs and actions](./interacting-with-runs): retrieve, monitor, and inspect runs and actions.
- [Run command options](./run-command-options): the full set of `flyte run` options.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps ===

# Apps

An app is a long-running service that Flyte keeps up, rather than a job that runs to completion. Use an app when something needs to stay reachable at a URL: a dashboard, a REST API, a model endpoint, or an MCP server.

Apps are declared the same way tasks are. An `AppEnvironment` names the image, the port, and the scaling behavior, and the code inside it is an ordinary web application.

```python
app = flyte.AppEnvironment(name="dashboard", image=..., port=8080)
```

Because apps and tasks live in the same project, an app can read what a task produced without moving data between systems, and a task can call an app it depends on.

Apps scale to zero when idle, so a rarely-used dashboard costs nothing between visits.

### **Apps > Configure apps**

Define `AppEnvironment`s with ports, autoscaling, custom domains, and authentication.

### **Apps > Build apps**

Build dashboards, REST APIs, and model endpoints with FastAPI, Streamlit, vLLM, and more.

### **Apps > Native app integrations**

Use pre-built environments for popular frameworks like Streamlit, FastAPI, vLLM, SGLang, and Ollama.

### **Apps > Serve and deploy apps**

Use `flyte serve` for fast iteration or `flyte deploy` for production deployments.

## Related

### **Tasks**

The batch workloads an app usually serves the results of.

### **Agents**

Agents are often deployed as apps so they can be reached over HTTP.

## Subpages
- **Apps > Configure apps**
- **Apps > Build apps**
- **Apps > Native app integrations**
- **Apps > Serve and deploy apps**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps ===

# Configure apps

`[[AppEnvironment]]`s allows you to configure the environment in which your app runs, including the container image, compute resources, secrets, domains, scaling behavior, and more.

Similar to `[[TaskEnvironment]]`, configuration can be set when creating the `[[AppEnvironment]]` object. Unlike tasks, apps are long-running services, so they have additional configuration options specific to web services:

- `port`: What port the app listens on
- `command` and `args`: How to start the app
- `scaling`: Autoscaling configuration for handling variable load
- `domain`: Custom domains and subdomains for your app
- `requires_auth`: Whether the app requires authentication to access
- `depends_on`: Other app or task environments that the app depends on

## Hello World example

Here's a complete example of deploying a simple Streamlit "hello world" app with a custom subdomain.

There are two ways to build apps in Flyte:

1. Defining `AppEnvironment(.., args=[...])` to run the app with the underlying `fserve` command.
2. Defining `@app_env.server` to run the app with a custom server function.

### Using fserve args

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1")
# {{/docs-fragment image}}

# {{docs-fragment app-env}}
app_env = flyte.app.AppEnvironment(
    name="hello-world-app",
    image=image,
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
    domain=flyte.app.Domain(subdomain="hello"),
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy the app
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app.py*

This example demonstrates:

- Creating a custom Docker image with Streamlit
- Setting the `args` to run the Streamlit hello app, which uses the underlying `fserve` command to run the app.
- Configuring the port
- Setting resource limits
- Disabling authentication (for public access)
- Using a custom subdomain

### Using @app_env.server

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1")
# {{/docs-fragment image}}

# {{docs-fragment app-env}}
app_env = flyte.app.AppEnvironment(
    name="hello-world-app-server",
    image=image,
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
    domain=flyte.app.Domain(subdomain="hello-server"),
)

@app_env.server
def server():
    import subprocess
    subprocess.run(["streamlit", "hello", "--server.port", "8080"], check=False)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy the app
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app-server.py*

This example demonstrates:

- Creating a custom Docker image with Streamlit
- Using the `@app_env.server` decorator to define a server function that runs the Streamlit hello app.
- Configuring the port
- Setting resource limits
- Disabling authentication (for public access)
- Using a custom subdomain

Once deployed, your app will be accessible at the generated URL or your custom subdomain.

## Differences from TaskEnvironment

While `AppEnvironment` inherits from `Environment` (the same base class as `TaskEnvironment`), it has several app-specific parameters:

| Parameter | AppEnvironment | TaskEnvironment | Description |
|-----------|----------------|-----------------|-------------|
| `type` | ✅ | ❌ | Type of app (e.g., "FastAPI", "Streamlit") |
| `port` | ✅ | ❌ | Port the app listens on |
| `args` | ✅ | ❌ | Arguments to pass to the app |
| `command` | ✅ | ❌ | Command to run the app |
| `requires_auth` | ✅ | ❌ | Whether app requires authentication |
| `scaling` | ✅ | ❌ | Autoscaling configuration |
| `domain` | ✅ | ❌ | Custom domain/subdomain |
| `links` | ✅ | ❌ | Links to include in the App UI page |
| `include` | ✅ | ❌ | Files to include in app |
| `parameters` | ✅ | ❌ | Parameters to pass to app |
| `cluster_pool` | ✅ | ❌ | Cluster pool for deployment |

Parameters like `image`, `resources`, `secrets`, `env_vars`, and `depends_on` are shared between both environment types. See the [task configuration](../../tasks/task-configuration/_index) docs for details on these shared parameters.

## Configuration topics

Learn more about configuring apps:

- **Apps > Configure apps > App environment settings**: Images, resources, secrets, and app-specific settings like `type`, `port`, `args`, `requires_auth`
- **Apps > Configure apps > App environment settings > App startup**: Understanding the difference between `args` and `command`
- **Apps > Configure apps > Including additional files**: How to include additional files needed by your app
- **Apps > Configure apps > Passing parameters into app environments**: Pass parameters to your app at deployment time
- **Apps > Configure apps > /// script**: Configure scaling up and down based on traffic with idle TTL
- **Apps > Configure apps > Apps depending on other environments**: Use `depends_on` to deploy dependent apps together

## Subpages
- **Apps > Configure apps > App environment settings**
- **Apps > Configure apps > Including additional files**
- **Apps > Configure apps > Passing parameters into app environments**
- **Apps > Configure apps > /// script**
- **Apps > Configure apps > Apps depending on other environments**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps/app-environment-settings ===

# App environment settings

`[[AppEnvironment]]`s control how your apps run in Flyte, including images, resources, secrets, startup behavior, and autoscaling.

## Shared environment settings

`[[AppEnvironment]]`s share many configuration options with `[[TaskEnvironment]]`s:

- **Images**: See [Container images](../../tasks/task-configuration/container-images) for details on creating and using container images
- **Resources**: See [Resources](../../tasks/task-configuration/resources) for CPU, memory, GPU, and storage configuration
- **Secrets**: See [Secrets](../../tasks/task-configuration/secrets) for injecting secrets into your app
- **Environment variables**: Set via the `env_vars` parameter (same as tasks)
- **Cluster pools**: Specify via the `cluster_pool` parameter

## App-specific environment settings

For complete parameter documentation, type signatures, and defaults, see the [`AppEnvironment` API reference](../../../api-reference/flyte-sdk/flyte.app/appenvironment).

Key app-specific parameters include `type`, `port`, `args`, `command`, `requires_auth`, `scaling`, `domain`, `links`, `include`, `parameters`, `cluster_pool`, and `timeouts`. See also:

- [Including additional files](./including-additional-files) in your app deployment
- [Passing parameters](./passing-parameters) to your app
- [Auto-scaling apps](./auto-scaling-apps)
- [App environment dependencies](./apps-depending-on-environments)

### Environment variable substitution in `args`

Environment variables are automatically substituted in `args` strings when they start with the `$` character. This works for both:

- Values from `env_vars`
- Secrets that are specified as environment variables (via `as_env_var` in `flyte.Secret`)

The `$VARIABLE_NAME` syntax will be replaced with the actual environment variable value at runtime:

```python
# Using env_vars
app_env = flyte.app.AppEnvironment(
    name="my-app",
    env_vars={"API_KEY": "secret-key-123"},
    args="--api-key $API_KEY",  # $API_KEY will be replaced with "secret-key-123"
    # ...
)

# Using secrets
app_env = flyte.app.AppEnvironment(
    name="my-app",
    secrets=flyte.Secret(key="AUTH_SECRET", as_env_var="AUTH_SECRET"),
    args=["--api-key", "$AUTH_SECRET"],  # $AUTH_SECRET will be replaced with the secret value
    # ...
)
```

This is particularly useful for passing API keys or other sensitive values to command-line arguments without hardcoding them in your code. The substitution happens at runtime, ensuring secrets are never exposed in your code or configuration files.

> [!TIP]
> For most `AppEnvironment`s, use `args` instead of `command` to specify the app startup command
> in the container. This is because `args` will use the `fserve` command to run the app, which
> enables features like local code bundling and file/directory mounting via parameter injection.

## App startup

There are two ways to start up an app in Flyte:

1. With a server function using `@app_env.server`
2. As a container command using `command` or `args`

### Server decorator via `@app_env.server`

The server function is a Python function that runs the app. It is defined using the `@app_env.server` decorator.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "uvicorn",
#    "flyte>=2.0.0b52",
# ]
# ///

import fastapi
import uvicorn

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment fastapi-app}}
app = fastapi.FastAPI()

env = FastAPIAppEnvironment(
    name="configure-fastapi-example",
    app=app,
    image=flyte.Image.from_uv_script(__file__, name="configure-fastapi-example"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    port=8080,
)

@env.server
def server():
    print("Starting server...")
    uvicorn.run(app, port=8080)

@app.get("/")
async def root() -> dict:
    return {"message": "Hello from FastAPI!"}
# {{/docs-fragment fastapi-app}}

# {{docs-fragment on-startup-decorator}}
state = {}

@env.on_startup
async def app_startup():
    print("App started up")
    state["data"] = ["Here's", "some", "data"]
# {{/docs-fragment on-startup-decorator}}

# {{docs-fragment on-shutdown-decorator}}
@env.on_shutdown
async def app_shutdown():
    print("App shut down")
    state.clear()  # clears the data
# {{/docs-fragment on-shutdown-decorator}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)
    deployed_app = flyte.serve(env)
    print(f"App served at: {deployed_app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/fastapi-server-example.py*

The `@app_env.server` decorator allows you to define a synchronous or asynchronous function that runs the app, either
with a server start command like `uvicorn.run`, [`HTTPServer.serve_forever`](https://docs.python.org/3/library/http.server.html), etc.

> [!NOTE]
> Generally the `[[FastAPIAppEnvironment]]` handles serving automatically under the hood,
> the example above just shows how the `@app_env.server` decorator can be used to define a server function
> that runs the app.

#### Startup hook

The server function is called after the app is started up, and before the app is shut down. It is defined using the `@app_env.on_startup` decorator. This is useful if you need to load any state or external connections needed to run the
app before it starts.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "uvicorn",
#    "flyte>=2.0.0b52",
# ]
# ///

import fastapi
import uvicorn

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment fastapi-app}}
app = fastapi.FastAPI()

env = FastAPIAppEnvironment(
    name="configure-fastapi-example",
    app=app,
    image=flyte.Image.from_uv_script(__file__, name="configure-fastapi-example"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    port=8080,
)

@env.server
def server():
    print("Starting server...")
    uvicorn.run(app, port=8080)

@app.get("/")
async def root() -> dict:
    return {"message": "Hello from FastAPI!"}
# {{/docs-fragment fastapi-app}}

# {{docs-fragment on-startup-decorator}}
state = {}

@env.on_startup
async def app_startup():
    print("App started up")
    state["data"] = ["Here's", "some", "data"]
# {{/docs-fragment on-startup-decorator}}

# {{docs-fragment on-shutdown-decorator}}
@env.on_shutdown
async def app_shutdown():
    print("App shut down")
    state.clear()  # clears the data
# {{/docs-fragment on-shutdown-decorator}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)
    deployed_app = flyte.serve(env)
    print(f"App served at: {deployed_app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/fastapi-server-example.py*

#### Shutdown hook

The server function is called before the app instance shuts down during scale down. It is defined using the
`@app_env.on_shutdown` decorator. This is useful if you need to clean up any state or external connections in the
container running the app.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "uvicorn",
#    "flyte>=2.0.0b52",
# ]
# ///

import fastapi
import uvicorn

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment fastapi-app}}
app = fastapi.FastAPI()

env = FastAPIAppEnvironment(
    name="configure-fastapi-example",
    app=app,
    image=flyte.Image.from_uv_script(__file__, name="configure-fastapi-example"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    port=8080,
)

@env.server
def server():
    print("Starting server...")
    uvicorn.run(app, port=8080)

@app.get("/")
async def root() -> dict:
    return {"message": "Hello from FastAPI!"}
# {{/docs-fragment fastapi-app}}

# {{docs-fragment on-startup-decorator}}
state = {}

@env.on_startup
async def app_startup():
    print("App started up")
    state["data"] = ["Here's", "some", "data"]
# {{/docs-fragment on-startup-decorator}}

# {{docs-fragment on-shutdown-decorator}}
@env.on_shutdown
async def app_shutdown():
    print("App shut down")
    state.clear()  # clears the data
# {{/docs-fragment on-shutdown-decorator}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)
    deployed_app = flyte.serve(env)
    print(f"App served at: {deployed_app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/fastapi-server-example.py*

### Container command via `command` vs `args`

The difference between `args` and `command` is crucial for properly configuring how your app starts.

- **`command`**: The full command to run your app, for example, `"streamlit hello --server.port 8080"`. For most use
  cases, you don't need to specify `command` as it's automatically configured, and uses the `fserve` executable to
  run the app. `fserve` does additional setup for you, like setting up the code bundle and loading [parameters](./passing-parameters) if provided, so it's highly recommended to use the default command.
- **`args`**: Arguments to pass to your app's command (used with the default Flyte command or your custom command). The
  `fserve` executable takes in additional arguments, which you can specify as the arguments needed to run your app, e.g.
  `uvicorn run main.py --server.port 8080`.

#### Default startup behavior

When you don't specify a `command`, Flyte generates a default command that uses `fserve` to run your app. This default command handles:

- Setting up the code bundle
- Configuring the version
- Setting up project/domain context
- Injecting parameters if provided

The default command looks like:

```bash
fserve --version <version> --project <project> --domain <domain> -- <args>
```

So if you specify `args`, they'll be appended after the `--` separator.

#### Using args with the default command

When you use `args` without specifying `command`, the args are passed to the default Flyte command:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment args-with-default-command}}
# Using args with default command
app_env = flyte.app.AppEnvironment(
    name="streamlit-app",
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py"],
    # command is None, so default Flyte command is used
)
# {{/docs-fragment args-with-default-command}}

# {{docs-fragment explicit-command}}
# Using explicit command
app_env2 = flyte.app.AppEnvironment(
    name="streamlit-hello",
    command="streamlit hello --server.port 8080",
    port=8080,
    # No args needed since command includes everything
)
# {{/docs-fragment explicit-command}}

# {{docs-fragment command-with-args}}
# Using command with args
app_env3 = flyte.app.AppEnvironment(
    name="custom-app",
    command="python -m myapp",
    args="--option1 value1 --option2 value2",
    # This runs: python -m myapp --option1 value1 --option2 value2
)
# {{/docs-fragment command-with-args}}

# {{docs-fragment fastapi-auto-command}}
# FastAPIAppEnvironment automatically sets command
from flyte.app.extras import FastAPIAppEnvironment
from fastapi import FastAPI

app = FastAPI()

env = FastAPIAppEnvironment(
    name="my-api",
    app=app,
    # You typically don't need to specify command or args, since the
    # FastAPIAppEnvironment automatically uses the bundled code to serve the
    # app via uvicorn.
)
# {{/docs-fragment fastapi-auto-command}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/app-startup-examples.py*

This effectively runs:

```bash
fserve --version ... --project ... --domain ... -- streamlit run main.py --server.port 8080
```

#### Using an explicit command

When you specify a `command`, it completely replaces the default command:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment args-with-default-command}}
# Using args with default command
app_env = flyte.app.AppEnvironment(
    name="streamlit-app",
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py"],
    # command is None, so default Flyte command is used
)
# {{/docs-fragment args-with-default-command}}

# {{docs-fragment explicit-command}}
# Using explicit command
app_env2 = flyte.app.AppEnvironment(
    name="streamlit-hello",
    command="streamlit hello --server.port 8080",
    port=8080,
    # No args needed since command includes everything
)
# {{/docs-fragment explicit-command}}

# {{docs-fragment command-with-args}}
# Using command with args
app_env3 = flyte.app.AppEnvironment(
    name="custom-app",
    command="python -m myapp",
    args="--option1 value1 --option2 value2",
    # This runs: python -m myapp --option1 value1 --option2 value2
)
# {{/docs-fragment command-with-args}}

# {{docs-fragment fastapi-auto-command}}
# FastAPIAppEnvironment automatically sets command
from flyte.app.extras import FastAPIAppEnvironment
from fastapi import FastAPI

app = FastAPI()

env = FastAPIAppEnvironment(
    name="my-api",
    app=app,
    # You typically don't need to specify command or args, since the
    # FastAPIAppEnvironment automatically uses the bundled code to serve the
    # app via uvicorn.
)
# {{/docs-fragment fastapi-auto-command}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/app-startup-examples.py*

This runs exactly:

```bash
streamlit hello --server.port 8080
```

#### Using a command with args

You can combine both, though this is less common:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment args-with-default-command}}
# Using args with default command
app_env = flyte.app.AppEnvironment(
    name="streamlit-app",
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py"],
    # command is None, so default Flyte command is used
)
# {{/docs-fragment args-with-default-command}}

# {{docs-fragment explicit-command}}
# Using explicit command
app_env2 = flyte.app.AppEnvironment(
    name="streamlit-hello",
    command="streamlit hello --server.port 8080",
    port=8080,
    # No args needed since command includes everything
)
# {{/docs-fragment explicit-command}}

# {{docs-fragment command-with-args}}
# Using command with args
app_env3 = flyte.app.AppEnvironment(
    name="custom-app",
    command="python -m myapp",
    args="--option1 value1 --option2 value2",
    # This runs: python -m myapp --option1 value1 --option2 value2
)
# {{/docs-fragment command-with-args}}

# {{docs-fragment fastapi-auto-command}}
# FastAPIAppEnvironment automatically sets command
from flyte.app.extras import FastAPIAppEnvironment
from fastapi import FastAPI

app = FastAPI()

env = FastAPIAppEnvironment(
    name="my-api",
    app=app,
    # You typically don't need to specify command or args, since the
    # FastAPIAppEnvironment automatically uses the bundled code to serve the
    # app via uvicorn.
)
# {{/docs-fragment fastapi-auto-command}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/app-startup-examples.py*

#### FastAPIAppEnvironment example

When using `FastAPIAppEnvironment`, the command is automatically configured to run uvicorn:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment args-with-default-command}}
# Using args with default command
app_env = flyte.app.AppEnvironment(
    name="streamlit-app",
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py"],
    # command is None, so default Flyte command is used
)
# {{/docs-fragment args-with-default-command}}

# {{docs-fragment explicit-command}}
# Using explicit command
app_env2 = flyte.app.AppEnvironment(
    name="streamlit-hello",
    command="streamlit hello --server.port 8080",
    port=8080,
    # No args needed since command includes everything
)
# {{/docs-fragment explicit-command}}

# {{docs-fragment command-with-args}}
# Using command with args
app_env3 = flyte.app.AppEnvironment(
    name="custom-app",
    command="python -m myapp",
    args="--option1 value1 --option2 value2",
    # This runs: python -m myapp --option1 value1 --option2 value2
)
# {{/docs-fragment command-with-args}}

# {{docs-fragment fastapi-auto-command}}
# FastAPIAppEnvironment automatically sets command
from flyte.app.extras import FastAPIAppEnvironment
from fastapi import FastAPI

app = FastAPI()

env = FastAPIAppEnvironment(
    name="my-api",
    app=app,
    # You typically don't need to specify command or args, since the
    # FastAPIAppEnvironment automatically uses the bundled code to serve the
    # app via uvicorn.
)
# {{/docs-fragment fastapi-auto-command}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/app-startup-examples.py*

The `FastAPIAppEnvironment` automatically:

1. Detects the module and variable name of your FastAPI app
2. Uses an internal server function to start the app via `uvicorn.run`.
3. Handles all the startup configuration for you

## Shared settings

For more details on shared settings like images, resources, and secrets, refer to the [task configuration](../../tasks/task-configuration/_index) documentation.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps/including-additional-files ===

# Including additional files

When your app needs additional files beyond the main script (like utility modules, configuration files, or data files), you can use the `include` parameter to specify which files to bundle with your app.

## How include works

The `include` parameter takes a list of file paths (relative to the directory containing your app definition). These files are bundled together and made available in the app container at runtime.

```python
include=["main.py", "utils.py", "config.yaml"]
```

## When to use include

Use `include` when:

- Your app spans multiple Python files (modules)
- You have configuration files that your app needs
- You have data files or templates your app uses
- You want to ensure specific files are available in the container

> [!NOTE]
> If you're using specialized app environments like `FastAPIAppEnvironment`, Flyte automatically detects and includes the necessary files, so you may not need to specify `include` explicitly.

## Examples

### Multi-file Streamlit app

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A custom Streamlit app with multiple files."""

import pathlib
import flyte
import flyte.app

# {{docs-fragment app-env}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "streamlit==1.41.1",
    "pandas==2.2.3",
    "numpy==2.2.3",
)

app_env = flyte.app.AppEnvironment(
    name="streamlit-multi-file-app",
    image=image,
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py", "utils.py"],  # Include your app files
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app = flyte.deploy(app_env)
    print(f"Deployed app: {app[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/multi_file_streamlit.py*

In this example:

- `main.py` is your main Streamlit app file
- `utils.py` contains helper functions used by `main.py`
- Both files are included in the app bundle

### Multi-file FastAPI app

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Multi-file FastAPI app example."""

from fastapi import FastAPI
from module import function  # Import from another file
import pathlib

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment app-definition}}
app = FastAPI(title="Multi-file FastAPI Demo")

app_env = FastAPIAppEnvironment(
    name="fastapi-multi-file",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    # FastAPIAppEnvironment automatically includes necessary files
    # But you can also specify explicitly:
    # include=["app.py", "module.py"],
)
# {{/docs-fragment app-definition}}

# {{docs-fragment endpoint}}
@app.get("/")
async def root():
    return function()  # Uses function from module.py
# {{/docs-fragment endpoint}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(app_env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/multi_file/app.py*

### App with configuration files

```python
include=["app.py", "config.yaml", "templates/"]
```

## File discovery

When using specialized app environments like `FastAPIAppEnvironment`, Flyte uses code introspection to automatically discover and include the necessary files. This means you often don't need to manually specify `include`.

However, if you have files that aren't automatically detected (like configuration files, data files, or templates), you should explicitly list them in `include`.

## Path resolution

Files in `include` are resolved relative to the directory containing your app definition file. For example:

```
project/
├── apps/
│   ├── app.py          # Your app definition
│   ├── utils.py        # Included file
│   └── config.yaml     # Included file
```

In `app.py`:

```python
include=["utils.py", "config.yaml"]  # Relative to apps/ directory
```

## Best practices

1. **Only include what you need**: Don't include unnecessary files as it increases bundle size
2. **Use relative paths**: Always use paths relative to your app definition file
3. **Include directories**: You can include entire directories, but be mindful of size
4. **Test locally**: Verify your includes work by testing locally before deploying
5. **Check automatic discovery**: Specialized app environments may already include files automatically

## Limitations

- Large files or directories can slow down deployment
- Binary files are supported but consider using data storage (S3, etc.) for very large files
- The bundle size is limited by your Flyte cluster configuration


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps/passing-parameters ===

# Passing parameters into app environments

`[[AppEnvironment]]`s support various parameter types that can be passed at deployment time. This includes primitive values, files, directories, and delayed values like `RunOutput` and `AppEndpoint`.

## Parameter types overview

There are several parameter types:

- **Primitive values**: Strings, numbers, booleans
- **Files**: `flyte.io.File` objects (use `flyte.io.File.from_existing_remote(...)` for remote files)
- **Directories**: `flyte.io.Dir` objects (use `flyte.io.Dir.from_existing_remote(...)` for remote directories)
- **Delayed values**: `RunOutput` (from task runs) or `AppEndpoint` (inject endpoint urls of other apps)

## Basic parameter types

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

> [!WARNING]
> Do not use `flyte.io.File(...)` or `flyte.io.Dir(...)` directly as parameter values.
> Use `flyte.io.File.from_existing_remote(...)` and `flyte.io.Dir.from_existing_remote(...)` to
> reference existing remote files and directories.

## Accessing parameters in your app

There are three ways to access parameter values at runtime: **mount paths**, **environment variables**, and the **`get_parameter` helper function**.

### `mount`

When `mount` is specified on a `Parameter`, the file or directory is downloaded to the given path.
If the mount path ends with a `/`, the file is placed inside that directory with its original
name. Otherwise, the file is downloaded to the exact path specified.

Your app reads directly from the mount path:

```python
with open("/tmp/data_file.txt", "rb") as fh:
    contents = fh.read()
```

### `env_var`

When `env_var` is specified, an environment variable is set containing the path to the downloaded
file or directory. This is useful when you want your app code to be decoupled from specific mount
paths:

```python
import os

data_path = os.environ["DATA_FILE_PATH"]
with open(data_path, "rb") as fh:
    contents = fh.read()
```

### `get_parameter`

The `get_parameter(name)` helper function from `flyte.app` returns the path to the downloaded
file or the string value of the parameter. This works regardless of whether `mount` or `env_var`
are specified and is the most flexible access method:

```python
from flyte.app import get_parameter

data_path = get_parameter("data")
with open(data_path, "rb") as fh:
    contents = fh.read()
```

When a parameter has no `mount` or `env_var` configured, `get_parameter` is the only way to
access its value at runtime.

### Full example

The following example defines a FastAPI app with parameters using all three access methods:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

## Delayed values

Delayed values are parameters whose actual values are materialized at deployment time.

### RunOutput

Use `RunOutput` to pass outputs from task runs as app parameters:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

The `type` argument is required and must be one of `string`, `file`, or `directory`.
When the app is deployed, it will make the remote calls needed to figure out the
actual value of the parameter.

### AppEndpoint

Use `AppEndpoint` to pass endpoints from other apps:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

The endpoint URL will be injected as the parameter value when the app starts.

This is particularly useful when you want to chain apps together (for example, a frontend app calling a backend app), without hardcoding URLs.

## Overriding parameters at serve time

You can override parameter values when serving apps using `parameter_values` in
`flyte.with_servecontext`. File and directory values must use `from_existing_remote`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

> [!NOTE]
> Parameter overrides are only available when using `flyte.with_servecontext().serve()`.
> The `flyte.deploy()` function does not support parameter overrides. Parameters must be specified in the `AppEnvironment` definition.

This is useful for:

- Testing different configurations during development
- Using different models or data sources for testing
- A/B testing different app configurations

## Example: FastAPI app with configurable model

Here's a complete example showing how to use parameters in a FastAPI app:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "fastapi",
#    "uvicorn",
#    "joblib",
#    "scikit-learn",
#    "flyte>=2.0.0b52",
# ]
# ///

from contextlib import asynccontextmanager
from pathlib import Path

import flyte
import flyte.app
import flyte.io
from flyte.app.extras import FastAPIAppEnvironment
from fastapi import FastAPI

# {{docs-fragment model-serving-api}}

image = flyte.Image.from_uv_script(__file__, name="app-parameters-fastapi-example")

task_env = flyte.TaskEnvironment(
    name="model_serving_task",
    image=image,
    resources=flyte.Resources(cpu=2, memory="1Gi"),
    cache="auto",
)

@task_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    import joblib
    import sklearn.ensemble
    import sklearn.datasets

    X, y = sklearn.datasets.make_classification(n_samples=1000, n_features=5, n_classes=2, random_state=42)
    model = sklearn.ensemble.RandomForestClassifier()
    model.fit(X, y)

    model_dir = Path("/tmp/model")
    model_dir.mkdir(parents=True, exist_ok=True)
    model_path = model_dir / "model.joblib"
    joblib.dump(model, model_path)
    return await flyte.io.File.from_local(model_path)

state = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    import joblib

    model = joblib.load("/root/models/model.joblib")
    state["model"] = model
    yield

app = FastAPI(lifespan=lifespan)

app_env = FastAPIAppEnvironment(
    name="model-serving-api",
    app=app,
    parameters=[
        flyte.app.Parameter(
            name="model_file",
            # this is a placeholder
            value=flyte.io.File.from_existing_remote("s3://bucket/models/default.pkl"),
            mount="/root/models/",
            download=True,
        ),
    ],
    image=image,
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    requires_auth=False,
)

@app.post("/predict")
async def predict(data: list[float]) -> dict[str, list[float]]:
    model = state["model"]
    return {"prediction": model.predict([data]).tolist()}

if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    run = flyte.run(train_model_task)
    print(f"Run: {run.url}")
    run.wait()

    model_file = run.outputs()[0]
    print(f"Model file: {model_file.path}")

    app = flyte.with_servecontext(
        parameter_values={
            "model-serving-api": {
                "model_file": flyte.io.File.from_existing_remote(model_file.path)
            }
        }
    ).serve(app_env)
    print(f"API URL: {app.url}")
# {{/docs-fragment model-serving-api}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/app-parameters-fastapi-example.py*

## Example: Using RunOutput for model serving

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Examples showing different ways to pass parameters into apps."""

import os

import flyte
import flyte.app
import flyte.io
from flyte.app import Parameter, get_parameter

# {{docs-fragment basic-parameter-types}}
# String parameters
app_env = flyte.app.AppEnvironment(
    name="configurable-app",
    parameters=[
        Parameter(name="environment", value="production"),
        Parameter(name="log_level", value="INFO"),
    ],
    # ...
)

# File parameters
app_env2 = flyte.app.AppEnvironment(
    name="app-with-model",
    parameters=[
        Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.pkl"),
            mount="/app/models/",
        ),
    ],
    # ...
)

# Directory parameters
app_env3 = flyte.app.AppEnvironment(
    name="app-with-data",
    parameters=[
        Parameter(
            name="data_dir",
            value=flyte.io.Dir.from_existing_remote("s3://bucket/data/"),
            mount="/app/data",
        ),
    ],
    # ...
)
# {{/docs-fragment basic-parameter-types}}

# {{docs-fragment parameter-access-methods}}
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment

DATA_MOUNT_PATH = "/tmp/data_file.txt"
DATA_ENV_VAR = "DATA_FILE_PATH"

demo_app = FastAPI()

demo_env = FastAPIAppEnvironment(
    name="parameter-access-demo",
    app=demo_app,
    parameters=[
        # With mount and env_var: the file is downloaded to the mount path,
        # and an environment variable is set to point to the file location.
        Parameter(
            name="data",
            type="file",
            mount=DATA_MOUNT_PATH,
            env_var=DATA_ENV_VAR,
        ),
        # With no mount or env_var: accessible only via get_parameter,
        # which returns the path to the downloaded file.
        Parameter(name="data_raw", type="file"),
    ],
    # ...
)

@demo_app.get("/from-mount")
def read_from_mount() -> dict:
    """Access the file directly at the mount path."""
    with open(DATA_MOUNT_PATH, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-env-var")
def read_from_env_var() -> dict:
    """Access the file through its environment variable."""
    data_path = os.environ[DATA_ENV_VAR]
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/from-get-parameter")
def read_from_get_parameter() -> dict:
    """Access the file using the get_parameter helper."""
    data_path = get_parameter("data")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}

@demo_app.get("/raw")
def read_raw() -> dict:
    """Access a parameter that has no mount or env_var."""
    data_path = get_parameter("data_raw")
    with open(data_path, "rb") as fh:
        return {"contents": fh.read().decode()}
# {{/docs-fragment parameter-access-methods}}

# {{docs-fragment parameter-serve-override}}
if __name__ == "__main__":
    import logging

    flyte.init_from_config(log_level=logging.DEBUG)

    app_handle = flyte.with_servecontext(
        parameter_values={
            demo_env.name: {
                "data": flyte.io.File.from_existing_remote("s3://bucket/data.txt"),
                "data_raw": flyte.io.File.from_existing_remote("s3://bucket/raw.txt"),
            }
        }
    ).serve(demo_env)
    print(f"Deployed app: {app_handle.url}")
# {{/docs-fragment parameter-serve-override}}

# {{docs-fragment runoutput-example}}
env = flyte.TaskEnvironment(name="training-env")

@env.task
async def train_model() -> flyte.io.File:
    # ... training logic ...
    return await flyte.io.File.from_local("/tmp/trained-model.pkl")

app_env4 = flyte.app.AppEnvironment(
    name="serving-app",
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="file", task_name="training-env.train_model"),
            mount="/app/model",
        ),
    ],
    # ...
)
# {{/docs-fragment runoutput-example}}

# {{docs-fragment appendpoint-example}}
app1_env = flyte.app.AppEnvironment(name="backend-api")

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    parameters=[
        Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
            env_var="BACKEND_URL",
        ),
    ],
    # ...
)
# {{/docs-fragment appendpoint-example}}

# {{docs-fragment runoutput-serving-example}}
import joblib
from sklearn.ensemble import RandomForestClassifier

training_env = flyte.TaskEnvironment(name="training-env")

@training_env.task
async def train_model_task() -> flyte.io.File:
    """Train a model and return it."""
    model = RandomForestClassifier()
    # ... training logic ...
    path = "./trained-model.pkl"
    joblib.dump(model, path)
    return await flyte.io.File.from_local(path)

serving_app = FastAPI()
model_serving_env = FastAPIAppEnvironment(
    name="model-serving-app",
    app=serving_app,
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(
                type="file",
                task_name="training-env.train_model_task",
            ),
            mount="/app/model",
            env_var="MODEL_PATH",
        ),
    ],
)
# {{/docs-fragment runoutput-serving-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/passing-parameters-examples.py*

## Best practices

1. **Use `from_existing_remote`**: Always use `flyte.io.File.from_existing_remote(...)` or `flyte.io.Dir.from_existing_remote(...)` to reference remote files and directories as parameter values.
2. **Use delayed parameters**: Use `RunOutput` and `AppEndpoint` to create app dependencies between tasks and apps, or app-to-app chains.
3. **Override for testing**: Use the `parameter_values` argument in `flyte.with_servecontext()` to test different configurations without changing code.
4. **Mount paths clearly**: Use descriptive mount paths for file/directory parameters so your app code is easy to understand.
5. **Use environment variables**: For paths that your app needs to reference dynamically, use `env_var` to inject values as environment variables.
6. **Use `get_parameter` for flexibility**: When you want to keep your parameter access decoupled from specific mount paths or env var names, use `get_parameter(name)`.

## Limitations

- Large files/directories can slow down app startup.
- Parameter overrides are only available when using `flyte.with_servecontext(...).serve(...)`.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps/auto-scaling-apps ===

## Autoscaling apps

Flyte apps support autoscaling, allowing them to scale up and down based on traffic. This helps optimize costs by scaling down when there's no traffic and scaling up when needed.

### Scaling configuration

The `scaling` parameter uses a `[[Scaling]]` object to configure autoscaling behavior:

```python
scaling=flyte.app.Scaling(
    replicas=(min_replicas, max_replicas),
    scaledown_after=idle_ttl_seconds,
)
```

#### Parameters

- **`replicas`**: A tuple `(min_replicas, max_replicas)` specifying the minimum and maximum number of replicas.
- **`scaledown_after`**: Time in seconds to wait before scaling down when idle (idle TTL).

### Basic scaling example

Here's a simple example with scaling from 0 to 1 replica:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment basic-scaling}}
# Basic example: scale from 0 to 1 replica
app_env = flyte.app.AppEnvironment(
    name="autoscaling-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale from 0 to 1 replica
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment basic-scaling}}

# {{docs-fragment always-on}}
# Always-on app
app_env2 = flyte.app.AppEnvironment(
    name="always-on-api",
    scaling=flyte.app.Scaling(
        replicas=(1, 1),  # Always keep 1 replica running
        # scaledown_after is ignored when min_replicas > 0
    ),
    # ...
)
# {{/docs-fragment always-on}}

# {{docs-fragment scale-to-zero}}
# Scale-to-zero app
app_env3 = flyte.app.AppEnvironment(
    name="scale-to-zero-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Can scale down to 0
        scaledown_after=600,  # Scale down after 10 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment scale-to-zero}}

# {{docs-fragment high-availability}}
# High-availability app
app_env4 = flyte.app.AppEnvironment(
    name="ha-api",
    scaling=flyte.app.Scaling(
        replicas=(2, 5),  # Keep at least 2, scale up to 5
        scaledown_after=300,  # Scale down after 5 minutes
    ),
    # ...
)
# {{/docs-fragment high-availability}}

# {{docs-fragment burstable}}
# Burstable app
app_env5 = flyte.app.AppEnvironment(
    name="bursty-app",
    scaling=flyte.app.Scaling(
        replicas=(1, 10),  # Start with 1, scale up to 10 under load
        scaledown_after=180,  # Scale down quickly after 3 minutes
    ),
    # ...
)
# {{/docs-fragment burstable}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/autoscaling-examples.py*

This configuration:

- Starts with 0 replicas (no running instances)
- Scales up to 1 replica when there's traffic
- Scales back down to 0 after 5 minutes (300 seconds) of no traffic

### Scaling patterns

#### Always-on app

For apps that need to always be running:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment basic-scaling}}
# Basic example: scale from 0 to 1 replica
app_env = flyte.app.AppEnvironment(
    name="autoscaling-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale from 0 to 1 replica
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment basic-scaling}}

# {{docs-fragment always-on}}
# Always-on app
app_env2 = flyte.app.AppEnvironment(
    name="always-on-api",
    scaling=flyte.app.Scaling(
        replicas=(1, 1),  # Always keep 1 replica running
        # scaledown_after is ignored when min_replicas > 0
    ),
    # ...
)
# {{/docs-fragment always-on}}

# {{docs-fragment scale-to-zero}}
# Scale-to-zero app
app_env3 = flyte.app.AppEnvironment(
    name="scale-to-zero-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Can scale down to 0
        scaledown_after=600,  # Scale down after 10 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment scale-to-zero}}

# {{docs-fragment high-availability}}
# High-availability app
app_env4 = flyte.app.AppEnvironment(
    name="ha-api",
    scaling=flyte.app.Scaling(
        replicas=(2, 5),  # Keep at least 2, scale up to 5
        scaledown_after=300,  # Scale down after 5 minutes
    ),
    # ...
)
# {{/docs-fragment high-availability}}

# {{docs-fragment burstable}}
# Burstable app
app_env5 = flyte.app.AppEnvironment(
    name="bursty-app",
    scaling=flyte.app.Scaling(
        replicas=(1, 10),  # Start with 1, scale up to 10 under load
        scaledown_after=180,  # Scale down quickly after 3 minutes
    ),
    # ...
)
# {{/docs-fragment burstable}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/autoscaling-examples.py*

#### Scale-to-zero app

For apps that can scale to zero when idle:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment basic-scaling}}
# Basic example: scale from 0 to 1 replica
app_env = flyte.app.AppEnvironment(
    name="autoscaling-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale from 0 to 1 replica
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment basic-scaling}}

# {{docs-fragment always-on}}
# Always-on app
app_env2 = flyte.app.AppEnvironment(
    name="always-on-api",
    scaling=flyte.app.Scaling(
        replicas=(1, 1),  # Always keep 1 replica running
        # scaledown_after is ignored when min_replicas > 0
    ),
    # ...
)
# {{/docs-fragment always-on}}

# {{docs-fragment scale-to-zero}}
# Scale-to-zero app
app_env3 = flyte.app.AppEnvironment(
    name="scale-to-zero-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Can scale down to 0
        scaledown_after=600,  # Scale down after 10 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment scale-to-zero}}

# {{docs-fragment high-availability}}
# High-availability app
app_env4 = flyte.app.AppEnvironment(
    name="ha-api",
    scaling=flyte.app.Scaling(
        replicas=(2, 5),  # Keep at least 2, scale up to 5
        scaledown_after=300,  # Scale down after 5 minutes
    ),
    # ...
)
# {{/docs-fragment high-availability}}

# {{docs-fragment burstable}}
# Burstable app
app_env5 = flyte.app.AppEnvironment(
    name="bursty-app",
    scaling=flyte.app.Scaling(
        replicas=(1, 10),  # Start with 1, scale up to 10 under load
        scaledown_after=180,  # Scale down quickly after 3 minutes
    ),
    # ...
)
# {{/docs-fragment burstable}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/autoscaling-examples.py*

#### High-availability app

For apps that need multiple replicas for availability:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment basic-scaling}}
# Basic example: scale from 0 to 1 replica
app_env = flyte.app.AppEnvironment(
    name="autoscaling-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale from 0 to 1 replica
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment basic-scaling}}

# {{docs-fragment always-on}}
# Always-on app
app_env2 = flyte.app.AppEnvironment(
    name="always-on-api",
    scaling=flyte.app.Scaling(
        replicas=(1, 1),  # Always keep 1 replica running
        # scaledown_after is ignored when min_replicas > 0
    ),
    # ...
)
# {{/docs-fragment always-on}}

# {{docs-fragment scale-to-zero}}
# Scale-to-zero app
app_env3 = flyte.app.AppEnvironment(
    name="scale-to-zero-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Can scale down to 0
        scaledown_after=600,  # Scale down after 10 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment scale-to-zero}}

# {{docs-fragment high-availability}}
# High-availability app
app_env4 = flyte.app.AppEnvironment(
    name="ha-api",
    scaling=flyte.app.Scaling(
        replicas=(2, 5),  # Keep at least 2, scale up to 5
        scaledown_after=300,  # Scale down after 5 minutes
    ),
    # ...
)
# {{/docs-fragment high-availability}}

# {{docs-fragment burstable}}
# Burstable app
app_env5 = flyte.app.AppEnvironment(
    name="bursty-app",
    scaling=flyte.app.Scaling(
        replicas=(1, 10),  # Start with 1, scale up to 10 under load
        scaledown_after=180,  # Scale down quickly after 3 minutes
    ),
    # ...
)
# {{/docs-fragment burstable}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/autoscaling-examples.py*

#### Burstable app

For apps with variable load:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

import flyte
import flyte.app

# {{docs-fragment basic-scaling}}
# Basic example: scale from 0 to 1 replica
app_env = flyte.app.AppEnvironment(
    name="autoscaling-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale from 0 to 1 replica
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment basic-scaling}}

# {{docs-fragment always-on}}
# Always-on app
app_env2 = flyte.app.AppEnvironment(
    name="always-on-api",
    scaling=flyte.app.Scaling(
        replicas=(1, 1),  # Always keep 1 replica running
        # scaledown_after is ignored when min_replicas > 0
    ),
    # ...
)
# {{/docs-fragment always-on}}

# {{docs-fragment scale-to-zero}}
# Scale-to-zero app
app_env3 = flyte.app.AppEnvironment(
    name="scale-to-zero-app",
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Can scale down to 0
        scaledown_after=600,  # Scale down after 10 minutes of inactivity
    ),
    # ...
)
# {{/docs-fragment scale-to-zero}}

# {{docs-fragment high-availability}}
# High-availability app
app_env4 = flyte.app.AppEnvironment(
    name="ha-api",
    scaling=flyte.app.Scaling(
        replicas=(2, 5),  # Keep at least 2, scale up to 5
        scaledown_after=300,  # Scale down after 5 minutes
    ),
    # ...
)
# {{/docs-fragment high-availability}}

# {{docs-fragment burstable}}
# Burstable app
app_env5 = flyte.app.AppEnvironment(
    name="bursty-app",
    scaling=flyte.app.Scaling(
        replicas=(1, 10),  # Start with 1, scale up to 10 under load
        scaledown_after=180,  # Scale down quickly after 3 minutes
    ),
    # ...
)
# {{/docs-fragment burstable}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/autoscaling-examples.py*

### Idle TTL (Time to live)

The `scaledown_after` parameter (idle TTL) determines how long an app instance can be idle before it's scaled down.

#### Considerations

- **Too short**: May cause frequent scale up/down cycles, leading to cold starts.
- **Too long**: Keeps resources running unnecessarily, increasing costs.
- **Optimal**: Balance between cost and user experience.

#### Common idle TTL values

- **Development/Testing**: 60-180 seconds (1-3 minutes) - quick scale down for cost savings.
- **Production APIs**: 300-600 seconds (5-10 minutes) - balance cost and responsiveness.
- **Batch processing**: 900-1800 seconds (15-30 minutes) - longer to handle bursts.
- **Always-on**: Set `min_replicas > 0` - never scale down.

### Autoscaling best practices

1. **Start conservative**: Begin with longer idle TTL values and adjust based on usage.
2. **Monitor cold starts**: Track how long it takes for your app to become ready after scaling up.
3. **Consider costs**: Balance idle TTL between cost savings and user experience.
4. **Use appropriate min replicas**: Set `min_replicas > 0` for critical apps that need to be always available.
5. **Test scaling behavior**: Verify your app handles scale up/down correctly (for example, state management and connections).

### Autoscaling limitations

- Scaling is based on traffic/request patterns, not CPU/memory utilization.
- Cold starts may occur when scaling from zero.
- Stateful apps need careful design to handle scaling (use external state stores).
- Maximum replicas are limited by your cluster capacity.

### Autoscaling troubleshooting

**App scales down too quickly:**

- Increase `scaledown_after` value.
- Set `min_replicas > 0` if the app needs to stay warm.

**App doesn't scale up fast enough:**

- Ensure your cluster has capacity.
- Check if there are resource constraints.

**Cold starts are too slow:**

- Pre-warm with `min_replicas = 1`.
- Optimize app startup time.
- Consider using faster storage for model loading.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/configure-apps/apps-depending-on-environments ===

# Apps depending on other environments

The `depends_on` parameter allows you to specify that one app depends on another app (or task environment). When you deploy an app with `depends_on`, Flyte ensures that all dependencies are deployed first.

## Basic usage

Use `depends_on` to specify a list of environments that this app depends on:

```python
app1_env = flyte.app.AppEnvironment(name="backend-api", ...)

app2_env = flyte.app.AppEnvironment(
    name="frontend-app",
    depends_on=[app1_env],  # Ensure backend-api is deployed first
    # ...
)
```

When you deploy `app2_env`, Flyte will:

1. First deploy `app1_env` (if not already deployed)
2. Then deploy `app2_env`
3. Make sure `app1_env` is available before `app2_env` starts

## Example: App calling another app

Here's a complete example where one FastAPI app calls another:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "httpx",
# ]
# ///

"""Example of one app calling another app."""

import httpx
from fastapi import FastAPI
import pathlib
import flyte
from flyte.app.extras import FastAPIAppEnvironment

image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi", "uvicorn", "httpx"
)

# {{docs-fragment backend-app}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment backend-app}}

# {{docs-fragment frontend-app}}
app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)

env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    depends_on=[env1],  # Depends on backend-api
)
# {{/docs-fragment frontend-app}}

# {{docs-fragment backend-endpoint}}
@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"
# {{/docs-fragment backend-endpoint}}

# {{docs-fragment frontend-endpoints}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint  # Access the backend endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str):
    """Proxy that calls the backend app."""
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        response.raise_for_status()
        return response.json()
# {{/docs-fragment frontend-endpoints}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    deployments = flyte.deploy(env2)
    print(f"Deployed FastAPI app: {deployments[0].env_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/app_calling_app.py*

When you deploy `env2`, Flyte will:

1. Deploy `env1` first (backend-api)
2. Wait for `env1` to be ready
3. Deploy `env2` (frontend-api)
4. `env2` can then access `env1.endpoint` to make requests

## Dependency chain

You can create chains of dependencies:

```python
app1_env = flyte.app.AppEnvironment(name="service-1", ...)
app2_env = flyte.app.AppEnvironment(name="service-2", depends_on=[app1_env], ...)
app3_env = flyte.app.AppEnvironment(name="service-3", depends_on=[app2_env], ...)

# Deploying app3_env will deploy in order: app1_env -> app2_env -> app3_env
```

## Multiple dependencies

An app can depend on multiple environments:

```python
backend_env = flyte.app.AppEnvironment(name="backend", ...)
database_env = flyte.app.AppEnvironment(name="database", ...)

api_env = flyte.app.AppEnvironment(
    name="api",
    depends_on=[backend_env, database_env],  # Depends on both
    # ...
)
```

When deploying `api_env`, both `backend_env` and `database_env` will be deployed first (they may be deployed in parallel if they don't depend on each other).

## Using AppEndpoint for dependency URLs

When one app depends on another, you can use `AppEndpoint` to get the URL:

```python
backend_env = flyte.app.AppEnvironment(name="backend-api", ...)

frontend_env = flyte.app.AppEnvironment(
    name="frontend-app",
    depends_on=[backend_env],
    parameters=[
        flyte.app.Parameter(
            name="backend_url",
            value=flyte.app.AppEndpoint(app_name="backend-api"),
        ),
    ],
    # ...
)
```

The `backend_url` parameter will be automatically set to the backend app's endpoint URL.
You can get this value in your app code using `flyte.app.get_input("backend_url")`.

## Deployment behavior

When deploying with `flyte.deploy()`:

```python
# Deploy the app (dependencies are automatically deployed)
deployments = flyte.deploy(env2)

# All dependencies are included in the deployment plan
for deployment in deployments:
    print(f"Deployed: {deployment.env.name}")
```

Flyte will:

1. Build a deployment plan that includes all dependencies
2. Deploy dependencies in the correct order
3. Ensure dependencies are ready before deploying dependent apps

## Task environment dependencies

You can also depend on task environments:

```python
task_env = flyte.TaskEnvironment(name="training-env", ...)

serving_env = flyte.app.AppEnvironment(
    name="serving-app",
    depends_on=[task_env],  # Can depend on task environments too
    # ...
)
```

This ensures the task environment is available when the app is deployed (useful if the app needs to call tasks in that environment).

## Best practices

1. **Explicit dependencies**: Always use `depends_on` to make app dependencies explicit
2. **Circular dependencies**: Avoid circular dependencies (app A depends on B, B depends on A)
3. **Dependency order**: Design your dependency graph to be a DAG (Directed Acyclic Graph)
4. **Endpoint access**: Use `AppEndpoint` to pass dependency URLs as inputs
5. **Document dependencies**: Make sure your app documentation explains its dependencies

## Example: A/B testing with dependencies

Here's an example of an A/B testing setup where a root app depends on two variant apps:

```python
app_a = FastAPI(title="Variant A")
app_b = FastAPI(title="Variant B")
root_app = FastAPI(title="Root App")

env_a = FastAPIAppEnvironment(name="app-a-variant", app=app_a, ...)
env_b = FastAPIAppEnvironment(name="app-b-variant", app=app_b, ...)

env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    depends_on=[env_a, env_b],  # Depends on both variants
    # ...
)
```

The root app can route traffic to either variant A or B based on A/B testing logic, and both variants will be deployed before the root app starts.

## Limitations

- Circular dependencies are not supported
- Dependencies must be in the same project/domain
- Dependency deployment order is deterministic but dependencies at the same level may deploy in parallel


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps ===

# Build apps

This section covers how to build different types of apps with Flyte, from single-script apps to multi-file projects, common usage patterns, and authentication.

> [!TIP]
> Go to **Get started > Core concepts > Apps** for an overview of apps and a quick example. For pre-built environments for popular frameworks like Streamlit, FastAPI, vLLM, SGLang, and Ollama, see [Native app integrations](../native-app-integrations/_index).

## App types

Flyte supports various types of apps:

- **UI dashboard apps**: Interactive web dashboards and data visualization tools like Streamlit and Gradio
- **Web API apps**: REST APIs, webhooks, and backend services like FastAPI and Flask
- **Model serving apps**: High-performance LLM serving with vLLM and SGLang, or lightweight serving with Ollama

For ready-to-use environments for these frameworks, see [Native app integrations](../native-app-integrations/_index).

## Usage patterns

Apps and tasks can interact in various ways: calling each other via HTTP, webhooks, WebSockets, or direct browser usage.

| Pattern | Use Case | Implementation |
|---------|----------|----------------|
| App | Stand-alone serving app | HTTP requests from arbitrary clients |
| App → App | Microservices, proxies, agent routers, LLM routers | HTTP requests between apps |
| App → Task | Webhooks, APIs triggering workflows | Flyte SDK in app |
| Task → App | Batch processing using inference services | HTTP requests from task |
| Browser app | User-facing dashboards (e.g. Streamlit, Gradio) | Direct browser access |

## Next steps

- **Apps > Build apps > Single-script apps**: The simplest way to build and deploy apps in a single Python script
- **Apps > Build apps > Multi-script apps**: Build FastAPI and Streamlit apps with multiple files
- **Apps > Build apps > Serving graphs**: Apps calling other apps for microservice architectures
- **Apps > Build apps > Hybrid app-task graphs**: Tasks calling apps and apps calling tasks (webhooks, APIs)
- **Apps > Build apps > WebSocket apps**: Real-time, bidirectional communication with WebSockets
- **Apps > Build apps > Browser apps**: User-facing dashboards and UIs
- **Apps > Build apps > Secret-based authentication**: Authenticate FastAPI apps using Flyte secrets

## Subpages
- **Apps > Build apps > Single-script apps**
- **Apps > Build apps > Multi-script apps**
- **Apps > Build apps > Serving graphs**
- **Apps > Build apps > Hybrid app-task graphs**
- **Apps > Build apps > WebSocket apps**
- **Apps > Build apps > Browser apps**
- **Apps > Build apps > Secret-based authentication**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/single-script-apps ===

# Single-script apps

The simplest way to build and deploy an app with Flyte is to write everything in a single Python script. This approach is perfect for:

- **Quick prototypes**: Rapidly test ideas and concepts
- **Simple services**: Basic HTTP servers, APIs, or dashboards
- **Learning**: Understanding how Flyte apps work without complexity
- **Minimal examples**: Demonstrating core functionality

All the code for your app (the application logic, the app environment configuration, and the deployment code) lives in one file. This makes it easy to understand, share, and deploy.

## Plain Python HTTP server

The simplest possible app is a plain Python HTTP server using Python's built-in `http.server` module. This requires no external dependencies beyond the Flyte SDK.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A plain Python HTTP server example - the simplest possible app."""

import flyte
import flyte.app
from pathlib import Path

# {{docs-fragment server-code}}
# Create a simple HTTP server handler
from http.server import HTTPServer, BaseHTTPRequestHandler

class SimpleHandler(BaseHTTPRequestHandler):
    """A simple HTTP server handler."""

    def do_GET(self):

        if self.path == "/":
            self.send_response(200)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(b"<h1>Hello from Plain Python Server!</h1>")

        elif self.path == "/health":
            self.send_response(200)
            self.send_header("Content-type", "application/json")
            self.end_headers()
            self.wfile.write(b'{"status": "healthy"}')

        else:
            self.send_response(404)
            self.end_headers()
# {{/docs-fragment server-code}}

# {{docs-fragment app-env}}
file_name = Path(__file__).name
app_env = flyte.app.AppEnvironment(
    name="plain-python-server",
    image=flyte.Image.from_debian_base(python_version=(3, 12)),
    args=["python", file_name, "--server"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import sys

    if "--server" in sys.argv:
        server = HTTPServer(("0.0.0.0", 8080), SimpleHandler)
        print("Server running on port 8080")
        server.serve_forever()
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"App URL: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/plain_python_server.py*

**Key points**

- **No external dependencies**: Uses only Python's standard library
- **Simple handler**: Define request handlers as Python classes
- **Basic command**: Run the server with a simple Python command
- **Minimal resources**: Requires only basic CPU and memory

## Streamlit app

Streamlit makes it easy to build interactive web dashboards. Here's a complete single-script Streamlit app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
# ]
# ///

"""A single-script Streamlit app example."""

import pathlib
import streamlit as st
import flyte
import flyte.app

# {{docs-fragment streamlit-app}}
def main():
    st.set_page_config(page_title="Simple Streamlit App", page_icon="🚀")

    st.title("Hello from Streamlit!")
    st.write("This is a simple single-script Streamlit app.")

    name = st.text_input("What's your name?", "World")
    st.write(f"Hello, {name}!")

    if st.button("Click me!"):
        st.balloons()
        st.success("Button clicked!")
# {{/docs-fragment streamlit-app}}

# {{docs-fragment app-env}}
file_name = pathlib.Path(__file__).name
app_env = flyte.app.AppEnvironment(
    name="streamlit-single-script",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit==1.41.1"
    ),
    args=["streamlit", "run", file_name, "--server.port", "8080", "--", "--server"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import sys

    if "--server" in sys.argv:
        main()
    else:
        flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"App URL: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit_single_script.py*

**Key points**

- **Interactive UI**: Streamlit provides widgets and visualizations out of the box
- **Single file**: All UI logic and deployment code in one script
- **Simple deployment**: Just specify the Streamlit command and port
- **Rich ecosystem**: Access to Streamlit's extensive component library

## FastAPI app

FastAPI is a modern, fast web framework for building APIs. Here's a minimal single-script FastAPI app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A single-script FastAPI app example - the simplest FastAPI app."""

from fastapi import FastAPI
import pathlib
import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment fastapi-app}}
app = FastAPI(
    title="Simple FastAPI App",
    description="A minimal single-script FastAPI application",
    version="1.0.0",
)

@app.get("/")
async def root():
    return {"message": "Hello, World!"}

@app.get("/health")
async def health():
    return {"status": "healthy"}
# {{/docs-fragment fastapi-app}}

# {{docs-fragment app-env}}
app_env = FastAPIAppEnvironment(
    name="fastapi-single-script",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.serve(app_env)
    print(f"Deployed: {app_deployment.url}")
    print(f"API docs: {app_deployment.url}/docs")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi_single_script.py*

**Key points**

- **FastAPIAppEnvironment**: Automatically configures uvicorn and FastAPI
- **Type hints**: FastAPI uses Python type hints for automatic validation
- **Auto docs**: Interactive API documentation at `/docs` endpoint
- **Async support**: Built-in support for async/await patterns

## Running single-script apps

To run any of these examples:

1. **Save the script** to a file (e.g., `my_app.py`)
2. **Ensure you have a config file** (`./.flyte/config.yaml` or `./config.yaml`)
3. **Run the script**:

```bash
python my_app.py
```

Or using `uv`:

```bash
uv run my_app.py
```

The script will:

- Initialize Flyte from your config
- Deploy the app to your Union/Flyte instance
- Print the app URL

## When to use single-script apps

**Use single-script apps when:**

- Building prototypes or proof-of-concepts
- Creating simple services with minimal logic
- Learning how Flyte apps work
- Sharing complete, runnable examples
- Building demos or tutorials

**Consider multi-script apps when:**

- Your app grows beyond a few hundred lines
- You need to organize code into modules
- You want to reuse components across apps
- You're building production applications

See [**Multi-script apps**](./multi-script-apps) for examples of organizing apps across multiple files.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/multi-script-apps ===

# Multi-script apps

Real-world applications often span multiple files. This page shows how to build FastAPI and Streamlit apps with multiple Python files.

## FastAPI multi-script app

### Project structure

```
project/
├── app.py          # Main FastAPI app file
└── module.py       # Helper module
```

### Example: Multi-file FastAPI app

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Multi-file FastAPI app example."""

from fastapi import FastAPI
from module import function  # Import from another file
import pathlib

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment app-definition}}
app = FastAPI(title="Multi-file FastAPI Demo")

app_env = FastAPIAppEnvironment(
    name="fastapi-multi-file",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    # FastAPIAppEnvironment automatically includes necessary files
    # But you can also specify explicitly:
    # include=["app.py", "module.py"],
)
# {{/docs-fragment app-definition}}

# {{docs-fragment endpoint}}
@app.get("/")
async def root():
    return function()  # Uses function from module.py
# {{/docs-fragment endpoint}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(app_env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/multi_file/app.py*

```
# {{docs-fragment helper-function}}
def function():
    """Helper function used by the FastAPI app."""
    return {"message": "Hello from module.py!"}
# {{/docs-fragment helper-function}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/multi_file/module.py*

### Automatic file discovery

`FastAPIAppEnvironment` automatically discovers and includes the necessary files by analyzing your imports. However, if you have files that aren't automatically detected (like configuration files or data files), you can explicitly include them:

```python
app_env = FastAPIAppEnvironment(
    name="fastapi-with-config",
    app=app,
    include=["app.py", "module.py", "config.yaml"],  # Explicit includes
    # ...
)
```

## Streamlit multi-script app

### Project structure

```
project/
├── main.py         # Main Streamlit app
├── utils.py        # Utility functions
└── components.py   # Reusable components
```

### Example: Multi-file Streamlit app

```
import os

import streamlit as st
from utils import generate_data

# {{docs-fragment streamlit-app}}
all_columns = ["Apples", "Orange", "Pineapple"]
with st.container(border=True):
    columns = st.multiselect("Columns", all_columns, default=all_columns)

all_data = st.cache_data(generate_data)(columns=all_columns, seed=101)

data = all_data[columns]

tab1, tab2 = st.tabs(["Chart", "Dataframe"])
tab1.line_chart(data, height=250)
tab2.dataframe(data, height=250, use_container_width=True)
st.write(f"Environment: {os.environ}")
# {{/docs-fragment streamlit-app}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/main.py*

```
import numpy as np
import pandas as pd

# {{docs-fragment utils-function}}
def generate_data(columns: list[str], seed: int = 42):
    rng = np.random.default_rng(seed)
    data = pd.DataFrame(rng.random(size=(20, len(columns))), columns=columns)
    return data
# {{/docs-fragment utils-function}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/utils.py*

### Deploying multi-file Streamlit app

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A custom Streamlit app with multiple files."""

import pathlib
import flyte
import flyte.app

# {{docs-fragment app-env}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "streamlit==1.41.1",
    "pandas==2.2.3",
    "numpy==2.2.3",
)

app_env = flyte.app.AppEnvironment(
    name="streamlit-multi-file-app",
    image=image,
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py", "utils.py"],  # Include your app files
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app = flyte.deploy(app_env)
    print(f"Deployed app: {app[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/multi_file_streamlit.py*

## Complex multi-file example

Here's a more complex example with multiple modules:

### Project structure

```
project/
├── app.py
├── models/
│   ├── __init__.py
│   └── user.py
├── services/
│   ├── __init__.py
│   └── auth.py
└── utils/
    ├── __init__.py
    └── helpers.py
```

### Example code

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Complex multi-file FastAPI app example."""

from pathlib import Path
from fastapi import FastAPI
from models.user import User
from services.auth import authenticate
from utils.helpers import format_response

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment complex-app}}
app = FastAPI(title="Complex Multi-file App")

@app.get("/users/{user_id}")
async def get_user(user_id: int):
    user = User(id=user_id, name="John Doe")
    return format_response(user)
# {{/docs-fragment complex-app}}

# {{docs-fragment complex-env}}
app_env = FastAPIAppEnvironment(
    name="complex-app",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "pydantic",
    ),
    # Include all necessary files
    include=[
        "app.py",
        "models/",
        "services/",
        "utils/",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment complex-env}}

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(app_env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/complex_multi_file/app.py*

```
# {{docs-fragment user-model}}
from pydantic import BaseModel

class User(BaseModel):
    id: int
    name: str
# {{/docs-fragment user-model}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/complex_multi_file/models/user.py*

```
# {{docs-fragment auth-service}}
def authenticate(token: str) -> bool:
    """Authenticate a user by token."""
    # ... authentication logic ...
    return True
# {{/docs-fragment auth-service}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/complex_multi_file/services/auth.py*

```
# {{docs-fragment helpers}}
def format_response(data):
    """Format a response with standard structure."""
    return {"data": data, "status": "success"}
# {{/docs-fragment helpers}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/complex_multi_file/utils/helpers.py*

### Deploying complex app

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Complex multi-file FastAPI app example."""

from pathlib import Path
from fastapi import FastAPI
from models.user import User
from services.auth import authenticate
from utils.helpers import format_response

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment complex-app}}
app = FastAPI(title="Complex Multi-file App")

@app.get("/users/{user_id}")
async def get_user(user_id: int):
    user = User(id=user_id, name="John Doe")
    return format_response(user)
# {{/docs-fragment complex-app}}

# {{docs-fragment complex-env}}
app_env = FastAPIAppEnvironment(
    name="complex-app",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "pydantic",
    ),
    # Include all necessary files
    include=[
        "app.py",
        "models/",
        "services/",
        "utils/",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment complex-env}}

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(app_env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/complex_multi_file/app.py*

## Best practices

1. **Use explicit includes**: For Streamlit apps, explicitly list all files in `include`
2. **Automatic discovery**: For FastAPI apps, `FastAPIAppEnvironment` handles most cases automatically
3. **Organize modules**: Use proper Python package structure with `__init__.py` files
4. **Test locally**: Test your multi-file app locally before deploying
5. **Include all dependencies**: Include all files that your app imports

## Troubleshooting

**Import errors:**

- Verify all files are included in the `include` parameter
- Check that file paths are correct (relative to app definition file)
- Ensure `__init__.py` files are included for packages

**Module not found:**

- Add missing files to the `include` list
- Check that import paths match the file structure
- Verify that the image includes all necessary packages

**File not found at runtime:**

- Ensure all referenced files are included
- Check mount paths for file/directory inputs
- Verify file paths are relative to the app root directory


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/serving-graphs ===

# Serving graphs

A *serving graph* is a set of Flyte apps that talk to each other inside the
cluster. Instead of putting every stage of a request into one process, you
split the work across multiple `AppEnvironment`s that you deploy together:
each one sized for its own bottleneck, with its own image and scaling policy.

This pattern is useful for:

- **Heterogeneous resource requirements**: CPU pre/postprocessing in front of a GPU forward pass
- **Microservice architectures**: Independent components with distinct lifecycles
- **A/B testing and canary rollouts**: A root app routes traffic across variant apps
- **Proxy / gateway patterns**: One app fronts several backends

## Core concepts: a minimal two-app chain

The simplest serving graph (`app2` proxies HTTP calls to `app1`) is enough
to introduce every core concept: deploying multiple apps together, discovering
an upstream app's endpoint, and sizing each app independently.

Both apps share an image and live in the same Python file:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

### Deploying multiple apps together with `depends_on`

The callee env is straightforward; it has no upstream dependencies of its
own:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

The caller declares `depends_on=[env1]`, which tells Flyte that `env1` must
be deployed alongside this one:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

Calling `flyte.serve(env2)` then deploys the whole dependency closure
transitively, so you only ever name the entry-point app:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

`depends_on` is about deployment co-scheduling, not request-time ordering:
at runtime each app is independent.

### Getting an upstream app's endpoint

There are two ways for one app to discover another app's URL. Both resolve
correctly across local, in-cluster, and external contexts.

**Pattern A: `env.endpoint` (Python property).** When both apps live in the
same Python module, the upstream env object is in scope and you can read
`env.endpoint` directly. The example above uses this pattern in `app2`'s
proxy endpoint:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

**Pattern B: `flyte.app.AppEndpoint` as a parameter.** When the upstream env
object isn't importable (different file, different process, looking it up by
name), declare it as a `flyte.app.Parameter` and have Flyte inject the
resolved URL via an environment variable. The `env2` definition above shows
this. `app1_url` becomes available as `os.getenv("APP1_URL")` at runtime:

```
import logging
import os
import pathlib
import typing

import httpx
from fastapi import FastAPI

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
# {{/docs-fragment image}}

# {{docs-fragment apps}}
app1 = FastAPI(
    title="App 1",
    description="A FastAPI app that runs some computations",
)

app2 = FastAPI(
    title="App 2",
    description="A FastAPI app that proxies requests to another FastAPI app",
)
# {{/docs-fragment apps}}

# {{docs-fragment env-direct}}
env1 = FastAPIAppEnvironment(
    name="app1-is-called-by-app2",
    app=app1,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment env-direct}}

# {{docs-fragment env-with-parameter}}
env2 = FastAPIAppEnvironment(
    name="app2-calls-app1",
    app=app2,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    parameters=[
        flyte.app.Parameter(
            name="app1_url",
            value=flyte.app.AppEndpoint(app_name="app1-is-called-by-app2"),
            env_var="APP1_URL",
        ),
    ],
    depends_on=[env1],
    env_vars={"LOG_LEVEL": "10"},
)
# {{/docs-fragment env-with-parameter}}

@app1.get("/greeting/{name}")
async def greeting(name: str) -> str:
    return f"Hello, {name}!"

# {{docs-fragment endpoint-property-pattern}}
@app2.get("/app1-endpoint")
async def get_app1_endpoint() -> str:
    return env1.endpoint

@app2.get("/greeting/{name}")
async def greeting_proxy(name: str) -> typing.Any:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{env1.endpoint}/greeting/{name}")
        return response.json()
# {{/docs-fragment endpoint-property-pattern}}

# {{docs-fragment endpoint-env-var-pattern}}
@app2.get("/app1-url")
async def get_app1_url() -> str:
    return os.getenv("APP1_URL")
# {{/docs-fragment endpoint-env-var-pattern}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.DEBUG,
    )
    app = flyte.serve(env2)
    print(f"Deployed FastAPI app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/two_app_chain.py*

### Sizing each node independently

Each `AppEnvironment` carries its own image, resources, and scaling. That's
the entire point of splitting: for example, the GPU side of an inference
graph can stay narrow with `scaling=Scaling(replicas=(1, 2))` while the CPU
side scales wide with `scaling=Scaling(replicas=(1, 8))`, with no shared
autoscaling policy between them. The next example shows this in practice.

## Example: CPU / GPU inference split

The canonical heterogeneous-resource pipeline: heavy CPU preprocessing in
front of a fast GPU forward pass, talking to each other over HTTP inside the
cluster.

```mermaid
flowchart LR
    client["client"] --> cpu["cpu_app (×N replicas)<br/>decode + resize<br/>+ softmax"]
    cpu --> gpu["gpu_app (×M replicas)<br/>ResNet18 forward only"]
    gpu --> cpu
    cpu --> client
```

In a typical vision/audio pipeline, the GPU forward pass takes milliseconds
but is sandwiched between slow CPU work (image decode, resize, normalization,
softmax, label lookup). If both stages share one process you pay for an idle
GPU during preprocessing. Splitting them lets each side scale independently:
cheap CPU wide, expensive GPU narrow.

### Disjoint images per node

The two apps share a small base image and add their own disjoint stacks. The
CPU app never imports `torch`; the GPU app never imports `PIL`:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

### GPU app: model.forward only

The GPU app loads the model once at startup using FastAPI's lifespan, so model
weights stay resident across requests:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

The inference endpoint speaks raw `float32` bytes over
`application/octet-stream`. For anything tensor-shaped this is the single
biggest perf knob. JSON-serializing a 19MB batch dominates end-to-end
latency:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

The GPU environment requests a GPU and keeps replicas narrow:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

### CPU app: pre/postprocess + call GPU

Preprocessing is deliberately CPU-bound (decode, denoise, resize, normalize):

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

The CPU app uses its lifespan to resolve the GPU endpoint via `gpu_env.endpoint`,
fetch labels once at startup, and build one persistent `httpx.AsyncClient` per
replica. Persistent clients avoid a TCP/TLS handshake per request, which
matters at high request rates:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

The `/classify` endpoint glues it all together. Heavy CPU work runs in this
process; the GPU forward pass is delegated over HTTP using the raw-bytes wire
format:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

The CPU environment scales wide and declares `depends_on=[gpu_env]` so both
sides deploy together:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

### Deploy

`flyte.serve(cpu_env)` deploys both apps. The CPU app is the public entry
point; the GPU app is reached only via the cluster-internal endpoint:

```
"""
Serving graph — CPU pre/post split from a GPU forward pass.

This example shows the canonical "two-app" inference graph: heavy CPU work
on one app, the GPU forward pass on another, talking to each other over HTTP
inside the cluster.

Why split? In a typical vision/audio/feature-engineering pipeline the GPU
forward pass is fast (millis) but is sandwiched between slow CPU work
(image decode, resize, denoise, NMS, label lookup, etc.). If you put both
stages in one process you pay for an idle GPU during preprocessing. Splitting
them lets each side scale independently:

    client ──► [cpu_app  x N replicas]  ──► [gpu_app x M replicas] ──► back
                preprocess + postprocess        model.forward only
                cheap CPU, scale wide           expensive GPU, scale narrow

Wire format between the two apps is raw float32 bytes (not JSON) — for
anything tensor-shaped this is the single biggest perf knob.
"""

import io
import ipaddress
import logging
import pathlib
import socket
from contextlib import asynccontextmanager

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Request, Response
from PIL import Image, ImageFilter
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment

# ---------------------------------------------------------------------------
# Images
# ---------------------------------------------------------------------------
# Shared base with the deps both apps need (HTTP server + numpy). The CPU and
# GPU images extend it with their own disjoint stacks — the CPU app never
# imports torch and the GPU app never imports PIL. Sharing the base layer
# means the registry only stores one copy of fastapi/uvicorn/numpy.

# {{docs-fragment images}}
base_image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "fastapi",
    "uvicorn",
    "numpy",
)

cpu_image = base_image.with_pip_packages(
    "httpx",
    "pillow",
)

gpu_image = base_image.with_pip_packages(
    "torch==2.7.1",
    "torchvision==0.22.1",
)
# {{/docs-fragment images}}

# ---------------------------------------------------------------------------
# Shared tensor layout
# ---------------------------------------------------------------------------

INPUT_C, INPUT_H, INPUT_W = 3, 224, 224
NUM_CLASSES = 1000
TENSOR_DTYPE = np.float32

# ===========================================================================
# GPU app — model.forward only
# ===========================================================================

# {{docs-fragment gpu-lifespan}}
@asynccontextmanager
async def _gpu_lifespan(app: FastAPI):
    # Imported lazily so the CPU app never has to import torch.
    import torch
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.IMAGENET1K_V1
    model = resnet18(weights=weights).eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        model = model.to("cuda")
    app.state.model = model
    app.state.device = device
    app.state.categories = list(weights.meta["categories"])
    logging.getLogger(__name__).info("model loaded on %s", device)
    yield

gpu_app = FastAPI(
    title="inference-gpu",
    description="ResNet18 forward pass.",
    lifespan=_gpu_lifespan,
)
# {{/docs-fragment gpu-lifespan}}

@gpu_app.get("/health")
async def gpu_health() -> dict:
    return {"status": "ok", "device": gpu_app.state.device}

@gpu_app.get("/labels")
async def labels() -> list[str]:
    # Exposed so the CPU side can fetch labels once at startup instead of
    # hard-coding the ImageNet class list.
    return gpu_app.state.categories

# {{docs-fragment gpu-infer}}
@gpu_app.post("/infer")
async def infer(request: Request) -> Response:
    """Run a batched forward pass.

    Request body:  raw float32 bytes, shape (B, 3, 224, 224), C-contiguous.
    Response body: raw float32 bytes, shape (B, 1000) — raw logits.

    We deliberately do NOT use JSON here. For a batch of 32 images the tensor
    is ~19MB; JSON-serializing that is the dominant cost end-to-end.
    """
    import torch

    raw = await request.body()
    arr = np.frombuffer(raw, dtype=TENSOR_DTYPE)
    if arr.size % (INPUT_C * INPUT_H * INPUT_W) != 0:
        raise HTTPException(400, "payload size is not a multiple of one image tensor")
    batch = arr.reshape(-1, INPUT_C, INPUT_H, INPUT_W)

    x = torch.from_numpy(batch).to(gpu_app.state.device)
    with torch.inference_mode():
        logits = gpu_app.state.model(x)
    out = logits.detach().to("cpu").numpy().astype(TENSOR_DTYPE, copy=False)
    return Response(content=out.tobytes(), media_type="application/octet-stream")
# {{/docs-fragment gpu-infer}}

# {{docs-fragment gpu-env}}
gpu_env = FastAPIAppEnvironment(
    name="serving-graph-gpu",
    app=gpu_app,
    image=gpu_image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu="A10G:1"),
    # GPU replicas are expensive; keep at least one warm so model weights stay
    # resident, and cap the max. Bump if a single replica saturates.
    scaling=flyte.app.Scaling(replicas=(1, 2)),
    requires_auth=True,
)
# {{/docs-fragment gpu-env}}

# ===========================================================================
# CPU app — pre/postprocess, calls the GPU app
# ===========================================================================

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=TENSOR_DTYPE).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=TENSOR_DTYPE).reshape(3, 1, 1)

class ClassifyRequest(BaseModel):
    image_url: str
    top_k: int = 5

class Prediction(BaseModel):
    label: str
    score: float

# {{docs-fragment cpu-preprocess}}
def _preprocess(img_bytes: bytes) -> np.ndarray:
    """Decode → denoise → resize → normalize. CPU-bound, deliberately so.

    Real preprocessing stacks (detection, OCR, audio) do substantially more
    than this — sliding window crops, color-space conversion, etc. The point
    is that none of it benefits from a GPU sitting next to it.
    """
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
    img = img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
    arr = np.asarray(img, dtype=TENSOR_DTYPE) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return np.ascontiguousarray(arr, dtype=TENSOR_DTYPE)
# {{/docs-fragment cpu-preprocess}}

def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# {{docs-fragment cpu-lifespan}}
@asynccontextmanager
async def _cpu_lifespan(app: FastAPI):
    # Resolved at serving time via the cluster-internal endpoint pattern,
    # so this stays correct across local/remote deploys without an env var.
    gpu_url = gpu_env.endpoint
    log = logging.getLogger(__name__)
    log.info("resolved GPU endpoint: %s", gpu_url)
    async with httpx.AsyncClient(timeout=30.0) as bootstrap:
        try:
            r = await bootstrap.get(f"{gpu_url}/labels")
            r.raise_for_status()
        except (httpx.HTTPError, OSError) as e:
            # Most common reason on a fresh deploy: GPU replica hasn't finished
            # pulling its image / loading weights yet. Crash-looping is fine —
            # the next attempt will likely succeed — but make the cause obvious.
            log.error("downstream GPU app at %s not ready: %s", gpu_url, e)
            raise
        app.state.labels = r.json()
    # One persistent client per replica — avoids TCP/TLS handshake per request,
    # which matters once you're doing 100s of req/s.
    async with httpx.AsyncClient(
        base_url=gpu_url,
        timeout=httpx.Timeout(30.0, connect=5.0),
        limits=httpx.Limits(max_connections=64, max_keepalive_connections=32),
    ) as client:
        app.state.client = client
        yield

cpu_app = FastAPI(
    title="inference-cpu",
    description="Pre/post around the GPU forward pass.",
    lifespan=_cpu_lifespan,
)
# {{/docs-fragment cpu-lifespan}}

@cpu_app.get("/health")
async def cpu_health() -> dict:
    return {"status": "ok", "labels_loaded": len(cpu_app.state.labels)}

# {{docs-fragment cpu-classify}}
async def validate_public_image_url(image_url: str) -> str:
     try:
         parsed = httpx.URL(image_url)
     except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid image_url.") from exc
     if parsed.scheme not in {"http", "https"}:
         raise HTTPException(status_code=400, detail="image_url must use http or https.")
     host = parsed.host
     if not host:
         raise HTTPException(status_code=400, detail="image_url must include a hostname.")
     try:
         addr_info = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
     except socket.gaierror as exc:
         raise HTTPException(status_code=400, detail="image_url host could not be resolved.") from exc
     for info in addr_info:
         ip_text = info[4][0]
         ip_obj = ipaddress.ip_address(ip_text)
         if not ip_obj.is_global:
             raise HTTPException(status_code=400, detail="image_url host resolves to a non-public address.")
     return str(parsed)

@cpu_app.post("/classify", response_model=list[Prediction])
async def classify(req: ClassifyRequest) -> list[Prediction]:
    async with httpx.AsyncClient(timeout=30.0) as client:
        img_resp = await client.get(await validate_public_image_url(req.image_url))
        img_resp.raise_for_status()

    tensor = _preprocess(img_resp.content)  # heavy CPU
    batch = tensor[np.newaxis, ...]  # add batch dim

    gpu_resp = await cpu_app.state.client.post(
        "/infer",
        content=batch.tobytes(),
        headers={"content-type": "application/octet-stream"},
    )
    gpu_resp.raise_for_status()
    logits = np.frombuffer(gpu_resp.content, dtype=TENSOR_DTYPE).reshape(1, NUM_CLASSES)

    probs = _softmax(logits, axis=-1)[0]  # back to CPU work
    top_idx = np.argsort(-probs)[: req.top_k]
    return [Prediction(label=cpu_app.state.labels[i], score=float(probs[i])) for i in top_idx]
# {{/docs-fragment cpu-classify}}

# {{docs-fragment cpu-env}}
cpu_env = FastAPIAppEnvironment(
    name="serving-graph-cpu",
    app=cpu_app,
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="4Gi"),
    # Cheap, so scale wide. Use scale-to-zero (replicas=(0, 8)) for bursty
    # traffic; keep replicas=(1, 8) here to avoid cold starts in the demo.
    scaling=flyte.app.Scaling(replicas=(1, 8)),
    requires_auth=True,
    depends_on=[gpu_env],
)
# {{/docs-fragment cpu-env}}

# ===========================================================================
# Deploy
# ===========================================================================

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(
        root_dir=pathlib.Path(__file__).parent,
        log_level=logging.INFO,
    )
    app = flyte.serve(cpu_env)
    print(f"Deployed serving graph; public CPU endpoint: {app.url}")
    print("Try: curl -X POST $URL/classify -H 'content-type: application/json' \\")
    print(
        '       -d \'{"image_url": "https://upload.wikimedia.org/wikipedia/commons/4/41/Sunflower_from_Silesia2.jpg"}\''
    )
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/image_classification.py*

## Example: A/B testing with Statsig

A serving graph also lets you shape traffic. A root app routes each incoming
request to one of two variant apps using a [Statsig](https://www.statsig.com/)
feature gate, with consistent per-user bucketing.

```mermaid
flowchart LR
    client["client"] --> root["root_app<br/>(check_gate)"]
    root -->|"gate off"| a["app_a<br/>fast-processing"]
    root -->|"gate on"| b["app_b<br/>enhanced-processing"]
```

### Statsig client singleton

The variant routing logic needs a single Statsig client per process. Wrap it
in a singleton so lifespan startup/shutdown is the only place that touches its
lifecycle:

```
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
CODE18
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
CODE19
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
CODE20
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
CODE21
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/ab_testing.py*

The root env declares `depends_on=[env_a, env_b]` so all three deploy
together, and pulls the Statsig API key from a Flyte secret:

```
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/ab_testing.py*

### Routing endpoint

The root app checks the `variant_b` feature gate against a user key and
proxies to the matching variant using its `endpoint` property:

```
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/ab_testing.py*

Use stable identifiers (user ID, session ID) for `user_key` so the same user
always lands in the same bucket. To swap `check_gate` for an experiment or
dynamic config:

CODE22

### Deploy

```
import os
import typing
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment statsig-client}}
class StatsigClient:
    """Singleton to manage Statsig client lifecycle."""

    _instance: "StatsigClient | None" = None
    _statsig = None

    @classmethod
    def initialize(cls, api_key: str):
        """Initialize Statsig client (call during lifespan startup)."""
        if cls._instance is None:
            cls._instance = cls()

        # Import statsig at runtime (only available in container)
        from statsig_python_core import Statsig

        cls._statsig = Statsig(api_key)
        cls._statsig.initialize().wait()

    @classmethod
    def get_client(cls):
        """Get the initialized Statsig instance."""
        if cls._statsig is None:
            raise RuntimeError("StatsigClient not initialized. Call initialize() first.")
        return cls._statsig

    @classmethod
    def shutdown(cls):
        """Shutdown Statsig client (call during lifespan shutdown)."""
        if cls._statsig is not None:
            cls._statsig.shutdown()
            cls._statsig = None
            cls._instance = None
# {{/docs-fragment statsig-client}}

# {{docs-fragment variant-apps}}
# Image with statsig-python-core for A/B testing
image = flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "httpx", "statsig-python-core")

# App A - First variant
app_a = FastAPI(
    title="App A",
    description="Variant A for A/B testing",
)

# App B - Second variant
app_b = FastAPI(
    title="App B",
    description="Variant B for A/B testing",
)
# {{/docs-fragment variant-apps}}

# {{docs-fragment root-lifespan}}
@asynccontextmanager
async def lifespan(_app: FastAPI):
    """Initialize and shutdown Statsig for A/B testing."""
    # Startup: Initialize Statsig using singleton
    api_key = os.getenv("STATSIG_API_KEY", None)
    if api_key is None:
        raise RuntimeError(f"StatsigClient API Key not set. ENV vars {os.environ}")
    StatsigClient.initialize(api_key)

    yield

    # Shutdown: Cleanup Statsig
    StatsigClient.shutdown()

# Root App - Performs A/B testing and routes to A or B
root_app = FastAPI(
    title="Root App - A/B Testing",
    description="Routes requests to App A or App B based on Statsig A/B test",
    lifespan=lifespan,
)
# {{/docs-fragment root-lifespan}}

# {{docs-fragment variant-envs}}
env_a = FastAPIAppEnvironment(
    name="app-a-variant",
    app=app_a,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

env_b = FastAPIAppEnvironment(
    name="app-b-variant",
    app=app_b,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)
# {{/docs-fragment variant-envs}}

# {{docs-fragment root-env}}
env_root = FastAPIAppEnvironment(
    name="root-ab-testing-app",
    app=root_app,
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[env_a, env_b],
    secrets=flyte.Secret("statsig-api-key", as_env_var="STATSIG_API_KEY"),
)
# {{/docs-fragment root-env}}

# {{docs-fragment variant-endpoints}}
# App A endpoints
@app_a.get("/process/{message}")
async def process_a(message: str) -> dict[str, str]:
    return {
        "variant": "A",
        "message": f"App A processed: {message}",
        "algorithm": "fast-processing",
    }

# App B endpoints
@app_b.get("/process/{message}")
async def process_b(message: str) -> dict[str, str]:
    return {
        "variant": "B",
        "message": f"App B processed: {message}",
        "algorithm": "enhanced-processing",
    }
# {{/docs-fragment variant-endpoints}}

# {{docs-fragment routing-endpoint}}
# Root app A/B testing endpoint
@root_app.get("/process/{message}")
async def process_with_ab_test(message: str, user_key: str) -> dict[str, typing.Any]:
    """
    Process a message using A/B testing to determine which app to call.

    Args:
        message: The message to process
        user_key: User identifier for A/B test bucketing (e.g., user_id, session_id)

    Returns:
        Response from either App A or App B, plus metadata about which variant was used
    """
    # Import StatsigUser at runtime (only available in container)
    from statsig_python_core import StatsigUser

    # Get statsig client from singleton
    statsig = StatsigClient.get_client()

    # Create Statsig user with the provided key
    user = StatsigUser(user_id=user_key)

    # Check the feature gate "variant_b" to determine which variant
    # If gate is enabled, use App B; otherwise use App A
    use_variant_b = statsig.check_gate(user, "variant_b")

    # Call the appropriate app based on A/B test result
    async with httpx.AsyncClient() as client:
        if use_variant_b:
            endpoint = f"{env_b.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()
        else:
            endpoint = f"{env_a.endpoint}/process/{message}"
            response = await client.get(endpoint)
            result = response.json()

    # Add A/B test metadata to response
    return {
        "ab_test_result": {
            "user_key": user_key,
            "selected_variant": "B" if use_variant_b else "A",
            "gate_name": "variant_b",
        },
        "response": result,
    }
# {{/docs-fragment routing-endpoint}}

@root_app.get("/endpoints")
async def get_endpoints() -> dict[str, str]:
    """Get the endpoints for App A and App B."""
    return {
        "app_a_endpoint": env_a.endpoint,
        "app_b_endpoint": env_b.endpoint,
    }

@root_app.get("/")
async def index():
    """Serve the A/B testing demo HTML page."""
    from fastapi.responses import HTMLResponse

    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>A/B Testing Demo - Statsig</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }

            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell,
                sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                display: flex;
                justify-content: center;
                align-items: center;
                padding: 20px;
            }

            .container {
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                padding: 40px;
                max-width: 600px;
                width: 100%;
            }

            h1 {
                color: #333;
                margin-bottom: 10px;
                font-size: 28px;
            }

            .subtitle {
                color: #666;
                margin-bottom: 30px;
                font-size: 14px;
            }

            .form-group {
                margin-bottom: 20px;
            }

            label {
                display: block;
                margin-bottom: 8px;
                color: #555;
                font-weight: 500;
                font-size: 14px;
            }

            input {
                width: 100%;
                padding: 12px 16px;
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                font-size: 14px;
                transition: border-color 0.3s;
            }

            input:focus {
                outline: none;
                border-color: #667eea;
            }

            button {
                width: 100%;
                padding: 14px;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                color: white;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: transform 0.2s, box-shadow 0.2s;
            }

            button:hover {
                transform: translateY(-2px);
                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
            }

            button:active {
                transform: translateY(0);
            }

            button:disabled {
                opacity: 0.6;
                cursor: not-allowed;
            }

            .result {
                margin-top: 30px;
                padding: 20px;
                border-radius: 12px;
                display: none;
            }

            .result.show {
                display: block;
            }

            .result.variant-a {
                background: #e3f2fd;
                border: 2px solid #2196f3;
            }

            .result.variant-b {
                background: #f3e5f5;
                border: 2px solid #9c27b0;
            }

            .result-header {
                font-size: 18px;
                font-weight: 600;
                margin-bottom: 15px;
                display: flex;
                align-items: center;
                gap: 10px;
            }

            .variant-badge {
                display: inline-block;
                padding: 4px 12px;
                border-radius: 12px;
                font-size: 12px;
                font-weight: 700;
            }

            .variant-a .variant-badge {
                background: #2196f3;
                color: white;
            }

            .variant-b .variant-badge {
                background: #9c27b0;
                color: white;
            }

            .result-content {
                margin-top: 10px;
            }

            .result-item {
                margin-bottom: 10px;
                padding: 10px;
                background: rgba(255, 255, 255, 0.8);
                border-radius: 6px;
            }

            .result-label {
                font-weight: 600;
                color: #555;
                font-size: 13px;
            }

            .result-value {
                color: #333;
                margin-top: 4px;
            }

            .error {
                background: #ffebee;
                border: 2px solid #f44336;
                color: #c62828;
                padding: 16px;
                border-radius: 8px;
                margin-top: 20px;
                display: none;
            }

            .error.show {
                display: block;
            }

            .info {
                background: #fff3e0;
                border-left: 4px solid #ff9800;
                padding: 12px 16px;
                margin-top: 20px;
                border-radius: 4px;
                font-size: 13px;
                color: #e65100;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎯 A/B Testing Demo</h1>
            <p class="subtitle">Test Statsig-powered variant selection</p>

            <form id="abTestForm">
                <div class="form-group">
                    <label for="message">Message to Process</label>
                    <input
                        type="text"
                        id="message"
                        name="message"
                        placeholder="e.g., hello, world, test"
                        required
                        value="hello"
                    >
                </div>

                <div class="form-group">
                    <label for="userKey">User Key (for A/B bucketing)</label>
                    <input
                        type="text"
                        id="userKey"
                        name="userKey"
                        placeholder="e.g., user123, session456"
                        required
                        value="user123"
                    >
                </div>

                <button type="submit" id="submitBtn">Run A/B Test</button>
            </form>

            <div id="result" class="result"></div>
            <div id="error" class="error"></div>

            <div class="info">
                💡 <strong>Tip:</strong> Try different user keys to see how Statsig routes to different variants.
                The same user key will always get the same variant (consistent bucketing).
            </div>
        </div>

        <script>
            const form = document.getElementById('abTestForm');
            const resultDiv = document.getElementById('result');
            const errorDiv = document.getElementById('error');
            const submitBtn = document.getElementById('submitBtn');

            form.addEventListener('submit', async (e) => {
                e.preventDefault();

                const message = document.getElementById('message').value;
                const userKey = document.getElementById('userKey').value;

                // Reset previous results
                resultDiv.classList.remove('show', 'variant-a', 'variant-b');
                errorDiv.classList.remove('show');
                submitBtn.disabled = true;
                submitBtn.textContent = 'Processing...';

                try {
                    const response =
                        await fetch(`/process/${encodeURIComponent(message)}?user_key=${encodeURIComponent(userKey)}`);

                    if (!response.ok) {
                        throw new Error(`HTTP error! status: ${response.status}`);
                    }

                    const data = await response.json();

                    // Display result
                    const variant = data.ab_test_result.selected_variant;
                    const variantClass = `variant-${variant.toLowerCase()}`;

                    resultDiv.className = `result show ${variantClass}`;
                    resultDiv.innerHTML = `
                        <div class="result-header">
                            <span>A/B Test Result</span>
                            <span class="variant-badge">Variant ${variant}</span>
                        </div>
                        <div class="result-content">
                            <div class="result-item">
                                <div class="result-label">User Key</div>
                                <div class="result-value">${data.ab_test_result.user_key}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Selected Variant</div>
                                <div class="result-value">Variant ${variant}
                                    (Gate: ${data.ab_test_result.gate_name})</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Response from App ${variant}</div>
                                <div class="result-value">${data.response.message}</div>
                            </div>
                            <div class="result-item">
                                <div class="result-label">Algorithm</div>
                                <div class="result-value">${data.response.algorithm}</div>
                            </div>
                        </div>
                    `;

                } catch (error) {
                    errorDiv.textContent = `Error: ${error.message}`;
                    errorDiv.classList.add('show');
                } finally {
                    submitBtn.disabled = false;
                    submitBtn.textContent = 'Run A/B Test';
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(env_root)
    print("Deployed A/B Testing Root App")
    print("\nUsage:")
    print("  Open your browser to '<endpoint>/' to access the interactive demo")
    print("  Or use curl: curl '<endpoint>/process/hello?user_key=user123'")
    print("\nNote: Set STATSIG_API_KEY secret to use real Statsig A/B testing.")
    print("      Create a feature gate named 'variant_b' in your Statsig dashboard.")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/serving_graphs/ab_testing.py*

**Setup before running:**

1. Get a Server Secret Key at [statsig.com](https://www.statsig.com/) → Settings → API Keys.
2. Create a feature gate named `variant_b` (e.g. 50% rollout).
3. Set the Flyte secret:

   CODE23

## When to split into a serving graph

Split when stages have:

- **Different bottlenecks**: CPU vs GPU vs memory
- **Different scaling needs**: bursty vs steady, wide vs narrow
- **Different lifecycles**: model weights you don't want to reload, expensive cold starts
- **Different routing concerns**: A/B, canary, proxy, gateway

Don't split just to separate code; a single app with a few endpoints is
simpler to operate.

## Best practices

1. **Use `depends_on`**: Always specify dependencies to ensure the dependency closure is deployed in one shot.
2. **Persistent HTTP clients**: Open one `httpx.AsyncClient` per replica in the app's lifespan rather than per request, to avoid TCP/TLS setup overhead.
3. **Pick the right wire format**: For tensor-shaped payloads, send raw bytes over `application/octet-stream` instead of JSON.
4. **Size each node independently**: GPU narrow, CPU wide; use scale-to-zero (`replicas=(0, N)`) for bursty downstream services.
5. **Authentication**: Use `requires_auth=True` on internal-only apps so they can't be reached from the public internet, and put public-facing auth on the entry-point app.
6. **Endpoint access**: Prefer `app_env.endpoint` for in-module references; use `flyte.app.AppEndpoint` parameters when the upstream env isn't importable.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/hybrid-graphs ===

# Hybrid app-task graphs

Apps and tasks can interact with each other: tasks can call apps via HTTP, and apps can trigger task execution via the Flyte SDK. This page covers both patterns.

## Call app from task

Tasks can call apps by making HTTP requests to the app's endpoint. This is useful when:

- You need to use a long-running service during task execution
- You want to call a model serving endpoint from a batch processing task
- You need to interact with an API from a workflow

### Example: FastAPI app called from a task

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "httpx",
# ]
# ///

"""Example of a task calling an app."""

import pathlib
import httpx
from fastapi import FastAPI
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(title="Add One", description="Adds one to the input", version="1.0.0")

image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")

# {{docs-fragment app-definition}}
app_env = FastAPIAppEnvironment(
    name="add-one-app",
    app=app,
    description="Adds one to the input",
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment app-definition}}

# {{docs-fragment task-env}}
task_env = flyte.TaskEnvironment(
    name="add_one_task_env",
    image=image,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    depends_on=[app_env],  # Ensure app is deployed before task runs
)
# {{/docs-fragment task-env}}

# {{docs-fragment app-endpoint}}
@app.get("/")
async def add_one(x: int) -> dict[str, int]:
    """Main endpoint for the add-one app."""
    return {"result": x + 1}
# {{/docs-fragment app-endpoint}}

# {{docs-fragment task}}
@task_env.task
async def add_one_task(x: int) -> int:
    print(f"Calling app at {app_env.endpoint}")
    async with httpx.AsyncClient() as client:
        response = await client.get(app_env.endpoint, params={"x": x})
        response.raise_for_status()
        return response.json()["result"]
# {{/docs-fragment task}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    deployments = flyte.deploy(task_env)
    print(f"Deployed task environment: {deployments}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/task_calling_app.py*

Key points:

- The task environment uses `depends_on=[app_env]` to ensure the app is deployed first
- Access the app endpoint via `app_env.endpoint`
- Use standard HTTP client libraries (like `httpx`) to make requests

### Example: Call a model inference service from a task

There are cases where you want to build a durable batch inference workflow that
calls to a reusable inference service. You can achieve this by creating a light-weight,
long-running [`AppEnvironment`](../../../api-reference/flyte-sdk/flyte.app/appenvironment) that the task calls via HTTP.

```mermaid
flowchart LR

    subgraph calls ["Batch inference task (CPU)"]
        D["Driver task<br/>fans out chunks<br/>(with concurrency cap)"]

        subgraph fanout
            T1["Inference task call 1"]
            T2["Inference task call 2"]
            T3["Inference task call N"]
        end

        D --> T1
        D --> T2
        D --> T3

        R["Aggregate batches"]
    end

    subgraph app ["Reusable inference service (GPU)"]
        FA["POST /generate"]
        B["Shared TokenBatcher"]
        M["vLLM model<br/>(loaded in lifespan)"]
        B --> FA
        M --> FA
    end

    T1 <--> FA
    T2 <--> FA
    T3 <--> FA
    fanout --> R
```

See [Batch inference](../../run-scaling/batch-inference) implementation details.

## Call task from app (webhooks / APIs)

Apps can trigger task execution using the Flyte SDK. This is useful for:

- Webhooks that trigger workflows
- APIs that need to run batch jobs
- Services that need to execute tasks asynchronously

Webhooks are HTTP endpoints that trigger actions in response to external events. Flyte apps can serve as webhook endpoints that trigger task runs, workflows, or other operations.

### Example: Basic webhook app

Here's a simple webhook that triggers Flyte tasks:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A webhook that triggers Flyte tasks."""

import pathlib
from fastapi import FastAPI, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from starlette import status
import os
from contextlib import asynccontextmanager
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment auth}}
WEBHOOK_API_KEY = os.getenv("WEBHOOK_API_KEY", "test-api-key")
security = HTTPBearer()

async def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> HTTPAuthorizationCredentials:
    """Verify the API key from the bearer token."""
    if credentials.credentials != WEBHOOK_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    return credentials
# {{/docs-fragment auth}}

# {{docs-fragment lifespan}}
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield
    # Cleanup if needed
# {{/docs-fragment lifespan}}

# {{docs-fragment app}}
app = FastAPI(
    title="Flyte Webhook Runner",
    description="A webhook service that triggers Flyte task runs",
    version="1.0.0",
    lifespan=lifespan,
)

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy"}
# {{/docs-fragment app}}

# {{docs-fragment webhook-endpoint}}
@app.post("/run-task/{project}/{domain}/{name}/{version}")
async def run_task(
    project: str,
    domain: str,
    name: str,
    version: str,
    inputs: dict,
    credentials: HTTPAuthorizationCredentials = Security(verify_token),
):
    """
    Trigger a Flyte task run via webhook.

    Returns information about the launched run.
    """
    # Fetch the task
    task = remote.Task.get(
        project=project,
        domain=domain,
        name=name,
        version=version,
    )

    # Run the task
    run = await flyte.run.aio(task, **inputs)

    return {
        "url": run.url,
        "id": run.id,
        "status": "started",
    }
# {{/docs-fragment webhook-endpoint}}

# {{docs-fragment env}}
env = FastAPIAppEnvironment(
    name="webhook-runner",
    app=app,
    description="A webhook service that triggers Flyte task runs",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,  # We handle auth in the app
    env_vars={"WEBHOOK_API_KEY": os.getenv("WEBHOOK_API_KEY", "test-api-key")},
)
# {{/docs-fragment env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed webhook: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/webhook/basic_webhook.py*

Once deployed, you can trigger tasks via HTTP POST:

```bash
curl -X POST "https://your-webhook-url/run-task/flytesnacks/development/my_task/v1" \
  -H "Authorization: Bearer test-api-key" \
  -H "Content-Type: application/json" \
  -d '{"input_key": "input_value"}'
```

Response:

```json
{
  "url": "https://console.union.ai/...",
  "id": "abc123",
  "status": "started"
}
```

### Advanced webhook patterns

**Webhook with validation**

Use Pydantic for input validation:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A webhook with Pydantic validation."""

import pathlib
from fastapi import FastAPI, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from starlette import status
import os
from contextlib import asynccontextmanager
from pydantic import BaseModel
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

WEBHOOK_API_KEY = os.getenv("WEBHOOK_API_KEY", "test-api-key")
security = HTTPBearer()

async def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> HTTPAuthorizationCredentials:
    """Verify the API key from the bearer token."""
    if credentials.credentials != WEBHOOK_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    return credentials

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield

app = FastAPI(
    title="Flyte Webhook Runner with Validation",
    description="A webhook service that triggers Flyte task runs with Pydantic validation",
    version="1.0.0",
    lifespan=lifespan,
)

# {{docs-fragment validation-model}}
class TaskInput(BaseModel):
    data: dict
    priority: int = 0
# {{/docs-fragment validation-model}}

# {{docs-fragment validated-webhook}}
@app.post("/run-task/{project}/{domain}/{name}/{version}")
async def run_task(
    project: str,
    domain: str,
    name: str,
    version: str,
    inputs: TaskInput,  # Validated input
    credentials: HTTPAuthorizationCredentials = Security(verify_token),
):
    task = remote.Task.get(
        project=project,
        domain=domain,
        name=name,
        version=version,
    )

    run = await flyte.run.aio(task, **inputs.model_dump())

    return {
        "run_id": run.id,
        "url": run.url,
    }
# {{/docs-fragment validated-webhook}}

env = FastAPIAppEnvironment(
    name="webhook-with-validation",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    env_vars={"WEBHOOK_API_KEY": os.getenv("WEBHOOK_API_KEY", "test-api-key")},
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed webhook: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/webhook_validation.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A webhook with Pydantic validation."""

import pathlib
from fastapi import FastAPI, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from starlette import status
import os
from contextlib import asynccontextmanager
from pydantic import BaseModel
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

WEBHOOK_API_KEY = os.getenv("WEBHOOK_API_KEY", "test-api-key")
security = HTTPBearer()

async def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> HTTPAuthorizationCredentials:
    """Verify the API key from the bearer token."""
    if credentials.credentials != WEBHOOK_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    return credentials

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield

app = FastAPI(
    title="Flyte Webhook Runner with Validation",
    description="A webhook service that triggers Flyte task runs with Pydantic validation",
    version="1.0.0",
    lifespan=lifespan,
)

# {{docs-fragment validation-model}}
class TaskInput(BaseModel):
    data: dict
    priority: int = 0
# {{/docs-fragment validation-model}}

# {{docs-fragment validated-webhook}}
@app.post("/run-task/{project}/{domain}/{name}/{version}")
async def run_task(
    project: str,
    domain: str,
    name: str,
    version: str,
    inputs: TaskInput,  # Validated input
    credentials: HTTPAuthorizationCredentials = Security(verify_token),
):
    task = remote.Task.get(
        project=project,
        domain=domain,
        name=name,
        version=version,
    )

    run = await flyte.run.aio(task, **inputs.model_dump())

    return {
        "run_id": run.id,
        "url": run.url,
    }
# {{/docs-fragment validated-webhook}}

env = FastAPIAppEnvironment(
    name="webhook-with-validation",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    env_vars={"WEBHOOK_API_KEY": os.getenv("WEBHOOK_API_KEY", "test-api-key")},
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed webhook: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/webhook_validation.py*

**Webhook with response waiting**

Wait for task completion:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A webhook that waits for task completion."""

import pathlib
from fastapi import FastAPI, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from starlette import status
import os
from contextlib import asynccontextmanager
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

WEBHOOK_API_KEY = os.getenv("WEBHOOK_API_KEY", "test-api-key")
security = HTTPBearer()

async def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> HTTPAuthorizationCredentials:
    """Verify the API key from the bearer token."""
    if credentials.credentials != WEBHOOK_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    return credentials

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield

app = FastAPI(
    title="Flyte Webhook Runner (Wait for Completion)",
    description="A webhook service that triggers Flyte task runs and waits for completion",
    version="1.0.0",
    lifespan=lifespan,
)

# {{docs-fragment wait-webhook}}
@app.post("/run-task-and-wait/{project}/{domain}/{name}/{version}")
async def run_task_and_wait(
    project: str,
    domain: str,
    name: str,
    version: str,
    inputs: dict,
    credentials: HTTPAuthorizationCredentials = Security(verify_token),
):
    task = remote.Task.get(
        project=project,
        domain=domain,
        name=name,
        version=version,
    )

    run = await flyte.run.aio(task, **inputs)
    run.wait()  # Wait for completion

    return {
        "run_id": run.id,
        "url": run.url,
        "status": run.status,
        "outputs": run.outputs(),
    }
# {{/docs-fragment wait-webhook}}

env = FastAPIAppEnvironment(
    name="webhook-wait-completion",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    env_vars={"WEBHOOK_API_KEY": os.getenv("WEBHOOK_API_KEY", "test-api-key")},
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed webhook: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/webhook_wait.py*

**Webhook with secret management**

Use Flyte secrets for API keys:

```python
env = FastAPIAppEnvironment(
    name="webhook-runner",
    app=app,
    secrets=flyte.Secret(key="webhook-api-key", as_env_var="WEBHOOK_API_KEY"),
    # ...
)
```

Then access in your app:

```python
WEBHOOK_API_KEY = os.getenv("WEBHOOK_API_KEY")
```

### Webhook security and best practices

- **Authentication**: Always secure webhooks with authentication (API keys, tokens, etc.).
- **Input validation**: Validate webhook inputs using Pydantic models.
- **Error handling**: Handle errors gracefully and return meaningful error messages.
- **Async operations**: Use async/await for I/O operations.
- **Health checks**: Include health check endpoints.
- **Logging**: Log webhook requests for debugging and auditing.
- **Rate limiting**: Consider implementing rate limiting for production.

Security considerations:

- Store API keys in Flyte secrets, not in code.
- Always use HTTPS in production.
- Validate all inputs to prevent injection attacks.
- Implement proper access control mechanisms.
- Log all webhook invocations for security auditing.

### Example: GitHub webhook

Here's an example webhook that triggers tasks based on GitHub events:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A GitHub webhook that triggers Flyte tasks based on GitHub events."""

import pathlib
import hmac
import hashlib
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, Header, HTTPException
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield

app = FastAPI(
    title="GitHub Webhook Handler",
    description="Triggers Flyte tasks based on GitHub events",
    version="1.0.0",
    lifespan=lifespan,
)

# {{docs-fragment github-webhook}}
@app.post("/github-webhook")
async def github_webhook(
    request: Request,
    x_hub_signature_256: str = Header(None),
):
    """Handle GitHub webhook events."""
    body = await request.body()

    # Verify signature
    secret = os.getenv("GITHUB_WEBHOOK_SECRET")
    signature = hmac.new(
        secret.encode(),
        body,
        hashlib.sha256
    ).hexdigest()

    expected_signature = f"sha256={signature}"
    if not hmac.compare_digest(x_hub_signature_256, expected_signature):
        raise HTTPException(status_code=403, detail="Invalid signature")

    # Process webhook
    event = await request.json()
    event_type = request.headers.get("X-GitHub-Event")

    if event_type == "push":
        # Trigger deployment task
        task = remote.Task.get(
            project="my-project",
            domain="development",
            name="deploy-task",
            version="v1",
        )
        run = await flyte.run.aio(task, commit=event["after"])
        return {"run_id": run.id, "url": run.url}

    return {"status": "ignored"}
# {{/docs-fragment github-webhook}}

# {{docs-fragment env}}
env = FastAPIAppEnvironment(
    name="github-webhook",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    secrets=flyte.Secret(key="GITHUB_WEBHOOK_SECRET", as_env_var="GITHUB_WEBHOOK_SECRET"),
)
# {{/docs-fragment env}}

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed GitHub webhook: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/github_webhook.py*

### Gradio agent UI

For AI agents, a Gradio app lets you build an interactive UI that kicks off agent runs. The app uses `flyte.with_runcontext()` to run the agent task either locally or on a remote cluster, controlled by an environment variable.

```python
import os
import flyte
import flyte.app
from research_agent import agent

RUN_MODE = os.getenv("RUN_MODE", "remote")

serving_env = flyte.app.AppEnvironment(
    name="research-agent-ui",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "gradio", "langchain-core", "langchain-openai", "langgraph",
    ),
    secrets=flyte.Secret(key="OPENAI_API_KEY", as_env_var="OPENAI_API_KEY"),
    port=7860,
)

def run_query(request: str):
    """Kick off the agent as a Flyte task."""
    result = flyte.with_runcontext(mode=RUN_MODE).run(agent, request=request)
    result.wait()
    return result.outputs()[0]

@serving_env.server
def app_server():
    create_demo().launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    create_demo().launch()
```

The `RUN_MODE` variable gives you a smooth development progression:

1. **Fully local**: `RUN_MODE=local python agent_app.py`. Everything runs in your local Python environment, great for rapid iteration.
2. **Local app, remote task**: `python agent_app.py`. The UI runs locally but the agent executes on the cluster with full compute resources.
3. **Full remote**: `flyte deploy agent_app.py serving_env`. Both the UI and agent run on the cluster.

## Best practices

1. **Use `depends_on`**: Always specify dependencies to ensure proper deployment order.
2. **Handle errors**: Implement proper error handling for HTTP requests.
3. **Use async clients**: Use async HTTP clients (`httpx.AsyncClient`) in async contexts.
4. **Initialize Flyte**: For apps calling tasks, initialize Flyte in the app's startup.
5. **Endpoint access**: Use `app_env.endpoint` or `AppEndpoint` parameter for accessing app URLs.
6. **Webhook security**: Secure webhooks with auth, validation, and HTTPS.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/websocket-apps ===

# WebSocket apps

WebSockets enable bidirectional, real-time communication between clients and servers. Flyte apps can serve WebSocket endpoints for real-time applications like chat, live updates, or streaming data.

## Example: Basic WebSocket app

Here's a simple FastAPI app with WebSocket support:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""A FastAPI app with WebSocket support."""

import pathlib
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import asyncio
import json
from datetime import UTC, datetime
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(
    title="Flyte WebSocket Demo",
    description="A FastAPI app with WebSocket support",
    version="1.0.0",
)

# {{docs-fragment connection-manager}}
class ConnectionManager:
    """Manages WebSocket connections."""

    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        """Accept and register a new WebSocket connection."""
        await websocket.accept()
        self.active_connections.append(websocket)
        print(f"Client connected. Total: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        """Remove a WebSocket connection."""
        self.active_connections.remove(websocket)
        print(f"Client disconnected. Total: {len(self.active_connections)}")

    async def send_personal_message(self, message: str, websocket: WebSocket):
        """Send a message to a specific WebSocket connection."""
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        """Broadcast a message to all active connections."""
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception as e:
                print(f"Error broadcasting: {e}")

manager = ConnectionManager()
# {{/docs-fragment connection-manager}}

# {{docs-fragment websocket-endpoint}}
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time communication."""
    await manager.connect(websocket)

    try:
        # Send welcome message
        await manager.send_personal_message(
            json.dumps({
                "type": "system",
                "message": "Welcome! You are connected.",
                "timestamp": datetime.now(UTC).isoformat(),
            }),
            websocket,
        )

        # Listen for messages
        while True:
            data = await websocket.receive_text()

            # Echo back to sender
            await manager.send_personal_message(
                json.dumps({
                    "type": "echo",
                    "message": f"Echo: {data}",
                    "timestamp": datetime.now(UTC).isoformat(),
                }),
                websocket,
            )

            # Broadcast to all clients
            await manager.broadcast(
                json.dumps({
                    "type": "broadcast",
                    "message": f"Broadcast: {data}",
                    "timestamp": datetime.now(UTC).isoformat(),
                    "connections": len(manager.active_connections),
                })
            )

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(
            json.dumps({
                "type": "system",
                "message": "A client disconnected",
                "connections": len(manager.active_connections),
            })
        )
# {{/docs-fragment websocket-endpoint}}

# {{docs-fragment env}}
env = FastAPIAppEnvironment(
    name="websocket-app",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed websocket app: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/basic_websocket.py*

## WebSocket patterns

**Echo server**

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""WebSocket patterns: echo, broadcast, streaming, and chat."""

import asyncio
import json
import random
from datetime import datetime, UTC
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(
    title="WebSocket Patterns Demo",
    description="Demonstrates various WebSocket patterns",
    version="1.0.0",
)

# {{docs-fragment echo-server}}
@app.websocket("/echo")
async def echo(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        pass
# {{/docs-fragment echo-server}}

# Connection manager for broadcast
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                pass

manager = ConnectionManager()

# {{docs-fragment broadcast-server}}
@app.websocket("/broadcast")
async def broadcast(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(data)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
# {{/docs-fragment broadcast-server}}

# {{docs-fragment streaming-server}}
@app.websocket("/stream")
async def stream_data(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Generate or fetch data
            data = {"timestamp": datetime.now(UTC).isoformat(), "value": random.random()}
            await websocket.send_json(data)
            await asyncio.sleep(1)  # Send update every second
    except WebSocketDisconnect:
        pass
# {{/docs-fragment streaming-server}}

# {{docs-fragment chat-room}}
class ChatRoom:
    def __init__(self, name: str):
        self.name = name
        self.connections: list[WebSocket] = []

    async def join(self, websocket: WebSocket):
        self.connections.append(websocket)

    async def leave(self, websocket: WebSocket):
        self.connections.remove(websocket)

    async def broadcast(self, message: str, sender: WebSocket):
        for connection in self.connections:
            if connection != sender:
                await connection.send_text(message)

rooms: dict[str, ChatRoom] = {}

@app.websocket("/chat/{room_name}")
async def chat(websocket: WebSocket, room_name: str):
    await websocket.accept()

    if room_name not in rooms:
        rooms[room_name] = ChatRoom(room_name)

    room = rooms[room_name]
    await room.join(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            await room.broadcast(data, websocket)
    except WebSocketDisconnect:
        await room.leave(websocket)
# {{/docs-fragment chat-room}}

env = FastAPIAppEnvironment(
    name="websocket-patterns",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed WebSocket patterns app: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/websocket_patterns.py*

**Broadcast server**

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""WebSocket patterns: echo, broadcast, streaming, and chat."""

import asyncio
import json
import random
from datetime import datetime, UTC
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(
    title="WebSocket Patterns Demo",
    description="Demonstrates various WebSocket patterns",
    version="1.0.0",
)

# {{docs-fragment echo-server}}
@app.websocket("/echo")
async def echo(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        pass
# {{/docs-fragment echo-server}}

# Connection manager for broadcast
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                pass

manager = ConnectionManager()

# {{docs-fragment broadcast-server}}
@app.websocket("/broadcast")
async def broadcast(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(data)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
# {{/docs-fragment broadcast-server}}

# {{docs-fragment streaming-server}}
@app.websocket("/stream")
async def stream_data(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Generate or fetch data
            data = {"timestamp": datetime.now(UTC).isoformat(), "value": random.random()}
            await websocket.send_json(data)
            await asyncio.sleep(1)  # Send update every second
    except WebSocketDisconnect:
        pass
# {{/docs-fragment streaming-server}}

# {{docs-fragment chat-room}}
class ChatRoom:
    def __init__(self, name: str):
        self.name = name
        self.connections: list[WebSocket] = []

    async def join(self, websocket: WebSocket):
        self.connections.append(websocket)

    async def leave(self, websocket: WebSocket):
        self.connections.remove(websocket)

    async def broadcast(self, message: str, sender: WebSocket):
        for connection in self.connections:
            if connection != sender:
                await connection.send_text(message)

rooms: dict[str, ChatRoom] = {}

@app.websocket("/chat/{room_name}")
async def chat(websocket: WebSocket, room_name: str):
    await websocket.accept()

    if room_name not in rooms:
        rooms[room_name] = ChatRoom(room_name)

    room = rooms[room_name]
    await room.join(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            await room.broadcast(data, websocket)
    except WebSocketDisconnect:
        await room.leave(websocket)
# {{/docs-fragment chat-room}}

env = FastAPIAppEnvironment(
    name="websocket-patterns",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed WebSocket patterns app: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/websocket_patterns.py*

**Real-time data streaming**

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""WebSocket patterns: echo, broadcast, streaming, and chat."""

import asyncio
import json
import random
from datetime import datetime, UTC
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(
    title="WebSocket Patterns Demo",
    description="Demonstrates various WebSocket patterns",
    version="1.0.0",
)

# {{docs-fragment echo-server}}
@app.websocket("/echo")
async def echo(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        pass
# {{/docs-fragment echo-server}}

# Connection manager for broadcast
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                pass

manager = ConnectionManager()

# {{docs-fragment broadcast-server}}
@app.websocket("/broadcast")
async def broadcast(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(data)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
# {{/docs-fragment broadcast-server}}

# {{docs-fragment streaming-server}}
@app.websocket("/stream")
async def stream_data(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Generate or fetch data
            data = {"timestamp": datetime.now(UTC).isoformat(), "value": random.random()}
            await websocket.send_json(data)
            await asyncio.sleep(1)  # Send update every second
    except WebSocketDisconnect:
        pass
# {{/docs-fragment streaming-server}}

# {{docs-fragment chat-room}}
class ChatRoom:
    def __init__(self, name: str):
        self.name = name
        self.connections: list[WebSocket] = []

    async def join(self, websocket: WebSocket):
        self.connections.append(websocket)

    async def leave(self, websocket: WebSocket):
        self.connections.remove(websocket)

    async def broadcast(self, message: str, sender: WebSocket):
        for connection in self.connections:
            if connection != sender:
                await connection.send_text(message)

rooms: dict[str, ChatRoom] = {}

@app.websocket("/chat/{room_name}")
async def chat(websocket: WebSocket, room_name: str):
    await websocket.accept()

    if room_name not in rooms:
        rooms[room_name] = ChatRoom(room_name)

    room = rooms[room_name]
    await room.join(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            await room.broadcast(data, websocket)
    except WebSocketDisconnect:
        await room.leave(websocket)
# {{/docs-fragment chat-room}}

env = FastAPIAppEnvironment(
    name="websocket-patterns",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed WebSocket patterns app: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/websocket_patterns.py*

**Chat application**

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""WebSocket patterns: echo, broadcast, streaming, and chat."""

import asyncio
import json
import random
from datetime import datetime, UTC
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(
    title="WebSocket Patterns Demo",
    description="Demonstrates various WebSocket patterns",
    version="1.0.0",
)

# {{docs-fragment echo-server}}
@app.websocket("/echo")
async def echo(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        pass
# {{/docs-fragment echo-server}}

# Connection manager for broadcast
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                pass

manager = ConnectionManager()

# {{docs-fragment broadcast-server}}
@app.websocket("/broadcast")
async def broadcast(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(data)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
# {{/docs-fragment broadcast-server}}

# {{docs-fragment streaming-server}}
@app.websocket("/stream")
async def stream_data(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Generate or fetch data
            data = {"timestamp": datetime.now(UTC).isoformat(), "value": random.random()}
            await websocket.send_json(data)
            await asyncio.sleep(1)  # Send update every second
    except WebSocketDisconnect:
        pass
# {{/docs-fragment streaming-server}}

# {{docs-fragment chat-room}}
class ChatRoom:
    def __init__(self, name: str):
        self.name = name
        self.connections: list[WebSocket] = []

    async def join(self, websocket: WebSocket):
        self.connections.append(websocket)

    async def leave(self, websocket: WebSocket):
        self.connections.remove(websocket)

    async def broadcast(self, message: str, sender: WebSocket):
        for connection in self.connections:
            if connection != sender:
                await connection.send_text(message)

rooms: dict[str, ChatRoom] = {}

@app.websocket("/chat/{room_name}")
async def chat(websocket: WebSocket, room_name: str):
    await websocket.accept()

    if room_name not in rooms:
        rooms[room_name] = ChatRoom(room_name)

    room = rooms[room_name]
    await room.join(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            await room.broadcast(data, websocket)
    except WebSocketDisconnect:
        await room.leave(websocket)
# {{/docs-fragment chat-room}}

env = FastAPIAppEnvironment(
    name="websocket-patterns",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed WebSocket patterns app: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/websocket_patterns.py*

## Using WebSockets with Flyte tasks

You can trigger Flyte tasks from WebSocket messages:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "websockets",
# ]
# ///

"""A WebSocket app that triggers Flyte tasks and streams updates."""

import json
from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import flyte
import flyte.remote as remote
from flyte.app.extras import FastAPIAppEnvironment

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize Flyte before accepting requests."""
    await flyte.init_in_cluster.aio()
    yield

app = FastAPI(
    title="WebSocket Task Runner",
    description="Triggers Flyte tasks via WebSocket and streams updates",
    version="1.0.0",
    lifespan=lifespan,
)

# {{docs-fragment task-runner-websocket}}
@app.websocket("/task-runner")
async def task_runner(websocket: WebSocket):
    await websocket.accept()

    try:
        while True:
            # Receive task request
            message = await websocket.receive_text()
            request = json.loads(message)

            # Trigger Flyte task
            task = remote.Task.get(
                project=request["project"],
                domain=request["domain"],
                name=request["task"],
                version=request["version"],
            )

            run = await flyte.run.aio(task, **request["inputs"])

            # Send run info back
            await websocket.send_json({
                "run_id": run.id,
                "url": run.url,
                "status": "started",
            })

            # Optionally stream updates
            async for update in run.stream():
                await websocket.send_json({
                    "status": update.status,
                    "message": update.message,
                })

    except WebSocketDisconnect:
        pass
# {{/docs-fragment task-runner-websocket}}

env = FastAPIAppEnvironment(
    name="task-runner-websocket",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "websockets",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed WebSocket task runner: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/websocket/task_runner_websocket.py*

## WebSocket client example

Connect from Python:

```python
import asyncio
import websockets
import json

async def client():
    uri = "ws://your-app-url/ws"
    async with websockets.connect(uri) as websocket:
        # Send message
        await websocket.send("Hello, Server!")
        
        # Receive message
        response = await websocket.recv()
        print(f"Received: {response}")

asyncio.run(client())
```

## Best practices

1. **Connection management**: Track active connections and handle disconnections gracefully.
2. **Heartbeats**: Implement ping/pong for connection health monitoring.
3. **Rate limiting**: Consider rate limiting for production deployments.
4. **Error handling**: Handle WebSocket errors and connection drops.
5. **Authentication**: Implement authentication for secure WebSocket connections.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/browser-apps ===

# Browser apps

For browser-based apps (like Streamlit, Gradio, or custom HTML/JS dashboards), users interact directly through the web interface. The app URL is accessible in a browser, and users interact with the UI directly: no API calls needed from other services.

## Accessing browser-based apps

To access a browser-based app:

1. Deploy the app using `flyte deploy` or `flyte serve`
2. Navigate to the app URL in a browser
3. Interact with the UI directly

## Common browser-based app types

### Streamlit apps

Streamlit is ideal for data dashboards and ML prototypes. See [Streamlit app](../native-app-integrations/streamlit-app) for details.

### Gradio apps

Gradio is great for ML model demos and interactive interfaces. You can deploy a Gradio app by building a custom [`AppEnvironment`](./single-script-apps) with the `gradio` package installed in your image.

### Custom HTML/JS apps

You can also serve custom HTML/JS applications using FastAPI's static file serving or any other web framework.

## Best practices

1. **Authentication**: For sensitive apps, enable authentication with `requires_auth=True`.
2. **Responsive design**: Design UIs that work on various screen sizes.
3. **Loading states**: Show loading indicators for long-running operations.
4. **Error handling**: Display user-friendly error messages.
5. **Resource management**: Configure appropriate CPU/memory resources based on expected usage.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/build-apps/secret-based-authentication ===

# Secret-based authentication

This guide deploys a FastAPI app that uses API key authentication with Flyte secrets. This allows you to invoke the endpoint from the public internet securely without exposing API keys in your code.

## Create the secret

Before defining and deploying the app, you need to create the `API_KEY` secret in Flyte. This secret will store your API key securely.

Create the secret using the Flyte CLI:

```bash
flyte create secret API_KEY <your-api-key-value>
```

For example:

```bash
flyte create secret API_KEY my-secret-api-key-12345
```

> [!NOTE]
> The secret name `API_KEY` must match the key specified in the `flyte.Secret()` call in your code. The secret will be available to your app as the environment variable specified in `as_env_var`.

## Define the FastAPI app

Here's a simple FastAPI app that uses `HTTPAuthorizationCredentials` to authenticate requests using a secret stored in Flyte:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Basic FastAPI authentication using dependency injection."""

from fastapi import FastAPI, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from starlette import status
import os
import pathlib
import flyte
from flyte.app.extras import FastAPIAppEnvironment

# Get API key from environment variable (loaded from Flyte secret)
# The secret must be created using: flyte create secret API_KEY <your-api-key-value>
API_KEY = os.getenv("API_KEY")
security = HTTPBearer()

async def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> HTTPAuthorizationCredentials:
    """Verify the API key from the bearer token."""
    if not API_KEY:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="API_KEY not configured",
        )
    if credentials.credentials != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    return credentials

app = FastAPI(title="Authenticated API")

@app.get("/public")
async def public_endpoint():
    """Public endpoint that doesn't require authentication."""
    return {"message": "This is public"}

@app.get("/protected")
async def protected_endpoint(
    credentials: HTTPAuthorizationCredentials = Security(verify_token),
):
    """Protected endpoint that requires authentication."""
    return {
        "message": "This is protected",
        "user": credentials.credentials,
    }

env = FastAPIAppEnvironment(
    name="authenticated-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,  # We handle auth in the app
    secrets=flyte.Secret(key="API_KEY", as_env_var="API_KEY"),
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/basic_auth.py*

As you can see, we:

1. Define a `FastAPI` app
2. Create a `verify_token` function that verifies the API key from the Bearer token
3. Define endpoints that use the `verify_token` function to authenticate requests
4. Configure the `FastAPIAppEnvironment` with:
   - `requires_auth=False` - This allows the endpoint to be reached without going through Flyte's authentication, since we're handling authentication ourselves using the `API_KEY` secret
   - `secrets=flyte.Secret(key="API_KEY", as_env_var="API_KEY")` - This injects the secret value into the `API_KEY` environment variable at runtime

The key difference from using `env_vars` is that secrets are stored securely in Flyte's secret store and injected at runtime, rather than being passed as plain environment variables.

## Deploy the FastAPI app

Once the secret is created, you can deploy the FastAPI app. Make sure your `config.yaml` file is in the same directory as your script, then run:

```bash
python basic_auth.py
```

Or use the Flyte CLI:

```bash
flyte serve basic_auth.py
```

Deploying the application will stream the status to the console and display the app URL:

```
✨ Deploying Application: authenticated-api
🔎 Console URL: https://<union-tenant>/console/projects/my-project/domains/development/apps/fastapi-with-auth
[Status] Pending: App is pending deployment
[Status] Started: Service is ready
🚀 Deployed Endpoint: https://rough-meadow-97cf5.apps.<union-tenant>
```

## Invoke the endpoint

Once deployed, you can invoke the authenticated endpoint using curl:

```bash
curl -X GET "https://rough-meadow-97cf5.apps.<union-tenant>/protected" \
  -H "Authorization: Bearer <your-api-key-value>"
```

Replace `<your-api-key-value>` with the actual API key value you used when creating the secret.

For example, if you created the secret with value `my-secret-api-key-12345`:

```bash
curl -X GET "https://rough-meadow-97cf5.apps.<union-tenant>/protected" \
  -H "Authorization: Bearer my-secret-api-key-12345"
```

You should receive a response:

```json
{
  "message": "This is protected",
  "user": "my-secret-api-key-12345"
}
```

## Authentication for vLLM and SGLang apps

Both vLLM and SGLang apps support API key authentication through their native `--api-key` argument. This allows you to secure your LLM endpoints while keeping them accessible from the public internet.

### Create the authentication secret

Create a secret to store your API key:

```bash
flyte create secret AUTH_SECRET <your-api-key-value>
```

For example:

```bash
flyte create secret AUTH_SECRET my-llm-api-key-12345
```

### Deploy vLLM app with authentication

Here's how to deploy a vLLM app with API key authentication:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b45",
# ]
# ///

"""vLLM app with API key authentication."""

import pathlib
from flyteplugins.vllm import VLLMAppEnvironment
import flyte

# The secret must be created using: flyte create secret AUTH_SECRET <your-api-key-value>
vllm_app = VLLMAppEnvironment(
    name="vllm-app-with-auth",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by vLLM
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # Disable Union's platform-level authentication so you can access the
    # endpoint from the public internet
    requires_auth=False,
    # Inject the secret as an environment variable
    secrets=flyte.Secret(key="AUTH_SECRET", as_env_var="AUTH_SECRET"),
    # Pass the API key to vLLM's --api-key argument
    # The $AUTH_SECRET will be replaced with the actual secret value at runtime
    extra_args=[
        "--api-key", "$AUTH_SECRET",
    ],
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app = flyte.serve(vllm_app)
    print(f"Deployed vLLM app: {app.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/vllm_with_auth.py*

Key points:

1. **`requires_auth=False`** - Disables Union's platform-level authentication so the endpoint can be accessed from the public internet
2. **`secrets=flyte.Secret(key="AUTH_SECRET", as_env_var="AUTH_SECRET")`** - Injects the secret as an environment variable
3. **`extra_args=["--api-key", "$AUTH_SECRET"]`** - Passes the API key to vLLM's `--api-key` argument. The `$AUTH_SECRET` will be replaced with the actual secret value at runtime

Deploy the app:

```bash
python vllm_with_auth.py
```

Or use the Flyte CLI:

```bash
flyte serve vllm_with_auth.py
```

### Deploy SGLang app with authentication

Here's how to deploy a SGLang app with API key authentication:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-sglang>=2.0.0b45",
# ]
# ///

"""SGLang app with API key authentication."""

import pathlib
from flyteplugins.sglang import SGLangAppEnvironment
import flyte

# The secret must be created using: flyte create secret AUTH_SECRET <your-api-key-value>
sglang_app = SGLangAppEnvironment(
    name="sglang-with-auth",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by SGLang
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    # Disable Union's platform-level authentication so you can access the
    # endpoint from the public internet
    requires_auth=False,
    # Inject the secret as an environment variable
    secrets=flyte.Secret(key="AUTH_SECRET", as_env_var="AUTH_SECRET"),
    # Pass the API key to SGLang's --api-key argument
    # The $AUTH_SECRET will be replaced with the actual secret value at runtime
    extra_args=[
        "--api-key", "$AUTH_SECRET",
    ],
)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app = flyte.serve(sglang_app)
    print(f"Deployed SGLang app: {app.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/sglang/sglang_with_auth.py*

The configuration is similar to vLLM:

1. **`requires_auth=False`** - Disables Union's platform-level authentication
2. **`secrets=flyte.Secret(key="AUTH_SECRET", as_env_var="AUTH_SECRET")`** - Injects the secret as an environment variable
3. **`extra_args=["--api-key", "$AUTH_SECRET"]`** - Passes the API key to SGLang's `--api-key` argument

Deploy the app:

```bash
python sglang_with_auth.py
```

Or use the Flyte CLI:

```bash
flyte serve sglang_with_auth.py
```

### Invoke authenticated LLM endpoints

Once deployed, you can invoke the authenticated endpoints using the OpenAI-compatible API format. Both vLLM and SGLang expose OpenAI-compatible endpoints.

For example, to make a chat completion request:

```bash
curl -X POST "https://your-app-url/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer <your-api-key-value>" \
  -d '{
    "model": "qwen3-0.6b",
    "messages": [
      {"role": "user", "content": "Hello, how are you?"}
    ]
  }'
```

Replace `<your-api-key-value>` with the actual API key value you used when creating the secret.

For example, if you created the secret with value `my-llm-api-key-12345`:

```bash
curl -X POST "https://your-app-url/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer my-llm-api-key-12345" \
  -d '{
    "model": "qwen3-0.6b",
    "messages": [
      {"role": "user", "content": "Hello, how are you?"}
    ]
  }'
```

You should receive a response with the model's completion.

> [!NOTE]
> The `$AUTH_SECRET` syntax in `extra_args` is automatically replaced with the actual secret value at runtime. This ensures the API key is never exposed in your code or configuration files.

## Accessing Swagger documentation

The app also includes a public health check endpoint and Swagger UI documentation:

- **Health check**: `https://your-app-url/health`
- **Swagger UI**: `https://your-app-url/docs`
- **ReDoc**: `https://your-app-url/redoc`

The Swagger UI will show an "Authorize" button where you can enter your Bearer token to test authenticated endpoints directly from the browser.

## Security best practices

1. **Use strong API keys**: Generate cryptographically secure random strings for your API keys
2. **Rotate keys regularly**: Periodically rotate your API keys for better security
3. **Scope secrets appropriately**: Use project/domain scoping when creating secrets if you want to limit access:

   ```bash
   flyte create secret --project my-project --domain development API_KEY my-secret-value
   ```

4. **Never commit secrets**: Always use Flyte secrets for API keys, never hardcode them in your code
5. **Use HTTPS**: Always use HTTPS in production (Flyte apps are served over HTTPS by default)

## Troubleshooting

**Authentication failing:**

- Verify the secret exists: `flyte get secret API_KEY`
- Check that the secret key name matches exactly (case-sensitive)
- Ensure you're using the correct Bearer token value
- Verify the `as_env_var` parameter matches the environment variable name in your code

**Secret not found:**

- Make sure you've created the secret before deploying the app
- Check the secret scope (organization vs project/domain) matches your app's project/domain
- Verify the secret name matches exactly (should be `API_KEY`)

**App not starting:**

- Check container logs for errors
- Verify all dependencies are installed in the image
- Ensure the secret is accessible in the app's project/domain

**LLM app authentication not working:**

- Verify the secret exists: `flyte get secret AUTH_SECRET`
- Check that `$AUTH_SECRET` is correctly specified in `extra_args` (note the `$` prefix)
- Ensure the secret name matches exactly (case-sensitive) in both the `flyte.Secret()` call and `extra_args`
- For vLLM, verify the `--api-key` argument is correctly passed
- For SGLang, verify the `--api-key` argument is correctly passed
- Check that `requires_auth=False` is set to allow public access

## Next steps

- Learn more about [managing secrets](../../tasks/task-configuration/secrets) in Flyte
- See [hybrid graphs](./hybrid-graphs) for webhook examples and authentication patterns
- Learn about [vLLM apps](../native-app-integrations/vllm-app) and [SGLang apps](../native-app-integrations/sglang-app) for serving LLMs


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations ===

# Native app integrations

Flyte ships with a set of pre-built [`AppEnvironment`](../build-apps/_index) integrations that wrap popular frameworks and serving runtimes, so you can deploy common app types without writing the integration glue yourself. Each integration provides a ready-to-use environment class: just configure your app, image, resources, and scaling, and Flyte handles the rest.

> [!TIP]
> If you're new to apps in Flyte, start with **Get started > Core concepts > Apps** for an overview, then see [Build apps](../build-apps/_index) to learn how to build custom app environments from scratch.

## When to use a native integration

Use a native integration when your app fits one of the supported frameworks and you want:

- **A minimal, opinionated setup**: sensible defaults for the framework, no boilerplate
- **First-class support**: features like model streaming, OpenAI-compatible APIs, and passthrough auth wired in for you
- **Faster time-to-deploy**: focus on your app logic, not on packaging and serving plumbing

For app types not covered here, build a custom [`AppEnvironment`](../build-apps/_index) using the patterns in the [Build apps](../build-apps/_index) section.

## Available integrations

| Integration | Framework | Typical use case |
|---|---|---|
| **Apps > Native app integrations > Streamlit app** | [Streamlit](https://streamlit.io/) | Interactive dashboards and data apps |
| **Apps > Native app integrations > FastAPI app** | [FastAPI](https://fastapi.tiangolo.com/) | REST APIs, webhooks, and backend services |
| **Apps > Native app integrations > vLLM app** | [vLLM](https://docs.vllm.ai/) | High-throughput LLM inference with an OpenAI-compatible API |
| **Apps > Native app integrations > SGLang app** | [SGLang](https://docs.sglang.io/) | Structured generation and LLM serving with an OpenAI-compatible API |
| **Apps > Native app integrations > Ollama app** | [Ollama](https://ollama.com/) | Lightweight local-style LLM serving with an OpenAI-compatible API |
| **Apps > Native app integrations > Flyte webhook** | [FastAPI](https://fastapi.tiangolo.com/) | Pre-built HTTP endpoints for common Flyte control plane operations |

## Next steps

- **Apps > Native app integrations > Streamlit app**: Build interactive Streamlit dashboards
- **Apps > Native app integrations > FastAPI app**: Create REST APIs and backend services
- **Apps > Native app integrations > vLLM app**: Serve large language models with vLLM
- **Apps > Native app integrations > SGLang app**: Serve LLMs with SGLang for structured generation
- **Apps > Native app integrations > Ollama app**: Serve lightweight LLMs with Ollama
- **Apps > Native app integrations > Flyte webhook**: Pre-built webhook for common Flyte operations

## Subpages
- **Apps > Native app integrations > Streamlit app**
- **Apps > Native app integrations > FastAPI app**
- **Apps > Native app integrations > vLLM app**
- **Apps > Native app integrations > SGLang app**
- **Apps > Native app integrations > Ollama app**
- **Apps > Native app integrations > Flyte webhook**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/streamlit-app ===

# Streamlit app

Streamlit is a popular framework for building interactive web applications and dashboards. Flyte makes it easy to deploy Streamlit apps as long-running services.

## Basic Streamlit app

The simplest way to deploy a Streamlit app is to use the built-in Streamlit "hello" demo:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A basic Streamlit app using the built-in hello demo."""

# {{docs-fragment app-definition}}
import flyte
import flyte.app

image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1")

app_env = flyte.app.AppEnvironment(
    name="streamlit-hello",
    image=image,
    args="streamlit hello --server.port 8080",
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.deploy(app_env)
    print(f"Deployed app: {app[0].summary_repr()}")
# {{/docs-fragment app-definition}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/basic_streamlit.py*

This just serves the built-in Streamlit "hello" demo.

## Single-file Streamlit app

For a single-file Streamlit app, you can wrap the app code in a function and use the `args` parameter to specify the command to run the app.
Note that the command is running the file itself, and uses the `--server` flag to start the server.

This is useful when you have a relatively small and simple app that you want to deploy as a single file.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
# ]
# ///

"""A single-script Streamlit app example."""

import sys
from pathlib import Path

import streamlit as st

import flyte
import flyte.app

# {{docs-fragment streamlit-app}}
def main():
    st.set_page_config(page_title="Simple Streamlit App", page_icon="🚀")

    st.title("Hello from Streamlit!")
    st.write("This is a simple single-script Streamlit app.")

    name = st.text_input("What's your name?", "World")
    st.write(f"Hello, {name}!")

    if st.button("Click me!"):
        st.balloons()
        st.success("Button clicked!")
# {{/docs-fragment streamlit-app}}

file_name = Path(__file__).name
# {{docs-fragment app-env}}
app_env = flyte.app.AppEnvironment(
    name="streamlit-single-script",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1"),
    args=[
        "streamlit",
        "run",
        file_name,
        "--server.port",
        "8080",
        "--",
        "--server",
    ],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    import logging
    import sys

    if "--server" in sys.argv:
        main()
    else:
        flyte.init_from_config(
            root_dir=Path(__file__).parent,
            log_level=logging.DEBUG,
        )
        app = flyte.serve(app_env)
        print(f"App URL: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/single_file_streamlit.py*

Note that the `if __name__ == "__main__"` block is used to both serve the `AppEnvironment` *and* run the app code via
the `streamlit run` command using the `--server` flag.

## Multi-file Streamlit app

When your streamlit application grows more complex, you may want to split your app into multiple files.
For a multi-file Streamlit app, use the `include` parameter to bundle your app files:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A custom Streamlit app with multiple files."""

import pathlib
import flyte
import flyte.app

# {{docs-fragment app-env}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "streamlit==1.41.1",
    "pandas==2.2.3",
    "numpy==2.2.3",
)

app_env = flyte.app.AppEnvironment(
    name="streamlit-multi-file-app",
    image=image,
    args="streamlit run main.py --server.port 8080",
    port=8080,
    include=["main.py", "utils.py"],  # Include your app files
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app = flyte.deploy(app_env)
    print(f"Deployed app: {app[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/multi_file_streamlit.py*

Where your project structure looks like this:

```
project/
├── main.py           # Main Streamlit app
├── utils.py          # Utility functions
└── components.py     # Reusable components
```

Your `main.py` file would contain your Streamlit app code:

```
import os

import streamlit as st
from utils import generate_data

# {{docs-fragment streamlit-app}}
all_columns = ["Apples", "Orange", "Pineapple"]
with st.container(border=True):
    columns = st.multiselect("Columns", all_columns, default=all_columns)

all_data = st.cache_data(generate_data)(columns=all_columns, seed=101)

data = all_data[columns]

tab1, tab2 = st.tabs(["Chart", "Dataframe"])
tab1.line_chart(data, height=250)
tab2.dataframe(data, height=250, use_container_width=True)
st.write(f"Environment: {os.environ}")
# {{/docs-fragment streamlit-app}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/main.py*

## Example: Data visualization dashboard

Here's a complete example of a Streamlit dashboard, all in a single file.

Define the streamlit app in the `main` function:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "pandas",
#    "numpy",
# ]
# ///

"""A data visualization dashboard example using Streamlit."""

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st

import flyte
import flyte.app

# {{docs-fragment streamlit-app}}
def main():
    st.set_page_config(page_title="Sales Dashboard", page_icon="📊")

    st.title("Sales Dashboard")

    # Load data
    @st.cache_data
    def load_data():
        return pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=100, freq="D"),
            "sales": np.random.randint(1000, 5000, 100),
        })

    data = load_data()

    # Sidebar filters
    st.sidebar.header("Filters")
    start_date = st.sidebar.date_input("Start date", value=data["date"].min())
    end_date = st.sidebar.date_input("End date", value=data["date"].max())

    # Filter data
    filtered_data = data[
        (data["date"] >= pd.Timestamp(start_date)) &
        (data["date"] <= pd.Timestamp(end_date))
    ]

    # Display metrics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Sales", f"${filtered_data['sales'].sum():,.0f}")
    with col2:
        st.metric("Average Sales", f"${filtered_data['sales'].mean():,.0f}")
    with col3:
        st.metric("Days", len(filtered_data))

    # Chart
    st.line_chart(filtered_data.set_index("date")["sales"])

# {{/docs-fragment streamlit-app}}

# {{docs-fragment app-env}}
file_name = Path(__file__).name
app_env = flyte.app.AppEnvironment(
    name="sales-dashboard",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit==1.41.1",
        "pandas==2.2.3",
        "numpy==2.2.3",
    ),
    args=["streamlit run", file_name, "--server.port", "8080", "--", "--server"],
    port=8080,
    resources=flyte.Resources(cpu="2", memory="2Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment serve}}
if __name__ == "__main__":
    import logging
    import sys

    if "--server" in sys.argv:
        main()
    else:
        flyte.init_from_config(
            root_dir=Path(__file__).parent,
            log_level=logging.DEBUG,
        )
        app = flyte.serve(app_env)
        print(f"Dashboard URL: {app.url}")
# {{/docs-fragment serve}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/data_visualization_dashboard.py*

Define the `AppEnvironment` to serve the app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "pandas",
#    "numpy",
# ]
# ///

"""A data visualization dashboard example using Streamlit."""

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st

import flyte
import flyte.app

# {{docs-fragment streamlit-app}}
def main():
    st.set_page_config(page_title="Sales Dashboard", page_icon="📊")

    st.title("Sales Dashboard")

    # Load data
    @st.cache_data
    def load_data():
        return pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=100, freq="D"),
            "sales": np.random.randint(1000, 5000, 100),
        })

    data = load_data()

    # Sidebar filters
    st.sidebar.header("Filters")
    start_date = st.sidebar.date_input("Start date", value=data["date"].min())
    end_date = st.sidebar.date_input("End date", value=data["date"].max())

    # Filter data
    filtered_data = data[
        (data["date"] >= pd.Timestamp(start_date)) &
        (data["date"] <= pd.Timestamp(end_date))
    ]

    # Display metrics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Sales", f"${filtered_data['sales'].sum():,.0f}")
    with col2:
        st.metric("Average Sales", f"${filtered_data['sales'].mean():,.0f}")
    with col3:
        st.metric("Days", len(filtered_data))

    # Chart
    st.line_chart(filtered_data.set_index("date")["sales"])

# {{/docs-fragment streamlit-app}}

# {{docs-fragment app-env}}
file_name = Path(__file__).name
app_env = flyte.app.AppEnvironment(
    name="sales-dashboard",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit==1.41.1",
        "pandas==2.2.3",
        "numpy==2.2.3",
    ),
    args=["streamlit run", file_name, "--server.port", "8080", "--", "--server"],
    port=8080,
    resources=flyte.Resources(cpu="2", memory="2Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment serve}}
if __name__ == "__main__":
    import logging
    import sys

    if "--server" in sys.argv:
        main()
    else:
        flyte.init_from_config(
            root_dir=Path(__file__).parent,
            log_level=logging.DEBUG,
        )
        app = flyte.serve(app_env)
        print(f"Dashboard URL: {app.url}")
# {{/docs-fragment serve}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/data_visualization_dashboard.py*

And finally the app serving logic:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "pandas",
#    "numpy",
# ]
# ///

"""A data visualization dashboard example using Streamlit."""

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st

import flyte
import flyte.app

# {{docs-fragment streamlit-app}}
def main():
    st.set_page_config(page_title="Sales Dashboard", page_icon="📊")

    st.title("Sales Dashboard")

    # Load data
    @st.cache_data
    def load_data():
        return pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=100, freq="D"),
            "sales": np.random.randint(1000, 5000, 100),
        })

    data = load_data()

    # Sidebar filters
    st.sidebar.header("Filters")
    start_date = st.sidebar.date_input("Start date", value=data["date"].min())
    end_date = st.sidebar.date_input("End date", value=data["date"].max())

    # Filter data
    filtered_data = data[
        (data["date"] >= pd.Timestamp(start_date)) &
        (data["date"] <= pd.Timestamp(end_date))
    ]

    # Display metrics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Sales", f"${filtered_data['sales'].sum():,.0f}")
    with col2:
        st.metric("Average Sales", f"${filtered_data['sales'].mean():,.0f}")
    with col3:
        st.metric("Days", len(filtered_data))

    # Chart
    st.line_chart(filtered_data.set_index("date")["sales"])

# {{/docs-fragment streamlit-app}}

# {{docs-fragment app-env}}
file_name = Path(__file__).name
app_env = flyte.app.AppEnvironment(
    name="sales-dashboard",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit==1.41.1",
        "pandas==2.2.3",
        "numpy==2.2.3",
    ),
    args=["streamlit run", file_name, "--server.port", "8080", "--", "--server"],
    port=8080,
    resources=flyte.Resources(cpu="2", memory="2Gi"),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment serve}}
if __name__ == "__main__":
    import logging
    import sys

    if "--server" in sys.argv:
        main()
    else:
        flyte.init_from_config(
            root_dir=Path(__file__).parent,
            log_level=logging.DEBUG,
        )
        app = flyte.serve(app_env)
        print(f"Dashboard URL: {app.url}")
# {{/docs-fragment serve}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/streamlit/data_visualization_dashboard.py*

## Best practices

1. **Use `include` for custom apps**: Always include your app files when deploying custom Streamlit code
2. **Set the port correctly**: Ensure your Streamlit app uses `--server.port 8080` (or match your `port` setting)
3. **Cache data**: Use `@st.cache_data` for expensive computations to improve performance
4. **Resource sizing**: Adjust resources based on your app's needs (data size, computations)
5. **Public vs private**: Set `requires_auth=False` for public dashboards, `True` for internal tools

## Troubleshooting

**App not loading:**

- Verify the port matches (use `--server.port 8080`)
- Check that all required files are included
- Review container logs for errors

**Missing dependencies:**

- Ensure all required packages are in your image's pip packages
- Check that file paths in `include` are correct

**Performance issues:**

- Increase CPU/memory resources
- Use Streamlit's caching features (`@st.cache_data`, `@st.cache_resource`)
- Optimize data processing


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/fastapi-app ===

# FastAPI app

FastAPI is a modern, fast web framework for building APIs. Flyte provides `FastAPIAppEnvironment` which makes it easy to deploy FastAPI applications.

## Basic FastAPI app

Here's a simple FastAPI app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""A basic FastAPI app example."""

from fastapi import FastAPI
import pathlib
import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment fastapi-app}}
app = FastAPI(
    title="My API",
    description="A simple FastAPI application",
    version="1.0.0",
)
# {{/docs-fragment fastapi-app}}

# {{docs-fragment fastapi-env}}
env = FastAPIAppEnvironment(
    name="my-fastapi-app",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment fastapi-env}}

# {{docs-fragment endpoints}}
@app.get("/")
async def root():
    return {"message": "Hello, World!"}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}
# {{/docs-fragment endpoints}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/basic_fastapi.py*

Once deployed, you can:

- Access the API at the generated URL
- View interactive API docs at `/docs` (Swagger UI)
- View alternative docs at `/redoc`

## Serving a machine learning model

Here's an example of serving a scikit-learn model:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
#    "scikit-learn",
#    "joblib",
# ]
# ///

"""Example of serving a machine learning model with FastAPI."""

import os
from contextlib import asynccontextmanager
from pathlib import Path

import joblib
import flyte
from fastapi import FastAPI
from flyte.app.extras import FastAPIAppEnvironment
from pydantic import BaseModel

# {{docs-fragment ml-model}}
app = FastAPI(title="ML Model API")

# Define request/response models
class PredictionRequest(BaseModel):
    feature1: float
    feature2: float
    feature3: float

class PredictionResponse(BaseModel):
    prediction: float
    probability: float

# Load model (you would typically load this from storage)
model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model_path = os.getenv("MODEL_PATH", "/app/models/model.joblib")
    # In production, load from your storage
    if os.path.exists(model_path):
        with open(model_path, "rb") as f:
            model = joblib.load(f)
    yield

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    # Make prediction
    # prediction = model.predict([[request.feature1, request.feature2, request.feature3]])

    # Dummy prediction for demo
    prediction = 0.85
    probability = 0.92

    return PredictionResponse(
        prediction=prediction,
        probability=probability,
    )

env = FastAPIAppEnvironment(
    name="ml-model-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "scikit-learn",
        "pydantic",
        "joblib",
    ),
    parameters=[
        flyte.app.Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.joblib"),
            mount="/app/models",
            env_var="MODEL_PATH",
        ),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    requires_auth=False,
)
# {{/docs-fragment ml-model}}

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"API URL: {app_deployment[0].url}")
    print(f"Swagger docs: {app_deployment[0].url}/docs")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/ml_model_serving.py*

## Accessing Swagger documentation

FastAPI automatically generates interactive API documentation. Once deployed:

- **Swagger UI**: Access at `{app_url}/docs`
- **ReDoc**: Access at `{app_url}/redoc`
- **OpenAPI JSON**: Access at `{app_url}/openapi.json`

The Swagger UI provides an interactive interface where you can:

- See all available endpoints
- Test API calls directly from the browser
- View request/response schemas
- See example payloads

## Example: REST API with multiple endpoints

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Example REST API with multiple endpoints."""

from pathlib import Path
from typing import List
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment rest-api}}
app = FastAPI(title="Product API")

# Data models
class Product(BaseModel):
    id: int
    name: str
    price: float

class ProductCreate(BaseModel):
    name: str
    price: float

# In-memory database (use real database in production)
products_db = []

@app.get("/products", response_model=List[Product])
async def get_products():
    return products_db

@app.get("/products/{product_id}", response_model=Product)
async def get_product(product_id: int):
    product = next((p for p in products_db if p["id"] == product_id), None)
    if not product:
        raise HTTPException(status_code=404, detail="Product not found")
    return product

@app.post("/products", response_model=Product)
async def create_product(product: ProductCreate):
    new_product = {
        "id": len(products_db) + 1,
        "name": product.name,
        "price": product.price,
    }
    products_db.append(new_product)
    return new_product

env = FastAPIAppEnvironment(
    name="product-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
)
# {{/docs-fragment rest-api}}

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app_deployment = flyte.deploy(env)
    print(f"API URL: {app_deployment[0].url}")
    print(f"Swagger docs: {app_deployment[0].url}/docs")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/rest_api.py*

## Multi-file FastAPI app

Here's an example of a multi-file FastAPI app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "fastapi",
# ]
# ///

"""Multi-file FastAPI app example."""

from fastapi import FastAPI
from module import function  # Import from another file
import pathlib

import flyte
from flyte.app.extras import FastAPIAppEnvironment

# {{docs-fragment app-definition}}
app = FastAPI(title="Multi-file FastAPI Demo")

app_env = FastAPIAppEnvironment(
    name="fastapi-multi-file",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
    ),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=False,
    # FastAPIAppEnvironment automatically includes necessary files
    # But you can also specify explicitly:
    # include=["app.py", "module.py"],
)
# {{/docs-fragment app-definition}}

# {{docs-fragment endpoint}}
@app.get("/")
async def root():
    return function()  # Uses function from module.py
# {{/docs-fragment endpoint}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_deployment = flyte.deploy(app_env)
    print(f"Deployed: {app_deployment[0].summary_repr()}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/multi_file/app.py*

The helper module:

```
# {{docs-fragment helper-function}}
def function():
    """Helper function used by the FastAPI app."""
    return {"message": "Hello from module.py!"}
# {{/docs-fragment helper-function}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/multi_file/module.py*

See [Multi-script apps](../build-apps/multi-script-apps) for more details on building FastAPI apps with multiple files.

## Local-to-remote model serving

A common ML pattern: train a model with a Flyte pipeline, then serve predictions from it. During local development, the app loads the model from a local file (e.g. `model.pt` saved by your training pipeline). When deployed remotely, Flyte's `Parameter` system automatically resolves the model from the latest training run output.

```python
from contextlib import asynccontextmanager
from pathlib import Path
import os

from fastapi import FastAPI
import flyte
from flyte.app import Parameter, RunOutput
from flyte.app.extras import FastAPIAppEnvironment

MODEL_PATH_ENV = "MODEL_PATH"

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model on startup, either local file or remote run output."""
    model_path = Path(os.environ.get(MODEL_PATH_ENV, "model.pt"))
    model = load_model(model_path)
    app.state.model = model
    yield

app = FastAPI(title="MNIST Predictor", lifespan=lifespan)

serving_env = FastAPIAppEnvironment(
    name="mnist-predictor",
    app=app,
    parameters=[
        # Remote: resolves model from the latest train run and sets MODEL_PATH
        Parameter(
            name="model",
            value=RunOutput(task_name="ml_pipeline.pipeline", type="file", getter=(1,)),
            download=True,
            env_var=MODEL_PATH_ENV,
        ),
    ],
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi", "uvicorn", "torch", "torchvision",
    ),
    resources=flyte.Resources(cpu=1, memory="4Gi"),
)

@app.get("/predict")
async def predict(index: int = 0) -> dict:
    return {"prediction": app.state.model(index)}

if __name__ == "__main__":
    # Local: skip RunOutput resolution, lifespan falls back to local model.pt
    serving_env.parameters = []
    local_app = flyte.with_servecontext(mode="local").serve(serving_env)
    local_app.activate(wait=True)
```

Locally, the app loads `model.pt` from disk:

```bash
python serve_model.py
```

Remotely, Flyte resolves the model from the latest training run:

```bash
flyte deploy serve_model.py serving_env
```

The key idea: `Parameter` with `RunOutput` bridges the gap between local and remote. Locally, the app falls back to a local file. Remotely, Flyte resolves the model artifact from the latest pipeline run automatically.

## Best practices

1. **Use Pydantic models**: Define request/response models for type safety and automatic validation
2. **Handle errors**: Use HTTPException for proper error responses
3. **Async operations**: Use async/await for I/O operations
4. **Environment variables**: Use environment variables for configuration
5. **Logging**: Add proper logging for debugging and monitoring
6. **Health checks**: Always include a `/health` endpoint
7. **API documentation**: FastAPI auto-generates docs, but add descriptions to your endpoints

## Advanced features

FastAPI supports many features that work with Flyte:

- **Dependencies**: Use FastAPI's dependency injection system
- **Background tasks**: Run background tasks with BackgroundTasks
- **WebSockets**: See [WebSocket apps](../build-apps/websocket-apps) for details
- **Authentication**: Add authentication middleware (see [secret-based authentication](../build-apps/secret-based-authentication))
- **CORS**: Configure CORS for cross-origin requests
- **Rate limiting**: Add rate limiting middleware

## Troubleshooting

**App not starting:**

- Check that uvicorn can find your app module
- Verify all dependencies are installed in the image
- Check container logs for startup errors

**Import errors:**

- Ensure all imported modules are available
- Use `include` parameter if you have custom modules
- Check that file paths are correct

**API not accessible:**

- Verify `requires_auth` setting
- Check that the app is listening on the correct port (8080)
- Review network/firewall settings


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/vllm-app ===

# vLLM app

vLLM is a high-performance library for serving large language models (LLMs). Flyte provides `VLLMAppEnvironment` for deploying vLLM model servers.

## Installation

First, install the vLLM plugin:

```bash
pip install flyteplugins-vllm
```

## Basic vLLM app

Here's a simple example serving a HuggingFace model:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b45",
# ]
# ///

"""A simple vLLM app example."""

from flyteplugins.vllm import VLLMAppEnvironment
import flyte

# {{docs-fragment basic-vllm-app}}
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by vLLM
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment basic-vllm-app}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(vllm_app)
    print(f"Deployed vLLM app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/basic_vllm.py*

## Using prefetched models

You can use models prefetched with `flyte.prefetch`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b45",
# ]
# override-dependencies = [
#    "cel-python; sys_platform == 'never'",
# ]
# ///

"""vLLM app using prefetched models."""

from flyteplugins.vllm import VLLMAppEnvironment
import flyte

# {{docs-fragment prefetch}}

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # this is a placeholder
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="10Gi"),
    stream_model=True,  # Stream model directly from blob store to GPU
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Prefetch the model first
    run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
    run.wait()

    # Use the prefetched model
    app = flyte.serve(
        vllm_app.clone_with(
            vllm_app.name,
            model_hf_path=None,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
        )
    )
    print(f"Deployed vLLM app: {app.url}")
# {{/docs-fragment prefetch}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/vllm_with_prefetch.py*

## Model streaming

`VLLMAppEnvironment` supports streaming models directly from blob storage to GPU memory, reducing startup time.
When `stream_model=True` and the `model_path` argument is provided with either a `flyte.io.Dir` or `RunOutput` pointing
to a path in object store:

- Model weights stream directly from storage to GPU
- Faster startup time (no full download required)
- Lower disk space requirements

> [!NOTE]
> The contents of the model directory must be compatible with the vLLM-supported formats, e.g. the HuggingFace model
> serialization format.

## Custom vLLM arguments

Use `extra_args` to pass additional arguments to vLLM:

```python
vllm_app = VLLMAppEnvironment(
    name="custom-vllm-app",
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    extra_args=[
        "--max-model-len", "8192",  # Maximum context length
        "--gpu-memory-utilization", "0.8",  # GPU memory utilization
        "--trust-remote-code",  # Trust remote code in models
    ],
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    # ...
)
```

See the [vLLM documentation](https://docs.vllm.ai/en/stable/configuration/engine_args.html) for all available arguments.

## Using the OpenAI-compatible API

Once deployed, your vLLM app exposes an OpenAI-compatible API:

```python
from openai import OpenAI

client = OpenAI(
    base_url="https://your-app-url/v1",  # vLLM endpoint
    api_key="your-api-key",  # If you passed an --api-key argument
)

response = client.chat.completions.create(
    model="qwen3-0.6b",  # Your model_id
    messages=[
        {"role": "user", "content": "Hello, how are you?"}
    ],
)

print(response.choices[0].message.content)
```

> [!TIP]
> If you passed an `--api-key` argument, you can use the `api_key` parameter to authenticate your requests.
> See [secret-based authentication](../build-apps/secret-based-authentication#deploy-vllm-app-with-authentication) for more details on how to pass auth secrets to your app.

## Multi-GPU inference (Tensor parallelism)

For larger models, use multiple GPUs with tensor parallelism:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b45",
# ]
# ///

"""vLLM app with multi-GPU tensor parallelism."""

from flyteplugins.vllm import VLLMAppEnvironment
import flyte

# {{docs-fragment multi-gpu}}
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    model_hf_path="meta-llama/Llama-2-70b-hf",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # 4 GPUs for tensor parallelism
        disk="100Gi",
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Use 4 GPUs
        "--max-model-len", "4096",
        "--gpu-memory-utilization", "0.9",
    ],
    requires_auth=False,
)
# {{/docs-fragment multi-gpu}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(vllm_app)
    print(f"Deployed vLLM app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/vllm_multi_gpu.py*

The `tensor-parallel-size` should match the number of GPUs specified in resources.

## Model sharding with prefetch

You can prefetch and shard models for multi-GPU inference:

```python
# Prefetch with sharding configuration
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    accelerator="L40s:4",
    shard_config=flyte.prefetch.ShardConfig(
        engine="vllm",
        args=flyte.prefetch.VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
)
run.wait()

# Use the sharded model
vllm_app = VLLMAppEnvironment(
    name="sharded-llm-app",
    model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
    model_id="llama-2-70b",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4", disk="100Gi"),
    extra_args=["--tensor-parallel-size", "4"],
    stream_model=True,
)
```

See [Prefetching models](../serve-and-deploy-apps/prefetching-models) for more details on sharding.

## Autoscaling

vLLM apps work well with autoscaling:

```python
vllm_app = VLLMAppEnvironment(
    name="autoscaling-llm-app",
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale to zero when idle
        scaledown_after=600,  # 10 minutes idle before scaling down
    ),
    # ...
)
```

## Best practices

1. **Use prefetching**: Prefetch models for faster deployment and better reproducibility
2. **Enable streaming**: Use `stream_model=True` to reduce startup time and disk usage
3. **Right-size GPUs**: Match GPU memory to model size
4. **Configure memory utilization**: Use `--gpu-memory-utilization` to control memory usage
5. **Use tensor parallelism**: For large models, use multiple GPUs with `tensor-parallel-size`
6. **Set autoscaling**: Use appropriate idle TTL to balance cost and performance
7. **Limit context length**: Use `--max-model-len` for smaller models to reduce memory usage

## Troubleshooting

**Model loading fails:**

- Verify GPU memory is sufficient for the model
- Check that the model path or HuggingFace path is correct
- Review container logs for detailed error messages

**Out of memory errors:**

- Reduce `--max-model-len`
- Lower `--gpu-memory-utilization`
- Use a smaller model or more GPUs

**Slow startup:**

- Enable `stream_model=True` for faster loading
- Prefetch models before deployment
- Use faster storage backends

## API reference

See the [vLLM API reference](../../../api-reference/integrations/vllm/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/sglang-app ===

# SGLang app

SGLang is a fast structured generation library for large language models (LLMs). Flyte provides `SGLangAppEnvironment` for deploying SGLang model servers.

## Installation

First, install the SGLang plugin:

```bash
pip install flyteplugins-sglang
```

## Basic SGLang app

Here's a simple example serving a HuggingFace model:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-sglang>=2.0.0b45",
# ]
# ///

"""A simple SGLang app example."""

from flyteplugins.sglang import SGLangAppEnvironment
import flyte

# {{docs-fragment basic-sglang-app}}
sglang_app = SGLangAppEnvironment(
    name="my-sglang-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by SGLang
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment basic-sglang-app}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(sglang_app)
    print(f"Deployed SGLang app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/sglang/basic_sglang.py*

## Using prefetched models

You can use models prefetched with `flyte.prefetch`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-sglang>=2.0.0b45",
# ]
# ///

"""SGLang app using prefetched models."""

from flyteplugins.sglang import SGLangAppEnvironment
import flyte

# {{docs-fragment prefetch}}

# Use the prefetched model
sglang_app = SGLangAppEnvironment(
    name="my-sglang-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # this is a placeholder
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="10Gi"),
    stream_model=True,  # Stream model directly from blob store to GPU
    requires_auth=False,
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Prefetch the model first
    run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
    run.wait()

    app = flyte.serve(
        sglang_app.clone_with(
            sglang_app.name,
            model_hf_path=None,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
        )
    )
    print(f"Deployed SGLang app: {app.url}")
# {{/docs-fragment prefetch}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/sglang/sglang_with_prefetch.py*

## Model streaming

`SGLangAppEnvironment` supports streaming models directly from blob storage to GPU memory, reducing startup time.
When `stream_model=True` and the `model_path` argument is provided with either a `flyte.io.Dir` or `RunOutput` pointing
to a path in object store:

- Model weights stream directly from storage to GPU
- Faster startup time (no full download required)
- Lower disk space requirements

> [!NOTE]
> The contents of the model directory must be compatible with the SGLang-supported formats, e.g. the HuggingFace model
> serialization format.

## Custom SGLang arguments

Use `extra_args` to pass additional arguments to SGLang:

```python
sglang_app = SGLangAppEnvironment(
    name="custom-sglang-app",
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    extra_args=[
        "--max-model-len", "8192",  # Maximum context length
        "--mem-fraction-static", "0.8",  # Memory fraction for static allocation
        "--trust-remote-code",  # Trust remote code in models
    ],
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    # ...
)
```

See the [SGLang server arguments documentation](https://docs.sglang.io/advanced_features/server_arguments.html) for all available options.

## Using the OpenAI-compatible API

Once deployed, your SGLang app exposes an OpenAI-compatible API:

```python
from openai import OpenAI

client = OpenAI(
    base_url="https://your-app-url/v1",  # SGLang endpoint
    api_key="your-api-key",  # If you passed an --api-key argument
)

response = client.chat.completions.create(
    model="qwen3-0.6b",  # Your model_id
    messages=[
        {"role": "user", "content": "Hello, how are you?"}
    ],
)

print(response.choices[0].message.content)
```

> [!TIP]
> If you passed an `--api-key` argument, you can use the `api_key` parameter to authenticate your requests.
> See [secret-based authentication](../build-apps/secret-based-authentication#deploy-sglang-app-with-authentication) for more details on how to pass auth secrets to your app.

## Multi-GPU inference (Tensor parallelism)

For larger models, use multiple GPUs with tensor parallelism:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-sglang>=2.0.0b45",
# ]
# ///

"""SGLang app with multi-GPU tensor parallelism."""

from flyteplugins.sglang import SGLangAppEnvironment
import flyte

# {{docs-fragment multi-gpu}}
sglang_app = SGLangAppEnvironment(
    name="multi-gpu-sglang-app",
    model_hf_path="meta-llama/Llama-2-70b-hf",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # 4 GPUs for tensor parallelism
        disk="100Gi",
    ),
    extra_args=[
        "--tp", "4",  # Tensor parallelism size (4 GPUs)
        "--max-model-len", "4096",
        "--mem-fraction-static", "0.9",
    ],
    requires_auth=False,
)
# {{/docs-fragment multi-gpu}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(sglang_app)
    print(f"Deployed SGLang app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/sglang/sglang_multi_gpu.py*

The tensor parallelism size (`--tp`) should match the number of GPUs specified in resources.

## Model sharding with prefetch

You can prefetch and shard models for multi-GPU inference using SGLang's sharding:

```python
# Prefetch with sharding configuration
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    accelerator="L40s:4",
    shard_config=flyte.prefetch.ShardConfig(
        engine="vllm",
        args=flyte.prefetch.VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
)
run.wait()

# Use the sharded model
sglang_app = SGLangAppEnvironment(
    name="sharded-sglang-app",
    model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
    model_id="llama-2-70b",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4", disk="100Gi"),
    extra_args=["--tp", "4"],
    stream_model=True,
)
```

See [Prefetching models](../serve-and-deploy-apps/prefetching-models) for more details on sharding.

## Autoscaling

SGLang apps work well with autoscaling:

```python
sglang_app = SGLangAppEnvironment(
    name="autoscaling-sglang-app",
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),  # Scale to zero when idle
        scaledown_after=600,  # 10 minutes idle before scaling down
    ),
    # ...
)
```

## Structured generation

SGLang is particularly well-suited for structured generation tasks. The deployed app supports standard OpenAI API calls, and you can use SGLang's advanced features through the API.

## Best practices

1. **Use prefetching**: Prefetch models for faster deployment and better reproducibility
2. **Enable streaming**: Use `stream_model=True` to reduce startup time and disk usage
3. **Right-size GPUs**: Match GPU memory to model size
4. **Use tensor parallelism**: For large models, use multiple GPUs with `--tp`
5. **Set autoscaling**: Use appropriate idle TTL to balance cost and performance
6. **Configure memory**: Use `--mem-fraction-static` to control memory allocation
7. **Limit context length**: Use `--max-model-len` for smaller models to reduce memory usage

## Troubleshooting

**Model loading fails:**

- Verify GPU memory is sufficient for the model
- Check that the model path or HuggingFace path is correct
- Review container logs for detailed error messages

**Out of memory errors:**

- Reduce `--max-model-len`
- Lower `--mem-fraction-static`
- Use a smaller model or more GPUs

**Slow startup:**

- Enable `stream_model=True` for faster loading
- Prefetch models before deployment
- Use faster storage backends

## API reference

See the [SGLang API reference](../../../api-reference/integrations/sglang/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/ollama-app ===

# Ollama app

[Ollama](https://ollama.com/) is a lightweight runtime for serving open large language models (LLMs) locally, with a built-in OpenAI-compatible API. It is a good fit for smaller models and quick local-style serving, and complements the higher-throughput [vLLM](./vllm-app) and [SGLang](./sglang-app) integrations.

Unlike vLLM and SGLang, Ollama has no dedicated `*AppEnvironment` plugin. Instead, you serve it with the generic [`AppEnvironment`](../build-apps/single-script-apps): an image that installs Ollama plus a small entrypoint that launches `ollama serve` and pulls the model on startup.

## Installation

Ollama needs no Flyte plugin — it is installed into the app image. Start from the default Flyte base image and add the Ollama binary with a build command:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A simple Ollama serving app example.

Unlike vLLM and SGLang, Ollama has no dedicated ``*AppEnvironment`` plugin, so it
is served with the generic ``flyte.app.AppEnvironment``: an image that installs
Ollama plus a small ``--server`` entrypoint that launches ``ollama serve`` and
pulls the model on startup. Ollama exposes an OpenAI-compatible API, so clients
call it exactly like the vLLM / SGLang apps.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

# Any tag from https://ollama.com/library. Small models run fine on CPU; larger
# ones benefit from the GPU requested below.
MODEL = "qwen3:0.6b"

# Bind Ollama to the app port so the platform can route to it.
PORT = 8080

file_name = Path(__file__).name

# {{docs-fragment ollama-image}}
# Install Ollama on top of the default Flyte base image. `install.sh` drops the
# `ollama` binary into /usr/local/bin (no systemd is needed inside a container).
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
)
# {{/docs-fragment ollama-image}}

# {{docs-fragment ollama-app}}
ollama_app = flyte.app.AppEnvironment(
    name="ollama-app",
    image=image,
    args=["python", file_name, "--server"],
    port=PORT,
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU accelerates inference; omit to run small models on CPU
        disk="20Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment ollama-app}}

# {{docs-fragment server}}
def serve() -> None:
    """Start `ollama serve`, wait for it, pull the model, then block."""
    # Bind to all interfaces so the platform can route to the server.
    server = subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{PORT}"},
    )

    # Wait for the server to accept connections.
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)

    # Pull the model so the OpenAI-compatible endpoint can serve it.
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{PORT}"},
        check=True,
    )
    print(f"Ollama serving '{MODEL}' on port {PORT}")

    server.wait()
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--server" in sys.argv:
        serve()
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(ollama_app)
        print(f"Deployed Ollama app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/ollama/basic_ollama.py*

The `install.sh` script drops the `ollama` binary into `/usr/local/bin`; no systemd service is needed inside a container.

## Basic Ollama app

Define the app with a GPU-backed `AppEnvironment`. The `args` run the script's `--server` entrypoint, which starts Ollama and pulls the model:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A simple Ollama serving app example.

Unlike vLLM and SGLang, Ollama has no dedicated ``*AppEnvironment`` plugin, so it
is served with the generic ``flyte.app.AppEnvironment``: an image that installs
Ollama plus a small ``--server`` entrypoint that launches ``ollama serve`` and
pulls the model on startup. Ollama exposes an OpenAI-compatible API, so clients
call it exactly like the vLLM / SGLang apps.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

# Any tag from https://ollama.com/library. Small models run fine on CPU; larger
# ones benefit from the GPU requested below.
MODEL = "qwen3:0.6b"

# Bind Ollama to the app port so the platform can route to it.
PORT = 8080

file_name = Path(__file__).name

# {{docs-fragment ollama-image}}
# Install Ollama on top of the default Flyte base image. `install.sh` drops the
# `ollama` binary into /usr/local/bin (no systemd is needed inside a container).
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
)
# {{/docs-fragment ollama-image}}

# {{docs-fragment ollama-app}}
ollama_app = flyte.app.AppEnvironment(
    name="ollama-app",
    image=image,
    args=["python", file_name, "--server"],
    port=PORT,
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU accelerates inference; omit to run small models on CPU
        disk="20Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment ollama-app}}

# {{docs-fragment server}}
def serve() -> None:
    """Start `ollama serve`, wait for it, pull the model, then block."""
    # Bind to all interfaces so the platform can route to the server.
    server = subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{PORT}"},
    )

    # Wait for the server to accept connections.
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)

    # Pull the model so the OpenAI-compatible endpoint can serve it.
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{PORT}"},
        check=True,
    )
    print(f"Ollama serving '{MODEL}' on port {PORT}")

    server.wait()
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--server" in sys.argv:
        serve()
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(ollama_app)
        print(f"Deployed Ollama app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/ollama/basic_ollama.py*

The `--server` entrypoint binds Ollama to the app port, waits for it to be ready, and pulls the model so the OpenAI-compatible endpoint can serve it:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A simple Ollama serving app example.

Unlike vLLM and SGLang, Ollama has no dedicated ``*AppEnvironment`` plugin, so it
is served with the generic ``flyte.app.AppEnvironment``: an image that installs
Ollama plus a small ``--server`` entrypoint that launches ``ollama serve`` and
pulls the model on startup. Ollama exposes an OpenAI-compatible API, so clients
call it exactly like the vLLM / SGLang apps.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

# Any tag from https://ollama.com/library. Small models run fine on CPU; larger
# ones benefit from the GPU requested below.
MODEL = "qwen3:0.6b"

# Bind Ollama to the app port so the platform can route to it.
PORT = 8080

file_name = Path(__file__).name

# {{docs-fragment ollama-image}}
# Install Ollama on top of the default Flyte base image. `install.sh` drops the
# `ollama` binary into /usr/local/bin (no systemd is needed inside a container).
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
)
# {{/docs-fragment ollama-image}}

# {{docs-fragment ollama-app}}
ollama_app = flyte.app.AppEnvironment(
    name="ollama-app",
    image=image,
    args=["python", file_name, "--server"],
    port=PORT,
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU accelerates inference; omit to run small models on CPU
        disk="20Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment ollama-app}}

# {{docs-fragment server}}
def serve() -> None:
    """Start `ollama serve`, wait for it, pull the model, then block."""
    # Bind to all interfaces so the platform can route to the server.
    server = subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{PORT}"},
    )

    # Wait for the server to accept connections.
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)

    # Pull the model so the OpenAI-compatible endpoint can serve it.
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{PORT}"},
        check=True,
    )
    print(f"Ollama serving '{MODEL}' on port {PORT}")

    server.wait()
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--server" in sys.argv:
        serve()
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(ollama_app)
        print(f"Deployed Ollama app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/ollama/basic_ollama.py*

Deploy it with `flyte.serve`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""A simple Ollama serving app example.

Unlike vLLM and SGLang, Ollama has no dedicated ``*AppEnvironment`` plugin, so it
is served with the generic ``flyte.app.AppEnvironment``: an image that installs
Ollama plus a small ``--server`` entrypoint that launches ``ollama serve`` and
pulls the model on startup. Ollama exposes an OpenAI-compatible API, so clients
call it exactly like the vLLM / SGLang apps.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

# Any tag from https://ollama.com/library. Small models run fine on CPU; larger
# ones benefit from the GPU requested below.
MODEL = "qwen3:0.6b"

# Bind Ollama to the app port so the platform can route to it.
PORT = 8080

file_name = Path(__file__).name

# {{docs-fragment ollama-image}}
# Install Ollama on top of the default Flyte base image. `install.sh` drops the
# `ollama` binary into /usr/local/bin (no systemd is needed inside a container).
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
)
# {{/docs-fragment ollama-image}}

# {{docs-fragment ollama-app}}
ollama_app = flyte.app.AppEnvironment(
    name="ollama-app",
    image=image,
    args=["python", file_name, "--server"],
    port=PORT,
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU accelerates inference; omit to run small models on CPU
        disk="20Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
# {{/docs-fragment ollama-app}}

# {{docs-fragment server}}
def serve() -> None:
    """Start `ollama serve`, wait for it, pull the model, then block."""
    # Bind to all interfaces so the platform can route to the server.
    server = subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{PORT}"},
    )

    # Wait for the server to accept connections.
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)

    # Pull the model so the OpenAI-compatible endpoint can serve it.
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{PORT}"},
        check=True,
    )
    print(f"Ollama serving '{MODEL}' on port {PORT}")

    server.wait()
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--server" in sys.argv:
        serve()
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(ollama_app)
        print(f"Deployed Ollama app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/ollama/basic_ollama.py*

## Choosing a model

`MODEL` can be any tag from the [Ollama library](https://ollama.com/library) — for example `qwen3:0.6b`, `llama3.2:1b`, or `gemma3:1b`. Larger models need more GPU memory and disk; size the `resources` accordingly.

> [!NOTE]
> Small models run comfortably on CPU. To run without a GPU, drop the `gpu` field from `resources`. A GPU is recommended for larger models or higher throughput.

## Using the OpenAI-compatible API

Once deployed, the Ollama app exposes an OpenAI-compatible API. Call it exactly like the vLLM or SGLang apps:

CODE0

## Chat UI with Streamlit

To put a browser UI in front of the model, run Ollama internally and expose a [Streamlit](./streamlit-app) chat interface instead of the raw API. Only the Streamlit port is exposed; the browser talks to Streamlit, and Streamlit talks to Ollama over localhost:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "openai",
# ]
# ///

"""An Ollama chat app fronted by a Streamlit UI.

Ollama runs internally as an OpenAI-compatible server; a Streamlit chat
interface fronts it. Only the Streamlit port is exposed to the platform — the
browser talks to Streamlit, and Streamlit talks to Ollama over localhost.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

MODEL = "qwen3:0.6b"
OLLAMA_PORT = 11434  # internal only, not exposed
APP_PORT = 8080  # Streamlit UI, exposed to the platform

file_name = Path(__file__).name

# {{docs-fragment app-env}}
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
    .with_pip_packages("streamlit==1.41.1", "openai")
)

app_env = flyte.app.AppEnvironment(
    name="ollama-streamlit",
    image=image,
    args=["python", file_name, "--server"],
    port=APP_PORT,
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="20Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=300),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment ui}}
def render_ui() -> None:
    """The Streamlit chat UI. Talks to the local Ollama OpenAI-compatible API."""
    import streamlit as st
    from openai import OpenAI

    st.set_page_config(page_title="Ollama Chat", page_icon="🦙")
    st.title("🦙 Ollama Chat")

    client = OpenAI(base_url=f"http://127.0.0.1:{OLLAMA_PORT}/v1", api_key="ollama")

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])

    if prompt := st.chat_input("Ask something..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)
        response = client.chat.completions.create(
            model=MODEL, messages=st.session_state.messages
        )
        answer = response.choices[0].message.content
        st.session_state.messages.append({"role": "assistant", "content": answer})
        st.chat_message("assistant").write(answer)
# {{/docs-fragment ui}}

# {{docs-fragment server}}
def start_ollama() -> None:
    """Launch `ollama serve` in the background and pull the model."""
    subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{OLLAMA_PORT}"},
    )
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{OLLAMA_PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{OLLAMA_PORT}"},
        check=True,
    )
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--ui" in sys.argv:
        # Re-entry: Streamlit runs this file to render the UI.
        render_ui()
    elif "--server" in sys.argv:
        # Container entrypoint: start Ollama, then hand the port to Streamlit.
        start_ollama()
        subprocess.run(
            [
                "streamlit", "run", file_name,
                "--server.port", str(APP_PORT),
                "--server.address", "0.0.0.0",
                "--", "--ui",
            ],
            check=True,
        )
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"Deployed app: {app.url}")
# {{/docs-fragment deploy}}
CODE1
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "openai",
# ]
# ///

"""An Ollama chat app fronted by a Streamlit UI.

Ollama runs internally as an OpenAI-compatible server; a Streamlit chat
interface fronts it. Only the Streamlit port is exposed to the platform — the
browser talks to Streamlit, and Streamlit talks to Ollama over localhost.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

MODEL = "qwen3:0.6b"
OLLAMA_PORT = 11434  # internal only, not exposed
APP_PORT = 8080  # Streamlit UI, exposed to the platform

file_name = Path(__file__).name

# {{docs-fragment app-env}}
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
    .with_pip_packages("streamlit==1.41.1", "openai")
)

app_env = flyte.app.AppEnvironment(
    name="ollama-streamlit",
    image=image,
    args=["python", file_name, "--server"],
    port=APP_PORT,
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="20Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=300),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment ui}}
def render_ui() -> None:
    """The Streamlit chat UI. Talks to the local Ollama OpenAI-compatible API."""
    import streamlit as st
    from openai import OpenAI

    st.set_page_config(page_title="Ollama Chat", page_icon="🦙")
    st.title("🦙 Ollama Chat")

    client = OpenAI(base_url=f"http://127.0.0.1:{OLLAMA_PORT}/v1", api_key="ollama")

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])

    if prompt := st.chat_input("Ask something..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)
        response = client.chat.completions.create(
            model=MODEL, messages=st.session_state.messages
        )
        answer = response.choices[0].message.content
        st.session_state.messages.append({"role": "assistant", "content": answer})
        st.chat_message("assistant").write(answer)
# {{/docs-fragment ui}}

# {{docs-fragment server}}
def start_ollama() -> None:
    """Launch `ollama serve` in the background and pull the model."""
    subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{OLLAMA_PORT}"},
    )
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{OLLAMA_PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{OLLAMA_PORT}"},
        check=True,
    )
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--ui" in sys.argv:
        # Re-entry: Streamlit runs this file to render the UI.
        render_ui()
    elif "--server" in sys.argv:
        # Container entrypoint: start Ollama, then hand the port to Streamlit.
        start_ollama()
        subprocess.run(
            [
                "streamlit", "run", file_name,
                "--server.port", str(APP_PORT),
                "--server.address", "0.0.0.0",
                "--", "--ui",
            ],
            check=True,
        )
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"Deployed app: {app.url}")
# {{/docs-fragment deploy}}
CODE2
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "openai",
# ]
# ///

"""An Ollama chat app fronted by a Streamlit UI.

Ollama runs internally as an OpenAI-compatible server; a Streamlit chat
interface fronts it. Only the Streamlit port is exposed to the platform — the
browser talks to Streamlit, and Streamlit talks to Ollama over localhost.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

MODEL = "qwen3:0.6b"
OLLAMA_PORT = 11434  # internal only, not exposed
APP_PORT = 8080  # Streamlit UI, exposed to the platform

file_name = Path(__file__).name

# {{docs-fragment app-env}}
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
    .with_pip_packages("streamlit==1.41.1", "openai")
)

app_env = flyte.app.AppEnvironment(
    name="ollama-streamlit",
    image=image,
    args=["python", file_name, "--server"],
    port=APP_PORT,
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="20Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=300),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment ui}}
def render_ui() -> None:
    """The Streamlit chat UI. Talks to the local Ollama OpenAI-compatible API."""
    import streamlit as st
    from openai import OpenAI

    st.set_page_config(page_title="Ollama Chat", page_icon="🦙")
    st.title("🦙 Ollama Chat")

    client = OpenAI(base_url=f"http://127.0.0.1:{OLLAMA_PORT}/v1", api_key="ollama")

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])

    if prompt := st.chat_input("Ask something..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)
        response = client.chat.completions.create(
            model=MODEL, messages=st.session_state.messages
        )
        answer = response.choices[0].message.content
        st.session_state.messages.append({"role": "assistant", "content": answer})
        st.chat_message("assistant").write(answer)
# {{/docs-fragment ui}}

# {{docs-fragment server}}
def start_ollama() -> None:
    """Launch `ollama serve` in the background and pull the model."""
    subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{OLLAMA_PORT}"},
    )
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{OLLAMA_PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{OLLAMA_PORT}"},
        check=True,
    )
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--ui" in sys.argv:
        # Re-entry: Streamlit runs this file to render the UI.
        render_ui()
    elif "--server" in sys.argv:
        # Container entrypoint: start Ollama, then hand the port to Streamlit.
        start_ollama()
        subprocess.run(
            [
                "streamlit", "run", file_name,
                "--server.port", str(APP_PORT),
                "--server.address", "0.0.0.0",
                "--", "--ui",
            ],
            check=True,
        )
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"Deployed app: {app.url}")
# {{/docs-fragment deploy}}
CODE3
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "streamlit",
#    "openai",
# ]
# ///

"""An Ollama chat app fronted by a Streamlit UI.

Ollama runs internally as an OpenAI-compatible server; a Streamlit chat
interface fronts it. Only the Streamlit port is exposed to the platform — the
browser talks to Streamlit, and Streamlit talks to Ollama over localhost.
"""

import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import flyte
import flyte.app

MODEL = "qwen3:0.6b"
OLLAMA_PORT = 11434  # internal only, not exposed
APP_PORT = 8080  # Streamlit UI, exposed to the platform

file_name = Path(__file__).name

# {{docs-fragment app-env}}
image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl")
    .with_commands(["curl -fsSL https://ollama.com/install.sh | sh"])
    .with_pip_packages("streamlit==1.41.1", "openai")
)

app_env = flyte.app.AppEnvironment(
    name="ollama-streamlit",
    image=image,
    args=["python", file_name, "--server"],
    port=APP_PORT,
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="20Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=300),
    requires_auth=False,
)
# {{/docs-fragment app-env}}

# {{docs-fragment ui}}
def render_ui() -> None:
    """The Streamlit chat UI. Talks to the local Ollama OpenAI-compatible API."""
    import streamlit as st
    from openai import OpenAI

    st.set_page_config(page_title="Ollama Chat", page_icon="🦙")
    st.title("🦙 Ollama Chat")

    client = OpenAI(base_url=f"http://127.0.0.1:{OLLAMA_PORT}/v1", api_key="ollama")

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])

    if prompt := st.chat_input("Ask something..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)
        response = client.chat.completions.create(
            model=MODEL, messages=st.session_state.messages
        )
        answer = response.choices[0].message.content
        st.session_state.messages.append({"role": "assistant", "content": answer})
        st.chat_message("assistant").write(answer)
# {{/docs-fragment ui}}

# {{docs-fragment server}}
def start_ollama() -> None:
    """Launch `ollama serve` in the background and pull the model."""
    subprocess.Popen(
        ["ollama", "serve"],
        env={**os.environ, "OLLAMA_HOST": f"0.0.0.0:{OLLAMA_PORT}"},
    )
    for _ in range(60):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{OLLAMA_PORT}/api/version", timeout=2)
            break
        except Exception:
            time.sleep(1)
    subprocess.run(
        ["ollama", "pull", MODEL],
        env={**os.environ, "OLLAMA_HOST": f"127.0.0.1:{OLLAMA_PORT}"},
        check=True,
    )
# {{/docs-fragment server}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    if "--ui" in sys.argv:
        # Re-entry: Streamlit runs this file to render the UI.
        render_ui()
    elif "--server" in sys.argv:
        # Container entrypoint: start Ollama, then hand the port to Streamlit.
        start_ollama()
        subprocess.run(
            [
                "streamlit", "run", file_name,
                "--server.port", str(APP_PORT),
                "--server.address", "0.0.0.0",
                "--", "--ui",
            ],
            check=True,
        )
    else:
        flyte.init_from_config(root_dir=Path(__file__).parent)
        app = flyte.serve(app_env)
        print(f"Deployed app: {app.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/ollama/ollama_streamlit.py*

## Authentication

The generic `AppEnvironment` uses Union's platform-level authentication. Leave `requires_auth=True` (the default) to require an authenticated caller, or set `requires_auth=False` for a public endpoint. Ollama has no built-in API-key argument of its own, so prefer platform auth over exposing a public endpoint. For app-managed authentication (verifying a Bearer token yourself with a Flyte secret), see [Secret-based authentication](../build-apps/secret-based-authentication).

## Autoscaling

Ollama apps work well with scale-to-zero, so an idle model server costs nothing:

CODE4

Because the model is pulled at startup, a larger `scaledown_after` avoids re-pulling on every cold start.

## Best practices

1. **Right-size resources**: Match GPU memory and disk to the model. Small models run on CPU; drop the `gpu` field to save cost.
2. **Bake big models into the image**: For faster, more reproducible cold starts, `ollama pull` the model in a build command instead of at startup.
3. **Use scale-to-zero**: Set an appropriate `scaledown_after` to balance cost against cold-start latency.
4. **Prefer platform auth**: Ollama has no native API-key auth, so rely on `requires_auth` rather than exposing a public endpoint.
5. **Pick the right runtime**: Use Ollama for lightweight or local-style serving; reach for [vLLM](./vllm-app) or [SGLang](./sglang-app) for high-throughput production inference.

## Troubleshooting

**Model pull fails or times out:**

- Verify the `MODEL` tag exists in the [Ollama library](https://ollama.com/library)
- Increase `disk` in `resources` for larger models
- Review container logs for the `ollama pull` output

**Server not reachable:**

- Confirm Ollama is bound to `0.0.0.0` on the app port via `OLLAMA_HOST`
- Check that the app `port` matches the port Ollama serves on

**Slow first response:**

- The model is pulled on startup; use a larger `scaledown_after` or bake the model into the image
- Use a smaller model, or add a GPU for faster inference


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/native-app-integrations/flyte-webhook ===

# Flyte webhook

`FlyteWebhookAppEnvironment` is a pre-built `FastAPIAppEnvironment` that exposes
HTTP endpoints for common Flyte operations. Instead of writing
your own FastAPI routes to interact with the control plane, you get a
ready-to-deploy webhook service with a single constructor call.

## Available endpoints

The webhook provides endpoints for the following operations:

| Group | Endpoints | Description |
|---|---|---|
| **core** | `GET /health`, `GET /me` | Health check and authenticated user info |
| **task** | `POST /run-task/{domain}/{project}/{name}`, `GET /task/{domain}/{project}/{name}` | Run tasks and retrieve task metadata |
| **run** | `GET /run/{name}`, `GET /run/{name}/io`, `POST /run/{name}/abort` | Get run status, inputs/outputs, and abort runs |
| **app** | `GET /app/{name}`, `POST /app/{name}/activate`, `POST /app/{name}/deactivate`, `POST /app/{name}/call` | Manage apps and call other app endpoints |
| **trigger** | `POST /trigger/{task_name}/{trigger_name}/activate`, `POST /trigger/{task_name}/{trigger_name}/deactivate` | Activate and deactivate triggers |
| **build** | `POST /build-image` | Build container images |
| **prefetch** | `POST /prefetch/hf-model`, `GET /prefetch/hf-model/{run_name}`, `GET /prefetch/hf-model/{run_name}/io`, `POST /prefetch/hf-model/{run_name}/abort` | Prefetch HuggingFace models |

All endpoints except `/health`, `/docs`, and `/openapi.json` use passthrough
authentication, forwarding the caller's credentials to the
Flyte control plane.

## Basic usage

Create a webhook with all endpoints enabled:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

Deploy and activate it:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

Once running, the webhook exposes OpenAPI docs at `{endpoint}/docs` (Swagger UI)
and `{endpoint}/redoc`.

## Filtering endpoints

You can restrict which endpoints the webhook exposes using either **endpoint
groups** or **individual endpoints**.

### Endpoint groups

Enable groups of related endpoints with `endpoint_groups`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

Available groups: `all`, `core`, `task`, `run`, `app`, `trigger`, `build`, `prefetch`.

### Individual endpoints

For finer control, specify exact endpoints with `endpoints`:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

> [!NOTE]
> You cannot specify both `endpoint_groups` and `endpoints` at the same time. Use
> one or the other.

## Allow-listing

Restrict which resources the webhook can access using allow-lists.

### Task allow-list

Limit which tasks can be run or queried through the webhook:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

Task identifiers support three formats:

- `domain/project/name`: exact match
- `project/name`: matches any domain
- `name`: matches any domain and project

### App allow-list

Limit which apps can be managed through the webhook:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

### Trigger allow-list

Limit which triggers can be activated or deactivated:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

Trigger identifiers support two formats:

- `task_name/trigger_name`: exact match
- `trigger_name`: matches any task

## Calling the webhook

Authenticate requests with a Flyte API key passed as a Bearer
token:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "fastapi",
#    "uvicorn",
#    "httpx",
# ]
# ///

"""Examples showing how to use FlyteWebhookAppEnvironment."""

import logging

import flyte
import flyte.app
from flyte.app.extras import FlyteWebhookAppEnvironment

# {{docs-fragment basic-webhook}}
webhook_env = FlyteWebhookAppEnvironment(
    name="my-webhook",
    title="My Flyte Webhook",
    description="A pre-built webhook service for Flyte operations",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    scaling=flyte.app.Scaling(replicas=1),
)
# {{/docs-fragment basic-webhook}}

# {{docs-fragment endpoint-groups}}
task_runner_webhook = FlyteWebhookAppEnvironment(
    name="task-runner-webhook",
    title="Task Runner Webhook",
    endpoint_groups=["core", "task", "run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment endpoint-groups}}

# {{docs-fragment individual-endpoints}}
minimal_webhook = FlyteWebhookAppEnvironment(
    name="minimal-webhook",
    title="Minimal Webhook",
    endpoints=["health", "run_task", "get_run"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment individual-endpoints}}

# {{docs-fragment task-allowlist}}
restricted_webhook = FlyteWebhookAppEnvironment(
    name="restricted-webhook",
    title="Restricted Webhook",
    endpoint_groups=["core", "task", "run"],
    task_allowlist=[
        "production/my-project/allowed-task",
        "my-project/another-task",
        "any-domain-task",
    ],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment task-allowlist}}

# {{docs-fragment app-allowlist}}
app_manager_webhook = FlyteWebhookAppEnvironment(
    name="app-manager-webhook",
    title="App Manager Webhook",
    endpoint_groups=["core", "app"],
    app_allowlist=["my-app", "another-app"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment app-allowlist}}

# {{docs-fragment trigger-allowlist}}
trigger_manager_webhook = FlyteWebhookAppEnvironment(
    name="trigger-manager-webhook",
    title="Trigger Manager Webhook",
    endpoint_groups=["core", "trigger"],
    trigger_allowlist=["my-task/my-trigger", "another-trigger"],
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
)
# {{/docs-fragment trigger-allowlist}}

# {{docs-fragment deploy-webhook}}
if __name__ == "__main__":
    import os

    import httpx

    flyte.init_from_config(log_level=logging.DEBUG)

    served_app = flyte.serve(webhook_env)
    url = served_app.url
    endpoint = served_app.endpoint
    print(f"Webhook is served on {url}")
    print(f"OpenAPI docs available at: {endpoint}/docs")

    served_app.activate(wait=True)
# {{/docs-fragment deploy-webhook}}

# {{docs-fragment call-webhook}}
    token = os.getenv("FLYTE_API_KEY")
    if not token:
        raise ValueError("FLYTE_API_KEY not set. Obtain with: flyte get api-key")

    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "flyte-webhook-client/1.0",
    }

    with httpx.Client(headers=headers) as client:
        # Health check (no auth required)
        health = client.get(f"{endpoint}/health")
        print(f"/health: {health.json()}")

        # Get current user info (requires auth)
        me = client.get(f"{endpoint}/me")
        print(f"/me: {me.json()}")

        # Run a task
        resp = client.post(
            f"{endpoint}/run-task/development/my-project/my-task",
            json={"x": 42, "y": "hello"},
        )
        result = resp.json()
        print(f"Run task: {result}")

        # Check run status
        run_name = result["name"]
        run = client.get(f"{endpoint}/run/{run_name}")
        print(f"Run status: {run.json()}")
# {{/docs-fragment call-webhook}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/flyte_webhook_examples.py*

## Authentication

`FlyteWebhookAppEnvironment` uses `FastAPIPassthroughAuthMiddleware`, which
extracts the caller's auth token from the `Authorization` header and sets up
a Flyte context so that every control plane call (e.g.
`remote.Task.get`, `flyte.run`) runs with the caller's identity.

The `/health`, `/docs`, `/openapi.json`, and `/redoc` endpoints are excluded
from authentication.

## Self-reference protection

App endpoints (`get_app`, `activate_app`, `deactivate_app`, `call_app`) prevent
the webhook from targeting itself. Attempting to activate, deactivate, or call
the webhook's own name returns a `400 Bad Request` error.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps ===

# Serve and deploy apps

Flyte provides two main ways to deploy apps: **serve** (for development) and **deploy** (for production). This section covers both methods and their differences.

## Serve vs deploy

### `flyte serve`

Serving is designed for development and iteration:

- **Dynamic parameter modification**: You can override app parameters when serving
- **Quick iteration**: Faster feedback loop for development
- **Interactive**: Better suited for testing and experimentation

### `flyte deploy`

Deployment is designed for production use:

- **Immutable**: Apps are deployed with fixed configurations
- **Production-ready**: Optimized for stability and reproducibility

## Using Python SDK

### Serve

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve and deploy examples for the _index.md documentation."""

import flyte
import flyte.app

# {{docs-fragment serve-example}}
app_env = flyte.app.AppEnvironment(
    name="my-app",
    image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"),
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"Served at: {app.url}")
# {{/docs-fragment serve-example}}

# {{docs-fragment deploy-example}}
app_env = flyte.app.AppEnvironment(
    name="my-app",
    image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"),
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)
    # Access deployed app URL from the deployment
    for deployed_env in deployments[0].envs.values():
        print(f"Deployed: {deployed_env.deployed_app.url}")
# {{/docs-fragment deploy-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_and_deploy_examples.py*

### Deploy

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve and deploy examples for the _index.md documentation."""

import flyte
import flyte.app

# {{docs-fragment serve-example}}
app_env = flyte.app.AppEnvironment(
    name="my-app",
    image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"),
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"Served at: {app.url}")
# {{/docs-fragment serve-example}}

# {{docs-fragment deploy-example}}
app_env = flyte.app.AppEnvironment(
    name="my-app",
    image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"),
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)
    # Access deployed app URL from the deployment
    for deployed_env in deployments[0].envs.values():
        print(f"Deployed: {deployed_env.deployed_app.url}")
# {{/docs-fragment deploy-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_and_deploy_examples.py*

## Using the CLI

### Serve

```bash
flyte serve path/to/app.py app_env
```

### Deploy

```bash
flyte deploy path/to/app.py app_env
```

## Next steps

- **Apps > Serve and deploy apps > How app serving works**: Understanding the serve process and configuration options
- **Apps > Serve and deploy apps > How app deployment works**: Understanding the deploy process and configuration options
- **Apps > Serve and deploy apps > Activating and deactivating apps**: Managing app lifecycle
- **Get started > Core concepts > Basic project: RAG**: Build a RAG embedding pipeline and semantic search app with Streamlit
- **Apps > Serve and deploy apps > Prefetching models**: Download and shard HuggingFace models for vLLM and SGLang

## Subpages
- **Apps > Serve and deploy apps > How app serving works**
- **Apps > Serve and deploy apps > How app custom domains work**
- **Apps > Serve and deploy apps > How app deployment works**
- **Apps > Serve and deploy apps > Activating and deactivating apps**
- **Apps > Serve and deploy apps > Prefetching models**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps/how-app-serving-works ===

# How app serving works

Serving is the recommended way to deploy apps during development. It provides a faster feedback loop and allows you to dynamically modify parameters.

## Overview

When you serve an app, the following happens:

1. **Code bundling**: Your app code is bundled and prepared
2. **Image building**: Container images are built (if needed)
3. **Deployment**: The app is deployed to your Flyte cluster
4. **Activation**: The app is automatically activated and ready to use
5. **URL generation**: A URL is generated for accessing the app

## Using the Python SDK

The simplest way to serve an app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve examples for the how-app-serving-works.md documentation."""

import logging
import flyte
import flyte.app

# {{docs-fragment basic-serve}}
app_env = flyte.app.AppEnvironment(
    name="my-dev-app",
    parameters=[flyte.app.Parameter(name="model_path", value="s3://bucket/models/model.pkl")],
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment basic-serve}}

# {{docs-fragment override-parameters}}
app = flyte.with_servecontext(
    input_values={
        "my-dev-app": {
            "model_path": "s3://bucket/models/test-model.pkl",
        }
    }
).serve(app_env)
# {{/docs-fragment override-parameters}}

# {{docs-fragment advanced-serving}}
app = flyte.with_servecontext(
    version="v1.0.0",
    project="my-project",
    domain="development",
    env_vars={"LOG_LEVEL": "DEBUG"},
    input_values={"app-name": {"input": "value"}},
    cluster_pool="dev-pool",
    log_level=logging.INFO,
    log_format="json",
    dry_run=False,
).serve(app_env)
# {{/docs-fragment advanced-serving}}

# {{docs-fragment return-value}}
app = flyte.serve(app_env)
print(f"URL: {app.url}")
print(f"Endpoint: {app.endpoint}")
print(f"Status: {app.deployment_status}")
# {{/docs-fragment return-value}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_examples.py*

## Overriding parameters

One key advantage of serving is the ability to override parameters dynamically:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve examples for the how-app-serving-works.md documentation."""

import logging
import flyte
import flyte.app

# {{docs-fragment basic-serve}}
app_env = flyte.app.AppEnvironment(
    name="my-dev-app",
    parameters=[flyte.app.Parameter(name="model_path", value="s3://bucket/models/model.pkl")],
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment basic-serve}}

# {{docs-fragment override-parameters}}
app = flyte.with_servecontext(
    input_values={
        "my-dev-app": {
            "model_path": "s3://bucket/models/test-model.pkl",
        }
    }
).serve(app_env)
# {{/docs-fragment override-parameters}}

# {{docs-fragment advanced-serving}}
app = flyte.with_servecontext(
    version="v1.0.0",
    project="my-project",
    domain="development",
    env_vars={"LOG_LEVEL": "DEBUG"},
    input_values={"app-name": {"input": "value"}},
    cluster_pool="dev-pool",
    log_level=logging.INFO,
    log_format="json",
    dry_run=False,
).serve(app_env)
# {{/docs-fragment advanced-serving}}

# {{docs-fragment return-value}}
app = flyte.serve(app_env)
print(f"URL: {app.url}")
print(f"Endpoint: {app.endpoint}")
print(f"Status: {app.deployment_status}")
# {{/docs-fragment return-value}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_examples.py*

This is useful for:

- Testing different configurations
- Using different models or data sources
- A/B testing during development

## Advanced serving options

Use `with_servecontext()` for more control over the serving process:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve examples for the how-app-serving-works.md documentation."""

import logging
import flyte
import flyte.app

# {{docs-fragment basic-serve}}
app_env = flyte.app.AppEnvironment(
    name="my-dev-app",
    parameters=[flyte.app.Parameter(name="model_path", value="s3://bucket/models/model.pkl")],
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment basic-serve}}

# {{docs-fragment override-parameters}}
app = flyte.with_servecontext(
    input_values={
        "my-dev-app": {
            "model_path": "s3://bucket/models/test-model.pkl",
        }
    }
).serve(app_env)
# {{/docs-fragment override-parameters}}

# {{docs-fragment advanced-serving}}
app = flyte.with_servecontext(
    version="v1.0.0",
    project="my-project",
    domain="development",
    env_vars={"LOG_LEVEL": "DEBUG"},
    input_values={"app-name": {"input": "value"}},
    cluster_pool="dev-pool",
    log_level=logging.INFO,
    log_format="json",
    dry_run=False,
).serve(app_env)
# {{/docs-fragment advanced-serving}}

# {{docs-fragment return-value}}
app = flyte.serve(app_env)
print(f"URL: {app.url}")
print(f"Endpoint: {app.endpoint}")
print(f"Status: {app.deployment_status}")
# {{/docs-fragment return-value}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_examples.py*

## Using CLI

You can also serve apps from the command line:

```bash
flyte serve path/to/app.py app
```

Where `app` is the variable name of the `AppEnvironment` object.

## Return value

`flyte.serve()` returns an `App` object with:

- `url`: The app's URL
- `endpoint`: The app's endpoint URL
- `deployment_status`: Current status of the app
- `name`: App name

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Serve examples for the how-app-serving-works.md documentation."""

import logging
import flyte
import flyte.app

# {{docs-fragment basic-serve}}
app_env = flyte.app.AppEnvironment(
    name="my-dev-app",
    parameters=[flyte.app.Parameter(name="model_path", value="s3://bucket/models/model.pkl")],
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    app = flyte.serve(app_env)
    print(f"App served at: {app.url}")
# {{/docs-fragment basic-serve}}

# {{docs-fragment override-parameters}}
app = flyte.with_servecontext(
    input_values={
        "my-dev-app": {
            "model_path": "s3://bucket/models/test-model.pkl",
        }
    }
).serve(app_env)
# {{/docs-fragment override-parameters}}

# {{docs-fragment advanced-serving}}
app = flyte.with_servecontext(
    version="v1.0.0",
    project="my-project",
    domain="development",
    env_vars={"LOG_LEVEL": "DEBUG"},
    input_values={"app-name": {"input": "value"}},
    cluster_pool="dev-pool",
    log_level=logging.INFO,
    log_format="json",
    dry_run=False,
).serve(app_env)
# {{/docs-fragment advanced-serving}}

# {{docs-fragment return-value}}
app = flyte.serve(app_env)
print(f"URL: {app.url}")
print(f"Endpoint: {app.endpoint}")
print(f"Status: {app.deployment_status}")
# {{/docs-fragment return-value}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_examples.py*

## Best practices

1. **Use for development**: App serving is ideal for development and testing.
2. **Override parameters**: Take advantage of parameter overrides for testing different configurations.
3. **Quick iteration**: Use `serve` for rapid development cycles.
4. **Switch to deploy**: Use [deploy](./how-app-deployment-works) for production deployments.

## Troubleshooting

**App not activating:**

- Check cluster connectivity
- Verify app configuration is correct
- Review container logs for errors

**Parameter overrides not working:**

- Verify parameter names match exactly
- Check that parameters are defined in the app environment
- Ensure you're using the `input_values` parameter correctly

**Slow serving:**

- Images may need to be built (first time is slower).
- Large code bundles can slow down deployment.
- Check network connectivity to the cluster.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps/how-app-custom-domain-works ===

# How app custom domains work

Below is a snippet of how to set a custom domain for your app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
# ]
# ///

"""A basic streamlit app that uses a custom domain."""

import flyte
import flyte.app

# {{docs-fragment custom-domain}}
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1")

# The `App` declaration.
# Uses the `ImageSpec` declared above.
# In this case we do not need to supply any app code
# as we are using the built-in Streamlit `hello` app.
app_env = flyte.app.AppEnvironment(
    name="streamlit-hello-custom-domain",
    image=image,
    args=["streamlit", "hello", "--server.port", "8080"],
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    domain=flyte.app.Domain(subdomain="custom-subdomain"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    d = flyte.deploy(app_env)
    print(d[0])
# {{/docs-fragment custom-domain}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/create_custom_domain.py*


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps/how-app-deployment-works ===

# How app deployment works

Deployment is the recommended way to deploy apps to production. It creates versioned, immutable app deployments.

## Overview

When you deploy an app, the following happens:

1. **Code bundling**: Your app code is bundled and prepared
2. **Image building**: Container images are built (if needed)
3. **Deployment**: The app is deployed to your Flyte cluster
4. **Activation**: The app is automatically activated and ready to use

## Using the Python SDK

Deploy an app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

`flyte.deploy()` returns a list of `Deployment` objects. Each `Deployment` contains a dictionary of `DeployedEnvironment` objects (one for each environment deployed, including environment dependencies). For apps, the `DeployedEnvironment` is a `DeployedAppEnvironment` which has a `deployed_app` property of type `App`.

## Deployment plan

Flyte automatically creates a deployment plan that includes:

- The app you're deploying
- All [app environment dependencies](../configure-apps/apps-depending-on-environments) (via `depends_on`)
- Proper deployment order

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

## Overriding app configuration at deployment time

If you need to override the app configuration at deployment time, you can use the `clone_with` method to create a new
app environment with the desired overrides.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

## Activation/deactivation

Unlike serving, deployment does not automatically activate apps. You need to activate them explicitly:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

See [Activating and deactivating apps](./activating-and-deactivating-apps) for more details.

## Using the CLI

Deploy from the command line:

```bash
flyte deploy path/to/app.py app
```

Where `app` is the variable name of the `AppEnvironment` object.

You can also specify the following options:

```bash
flyte deploy path/to/app.py app \
    --version v1.0.0 \
    --project my-project \
    --domain production \
    --dry-run
```

## Example: Full deployment configuration

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

## Best practices

1. **Use for production**: Deploy is designed for production use.
2. **Version everything**: Always specify versions for reproducibility.
3. **Test first**: Test with serve before deploying to production.
4. **Manage dependencies**: Use `depends_on` to manage app dependencies.
5. **Activation strategy**: Have a strategy for activating/deactivating apps.
7. **Use dry-run**: Test deployments with `dry_run=True` first.
8. **Separate environments**: Use different projects/domains for different environments.
9. **Parameter management**: Consider using environment-specific parameter values.

## Deployment status and return value

`flyte.deploy()` returns a list of `Deployment` objects. Each `Deployment` contains a dictionary of `DeployedEnvironment` objects:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Deploy examples for the how-app-deployment-works.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

# {{docs-fragment basic-deploy}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ...
)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment basic-deploy}}

# {{docs-fragment deployment-plan}}
app1_env = flyte.app.AppEnvironment(name="backend", ...)
app2_env = flyte.app.AppEnvironment(name="frontend", depends_on=[app1_env], ...)

# Deploying app2_env will also deploy app1_env
deployments = flyte.deploy(app2_env)

# deployments contains both app1_env and app2_env
assert len(deployments) == 2
# {{/docs-fragment deployment-plan}}

# {{docs-fragment clone-with}}
app_env = flyte.app.AppEnvironment(name="my-app", ...)

if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(
        app_env.clone_with(app_env.name, resources=flyte.Resources(cpu="2", memory="2Gi"))
    )
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {deployed_env.deployed_app.url}")
# {{/docs-fragment clone-with}}

# {{docs-fragment activation-deactivation}}
if __name__ == "__main__":
    flyte.init_from_config()
    deployments = flyte.deploy(app_env)

    app = App.get(name=app_env.name)

    # deactivate the app
    app.deactivate()

    # activate the app
    app.activate()
# {{/docs-fragment activation-deactivation}}

# {{docs-fragment full-deployment}}
if __name__ == "__main__":
    flyte.init_from_config()

    deployments = flyte.deploy(
        app_env,
        dryrun=False,
        version="v1.0.0",
        interactive_mode=False,
        copy_style="loaded_modules",
    )

    # Access deployed apps from deployments
    for deployment in deployments:
        for deployed_env in deployment.envs.values():
            app = deployed_env.deployed_app
            print(f"Deployed: {deployed_env.env.name}")
            print(f"URL: {app.url}")

            # Activate the app
            app.activate()
            print(f"Activated: {app.name}")
# {{/docs-fragment full-deployment}}

# {{docs-fragment deployment-status}}
deployments = flyte.deploy(app_env)

for deployment in deployments:
    for deployed_env in deployment.envs.values():
        if hasattr(deployed_env, 'deployed_app'):
            # Access deployed environment
            env = deployed_env.env
            app = deployed_env.deployed_app

            # Access deployment info
            print(f"Name: {env.name}")
            print(f"URL: {app.url}")
            print(f"Status: {app.deployment_status}")
# {{/docs-fragment deployment-status}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/deploy_examples.py*

For apps, each `DeployedAppEnvironment` includes:

- `env`: The `AppEnvironment` that was deployed
- `deployed_app`: The `App` object with properties like `url`, `endpoint`, `name`, and `deployment_status`

## Troubleshooting

**Deployment fails:**

- Check that all dependencies are available
- Verify image builds succeed
- Review deployment logs

**App not accessible:**

- Ensure the app is activated
- Check cluster connectivity
- Verify app configuration

**Version conflicts:**

- Use unique versions for each deployment
- Check existing app versions
- Clean up old versions if needed


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps/activating-and-deactivating-apps ===

# Activating and deactivating apps

Apps deployed with `flyte.deploy()` need to be explicitly activated before they can serve traffic. Apps served with `flyte.serve()` are automatically activated.

## Activation

### Activate after deployment

After deploying an app, activate it:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

### Activate an app

When you get an app by name, you get the current app instance:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

### Check activation status

Check if an app is active:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

## Deactivation

Deactivate an app when you no longer need it:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

## Lifecycle management

### Typical deployment workflow

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

### Blue-green deployment

For zero-downtime deployments:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

## Using CLI

### Activate

```bash
flyte update app --activate my-app
```

### Deactivate

```bash
flyte update app --deactivate my-app
```

### Check status

```bash
flyte get app my-app
```

Use `--project` and `--domain` to target a specific [project-domain pair](../../get-started/core-concepts/projects-and-domains).
For all available options, see the [CLI reference](../../../api-reference/flyte-cli).

## Best practices

1. **Activate after testing**: Test deployed apps before activating
2. **Version management**: Keep track of which version is active
4. **Blue-green deployments**: Use blue-green for zero-downtime
5. **Monitor**: Monitor apps after activation
6. **Cleanup**: Deactivate and remove old versions periodically

## Automatic activation with serve

Apps served with `flyte.serve()` are automatically activated:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

This is convenient for development but less suitable for production where you want explicit control over activation.

## Example: Complete deployment and activation

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# ///

"""Activation examples for the activating-and-deactivating-apps.md documentation."""

import flyte
import flyte.app
from flyte.remote import App

app_env = flyte.app.AppEnvironment(
    name="my-app",
    # ...
)

# {{docs-fragment activate-after-deployment}}
# Deploy the app
deployments = flyte.deploy(app_env)

# Activate the app
app = App.get(name=app_env.name)
app.activate()

print(f"Activated app: {app.name}")
print(f"URL: {app.url}")
# {{/docs-fragment activate-after-deployment}}

# {{docs-fragment activate-app}}
app = App.get(name="my-app")
app.activate()
# {{/docs-fragment activate-app}}

# {{docs-fragment check-activation-status}}
app = App.get(name="my-app")
print(f"Active: {app.is_active()}")
print(f"Revision: {app.revision}")
# {{/docs-fragment check-activation-status}}

# {{docs-fragment deactivation}}
app = App.get(name="my-app")
app.deactivate()

print(f"Deactivated app: {app.name}")
# {{/docs-fragment deactivation}}

# {{docs-fragment typical-deployment-workflow}}
# 1. Deploy new version
deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

# 2. Get the deployed app
new_app = App.get(name="my-app")
# Test endpoints, etc.

# 3. Activate the new version
new_app.activate()

print(f"Deployed and activated version {new_app.revision}")
# {{/docs-fragment typical-deployment-workflow}}

# {{docs-fragment blue-green-deployment}}
# Deploy new version without deactivating old
new_deployments = flyte.deploy(
    app_env,
    version="v2.0.0",
)

new_app = App.get(name="my-app")

# Test new version
# ... testing ...

# Switch traffic to new version
new_app.activate()

print(f"Activated revision {new_app.revision}")
# {{/docs-fragment blue-green-deployment}}

# {{docs-fragment automatic-activation}}
# Automatically activated
app = flyte.serve(app_env)
print(f"Active: {app.is_active()}")  # True
# {{/docs-fragment automatic-activation}}

# {{docs-fragment complete-example}}
app_env = flyte.app.AppEnvironment(
    name="my-prod-app",
    # ... configuration ...
)

if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy
    deployments = flyte.deploy(
        app_env,
        version="v1.0.0",
        project="my-project",
        domain="production",
    )

    # Get the deployed app
    app = App.get(name="my-prod-app")

    # Activate
    app.activate()

    print(f"Deployed and activated: {app.name}")
    print(f"Revision: {app.revision}")
    print(f"URL: {app.url}")
    print(f"Active: {app.is_active()}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/activation_examples.py*

## Troubleshooting

**App not accessible after activation:**

- Verify activation succeeded
- Check app logs for startup errors
- Verify cluster connectivity
- Check that the app is listening on the correct port

**Activation fails:**

- Check that the app was deployed successfully
- Verify app configuration is correct
- Check cluster resources
- Review deployment logs

**Cannot deactivate:**

- Ensure you have proper permissions
- Check if there are dependencies preventing deactivation
- Verify the app name and version


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/apps/serve-and-deploy-apps/prefetching-models ===

# Prefetching models

Prefetching allows you to download and prepare HuggingFace models (including sharding for multi-GPU inference) before
deploying [vLLM](../native-app-integrations/vllm-app) or [SGLang](../native-app-integrations/sglang-app) apps. This speeds up deployment and ensures models are ready when your app starts.

## Why prefetch?

Prefetching models provides several benefits:

- **Faster deployment**: Models are pre-downloaded, so apps start faster
- **Reproducibility**: Models are versioned and stored in Flyte's object store
- **Sharding support**: Pre-shard models for multi-GPU tensor parallelism
- **Cost efficiency**: Download once, use many times
- **Offline support**: Models are cached in your storage backend

## Basic prefetch

### Using Python SDK

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

### Using CLI

```bash
flyte prefetch hf-model Qwen/Qwen3-0.6B
```

Wait for completion:

```bash
flyte prefetch hf-model Qwen/Qwen3-0.6B --wait
```

## Using prefetched models

Use the prefetched model in your vLLM or SGLang app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

> [!TIP]
> You can also use prefetched models as parameters to your generic `[[AppEnvironment]]`s or `FastAPIAppEnvironment`s.

## Prefetch options

### Custom artifact name

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

### With HuggingFace token

If the model requires authentication:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

The default value for `hf_token_key` is `HF_TOKEN`, where `HF_TOKEN` is the name of the Flyte secret containing your
HuggingFace token. If this secret doesn't exist, you can create a secret using the [flyte create secret CLI](../../tasks/task-configuration/secrets).

### With resources

By default, the prefetch task uses minimal resources (2 CPUs, 8GB of memory, 50Gi of disk storage), using
filestreaming logic to move the model weights from HuggingFace to your storage backend directly.

In some cases, the HuggingFace model may not support filestreaming, in which case the prefetch task will fallback to
downloading the model weights to the task pod's disk storage first, then uploading them to your storage backend. In this
case, you can specify custom resources for the prefetch task to override the default resources.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

## Sharding models for multi-GPU

### vLLM sharding

Shard a model for tensor parallelism:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

Currently, the `flyte.prefetch.hf_model` function only supports sharding models
using the `vllm` engine. Once sharded, these models can be loaded with other
frameworks such as `transformers`, `torch`, or `sglang`.

### Using shard config via CLI

You can also use a YAML file for sharding configuration to use with the
`flyte prefetch hf-model` CLI command:

```yaml
# shard_config.yaml
engine: vllm
args:
  tensor_parallel_size: 8
  dtype: auto
  trust_remote_code: true
```

Then run the CLI command:

```bash
flyte prefetch hf-model meta-llama/Llama-2-70b-hf \
    --shard-config shard_config.yaml \
    --accelerator L40s:8 \
    --hf-token-key HF_TOKEN
```

## Using prefetched sharded models

After prefetching and sharding, serve the model in your app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

## CLI options

Complete CLI usage:

```bash
flyte prefetch hf-model <repo> \
    --artifact-name <name> \
    --architecture <arch> \
    --task <task> \
    --modality text \
    --format safetensors \
    --model-type transformer \
    --short-description "Description" \
    --force 0 \
    --wait \
    --hf-token-key HF_TOKEN \
    --cpu 4 \
    --mem 16Gi \
    --ephemeral-storage 100Gi \
    --accelerator L40s:4 \
    --shard-config shard_config.yaml
```

## Complete example

Here's a complete example of prefetching and using a model:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-vllm>=2.0.0b49",
# ]
# ///

"""Prefetch examples for the prefetching-models.md documentation."""

import flyte
from flyte.prefetch import ShardConfig, VLLMShardArgs
from flyteplugins.vllm import VLLMAppEnvironment

# {{docs-fragment basic-prefetch}}
# Prefetch a HuggingFace model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")

# Wait for prefetch to complete
run.wait()

# Get the model path
model_path = run.outputs()[0].path
print(f"Model prefetched to: {model_path}")
# {{/docs-fragment basic-prefetch}}

# {{docs-fragment using-prefetched-models}}
# Prefetch the model
run = flyte.prefetch.hf_model(repo="Qwen/Qwen3-0.6B")
run.wait()

# Use the prefetched model
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_path=flyte.app.RunOutput(
        type="directory",
        run_name=run.name,
    ),
    model_id="qwen3-0.6b",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
    stream_model=True,
)

app = flyte.serve(vllm_app)
# {{/docs-fragment using-prefetched-models}}

# {{docs-fragment custom-artifact-name}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    artifact_name="qwen-0.6b-model",  # Custom name for the stored model
)
# {{/docs-fragment custom-artifact-name}}

# {{docs-fragment hf-token}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-7b-hf",
    hf_token_key="HF_TOKEN",  # Name of Flyte secret containing HF token
)
# {{/docs-fragment hf-token}}

# {{docs-fragment with-resources}}
run = flyte.prefetch.hf_model(
    repo="Qwen/Qwen3-0.6B",
    cpu="4",
    mem="16Gi",
    ephemeral_storage="100Gi",
)
# {{/docs-fragment with-resources}}

# {{docs-fragment vllm-sharding}}
run = flyte.prefetch.hf_model(
    repo="meta-llama/Llama-2-70b-hf",
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L40s:4"),
    shard_config=ShardConfig(
        engine="vllm",
        args=VLLMShardArgs(
            tensor_parallel_size=4,
            dtype="auto",
            trust_remote_code=True,
        ),
    ),
    hf_token_key="HF_TOKEN",
)

run.wait()
# {{/docs-fragment vllm-sharding}}

# {{docs-fragment using-sharded-models}}
# Use in vLLM app
vllm_app = VLLMAppEnvironment(
    name="multi-gpu-llm-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="llama-2-70b",
    resources=flyte.Resources(
        cpu="8",
        memory="32Gi",
        gpu="L40s:4",  # Match the number of GPUs used for sharding
    ),
    extra_args=[
        "--tensor-parallel-size", "4",  # Match sharding config
    ],
)

if __name__ == "__main__":
    # Prefetch with sharding
    run = flyte.prefetch.hf_model(
        repo="meta-llama/Llama-2-70b-hf",
        accelerator="L40s:4",
        shard_config=ShardConfig(
            engine="vllm",
            args=VLLMShardArgs(tensor_parallel_size=4),
        ),
    )
    run.wait()

    flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            # override the model path to use the prefetched model
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            # set the hf_model_path to None
            hf_model_path=None,
            # stream the model from flyte object store directly to the GPU
            stream_model=True,
        )
    )
# {{/docs-fragment using-sharded-models}}

# {{docs-fragment complete-example}}
# define the app environment
vllm_app = VLLMAppEnvironment(
    name="qwen-serving-app",
    # this will download the model from HuggingFace into the app container's filesystem
    model_hf_path="Qwen/Qwen3-0.6B",
    model_id="qwen3-0.6b",
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=600,
    ),
    requires_auth=False,
)

if __name__ == "__main__":
    # prefetch the model
    print("Prefetching model...")
    run = flyte.prefetch.hf_model(
        repo="Qwen/Qwen3-0.6B",
        artifact_name="qwen-0.6b",
        cpu="4",
        mem="16Gi",
        ephemeral_storage="50Gi",
    )

    # wait for completion
    print("Waiting for prefetch to complete...")
    run.wait()
    print(f"Model prefetched: {run.outputs()[0].path}")

    # deploy the app
    print("Deploying app...")
    flyte.init_from_config()
    app = flyte.serve(
        vllm_app.clone_with(
            name=vllm_app.name,
            model_path=flyte.app.RunOutput(type="directory", run_name=run.name),
            hf_model_path=None,
            stream_model=True,
        )
    )
    print(f"App deployed: {app.url}")
# {{/docs-fragment complete-example}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/prefetch_examples.py*

## Best practices

1. **Prefetch before deployment**: Prefetch models before deploying apps for faster startup
2. **Version models**: Use meaningful artifact names to easily identify the model in object store paths
3. **Shard appropriately**: Shard models for the GPU configuration you'll use for inference
4. **Cache prefetched models**: Once prefetched, models are cached in your storage backend for faster serving

## Troubleshooting

**Prefetch fails:**

- Check HuggingFace token (if required)
- Verify model repo exists and is accessible
- Check resource availability
- Review prefetch task logs

**Sharding fails:**

- Ensure accelerator matches shard config
- Check GPU memory is sufficient
- Verify `tensor_parallel_size` matches GPU count
- Review prefetch task logs for sharding-related errors

**Model not found in app:**

- Verify RunOutput references correct run name
- Check that prefetch completed successfully
- Ensure model_path is set correctly
- Review app startup logs


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents ===

# Agents

An agent is a program that decides what to do next by calling a model in a loop. On Flyte, each step of that loop is a task, which is what makes the agent durable: every model call, tool call, and intermediate result is recorded, so a run that fails partway through can be inspected and resumed rather than restarted.

This matters more for agents than for batch jobs. An agent's control flow is decided at runtime by the model, so you cannot know in advance which path a run took. Recording each step is what makes the run explainable afterward.

```python
@env.task
def step(state: State) -> State:
    ...  # one model call plus its tool calls
```

Nothing here is a separate agent framework. An agent is tasks for the reasoning steps, an app when it needs to be reachable over HTTP, and a sandbox when it executes code the model wrote.

### **Agents > Build an agent**

Implement ReAct, Plan-and-Execute, and other agent patterns with full observability.

### **Agents > Build an MCP**

Serve Model Context Protocol servers for AI assistants to interact with, hosted on Flyte.

### **Agents > Sandboxing**

Safely execute LLM-generated code with workflow sandboxes or ephemeral containers.

## Related

### **Agent frameworks**

Run agents from OpenAI, Claude, LangGraph, CrewAI, and more as durable Flyte tasks.

### **Tasks**

The unit each reasoning step is built from.

## Subpages
- **Agents > Build an agent**
- **Agents > Build an MCP**
- **Agents > Sandboxing**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent ===

# Build an agent

This section covers how to build, deploy, and run agentic AI applications on Flyte.

Building an agent on Flyte breaks down into two **orthogonal** choices:

1. **How you build the agent loop**: plain Python, the built-in `flyte.ai.agents.Agent` harness, or a third-party framework (LangGraph, PydanticAI, OpenAI Agents SDK).
2. **How you deploy and run it**: as a task you invoke on demand, as a scheduled task driven by a `flyte.Trigger`, or as a long-running app (e.g. a webhook or chat UI).

Any agent from axis (1) can be deployed via any pattern in axis (2). The two are independent, so you can start with a pure-Python loop run on demand and later move it behind a schedule or a webhook without rewriting the agent.

## How Flyte maps to the agentic world

- **`TaskEnvironment`**: The sandboxed execution environment for your agent steps. It configures the container image, hardware resources (CPU, GPU), and secrets (API keys). Think of it as defining "where this code runs."
- **`@env.task`**: Turns any Python function into a remotely-executed step. Each task runs in its own container with the resources you specified. This is the equivalent of a node in LangGraph or n8n.
- **Tasks calling tasks**: A task can `await` other tasks, and each called task gets its own container automatically. No separate workflow decorator needed. The calling task IS your workflow, this is how you build multi-step agentic pipelines.
- **`@flyte.trace`**: Marks helper functions inside a task for fine-grained observability and caching. Each traced call appears as a span in the Flyte dashboard, with its inputs and outputs captured and checkpointed. Use this on your LLM calls, tool executions, and routing decisions to get full visibility into every turn of the agent loop.

> [!TIP]
> See the **Get started > Quickstart** for a hands-on walkthrough.

## Ways to build an agent

| Approach | When to use it | Guide |
|----------|----------------|-------|
| **Pure Python** | You want full control over the loop and the lightest possible dependency footprint | **Agents > Build an agent > Pure Python agents** |
| **The `Agent` harness** | You want a batteries-included tool-use loop with tools, MCP servers, memory, and HITL built in | **Agents > Build an agent > Flyte-native agents** |
| **Third-party frameworks** | You already have agents in LangGraph, CrewAI, OpenAI Agents SDK, Pydantic AI, and more | [Agent frameworks](../../../integrations/agents/_index) |
| **An unsupported framework** | Your framework has no first-party plugin | **Agents > Build an agent > Bring your own framework** |

The `Agent` harness has a few dedicated guides of its own:

- **Agents > Build an agent > Flyte-native agents > Extending the agent class**: customize the loop by overriding `run`.
- **Agents > Build an agent > Agent memory**: persist conversation history and artifacts across runs with `MemoryStore`.
- **Agents > Build an agent > Agent chat UI**: give any agent a hosted chat interface.

## Deploying an agent

Once you've built an agent, **Agents > Build an agent > Deploy an agent as a service** covers running it as a task, on a schedule via `flyte.Trigger`, and behind an `AppEnvironment` webhook.

## Related

- [**Grafana Agent Observability**](../../../integrations/grafana-agent-observability/_index): export generations, tool calls, token usage, and cost, grouped by run.
- [**OpenTelemetry**](../../../integrations/opentelemetry/_index): export tasks and traced steps as spans, with durable runs arriving as one trace.
- [**Sandboxing**](../sandboxing/_index): safely execute LLM-generated code.
- [**Build an MCP server**](../build-mcp/_index): serve Model Context Protocol servers for AI assistants to interact with Flyte.

## Subpages
- **Agents > Build an agent > Pure Python agents**
- **Agents > Build an agent > Flyte-native agents**
- **Agents > Build an agent > Bring your own framework**
- **Agents > Build an agent > Agent memory**
- **Agents > Build an agent > Agent chat UI**
- **Agents > Build an agent > Deploy an agent as a service**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/python-agents ===

# Pure Python agents

The lightest way to build an agent on Flyte is to write the loop yourself in plain Python. Flyte is framework-agnostic: use any Python LLM library (OpenAI SDK, Anthropic SDK, LiteLLM, etc.) inside your tasks. The platform provides the production infrastructure layer: sandboxed execution, parallel fan-out, durable checkpointing, and observability for every step of the agent loop.

This approach gives you full control over the loop and the smallest possible dependency footprint. If you'd rather not hand-roll the tool-call loop, see [The Flyte Agent harness](./flyte-agents), which provides a batteries-included loop with tools, MCP servers, memory, and HITL. If you already have agents written in a third-party framework, see [Agent frameworks](../../../integrations/agents/_index) for the ten supported plugins, or [Bring your own framework](./bring-your-own-framework) for the framework-agnostic pattern.

Two decorators are all you need:

| Decorator | What it does | Think of it as... |
|-----------|-------------|-------------------|
| **`@env.task`** | Runs a function in its own container on Flyte with dedicated resources, dependencies, and secrets | A sandboxed agent step with its own execution environment |
| **`@flyte.trace`** | Marks a helper function for observability, where each call appears as a span in the Flyte dashboard with captured I/O | An observability hook on your LLM calls, tool executions, and routing decisions |

## ReAct pattern: Reason, Act, Observe (no framework needed)

The [ReAct pattern](https://arxiv.org/abs/2210.03629) is the most common agent architecture: the LLM reasons about what to do, calls a tool, observes the result, and repeats until done. This example is implemented directly with flyte:

```
Thought → Action → Observation → repeat until done
```

```
import json

from openai import AsyncOpenAI
from pydantic import BaseModel

import flyte

env = flyte.TaskEnvironment(
    name="agent_env",
    image=flyte.Image.from_debian_base(python_version=(3, 13)).with_pip_packages("openai"),
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    secrets=[flyte.Secret(key="OPENAI_API_KEY")],
)

TOOLS = {"add": lambda a, b: a + b, "multiply": lambda a, b: a * b}

@flyte.trace  # each call = a span in the dashboard
async def reason(goal: str, history: str) -> dict:
    """LLM picks a tool or returns a final answer."""
    r = await AsyncOpenAI().chat.completions.create(
        model="gpt-4.1-nano",
        response_format={"type": "json_object"},
        messages=[
            {
                "role": "system",
                "content": f"Tools: {list(TOOLS)}. Respond JSON: "
                '{"thought":..,"tool":..,"args":{}} or {"thought":..,"done":true,"answer":..}',
            },
            {"role": "user", "content": f"Goal: {goal}\n\n{history}\nWhat next?"},
        ],
    )
    return json.loads(r.choices[0].message.content)

@flyte.trace
async def act(tool: str, args: dict) -> str:
    """Execute the chosen tool."""
    return str(TOOLS[tool](**args))

class AgentResult(BaseModel):
    answer: str
    steps: int

@env.task  # runs in its own container
async def react_agent(goal: str, max_steps: int = 10) -> AgentResult:
    history = ""
    for step in range(1, max_steps + 1):  # the agent loop
        decision = await reason(goal, history)  # Thought
        if decision.get("done"):
            return AgentResult(answer=str(decision["answer"]), steps=step)
        result = await act(decision["tool"], decision["args"])  # Action
        # Observation
        history += f"Step {step}: {decision['thought']} -> {decision['tool']}({decision['args']}) = {result}\n"
    return AgentResult(answer="Max steps reached", steps=max_steps)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/building-agents/react_agent.py*

```bash
flyte run agent.py react_agent --goal "What is (12 + 8) * 3?"
# => AgentResult(answer='60', steps=3)
```

**What's happening under the hood:**

- `react_agent` runs in a container with only `openai` installed and 2 CPU / 2GB RAM
- Each `reason()` and `act()` call is traced, so you see every LLM call, every tool invocation, and every intermediate result in the Flyte dashboard
- The agent's inputs and final output are durably persisted, letting you inspect any past run end-to-end
- Swap in your own tools (web search, database queries, API calls) by adding to the `TOOLS` dict

> [!TIP]
> See the [Agentic Refinement docs](../../advanced-project/agentic-refinement), [Traces docs](../../tasks/task-programming/traces), and [more patterns (planner, debate, etc.)](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows).

## Plan-and-Execute with parallel fan-out

The [Plan-and-Execute pattern](https://blog.langchain.com/plan-and-execute-agents/) splits a complex query into sub-tasks, fans them out in parallel, then synthesizes the results. With Flyte the fan-out is just `asyncio.gather()`, and each sub-task gets its own container, giving you true parallelism on separate hardware.

```python
# workflow.py
import os, json, asyncio, flyte
from openai import AsyncOpenAI

env = flyte.TaskEnvironment(
    name="research_env",
    image=flyte.Image.from_debian_base(python_version=(3, 13)).with_pip_packages("openai"),
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    secrets=[flyte.Secret(key="OPENAI_API_KEY")],
)

@flyte.trace
async def llm(prompt: str) -> str:
    r = await AsyncOpenAI().chat.completions.create(
        model="gpt-4.1-nano",
        messages=[{"role": "user", "content": prompt}],
    )
    return r.choices[0].message.content

@env.task
async def plan(query: str, n: int = 3) -> list[str]:
    """Split the query into sub-topics."""
    raw = await llm(
        f'Break this into exactly {n} sub-topics. Return ONLY a JSON array of strings.\n\n{query}'
    )
    return json.loads(raw)[:n]

@env.task
async def research(topic: str) -> str:
    """Research one sub-topic (each call = its own container)."""
    return await llm(f"Write a short, factual report on: {topic}")

@env.task
async def synthesize(query: str, reports: list[str]) -> str:
    """Combine the sub-reports into a final answer."""
    sections = "\n\n".join(reports)
    return await llm(f"Synthesize a final answer to '{query}' from:\n\n{sections}")

@env.task
async def research_workflow(query: str, num_topics: int = 3) -> str:
    topics = await plan(query, num_topics)
    reports = list(await asyncio.gather(*[research(t) for t in topics]))  # parallel fan-out
    return await synthesize(query, reports)
```

```bash
flyte run workflow.py research_workflow --query "Impact of storms on travel insurance payouts"
```

**What's happening under the hood:**

```
research_workflow (orchestrator)
  ├── plan          → LLM breaks query into N sub-topics      [container 1]
  ├── research(t1)  → researches one sub-topic                [container 2]  ┐
  ├── research(t2)  → researches one sub-topic                [container 3]  ├ parallel
  ├── research(t3)  → researches one sub-topic                [container 4]  ┘
  └── synthesize    → LLM combines reports into final answer  [container 5]
```

- **Fan-out:** `asyncio.gather()` launches all research tasks in parallel, each in its own container
- **Observability:** `@flyte.trace` on each LLM call means every prompt and response is visible as a span in the Flyte dashboard
- **Durable checkpointing:** Each task's output is persisted. If `synthesize` fails, re-running skips the completed `plan` and `research` steps (with caching enabled)

> [!TIP]
> The same fan-out works with any framework inside the `research` task. See [Agent frameworks](../../../integrations/agents/_index) for versions that run a LangGraph, CrewAI or OpenAI Agents SDK researcher inside each parallel container, with the tool calls durable on their own.

## More agentic patterns

Flyte is framework-agnostic, so these patterns work with any LLM library. Each maps to well-known agent architectures:

| Pattern | What it does | When to use it | Link |
|---------|-------------|----------------|------|
| **ReAct** | Reason → Act → Observe loop with tool calling | Single-agent tasks with tools (API calls, search, code execution) | [multi-agent-workflows/react](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows) |
| **Plan-and-Execute** | LLM creates a plan, independent steps fan out in parallel, results are synthesized | Complex queries that decompose into parallel sub-tasks | [multi-agent-workflows/planner](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows) |
| **Evaluator-Optimizer (Reflection)** | Generate → Critique → Refine loop until quality threshold met | Content generation, code generation, any task with clear quality criteria | [Agentic Refinement docs](../../advanced-project/agentic-refinement) |
| **Orchestrator-Workers (Manager)** | Supervisor agent delegates to specialist worker agents, reviews quality, requests revisions | Multi-agent systems where sub-tasks require different expertise | [multi-agent-workflows/manager](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows) |
| **Debate** | Multiple agents solve independently, then debate to consensus | High-stakes decisions where diverse reasoning improves accuracy | [multi-agent-workflows/debate](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows) |
| **Sequential (Prompt Chaining)** | Static pipeline of LLM calls, no dynamic routing | Predictable multi-step transformations (extract → validate → format) | [multi-agent-workflows/sequential](https://github.com/unionai/workshops/tree/main/tutorials/multi-agent-workflows) |

## How Flyte's primitives map to the agent stack

If you're coming from LangGraph, CrewAI, OpenAI Agents SDK, or similar frameworks, here's how the concepts you already know translate:

**Your agent loop** is a Python `for`/`while` loop inside an `@env.task`. Each iteration calls `@flyte.trace`-decorated functions for reasoning and tool execution. Flyte doesn't impose a loop structure; you write it in plain Python, which means any pattern (ReAct, reflection, plan-and-execute) works naturally.

**Tool calling** is just calling Python functions. Define your tools as regular functions, decorate them with `@flyte.trace` for observability, and call them from within the agent loop. Use any tool-calling mechanism your LLM SDK provides (OpenAI function calling, Anthropic tool use, LangChain `bind_tools()`). MCP servers can be accessed from within tasks using the MCP Python SDK.

**Parallel fan-out** (LangGraph's `Send()`, n8n's Split in Batches) is `asyncio.gather()`. Each awaited task gets its own container, giving you true parallelism on separate hardware, not just concurrent coroutines.

**State and checkpointing** (LangGraph's Checkpointers, Threads) is automatic. Every task's inputs and outputs are durably persisted. `@flyte.trace` adds sub-step checkpoints within a task. Re-running with caching enabled skips completed steps, Flyte's equivalent of replaying from a checkpoint.

**Routing and conditional logic** (LangGraph's `add_conditional_edges`, n8n's If/Switch nodes) is Python `if/else`. No special API needed.

**Environment isolation** (different dependencies per step) is `TaskEnvironment`. Your LLM step can use `langchain==0.3`; your data step can use `pandas` + GPU. Each gets its own container image.

**Guardrails and validation** are Python code between steps: `if/else` checks, Pydantic validation, structured output parsing, or libraries like NeMo Guardrails. Raise an exception to fail a step and trigger retries.

**Observability:** The Flyte dashboard shows the full execution tree with per-step inputs, outputs, logs, resource usage, and timing. `@flyte.trace` adds spans within a task for fine-grained visibility into individual LLM calls and tool invocations. For LLM-specific metrics (token usage, cost per call), integrate with Langfuse or LangSmith from within your tasks.

## Next steps

- [The Flyte Agent harness](./flyte-agents): skip the boilerplate with a built-in tool-use loop.
- [Deploy an agent as a service](./deploy-agent-as-service): run this agent on a schedule or behind a webhook.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/flyte-agents ===

# Flyte-native agents

`flyte.ai.agents.Agent` is a flyte-native, batteries-included agent harness. Instead of hand-rolling the tool-call loop (as in [Build an agent with pure Python](./python-agents)), you declare a set of tools and instructions, and the harness drives a robust LLM ↔ tool loop for you.

The harness deeply integrates with Flyte:

- **Tools** can be plain Python callables, `@flyte.trace` helpers, `@env.task` durable tasks, `LazyEntity` remote-task references, or pre-built `AgentTool` instances.
- **MCP servers** (Slack, GitHub, Linear, filesystem, …) are first-class: pass a `MCPServerSpec` and their tools are loaded into the catalog automatically.
- **Memory** persists across runs via `flyte.io.Dir`. See [Agent memory](./agent-memory).
- **HITL** support pauses the loop and asks a human for approval before sensitive tools execute.

## How it works

`Agent(...)` collapses heterogeneous tool sources into a single tool registry plus an auto-generated system prompt. `agent.run(message)` then drives an LLM ↔ tool-call loop:

1. Send the conversation and tool catalog to the LLM.
2. If the assistant returns tool calls, execute each one (sequentially or concurrently), append the results back into the message history, and loop.
3. Stop when the assistant returns a plain-text reply (no tool calls) or when `max_turns` is reached.

```mermaid
flowchart TB
    inputs["tools (fn / task / MCP),<br/> skills, memory, and instructions"]
    inputs --> agent[["Agent<br/>(tool registry, skills, system prompt)"]]

    subgraph loop["agent.run(message)  ·  agent.run.aio(message) in async code"]
      direction TB
      llm["call_llm"] --loop max_turns times--> branch{"tool_calls?"}
      branch -- yes --> exec["execute tools<br/>(optional HITL approval)"]
      exec --> llm
      branch -- no --> done["final reply"]
    end

    user(["user message"]) --> llm
    agent --> llm
    done --> result(["AgentResult<br/>+ updated memory"])
```

The call returns an `AgentResult` with the final `summary`, an `error` string (empty on success), and the number of `attempts` (turns) taken.

### Sync vs async

`agent.run` is synchronous by default. Inside `async def` code (Flyte tasks, FastAPI handlers, etc.) use the `.aio(...)` companion instead.

| Context | Call |
|---------|------|
| Scripts, notebooks, sync code | `result = agent.run(message)` |
| `async def` tasks / handlers | `result = await agent.run.aio(message)` |

## A minimal agent

Declare a few tools as plain async functions, build an `Agent`, and call `run`. The harness reads each tool's signature and docstring to build the JSON schema and description the LLM sees, so well-documented tools work best.

```
from flyte.ai.agents import Agent

async def add(x: float, y: float) -> float:
    """Add two numbers and return their sum."""
    return x + y

async def multiply(x: float, y: float) -> float:
    """Multiply two numbers and return their product."""
    return x * y

async def get_weather(city: str) -> dict[str, str | float]:
    """Return a weather snapshot for `city`.

    In a real agent, replace this stub with a call to a weather API (and
    promote it to a ``@env.task`` for durable, retryable execution).
    """
    fake = {
        "new york": {"temperature_f": 68.4, "conditions": "partly cloudy"},
        "san francisco": {"temperature_f": 61.0, "conditions": "foggy"},
        "tokyo": {"temperature_f": 74.2, "conditions": "sunny"},
    }
    return fake.get(city.lower(), {"temperature_f": 70.0, "conditions": "clear"})

agent = Agent(
    name="basic-helper",
    instructions=(
        "You are a friendly assistant. Use the available tools to look up "
        "weather and compute math. Reply with a single sentence summary."
    ),
    model="claude-haiku-4-5",
    tools=[add, multiply, get_weather],
    max_turns=6,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/agent-harness/basic_agent.py*

Call it synchronously, or with `await agent.run.aio(message)` inside async code:

CODE1

## Tools

The `tools=` argument accepts a sequence (or a `{name: tool}` mapping) of any mix of:

| Tool source | What it is | When to use it |
|-------------|------------|----------------|
| Plain callable | A sync or async Python function | Lightweight, in-process helpers |
| `@flyte.trace` helper | A traced function | In-process helpers you want as spans in the dashboard |
| `@env.task` template | A durable Flyte task | Heavy compute / IO that should run on-cluster, be retryable, and observable |
| `LazyEntity` | A reference to a remote deployed task | Calling already-deployed tasks by name |
| `AgentTool` | A pre-built tool descriptor | Renaming, custom schema, or HITL gating |

Pass a mapping to expose a tool under a different name to the LLM:

CODE2

When a tool is an `@env.task`, the harness invokes it with `task.aio(...)`, so each tool call executes durably on the cluster and shows up in the Flyte dashboard.

### Customizing a tool with `tool(...)`

Use the `tool` decorator/wrapper to rename a tool, override its description, or gate it behind human approval, without writing an `AgentTool` by hand:

CODE3

When the LLM tries to call a tool marked `requires_approval=True`, the harness invokes the agent's `approval_callback` and waits for a boolean decision before executing. The default callback raises a human-input request via the `flyteplugins-hitl` plugin and blocks until a human approves or denies. If denied, the agent receives a synthetic tool message explaining the rejection so it can recover gracefully.

Pass `call_handler` to intercept *how* a tool is invoked. The handler is an async callback `(call_llm, tool_fn, **kwargs) -> result` that runs in place of the default execution. Await `tool_fn` to run the default behavior, or reach into `tool_fn.target` (the underlying task / callable) and `call_llm` (the agent's LLM callback) to do something custom. For example, ask the LLM how to size compute, then run the task with overridden resources and retry on OOM:

CODE4

## MCP integration

The harness can connect to one or more [Model Context Protocol (MCP)](https://modelcontextprotocol.io) servers and surface their tools transparently. On the first `run` call, the harness connects to each server, lists its tools, and registers them in the catalog.

Declare servers with `MCPServerSpec`: either an HTTP(S) `url` (for streamable-http / SSE transports) or a `command` (for stdio servers):

CODE5

Useful `MCPServerSpec` knobs:

- `tool_prefix`: prepend a prefix to every tool name from this server to avoid collisions.
- `tool_filter`: an allowlist of tool names to expose to the LLM (`None` exposes all).
- `headers`: HTTP headers (e.g. `Authorization`) for authenticated servers.

MCP support requires the `mcp` package: `pip install 'flyte[mcp]'`. To serve your own MCP servers on Flyte, see [Build an MCP](../build-mcp/_index).

## Skills

Pass extra context to append to the system prompt via `skills=`. Each entry is either a literal string or a `pathlib.Path` to a local text file:

CODE6

## Observability

Every step of the loop emits a typed `AgentEvent` (`agent_start`, `turn_start`, `tool_start`, `tool_end`, `approval_request`, …). Subscribe by setting the `agent_progress_cb` context variable to forward events to logs, NDJSON streams, websockets, or Flyte reports. The built-in chat UI uses this hook to stream progress; see [Add a chat UI](./agent-chat-ui).

## Extending the agent class

The default loop is robust, but sometimes you need custom behavior around it: input guardrails, output post-processing, a different control flow, or extra bookkeeping. The cleanest way to do this is to subclass `Agent` and override its `run` method.

`Agent` is a [dataclass](https://docs.python.org/3/library/dataclasses.html), and `run` is the single public entry point that drives the loop. There are two common strategies:

1. **Wrap the built-in loop**: add logic before and after `super().run(...)`. Best when you mostly want the default behavior plus pre/post steps.
2. **Replace the loop entirely**: implement `run` (and `tool_descriptions`) yourself. Best when you need a fundamentally different control flow but still want to plug into the rest of the ecosystem (e.g. the chat UI).

### `run` is sync-by-default

`Agent.run` is wrapped with `@syncify`, which means callers can use it synchronously (`agent.run(...)`) or await the async companion (`await agent.run.aio(...)`). When you override `run`, decorate your async implementation with `@syncify` to keep the same dual interface, and call the parent loop via `await super().run.aio(...)`.

### Strategy 1: wrap the built-in loop

Subclass `Agent` as a dataclass so you can add your own fields, then override `run` to add an input guardrail and post-process the answer:

CODE7

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/extending-the-agent/guarded_agent.py*

Instantiate and call it just like a regular `Agent`:

CODE8

Because `GuardedAgent` still subclasses `Agent`, every other feature (tools, MCP servers, memory, HITL) keeps working unchanged.

### Strategy 2: implement `run` from scratch

If you want a completely custom loop, implement the `AgentProtocol`: a class exposing `run(message, memory) -> AgentResult` and `tool_descriptions() -> list[dict]`. Any object satisfying this protocol can be used anywhere the harness is accepted, including the [chat UI](./agent-chat-ui). `memory` may be a `list[dict]` of prior messages (a chat history) or a `MemoryStore`.

CODE9

> [!NOTE]
> `AgentResult` carries `summary`, `error`, `attempts`, and (for code-generating agents) `code` and `charts`. Populate the fields relevant to your loop; downstream consumers like the chat UI read `summary` and `error`.

### Choosing between subclassing and composition

Subclassing is the right tool when you need to change *how the loop runs*. If you only need to change *what happens around a run* (for example, looping the agent until a condition is met, or combining several agents) prefer plain composition: call `agent.run.aio(...)` from inside your own `@env.task`. This keeps the harness untouched and your orchestration logic explicit and observable in the dashboard.

## Next steps

- [Agent memory](./agent-memory): persist transcript and artifacts across runs.
- [Add a chat UI](./agent-chat-ui): wrap the agent in a hosted chat interface.
- [Deploy an agent as a service](./deploy-agent-as-service): run on a schedule or behind a webhook.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/bring-your-own-framework ===

# Bring your own framework

The [supported plugins](../../../integrations/agents/_index) are worked examples of one underlying idea: **Flyte is the runtime, your framework is the loop.** If your agent library is written in Python, it runs on Flyte with no special plugin.

This page is a framework-agnostic template for the frameworks that don't have one. Drop in any library ([AutoGen](https://microsoft.github.io/autogen/), [smolagents](https://github.com/huggingface/smolagents), [Atomic Agents](https://github.com/BrainBlend-AI/atomic-agents), a raw provider SDK, or your own homegrown loop) wherever the comments say so.

> [!NOTE] Check for a plugin first
> If your framework is one of the ten with a first-party adapter, use that instead. The plugins give you per-tool containerization, model-turn replay and cross-run memory that this template does not. See [Agent frameworks](../../../integrations/agents/_index).

## How much control does the framework give you?

Frameworks differ in how much of the agent loop they own, which shapes how you integrate them:

| Level of control | What it means | Integration pattern |
|------------------|---------------|---------------------|
| **You own the loop** | The framework gives you primitives (graph nodes, tools) and you wire the control flow | Decorate nodes with `@flyte.trace` and run the compiled graph inside a task |
| **The framework owns the loop, you own the tools** | The framework runs the tool-calling loop; you provide tools as plain functions | Have tools delegate to durable `@env.task`s |
| **The framework owns everything** | The framework builds and runs the agent from configuration | Wrap the whole run in a task and trace the seam below the loop |

Whichever model your framework uses, the integration is the same in spirit: the framework decides *what* the agent does next, and Flyte decides *where and how durably* each step runs.

## The core pattern

Put your framework's agent invocation inside an `@env.task`. The task gives you a container, durable inputs/outputs, retries, and a span in the dashboard. Everything inside the task is ordinary Python, so the framework behaves exactly as it does locally.

```python
import flyte

# 1. Declare the runtime: image (with YOUR framework's deps), resources, secrets.
env = flyte.TaskEnvironment(
    name="my-agent",
    image=flyte.Image.from_debian_base(python_version=(3, 13)).with_pip_packages(
        # --> your agent framework + its provider packages go here, e.g.:
        # "crewai", "smolagents", "autogen-agentchat", ...
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    secrets=[flyte.Secret(key="ANTHROPIC_API_KEY")],  # --> your model provider key(s)
)

@env.task(report=True)
async def run_agent(prompt: str) -> str:
    # 2. Build/configure your framework's agent exactly as you would locally.
    #    --> your framework setup goes here
    # agent = MyFramework.Agent(model=..., tools=[...], instructions=...)

    # 3. Run it. Use the framework's own (sync or async) entry point.
    #    --> invoke your framework here
    # result = await agent.run(prompt)

    # 4. Return a serializable value (str, pydantic model, dataclass, ...).
    # return result.output
    ...

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(run_agent, prompt="...")  # --> your prompt / inputs
    print(run.url)
```

That is the whole integration. The remaining sections are optional enhancements that make the framework more durable and observable.

## Make tools durable

Most frameworks let a tool be any Python callable. To make a tool durable, retryable, and independently observable, have the framework's tool delegate to an `@env.task`. The framework still "owns" the tool; the heavy lifting runs on-cluster.

```python
# A durable task that does the real work (IO, compute, external calls).
@env.task
async def fetch_data(source: str) -> dict:
    # --> your real tool implementation (API call, DB query, scrape, ...)
    ...

# Register it with your framework using whatever tool API it exposes.
# The body just awaits the durable task.
#
#   @my_framework.tool                  # --> your framework's tool decorator/registration
#   async def get_data(source: str) -> dict:
#       """Tool description the LLM sees."""
#       return await fetch_data(source)   # runs as a Flyte task, durable + traced
```

> [!TIP]
> Reach for an `@env.task` when a tool does real work you want retried, cached, or run on its own hardware (GPU, more memory). For lightweight in-process helpers, a plain `@flyte.trace` function (below) is enough.

## Trace the framework's internals

If your framework exposes hooks, callbacks, or lets you wrap its node/step functions, decorate those with `@flyte.trace` to turn each LLM call, tool call, and routing decision into a span, with inputs and outputs captured and checkpointed.

```python
@flyte.trace
async def call_model(messages: list[dict]) -> str:
    # --> wrap the framework's model call (or pass this as the framework's LLM hook)
    ...

@flyte.trace
async def route(state) -> str:
    # --> wrap a routing / decision function so the branch is visible in the dashboard
    ...
```

For frameworks that don't expose hooks, wrap the whole run in `flyte.group(...)` to keep its trace tidy:

```python
@env.task(report=True)
async def run_agent(prompt: str) -> str:
    with flyte.group("my-framework-run"):  # groups everything below under one span
        # --> your framework invocation
        ...
```

## Fan out across containers

Run many independent agents in parallel, each in its own container, with `asyncio.gather()`. This works for any framework because each call is just an awaited task.

```python
import asyncio

@env.task
async def run_one(task_input: str) -> str:
    # --> one self-contained agent run for a single input
    ...

@env.task
async def run_many(inputs: list[str]) -> list[str]:
    # Each run_one call lands in its own container.
    results = await asyncio.gather(*[run_one(i) for i in inputs])
    return list(results)
```

## Checklist

To bring any Python agent framework onto Flyte:

1. **Wrap the run**: call the framework's entry point inside an `@env.task`.
2. **Declare deps**: add the framework + provider packages to the task's `image`.
3. **Supply secrets**: mount model-provider API keys via `flyte.Secret`.
4. **(Optional) Durable tools**: have tools delegate to `@env.task`s.
5. **(Optional) Observe**: decorate hooks/steps with `@flyte.trace`, or wrap in `flyte.group(...)`.
6. **(Optional) Scale**: fan out with `asyncio.gather()` for parallel, per-container runs.

## Next steps

- [Agent frameworks](../../../integrations/agents/_index): the ten frameworks with a first-party plugin, if yours is one of them.
- [The Flyte Agent harness](./flyte-agents): a built-in, batteries-included loop if you'd rather not bring a framework.
- [Build an agent with pure Python](./python-agents): hand-roll the loop with no framework at all.
- [Deploy an agent as a service](./deploy-agent-as-service): run your agent on a schedule or behind a webhook.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/agent-memory ===

# Agent memory

By default, an [`Agent`](./flyte-agents) is stateless: each `run` starts from a blank conversation. `MemoryStore` gives an agent continuity across runs by persisting both the conversation transcript and arbitrary path-addressed artifacts to a `flyte.io.Dir`. This is what lets a scheduled or webhook-driven agent remember what it did last time.

Use cases:

- An "inbox triage" agent that recalls which threads it has already responded to.
- A research agent that builds up a scratchpad over many days.
- Any sleep/wake pattern where the agent wakes on a schedule and resumes prior context.

## What a `MemoryStore` holds

A `MemoryStore` combines two complementary stores backed by a working directory:

- **`messages`**: the live LLM conversation transcript, managed by the agent. Mutate it only via `append()` / `extend()`.
- **Path-addressed files**: arbitrary named blobs under the working root. Read and write them with `read_text` / `write_text` / `read_json` / `write_json` / `list_paths`.

The on-disk layout under the root looks like:

```
<root>/messages.json                           # transcript
<root>/<your/path>.{txt,json,…}                # path-addressed entries
<root>/meta/<encoded_path>.json                # per-entry metadata (sha, actor, ts)
<root>/audit/log.jsonl                         # opt-in audit trail
<root>/versions/<encoded_path>/<ts>_<sha>.txt  # opt-in version history
```

## Sync vs async

The path-addressed I/O methods (`read_text`, `read_json`, `write_text`, `write_json`, `get_meta`, `current_sha`) and the lifecycle methods (`create`, `get_or_create`, `save`) are sync-by-default with an `.aio(...)` companion. Inside `async def` tasks, use the `.aio` form.

## Keyed stores: the easy path

For durable agent memory, use a **keyed store**. `MemoryStore.get_or_create(key=...)` loads an existing store if present, otherwise creates a new one, saving to a deterministic blob-store namespace under the active Flyte raw-data root:

```
{storage_root}/agents/memory-store/v0/{org}/{project}/{domain}/{key}
```

First, define the agent. Here it's a small research assistant with a single, **stateless** `web_search` tool. Its continuity comes from memory, not from the tool:

```
@env.task
async def web_search(query: str, max_results: int = 3) -> list[dict[str, str]]:
    """Search the web for `query` and return the top matching results.

    A stateless tool — it knows nothing about the agent's memory. But because the
    results it returns are recorded in the conversation transcript, the agent can
    recall or build on them in a later run without searching again.

    This stub returns canned results so the example runs offline. In a real
    agent, replace it with a call to a search API (Tavily, Brave, SerpAPI, …);
    keeping it an `@env.task` makes each search durable, retryable, and
    observable in the dashboard.
    """
    return [
        {
            "title": f"{query.title()} — overview ({i + 1})",
            "url": f"https://example.com/?q={query.replace(' ', '+')}&r={i + 1}",
            "snippet": f"Key point #{i + 1} about {query}.",
        }
        for i in range(max_results)
    ]

agent = Agent(
    name="memory-assistant",
    instructions=(
        "You are a personal research assistant with long-term memory. You "
        "remember what the user is working on and the facts they share, because "
        "your prior conversation transcript is always available. Use web_search "
        "to look things up, and reuse earlier findings from the conversation "
        "instead of searching again when you already have the answer."
    ),
    model="claude-haiku-4-5",
    tools=[web_search],
    max_turns=12,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/agent-memory/agent_with_memory.py*

Reuse the same `key` across runs to keep continuity. The chat task below picks up where the previous run left off (see the [full example](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/build-agent/agent-memory/agent_with_memory.py) for the `TaskEnvironment` setup):

CODE2

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/agent-memory/agent_with_memory.py*

The agent has no note-taking tools. Continuity comes entirely from the persisted transcript, and it remembers two kinds of things for free: the **facts the user shares** and the **results its tools return**. The first run records both in `messages.json`; a later run with the same `memory_key` reloads and prepends them, so the agent recalls earlier context, and reuses prior `web_search` findings instead of searching again. That is the core value of `MemoryStore`: no extra plumbing required.

## Working with a MemoryStore independently

Beyond the transcript, you can persist structured artifacts under arbitrary paths in the same store. This is optional (most agents get all the continuity they need from the transcript above) but it's useful when you want durable, queryable state (a scratchpad, a dedupe ledger, intermediate results).

A flyte task can commit its own artifact by loading the keyed store, read-modify-writing a path-addressed file, and calling `save()`. Every write is recorded in a metadata sidecar (sha256, actor, timestamp) and, by default, appended to an audit log:

CODE3

> [!NOTE] Coordinating tool writes with the transcript
> Artifacts live on independent paths (e.g. `notes/notes.json`) from the transcript (`messages.json`), so they never collide. But when a tool writes to the same keyed store that the orchestrator also saves, the orchestrator's working copy goes stale mid-run. Reload the store with `get_or_create` after `agent.run`, carry over the updated transcript (`reloaded.messages = result.memory.messages`), and save once. Otherwise the orchestrator's final save re-uploads a stale copy and clobbers the tool's artifact.

## Optimistic concurrency

When several tasks or agents share one keyed store (e.g. parallel tool calls, or a sleep/wake pattern), guard against lost updates by passing `expected_sha=`. The write succeeds only if the current content still matches; otherwise it raises `ConcurrencyError`:

CODE4

## Optional capabilities

`MemoryStore` (and `create` / `get_or_create`) accept a few flags:

| Option | Default | What it does |
|--------|---------|--------------|
| `audit` | `True` | Append every successful write to `audit/log.jsonl`. Inspect with `audit_tail(n)`. |
| `keep_versions` | `False` | Snapshot every write under `versions/` for full history (≈ 2× storage per write). |
| `read_only_prefixes` | `()` | Reject direct writes into the given prefixes (e.g. `("memory/",)`), so the agent must stage proposals elsewhere and a trusted pipeline promotes them. |

The internal `audit/`, `meta/`, and `versions/` prefixes and `messages.json` are reserved: writes to them are rejected, and they're excluded from `list_paths`.

## Passing memory to the agent

Memory is not attached to the agent: it is passed in per call and returned on the result. `agent.run(message, memory=store)` prepends the store's prior transcript, runs the loop, and appends the new turn back onto the store. Persisting is explicit: `run` never writes on its own, so call `memory.save()` (or `await memory.save.aio()`) yourself afterward.

CODE5

You can also pass a plain `list[dict]` of prior messages as `memory` for a stateless, single-shot history (nothing is persisted in that case).

## Lower-level usage

Every `MemoryStore` is **keyed**: there is no unkeyed/ephemeral store. You normally obtain one via `MemoryStore.create(key=...)` or `MemoryStore.get_or_create(key=...)`, but direct construction is supported for advanced use and serialization (`MemoryStore` is a Flyte I/O type, so it can be passed as a task input/output). `save()` takes no arguments: it always uploads the working root to the deterministic keyed `remote_path`. When `root` is omitted, a temporary working directory is created and cleaned up automatically.

## Next steps

- [The Flyte Agent harness](./flyte-agents): how the agent loop uses `memory`.
- [Deploy an agent as a service](./deploy-agent-as-service): schedule a memory-backed agent so it resumes context on each wakeup.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/agent-chat-ui ===

# Agent chat UI

A useful way to interact with an agent is through a chat interface. Because Flyte can [host apps](../../apps/build-apps/_index) behind a URL, you can serve a chat UI for your agent with no separate infrastructure. There are two approaches:

1. **`AgentChatAppEnvironment`**: the fastest path. Any agent that implements the `AgentProtocol` (including the built-in `Agent`, in tool-use or `code_mode`) gets a hosted chat shell, tool sidebar, and streaming for free.
2. **A custom FastAPI app**: full control over the UI. Wrap the agent in a `FastAPIAppEnvironment` and serve your own HTML/CSS/JS.

Both reuse the same agent object, so you can start with the built-in shell and graduate to a custom UI later.

## Option 1: the built-in chat UI

`flyte.ai.chat.AgentChatAppEnvironment` wraps an agent in a hosted chat app. Since `Agent` implements the `AgentProtocol`, it plugs straight in:

```
import flyte
from flyte.ai.agents import Agent
from flyte.ai.chat import AgentChatAppEnvironment, CustomTheme

task_env = flyte.TaskEnvironment(
    name="chat-agent-tools",
    image=flyte.Image.from_debian_base().with_pip_packages("litellm", "httpx"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
)

@task_env.task
async def search_docs(query: str, max_results: int = 3) -> list[dict[str, str]]:
    """Search internal documentation (stub) and return matching snippets."""
    corpus = [
        {"title": "Tasks", "body": "Define a task by decorating an async function with @env.task."},
        {"title": "Triggers", "body": "Schedule a task by attaching a flyte.Trigger with a flyte.Cron automation."},
        {"title": "Secrets", "body": "Mount cluster-managed secrets into a task with flyte.Secret(...)."},
    ]
    needle = query.lower()
    matches = [d for d in corpus if needle in d["body"].lower() or needle in d["title"].lower()]
    return matches[:max_results]

agent = Agent(
    name="docs-helper",
    instructions=(
        "You are a friendly internal docs assistant. Use search_docs to find "
        "relevant snippets. Always cite the doc title in your final answer."
    ),
    model="claude-haiku-4-5",
    tools=[search_docs],
    max_turns=8,
)

@task_env.task(report=True)
async def chat_entrypoint(message: str, memory: list[dict]) -> dict:
    """Parent task that owns the agent loop and the nested tool tasks."""
    result = await agent.run.aio(message, memory=memory)
    return {
        "summary": result.summary,
        "error": result.error,
        "attempts": result.attempts,
        "charts": [],
        "code": "",
    }

env = AgentChatAppEnvironment(
    name="docs-agent-chat-ui",
    agent=agent,
    task_entrypoint=chat_entrypoint,
    title="Internal docs assistant",
    subtitle="Backed by a flyte.ai.agents.Agent + durable Flyte task tools.",
    theme=CustomTheme(accent_color="#6F2AEF", accent_hover_color="#8B52F2"),
    prompt_nudges=[
        {"label": "Basics", "prompt": "Can you show me a hello world example?"},
        {"label": "Triggers", "prompt": "How do I schedule a task?"},
    ],
    depends_on=[task_env],
    image=flyte.Image.from_debian_base().with_pip_packages("litellm", "fastapi", "uvicorn"),
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    secrets=flyte.Secret("internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/agent-chat-ui/agent_chat_ui.py*

The `task_entrypoint` is a parent task that owns the agent loop, so the nested durable tool tasks run correctly under it. The chat shell streams progress by subscribing to the agent's `agent_progress_cb` events.

## Option 2: a custom FastAPI chat app

When you want to control the look and feel, wrap any `AgentProtocol`-compatible agent in a `FastAPIAppEnvironment` and serve your own UI. A natural fit is an `Agent` in **code mode** (`code_mode=True`): each turn the LLM writes Python that runs in a [sandbox](../sandboxing/_index) with the tools exposed as functions, returning code + a summary (and any charts you choose to surface), all behind a conversational web interface.

The architecture is small:

```
Browser (Chat UI)
  ├── GET  /            -> embedded HTML/CSS/JS chat interface
  ├── GET  /api/tools   -> JSON list of available tool descriptions
  └── POST /api/chat    -> { message, memory } -> { code, charts, summary, error }
           └── Agent.run(message, memory)
```

The app itself is just a FastAPI server. The endpoints call the agent's `run.aio` and `tool_descriptions` methods:

```python
import pathlib

from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

import flyte
from flyte.ai.agents import Agent
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(title="Chat Data Analytics Agent")

env = FastAPIAppEnvironment(
    name="chat-analytics-agent",
    app=app,
    image=flyte.Image.from_debian_base().with_pip_packages(
        "fastapi", "uvicorn", "httpx", "pydantic-monty", "litellm",
    ),
    secrets=flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    scaling=flyte.app.Scaling(replicas=1),
)

agent = Agent(
    name="analytics",
    instructions="You are a data analyst. Use the tools to fetch, aggregate, and chart data.",
    tools=ALL_TOOLS,
    code_mode=True,
    max_turns=15,
)

class ChatRequest(BaseModel):
    message: str
    memory: list[dict] = []

class ChatResponse(BaseModel):
    code: str = ""
    charts: list[str] = []
    summary: str = ""
    error: str = ""

@app.get("/api/tools")
async def get_tools() -> list[dict]:
    """Return JSON descriptions of available tool functions (for the sidebar)."""
    return agent.tool_descriptions()

@app.post("/api/chat")
async def chat(req: ChatRequest) -> ChatResponse:
    """Generate code, run it in the sandbox, and return results."""
    result = await agent.run.aio(req.message, memory=req.memory)
    return ChatResponse(code=result.code, charts=result.charts,
                        summary=result.summary, error=result.error)

@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    """Serve the embedded chat UI."""
    return HTMLResponse(content=CHAT_HTML)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_handle = flyte.serve(env)
    print(f"Deployed Chat Analytics Agent: {app_handle.url}")
```

`CHAT_HTML` is the embedded front-end (a single HTML string with the chat markup, styles, and a small fetch-based client that POSTs to `/api/chat` and renders the returned charts and summary). `ALL_TOOLS` is the agent's tool registry. Keeping both in their own modules means adding a tool is the only change required; the agent auto-generates its system prompt from each tool's signature and docstring.

Run it locally during development, then deploy with one command:

```bash
# Local development
python chat_app.py

# Deploy to Flyte
flyte deploy chat_app.py env
```

Flyte assigns a URL, handles TLS, and auto-scales the app.

> [!TIP]
> Drop `code_mode=True` to serve a standard tool-use [`Agent`](./flyte-agents) (or plug in any object implementing the `AgentProtocol`) behind the same UI. The endpoints only depend on `run.aio` and `tool_descriptions`.

## Next steps

- [The Flyte Agent harness](./flyte-agents): the agent powering the chat UI.
- [Sandboxing](../sandboxing/_index): how an `Agent` in code mode safely executes generated code.
- [Deploy an agent as a service](./deploy-agent-as-service): other ways to run an agent (task, schedule, webhook).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-agent/deploy-agent-as-service ===

# Deploy an agent as a service

Once you've built an agent (with [pure Python](./python-agents), the [`Agent` harness](./flyte-agents), or a [third-party framework](../../../integrations/agents/_index)), *how* you run it is an independent choice. The same agent object can be deployed in several ways:

| Pattern | When to use it | What invokes the agent |
|---------|----------------|------------------------|
| **As a task** | On-demand runs from the CLI, a notebook, or another service | `flyte.run(...)` |
| **As a scheduled task** | Recurring autonomous wakeups (triage, monitoring, reports) | A `flyte.Trigger` (cron or fixed-rate) |
| **Behind a webhook** | React to external events (GitHub, paging tools, CI) | An HTTP `POST` to an `AppEnvironment` |

All three wrap the agent loop in a regular Flyte task, so every run is durable, retryable, and observable in the Flyte dashboard. The examples below use the `Agent` harness, but the pattern is identical for any agent: just call your agent's entry point inside the task.

## As a task

The simplest deployment: put the agent loop in an `@env.task` and invoke it on demand. This works for any agent.

```python
import flyte
from flyte.ai.agents import Agent

env = flyte.TaskEnvironment(
    name="concierge-agent",
    image=flyte.Image.from_debian_base().with_pip_packages("litellm"),
    secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
)

agent = Agent(
    name="customer-concierge",
    instructions="You are a customer-service concierge.",
    tools=[...],
)

@env.task(report=True)
async def concierge(request: str) -> str:
    """Run the agent for a single request."""
    result = await agent.run.aio(request)
    return result.summary or result.error
```

Run it on demand:

```bash
flyte run agent.py concierge --request "Refund order #12345 to the customer."
```

Or from Python with `flyte.run(concierge, request="...")`. To register a stable, deployed version of the task, use `flyte deploy agent.py env`.

## As a scheduled task (via `Trigger`)

To run an agent autonomously on a schedule, attach a `flyte.Trigger` to the task. The "wakeup" is a regular Flyte task: the agent loop runs inside it, so every tool call is durable, observable, and retryable. Pair this with [agent memory](./agent-memory) so the agent resumes prior context on each wakeup.

```
agent = Agent(
    name="github-triage",
    instructions=(
        "You are a GitHub issue triager. For each wakeup: list open issues for "
        "the configured repo, classify each one, group them by severity, and "
        "post a concise digest to the team channel. Always end by calling post_digest."
    ),
    model="claude-haiku-4-5",
    tools=[list_open_issues, classify_issue, post_digest],
    max_turns=20,
)

@env.task(
    triggers=flyte.Trigger(
        "daily-triage",
        flyte.Cron("0 9 * * *"),  # every day at 09:00
        inputs={"trigger_time": flyte.TriggerTime, "repo": "flyteorg/flyte", "channel": "#flyte-triage"},
    ),
    report=True,
)
async def triage_repo(trigger_time: datetime, repo: str, channel: str) -> str:
    """Scheduled wakeup that runs the triage agent end-to-end."""
    message = f"It is {trigger_time.isoformat()}. Triage the open issues in {repo} and post a digest to {channel}."
    with flyte.group("triage-loop"):
        result = await agent.run.aio(message)
    return result.summary or f"[triage failed] {result.error}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/deploy/scheduled_triage_agent.py*

The agent's tools (`list_open_issues`, `classify_issue`, `post_digest`) are durable `@env.task`s; see the [full example](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/build-agent/deploy/scheduled_triage_agent.py) for their definitions.

Deploying the task registers the trigger; from then on Flyte wakes the agent on schedule. Use `flyte.Cron(...)` for calendar schedules or `flyte.FixedRate(...)` for fixed intervals. The `flyte.TriggerTime` input is filled with the scheduled fire time. See [Triggers](../../tasks/task-configuration/triggers) for the full schedule reference.

## Behind a webhook (`AppEnvironment`)

To kick off an agent run in response to an external event, deploy a small FastAPI app via an `AppEnvironment` that exposes an HTTP endpoint. The endpoint launches the agent task with `flyte.run.aio(...)`, so the long-running agent loop executes durably in the background while the webhook returns immediately with a run URL.

```
@tool_env.task(report=True)
async def review_pr(repo: str, pr_number: int, event: str) -> str:
    """Durable task that runs the agent for a single webhook event."""
    message = f"GitHub webhook fired for {repo}#{pr_number} (event={event}). Review the PR."
    result = await agent.run.aio(message)
    return result.summary or result.error

def _build_app():
    from fastapi import FastAPI

    api = FastAPI(title="flyte-agent-webhook")

    @api.post("/trigger")
    async def trigger(payload: dict) -> dict[str, str]:
        repo = payload.get("repository")
        pr_number = int(payload.get("pull_request", {}).get("number", 0))
        event = payload.get("action")
        run = await flyte.run.aio(review_pr, repo=repo, pr_number=pr_number, event=event)
        return {"run_url": run.url, "name": run.name}

    return api

webhook_env = flyte.app.AppEnvironment(
    name="flyte-agent-webhook",
    image=flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "litellm"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    requires_auth=True,
    depends_on=[tool_env],
)

@webhook_env.server
async def serve():
    import uvicorn

    config = uvicorn.Config(_build_app(), host="0.0.0.0", port=webhook_env.get_port().port)
    await uvicorn.Server(config).serve()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-agent/deploy/webhook_agent.py*

The agent and its tools (`fetch_pr`, `post_comment`) are defined in the [full example](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/build-agent/deploy/webhook_agent.py).

Once deployed, point your external system at the `/trigger` URL:

```bash
curl -X POST -H "Content-Type: application/json" \
    -d '{"repository": "flyteorg/flyte", "pull_request": {"number": 123}, "action": "opened"}' \
    https://<subdomain>.apps.<endpoint>/trigger
```

> [!NOTE]
> When the webhook app submits runs on behalf of incoming requests, it needs valid Flyte credentials. Use passthrough auth (a `FastAPIPassthroughAuthMiddleware` and `flyte.init_passthrough`) so the run is submitted with the caller's identity. See [FastAPI apps](../../apps/native-app-integrations/fastapi-app).

## Chat and other app patterns

- **Chat UI:** To let users converse with the agent in a browser, serve it behind a chat interface. See [Add a chat UI](./agent-chat-ui).
- **FastAPI endpoint:** For API-first agents, expose your agent behind a REST endpoint with `FastAPIAppEnvironment` so other services or agents can call it programmatically.
- **Model serving:** [Serve open-weight LLMs](../../apps/native-app-integrations/vllm-app) on GPUs behind an OpenAI-compatible API with `VLLMAppEnvironment` or `SGLangAppEnvironment`.

> [!TIP]
> See [Build Apps](../../apps/build-apps/_index) and [Configure Apps](../../apps/configure-apps/_index) for more details on hosting services on Flyte.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-mcp ===

# Build an MCP

Flyte supports serving [Model Context Protocol (MCP)](https://modelcontextprotocol.io) servers over HTTP.

There are two main MCP environment types:

| Environment | Purpose |
|-------------|---------|
| **`MCPAppEnvironment`** | Serve any FastMCP instance with custom tools |
| **`FlyteMCPAppEnvironment`** | Flyte-specific server that exposes Flyte operations as tools |

See the sub-pages for detailed guides:

- **Agents > Build an MCP > User-defined MCP server**: Build and deploy your own FastMCP instances
- **Agents > Build an MCP > Flyte MCP server**: Use Flyte-specific tools to interact with your cluster

## HTTP layout

All MCP app environments expose the same HTTP endpoints:

- `GET /health`: Liveness/readiness check (`{"status": "healthy"}`)
- `POST {mcp_mount_path}/mcp` or `{mcp_mount_path}/sse`: MCP protocol endpoint (default: `/mcp` for generic, `/flyte-mcp` for Flyte)

## Quickstart

The fastest way to try Flyte MCP is locally (no deployment needed):

```bash
uvx --from "flyte[mcp]>=2.5.18" flyte-mcp --transport stdio
```

`--transport stdio` is required when a client launches the server as a subprocess. The CLI defaults to `streamable-http`, which starts an HTTP listener instead.

For client setup, tool selection, allowlists, and remote deployment, see **Agents > Build an MCP > Flyte MCP server**.

## Subpages
- **Agents > Build an MCP > User-defined MCP server**
- **Agents > Build an MCP > Flyte MCP server**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-mcp/mcp_server ===

# User-defined MCP server

`MCPAppEnvironment` deploys any [FastMCP](https://modelcontextprotocol.io) instance as a long-running Flyte app, serving its tools over HTTP. Use it when you want to expose your *own* custom tools to AI assistants and LLM clients, rather than the built-in Flyte operations covered in [Flyte MCP server](./flyte_mcp_server).

## When to use it

Reach for `MCPAppEnvironment` when:

- You have domain-specific logic (database lookups, internal APIs, business rules, retrieval over your own corpus) that you want to package as MCP tools.
- You want an LLM client like Claude Code, Claude Desktop, or OpenCode to call those tools over a stable, authenticated HTTP endpoint.
- You want the tool server to run on Flyte infrastructure, with the same image, secrets, resources, and autoscaling story as any other app.

If instead you want assistants to *operate your Flyte cluster* (run tasks, inspect runs, build images, search docs), use [`FlyteMCPAppEnvironment`](./flyte_mcp_server). It ships those tools for you. The two can also run side by side as separate apps.

## How it works

You build a `FastMCP` instance and register tools on it with the `@mcp.tool()` decorator. `MCPAppEnvironment` wraps that instance in a Starlette + Uvicorn server and deploys it as an app. Every tool you defined on the instance is exposed automatically; there is no extra registration step.

## Basic example

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "mcp",
#    "starlette",
#    "uvicorn",
# ]
# ///

"""A basic MCP server app that serves a custom FastMCP instance.

This example shows how to deploy any FastMCP server as a Flyte app using
``MCPAppEnvironment``. The server exposes tools via the Model Context Protocol
(MCP) over HTTP.
"""

import flyte
from flyte.ai.mcp import MCPAppEnvironment
from mcp.server.fastmcp import FastMCP

# {{docs-fragment basic-mcp}}
mcp = FastMCP(name="demo-generic-mcp")

@mcp.tool()
def ping() -> str:
    """Health-style echo for demos."""
    return "pong"

@mcp.tool()
def add(a: int, b: int) -> int:
    """Add two numbers together."""
    return a + b

env = MCPAppEnvironment(
    name="generic-mcp-demo",
    mcp=mcp,
    transport="streamable-http",
    mcp_mount_path="/mcp",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

if __name__ == "__main__":
    flyte.init_from_config()
    handle = flyte.serve(env)
    handle.activate(wait=True)
    print(f"App is ready at {handle.endpoint}")
# {{/docs-fragment basic-mcp}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-mcp/basic_mcp_app.py*

Deploying with `flyte.serve(env)` activates the app and prints its public endpoint. See [Serve and deploy apps](../../apps/serve-and-deploy-apps/_index) for how deployment, activation, and scaling work in general.

## HTTP layout

The MCP ASGI app is mounted at `mcp_mount_path` (default `/mcp`). The resulting endpoints are:

CODE0

The session path is `{mcp_mount_path}/mcp` for `streamable-http` and `{mcp_mount_path}/sse` for `sse`. To get a cleaner URL such as `/mcp`, set `mcp_mount_path="/"`.

## Choosing a transport

| Transport | Session endpoint | When to use |
|-----------|------------------|-------------|
| `streamable-http` (default) | `{mcp_mount_path}/mcp` | The right choice for almost all remote deployments. Works with Claude Code, Claude Desktop, and OpenCode. |
| `sse` | `{mcp_mount_path}/sse` | Only when a client specifically requires a Server-Sent Events stream. Being phased out across the MCP ecosystem. |

CODE1

## Connecting a client

User-defined MCP servers are always deployed as remote HTTP apps; there is no local CLI equivalent. Once `flyte.serve(env)` prints your endpoint, register it with your client.

The session URL depends on your `mcp_mount_path` (default `/mcp`) and transport (default `streamable-http`), so the full endpoint is `https://<YOUR_HOST>/mcp/mcp`.

### Claude Code

CODE2

### OpenCode

CODE3

Replace `<YOUR_HOST>` with the hostname from `handle.endpoint`. If you deployed with a custom `mcp_mount_path`, adjust the path accordingly.

## Configuration tips

- **Resources**: Tool servers are typically I/O-bound and lightweight. `flyte.Resources(cpu=1, memory="512Mi")` is a good starting point; raise it only if a tool does heavy in-process work.
- **Secrets**: If your tools call external APIs, pass credentials with `secrets=...` rather than baking them into the image. See [secret-based authentication](../../apps/build-apps/secret-based-authentication).
- **Custom dependencies**: Add the libraries your tools need to the app `image`. Remember to include `mcp`, `starlette`, and `uvicorn` (these come from `pip install 'flyte[mcp]'`).
- **Extra files**: If your tools import local helper modules, use the `include` parameter so they ship with the app. See [including additional files](../../apps/configure-apps/including-additional-files).

## Best practices

1. **Write clear docstrings**: The docstring on each `@mcp.tool()` function is what the LLM sees when deciding whether and how to call it. Treat it as the tool's API contract: describe the purpose, arguments, and return value.
2. **Use precise type hints**: FastMCP derives the tool's input schema from your function signature, so accurate types lead to better-formed tool calls.
3. **Keep tools focused**: Prefer several small, single-purpose tools over one tool with many modes. LLMs select narrowly-scoped tools more reliably.
4. **Require auth in production**: Keep `requires_auth=True` (the default) so only authenticated clients can reach your tools.
5. **Right-size resources**: Start small (1 CPU, 512Mi) and scale up only if profiling shows you need to.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/build-mcp/flyte_mcp_server ===

# Flyte MCP server

`FlyteMCPAppEnvironment` exposes Flyte operations as standardized [MCP](https://modelcontextprotocol.io) tools, so AI assistants and LLM clients can drive your cluster programmatically: running tasks, monitoring runs, reading logs, inspecting actions, managing apps and triggers, and searching the SDK and docs.

Unlike [`MCPAppEnvironment`](./mcp_server), where you supply your own tools, this environment ships a curated set of Flyte tools out of the box. You decide which of them to expose and what they're allowed to touch.

> [!TIP] There is a prebuilt plugin
> [`flyte-agent-plugins`](https://github.com/flyteorg/flyte-agent-plugins) ships these same
> control-plane tools as a `flyte-cluster` MCP server, plus a hosted `flyte-docs` search
> server. Claude Code and Codex wire up both for you; Hermes, opencode, and pi take a few
> lines of config. See [Flyte agent plugins](../../../api-reference/agent-plugins).
>
> Build your own `FlyteMCPAppEnvironment` below when you need to scope the tools,
> allowlist resources, or deploy a shared server.

## When to use it

Use a Flyte MCP server when you want an assistant to *act on your cluster on your behalf*. Common scenarios:

- **Agentic development loops**: Let Claude Code or OpenCode run a task, wait for it, read its outputs, and iterate, without you copy-pasting commands.
- **Conversational operations**: Ask an assistant to list recent runs, check a run's status, or abort a stuck run.
- **Docs- and example-aware coding**: Enable the `search` tools so an assistant can ground its answers in the Flyte SDK examples and Union documentation.
- **Self-service automation**: Give a trusted internal agent a tightly-scoped server (a few allowlisted tasks, no destructive tools) to perform a narrow job.

If you only need to expose custom, non-Flyte tools, use [`MCPAppEnvironment`](./mcp_server) instead.

## How to run it

There are two ways to run a Flyte MCP server, suited to different stages:

| Mode | When to use |
|------|-------------|
| **Local (stdio)** | One user on one machine. The client launches the server as a subprocess, using your local Flyte config and your existing login. Nothing is deployed, and no data leaves your machine. |
| **Remote (HTTP)** | A whole team, or a client that cannot launch a subprocess (a browser-based assistant, for example). The server runs as a deployed app with a stable, authenticated URL. |

Prefer stdio unless you need one of the two things only HTTP gives you: a single shared server that nobody has to install, or a URL for a client that has no local process.

> [!NOTE] Remote deployment is not available here
> Deploying the server as a long-running app requires Union.ai apps, which open-source Flyte does not provide. Use the local stdio mode below. Everything else on this page — tool groups, individual tools, and allowlists — applies to both modes.

### Running locally with `uvx`

The `flyte[mcp]` extra ships a `flyte-mcp` CLI entrypoint. Run it with [`uvx`](https://docs.astral.sh/uv/guides/tools/) (no global install required):

```bash
uvx --from "flyte[mcp]>=2.5.18" flyte-mcp --transport stdio
```

`uvx` downloads `flyte[mcp]` into an isolated environment, runs `flyte-mcp`, and exits cleanly when you're done. The server reads your active Flyte config (the same one used by the `flyte` CLI or `flyte.init_from_config()`), so whichever project and cluster you're pointed at is what the tools operate on.

Two parts of that command matter:

- **`--transport stdio` is required.** The CLI defaults to `streamable-http`, which starts an HTTP listener instead of speaking JSON-RPC on stdin and stdout. A client that launches the process expecting stdio will not connect without this flag.
- **`>=2.5.18` is the minimum version.** It is the first release that constrains its `mcp` dependency below 2.0. Earlier versions resolve `mcp` 2.0.0, which removed the module the server imports, and the server exits at startup reporting `mcp is not installed`.

> [!TIP]
> Pin to an exact version instead of a floor: `uvx --from "flyte[mcp]==2.5.18" flyte-mcp --transport stdio`

The server starts even when no Flyte config is present. In that case the tools fail when the assistant calls them, rather than the server failing to start.

> [!NOTE] Skip the search corpus
> The three `search_*` tools grep a local copy of the Flyte SDK examples, the docs examples, and `llms.txt`. Enabling them makes the CLI clone roughly 120 MB into `~/.flyte/mcp` on first launch. Pass `--tool-groups` without `search` to skip it:
> ```bash
> uvx --from "flyte[mcp]>=2.5.18" flyte-mcp --transport stdio \
>   --tool-groups task,run,action,logs,app,trigger,project,secret,condition,identity
> ```

Once running, register it with your client as a **stdio** transport: the client manages the process lifetime. See the **Agents > Build an MCP > Flyte MCP server > Connecting a client** section below.

## Basic example

A server with **all** tools enabled. The environment definition is the same for both modes — only the last step differs, since `flyte.serve` deploys it as an app:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0",
#    "mcp",
#    "starlette",
#    "uvicorn",
# ]
# ///

"""A Flyte MCP server app that exposes Flyte operations as MCP tools.

This example deploys an MCP (Model Context Protocol) server that allows AI
assistants and LLM-based clients to interact with the Flyte control plane
using the standardized MCP protocol.

The server exposes tools for running tasks, monitoring runs, managing apps
and triggers, building container images, building and running UV scripts
remotely, and searching Flyte SDK/docs examples.

Requirements:
    pip install 'flyte[mcp]'

Usage:

    Deploy all tools

    $ python v2/user-guide/build-mcp/flyte_mcp_app.py

    Or serve locally for development (recommended: `uvx`)

    $ uvx --from "flyte[mcp]" flyte-mcp
"""

import flyte
from flyte.ai.mcp import FlyteMCPAppEnvironment

# {{docs-fragment flyte-mcp-all-tools}}
# Deploy an MCP server with all tools enabled
mcp_env = FlyteMCPAppEnvironment(
    name="flyte-mcp-server",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    transport="streamable-http",
    instructions=(
        "This MCP server provides tools to interact with the Flyte control plane. "
        "Use the available tools to run tasks, monitor runs, manage apps, build images, "
        "build and run UV scripts remotely, and search SDK/docs examples."
    ),
)

if __name__ == "__main__":
    flyte.init_from_config()
    app_handle = flyte.serve(mcp_env)
    app_handle.activate(wait=True)
    print(f"App is ready at {app_handle.endpoint}")
# {{/docs-fragment flyte-mcp-all-tools}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-mcp/flyte_mcp_app.py*

The default mount path is `/flyte-mcp`, so with the default `streamable-http` transport the MCP endpoint is `/flyte-mcp/mcp`.

> [!TIP] Set `instructions`
> The `instructions` string is sent to the LLM as guidance on what the server is for and how to use its tools. A clear, specific instruction string measurably improves how reliably the assistant picks the right tool.

## Scoping the server

A server with every tool enabled and no restrictions is convenient for trusted local use, but for anything shared you should narrow it down. There are three layers of control, from coarse to fine.

### 1. Tool groups

Tools are organized into groups. Pass `tool_groups` to enable only the groups you need:

CODE2

| Group | Tools | Typical use |
|-------|-------|-------------|
| `all` | All tools (default when both `tool_groups` and `tools` are omitted) | Trusted local development |
| `core` | No tools (only HTTP routes) | Health-check-only / building up explicitly |
| `task` | `run_task`, `get_task`, `list_tasks` | Launching and inspecting tasks |
| `run` | `get_run`, `get_run_io`, `abort_run`, `list_runs`, `wait_for_run`, `rerun_run` | Monitoring and controlling runs |
| `action` | `list_actions`, `get_action`, `abort_action` | Debugging a run step by step: phases, attempts, failure details, timing |
| `logs` | `get_logs` | Reading task logs |
| `app` | `get_app`, `list_apps`, `activate_app`, `deactivate_app` | Managing deployed apps |
| `trigger` | `list_triggers`, `get_trigger`, `activate_trigger`, `deactivate_trigger` | Managing triggers |
| `project` | `list_projects`, `get_project` | Discovering projects and domains |
| `secret` | `list_secrets`, `create_secret`, `delete_secret` | Managing secrets (names only — values are never returned) |
| `condition` | `list_conditions`, `signal_condition` | Answering human-in-the-loop gates |
| `identity` | `whoami` | Confirming which identity and org the server is acting as |
| `search` | `search_flyte_sdk_examples`, `search_flyte_docs_examples`, `search_full_docs` | Grounding answers in SDK/docs |

### 2. Individual tools

For the tightest control, pass `tools` with an explicit list of tool names instead of `tool_groups` (pass one or the other, not both):

CODE3

This is the way to build, for example, a strictly read-only server.

### 3. Allowlists

Even with a tool enabled, you can restrict *which resources* it may target. Allowlists are the safest way to expose `run_task` or app/trigger management to an agent:

CODE4

When an allowlist is set, calls targeting anything outside it are rejected. Omitting an allowlist leaves that resource type unrestricted.

## Enabling the search tools

The `search` tools need a corpus to scan, so you must point them at filesystem paths that exist *inside the app image*:

- `sdk_examples_path`: Flyte SDK examples (powers `search_flyte_sdk_examples`)
- `docs_examples_path`: Union examples (powers `search_flyte_docs_examples`)
- `full_docs_path`: the docs `llms.txt` index (powers `search_full_docs`)

The default image already clones the flyte-sdk and unionai-examples repos and downloads `llms.txt` into `/root` for you. If you supply a custom `image`, bake the corpora in yourself and pass matching paths:

CODE5

## Putting it together: a filtered server

This example combines tool groups, an allowlist, search paths, and instructions to build a scoped, production-ready server:

CODE6

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-mcp/flyte_mcp_app_filtered.py*

## Connecting a client

### Claude Code: local (stdio)

Registers the `flyte-mcp` process as a locally managed stdio server. Claude Code starts and stops the `uvx` process automatically:

CODE7

`--transport stdio` appears twice on purpose. The first tells Claude Code how to talk to the server; the second tells the server which transport to serve. Without the second, the server starts an HTTP listener and the connection fails.

### OpenCode: local

OpenCode spawns the `uvx` command for you:

CODE8

## Best practices

1. **Start broad locally, scope down for sharing**: Run with `all` tools via `uvx` while exploring, then enable only the groups or tools you need before deploying a shared server.
2. **Always allowlist mutating tools**: If `run_task`, `abort_run`, or app/trigger management is enabled on a shared server, set the corresponding allowlist so an agent can't touch arbitrary resources.
3. **Prefer `tools` over `tool_groups` for read-only servers**: An explicit allowlist of read tools is the clearest way to guarantee an agent can't change anything.
4. **Write specific `instructions`**: Describe what the server does and any constraints (e.g. "only allowlisted tasks can be run"). This guides tool selection and reduces wasted calls.
5. **Keep auth on**: Leave `requires_auth=True` so only authenticated clients can reach a deployed server.

## MCP tools reference

| Tool | Group | Description |
|------|-------|-------------|
| `run_task` | `task` | Run a task |
| `get_task` | `task` | Get task details |
| `list_tasks` | `task` | List tasks |
| `get_run` | `run` | Get a run |
| `get_run_io` | `run` | Get run inputs and outputs |
| `abort_run` | `run` | Abort a run |
| `list_runs` | `run` | List runs |
| `wait_for_run` | `run` | Wait for a run to finish |
| `rerun_run` | `run` | Re-run a prior run |
| `list_actions` | `action` | List the actions of a run |
| `get_action` | `action` | Get action details and timing |
| `abort_action` | `action` | Abort a single action |
| `get_logs` | `logs` | Read action logs |
| `get_app` | `app` | Get an app |
| `list_apps` | `app` | List apps |
| `activate_app` | `app` | Activate an app |
| `deactivate_app` | `app` | Deactivate an app |
| `list_triggers` | `trigger` | List triggers |
| `get_trigger` | `trigger` | Get a trigger |
| `activate_trigger` | `trigger` | Activate a trigger |
| `deactivate_trigger` | `trigger` | Deactivate a trigger |
| `list_projects` | `project` | List projects |
| `get_project` | `project` | Get a project |
| `list_secrets` | `secret` | List secret names |
| `create_secret` | `secret` | Create a secret |
| `delete_secret` | `secret` | Delete a secret |
| `list_conditions` | `condition` | List the conditions of a run |
| `signal_condition` | `condition` | Signal a waiting condition |
| `whoami` | `identity` | Show the caller's identity |
| `search_flyte_sdk_examples` | `search` | Search Flyte SDK examples |
| `search_flyte_docs_examples` | `search` | Search Flyte docs examples |
| `search_full_docs` | `search` | Search the full Flyte docs |


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/sandboxing ===

# Sandboxing

A **sandbox** is an isolated, secure environment where code can run without affecting the host system.
Sandboxes restrict what the executing code can do (limiting filesystem access, blocking network calls, and preventing arbitrary system operations) so that even malicious or buggy code cannot cause harm.

The exact restrictions depend on the sandboxing approach: some sandboxes eliminate dangerous operations entirely, while others provide full capabilities within an isolated, disposable container.

## Why sandboxing matters for AI

LLM-generated code is inherently untrusted.
The model may produce code that is correct and useful, but it can also produce code that is dangerous, and it does so without intent or awareness.

| Risk | Example |
|------|---------|
| Data destruction | `DELETE FROM orders WHERE 1=1`: wipes an entire table |
| Credential exfiltration | Reads environment variables and sends API keys to an external endpoint |
| Infinite loops | `while True: pass` consumes CPU indefinitely |
| Resource abuse | Spawns thousands of threads or allocates unbounded memory |
| Filesystem damage | `rm -rf /` or overwrites critical configuration files |
| Network abuse | Makes unauthorized API calls, sends spam, or joins a botnet |

Running LLM-generated code without a sandbox means trusting the model to never make these mistakes.
Sandboxing eliminates this trust requirement by making dangerous operations structurally impossible.

## Types of sandboxes

There are three broad approaches to sandboxing LLM-generated code, each with different tradeoffs:

| Type | How it works | Tradeoffs | Examples |
|------|-------------|-----------|----------|
| **One-shot execution** | Code runs to completion in a disposable container, then the container is discarded. Stdout, stderr, and outputs are captured. | Simple, no state reuse. Good for single-turn tasks. | Container tasks, serverless functions |
| **Interactive sessions** | A persistent VM or container where you send commands incrementally and observe results between steps. Sessions last for the lifetime of the VM. | Flexible and multi-turn, but heavier to provision and manage. | E2B, Daytona, fly.io |
| **Programmatic tool calling** | The LLM generates orchestration code that calls a predefined set of tools. The orchestration code runs in a sandbox while the tools run in full containers. | Durable, observable, and secure. Tools are known ahead of time. | Flyte workflow sandboxing |

## What Flyte offers

Flyte provides two complementary sandboxing approaches:

### Workflow sandbox (Monty)

A **sandboxed orchestrator** built on [Monty](https://github.com/pydantic/monty), a Rust-based sandboxed Python interpreter.
The sandbox starts in microseconds, runs pure Python control flow, and dispatches heavy work to full container tasks through the Flyte controller.

This enables the **programmatic tool calling** pattern (also known as code mode): LLMs generate Python orchestration code that invokes registered tools, and Flyte executes it safely with full durability, observability, and type checking.

### Code sandbox (container)

A **stateless code sandbox** that runs arbitrary Python scripts or shell commands inside an ephemeral Docker container.
The container is built on demand from declared dependencies, executed once, and discarded.

This is the right choice when you need full Python capabilities: third-party packages, file I/O, shell commands, or any computation that goes beyond pure control flow.

### When to use which

| | Workflow sandbox | Code sandbox |
|---|---|---|
| **Runtime** | Monty (Rust-based Python interpreter) | Ephemeral Docker container |
| **Startup** | Microseconds | Seconds (image build + container spin-up) |
| **Capabilities** | Pure Python control flow only: no imports, no I/O, no network | Full Python environment: any package, any library, full I/O |
| **Use case** | LLM-generated orchestration logic that calls registered tools | Arbitrary computation: data processing, test execution, ETL, shell pipelines |
| **State** | Runs within a worker container process | Stateless: fresh container per invocation |
| **Security model** | Dangerous operations are structurally impossible | Isolated container |

- Use the **workflow sandbox** when you need to run untrusted control flow (loops, conditionals, routing) that dispatches work to known tasks. It starts in microseconds and provides the strongest isolation guarantees.
- Use the **code sandbox** when you need full Python capabilities: third-party packages, file I/O, shell commands, or any computation that goes beyond pure control flow.

### Learn more

- **Agents > Sandboxing > Workflow sandboxing in Flyte**: How the Monty-based sandboxed orchestrator works, with examples
- **Agents > Sandboxing > Programmatic tool calling for agents**: The concept behind programmatic tool calling and how to build agents that use it
- **Agents > Sandboxing > Code sandboxing**: Running arbitrary code and commands in ephemeral containers with `flyte.sandbox.create()`

## Subpages
- **Agents > Sandboxing > Workflow sandboxing in Flyte**
- **Agents > Sandboxing > Programmatic tool calling for agents**
- **Agents > Sandboxing > Code sandboxing**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/sandboxing/workflow-sandboxing-flyte ===

# Workflow sandboxing in Flyte

Flyte provides a sandboxed orchestrator that lets you run pure Python control flow in a secure sandbox while dispatching heavy work to full container tasks.
This enables patterns where LLMs generate orchestration code dynamically, and Flyte executes it safely with full durability and observability.

## Why workflow sandboxing?

Three properties of Flyte make it a natural fit for sandboxed code execution:

1. **Infrastructure on demand**: Flyte spins up containers with specific permissions, secrets, and resources for each task.
2. **LLMs are great at Python**: Models trained on billions of lines of code can reliably generate Python orchestration logic.
3. **Microsecond startup**: The sandbox is powered by [Monty](https://github.com/pydantic/monty) (Pydantic's Rust-based Python interpreter), which starts in microseconds without the overhead of VMs or containers.

The result: LLMs generate the orchestration code (control flow, conditionals, loops), and Flyte tasks handle the heavy lifting (data access, computation, external APIs) in full containers.

## How it works

Your generated code runs inside one or more **Monty sandboxes**: lightweight Python interpreters embedded within a **worker container**. Each sandbox can execute pure Python (variables, loops, conditionals, function calls) but has no access to the filesystem, network, imports, or OS. A **bridge layer** acts as a hypervisor between the worker container and the sandboxes, handling opaque IO and routing callable tasks. When your code calls an external task, the bridge dispatches it, either as a method in the outer Python process or as a durable remote call through the Flyte controller (via the Queue Service):

```mermaid
flowchart TB
    subgraph worker["Worker Container"]
        subgraph bridge["Bridge / Hypervisor"]
            IO["Opaque IO: File, Dir, DataFrame"]
            subgraph sandbox1["Monty Sandbox 1"]
                A1["Your code: loops, variables, conditionals"]
                B1["result = add(x, y)"]
            end
            subgraph sandbox2["Monty Sandbox 2"]
                A2["More sandboxed code"]
            end
        end
    end

    A1 --> B1
    B1 -- "callable task" --> bridge
    bridge -- "result" --> B1
    IO -. "routed to tasks" .-> bridge
    bridge -- "external call" --> QS["Queue Service"]
    QS -- "completion" --> bridge
```

Each sandbox sees external tasks as opaque function calls. When your code hits one, Monty **pauses**, and the bridge layer dispatches the task, either directly in the outer Python process or as a remote durable call through the Flyte controller system (Queue Service). Once the call completes, Monty **resumes** with the result. Your code never knows the difference: it just looks like a regular function call that returns a value. Multiple Monty sandboxes can run within the same worker container, each isolated like a lightweight VM.

**Opaque IO types** like `File`, `Dir`, and `DataFrame` are managed by the bridge layer and pass through the sandbox without inspection. Your code can route them between tasks but cannot read or modify their contents.

## Example: sandboxed orchestrator

Use `@env.sandbox.orchestrator` to define a sandboxed task that calls regular worker tasks.
The orchestrator contains only pure Python control flow; all heavy computation runs in worker containers.

```python
import flyte

env = flyte.TaskEnvironment(name="sandboxed-demo")

# Worker tasks — run in their own containers
@env.task
def add(x: int, y: int) -> int:
    return x + y

@env.task
def multiply(x: int, y: int) -> int:
    return x * y

@env.task
def fib(n: int) -> int:
    """Compute the nth Fibonacci number iteratively."""
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
    return a

# Sandboxed orchestrator — pure Python control flow
@env.sandbox.orchestrator
def pipeline(n: int) -> dict[str, int]:
    fib_result = fib(n)
    linear_result = add(multiply(n, 2), 5)
    total = add(fib_result, linear_result)

    return {
        "fib": fib_result,
        "linear": linear_result,
        "total": total,
    }
```

When `pipeline` runs, Monty executes the control flow in the sandbox. Each call to `fib`, `multiply`, and `add` pauses the sandbox, runs the worker task in a container, and resumes with the result.

Both `def` and `async def` orchestrators are supported; Monty natively handles `await` expressions.

## Example: dynamic code execution

For cases where the code itself is generated at runtime (from templates, user input, or LLM output), use `orchestrator_from_str()` and `orchestrate_local()`.

### Reusable task from a code string

`orchestrator_from_str()` creates a reusable task template from a Python code string.
The value of the **last expression** becomes the return value.

```python
import flyte
import flyte.sandbox

env = flyte.TaskEnvironment(name="code-string-demo")

@env.task
def add(x: int, y: int) -> int:
    return x + y

@env.task
def multiply(x: int, y: int) -> int:
    return x * y

# Create a reusable task from a code string
compute_pipeline = flyte.sandbox.orchestrator_from_str(
    """
    partial = add(x, y)
    multiply(partial, scale)
    """,
    inputs={"x": int, "y": int, "scale": int},
    output=int,
    tasks=[add, multiply],
    name="compute-pipeline",
)
# flyte.run(compute_pipeline, x=2, y=3, scale=4)  → 20
```

### One-shot local execution

`orchestrate_local()` executes a code string and returns the result directly: no task template, no controller.
Use it for quick one-off computations.

```python
result = await flyte.sandbox.orchestrate_local(
    "add(x, y) * 2",
    inputs={"x": 1, "y": 2},
    tasks=[add],
)
# result → 6
```

### Parameterized code generation

Because the code is a string, you can generate it programmatically:

```python
def make_reducer(operation: str) -> flyte.sandbox.CodeTaskTemplate:
    """Create a sandboxed task that reduces a list using the given operation."""
    if operation == "sum":
        body = """
            acc = 0
            for v in values:
                acc = acc + v
            acc
        """
    elif operation == "product":
        body = """
            acc = 1
            for v in values:
                acc = acc * v
            acc
        """
    else:
        raise ValueError(f"Unknown operation: {operation}")

    return flyte.sandbox.orchestrator_from_str(
        body,
        inputs={"values": list},
        output=int,
        name=f"reduce-{operation}",
    )

sum_task = make_reducer("sum")
product_task = make_reducer("product")
```

## Building agents with programmatic tool calling

The sandboxed orchestrator and `orchestrate_local()` are the foundation for building agents that use **programmatic tool calling**: systems where an LLM generates Python orchestration code, and the sandbox executes it with registered tools.

Because `orchestrate_local()` accepts a plain code string and a list of tool functions, you can wire it into an LLM generate-execute-retry loop: the model writes code, the sandbox runs it, and on failure the error feeds back to the model for correction.

See [Programmatic tool calling for agents](./code-mode) for the full concept, agent implementation patterns, and end-to-end examples.

## Syntax restrictions

Monty enforces strict syntax restrictions to guarantee sandbox safety.
These restrictions are a feature, not a limitation: they ensure that sandboxed code is deterministic and side-effect free.

### Allowed

| Feature | Notes |
|---------|-------|
| Variables and assignment | `x = 1` |
| Arithmetic and comparisons | `x + y`, `x > y` |
| String operations | Concatenation, formatting |
| `if`/`elif`/`else` | Conditional logic |
| `for` loops | Iteration over lists, ranges, dicts |
| `while` loops | Condition-based loops |
| Function definitions (`def`) | Local helper functions |
| `async def` and `await` | Async orchestrators |
| List/dict/tuple literals | `[1, 2, 3]`, `{"key": "value"}` |
| List comprehensions | `[x * 2 for x in items]` |
| `.append()` on lists | Building lists incrementally |
| Subscript reading | `x = d["key"]`, `x = l[0]` |
| External task calls | Calling registered `@env.task` workers |
| `raise` | Raising exceptions |

### Not allowed

| Feature | Workaround |
|---------|------------|
| `import` | All available functions are provided directly |
| Subscript assignment (`d[k] = v`, `l[i] = v`) | Build dicts as literals; use `.append()` for lists |
| Augmented assignment (`x += 1`) | Use `x = x + 1` |
| `class` definitions | Use dicts or tuples |
| `with` statements | Not needed: no resource management in sandbox |
| `try`/`except` | Errors propagate to the controller |
| Walrus operator (`:=`) | Use separate assignment |
| `yield`/`yield from` | Not supported |
| `global`/`nonlocal` | Not supported |
| Set literals/comprehensions | Use lists |
| `del` statements | Not supported |
| `assert` statements | Use `if` + `raise` |

### Type restrictions

- **Primitive types**: `int`, `float`, `str`, `bool`, `bytes`, `None`
- **Collection types**: `list`, `dict`, `tuple` (including generic forms like `list[int]`, `dict[str, float]`)
- **Opaque IO handles**: `File`, `Dir`, `DataFrame`; pass-through only, cannot be inspected in the sandbox
- **Union types**: `Optional[T]` and `Union` of allowed types
- **Not allowed**: Custom classes, dataclasses, Pydantic models, or any user-defined types

## Security model

The sandboxed orchestrator provides security through restriction, not trust:

- **No filesystem access**: Cannot read, write, or list files
- **No network access**: Cannot make HTTP requests, open sockets, or resolve DNS
- **No OS access**: Cannot spawn processes, read environment variables, or access system resources
- **No imports**: Cannot load any Python modules
- **Opaque IO**: `File`, `Dir`, and `DataFrame` values pass through the sandbox without inspection; the sandbox can route them between tasks but cannot read their contents
- **Type-checked boundaries**: Inputs and outputs are validated against declared types at the sandbox boundary
- **Deterministic execution**: The same inputs always produce the same outputs (excluding external task results)

The sandbox runs untrusted code safely because dangerous operations are not just discouraged: they are structurally impossible in the Monty runtime.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/sandboxing/code-mode ===

# Programmatic tool calling for agents

**Programmatic tool calling** (also known as **code mode**) is a pattern where LLMs write executable code instead of making individual tool calls.
Rather than the model emitting a sequence of JSON tool-call objects and the system routing each one, the model generates a single block of code that calls multiple tools, transforms data, and applies logic, all executed in a sandbox.

The key insight: LLMs are trained on billions of lines of code, but only a small amount of synthetic tool-call data.
Code generation is a more natural and reliable output modality for models than structured tool-call schemas.

## Programmatic tool calling vs sequential tool calling

In sequential tool calling, every intermediate result passes through the model's context window.
The model calls one tool, reads the result, decides what to do next, calls another tool, and so on.
Each round-trip costs tokens and latency.

With programmatic tool calling, the model generates a complete program upfront.
The sandbox executes it, and only the final result returns to the model.

| Aspect | Sequential tool calling | Programmatic tool calling |
|--------|-------------|-----------|
| **Output format** | JSON tool-call objects, one at a time | A single block of executable code |
| **Data flow** | Every intermediate result passes through the model | Intermediate results stay in the sandbox |
| **Context overhead** | Grows with each tool call (all results in context) | Fixed: only tool signatures in context |
| **Multi-step logic** | Model re-invoked at every step | Sandbox executes loops, conditionals, transforms |
| **Scaling with tools** | Context grows linearly with number of tool definitions | Tools discovered progressively or loaded on demand |

## Why programmatic tool calling is powerful

### Token efficiency

Sequential tool calling loads all tool definitions into the context window upfront and passes every intermediate result through the model.
Programmatic tool calling reduces this dramatically:

- **98%+ context reduction** reported by Anthropic when using code execution with MCP servers: from 150,000 tokens down to 2,000 tokens for the same task.
- **99.9% reduction** reported by Cloudflare for large APIs: approximately 1,000 tokens with programmatic tool calling versus 1.17 million tokens when exposing each API endpoint as a separate tool.

### Performance

By eliminating round-trips through the model for intermediate steps, programmatic tool calling achieves significant speed improvements.
The sandbox evaluates conditionals, loops, and data transformations locally: no "time to first token" delay for each step.

### Natural programming patterns

Code naturally expresses patterns that are awkward or impossible in tool-call sequences:

- **Loops**: Process a list of items without the model deciding "call this tool again" for each one
- **Conditionals**: Branch on intermediate results without another model invocation
- **Data transformation**: Filter, map, and aggregate data before passing it to the next tool
- **Variable reuse**: Store intermediate results and reference them later

### Progressive tool discovery

Instead of loading hundreds of tool definitions into the context window, programmatic tool calling supports progressive discovery.
The model can search for relevant tools, load only what it needs, and compose them in code.

### Data privacy

Intermediate results stay in the sandbox execution environment.
They never re-enter the model's context window, which means sensitive data (PII, credentials, financial records) can be processed without the model seeing it.

## Example: sequential vs programmatic tool calling

Consider a task: "Analyze sales data, filter for Q4, calculate statistics, and create a chart."

### Sequential tool calling approach

The model makes serial tool calls, with each result passing through the context window:

```
Step 1: Model → tool_call: fetch_data("sales_2024")
        Result: [150KB of sales data] → back into model context

Step 2: Model → tool_call: filter_data(data, "month", ">=", "Oct")
        Result: [40KB of filtered data] → back into model context

Step 3: Model → tool_call: calculate_statistics(filtered, "revenue")
        Result: {"mean": 112000, ...} → back into model context

Step 4: Model → tool_call: create_chart("bar", "Q4 Revenue", ...)
        Result: "<canvas>...</canvas>" → back into model context
```

Four round-trips through the model.
The 150KB dataset enters the context window and stays there.

### Programmatic tool calling approach

The model generates a single code block:

```python
data = fetch_data("sales_2024")
q4_months = ["Oct", "Nov", "Dec"]
q4_data = [row for row in data if row["month"] in q4_months]
stats = calculate_statistics(q4_data, "revenue")

months = []
revenues = []
for row in q4_data:
    if row["month"] not in months:
        months.append(row["month"])
for month in months:
    total = 0
    for row in q4_data:
        if row["month"] == month:
            total = total + row["revenue"]
    revenues.append(total)

chart = create_chart("bar", "Q4 Revenue by Month", months, revenues)
{"charts": [chart], "summary": "Q4 stats: " + str(stats)}
```

One model invocation.
The data never re-enters the model's context window.
The sandbox handles the filtering, aggregation, and chart creation locally.

## Example: defining tools

Tools are plain Python functions with type annotations and docstrings.
The agent auto-generates its system prompt from these signatures, so adding a tool requires no other changes.

```python
async def fetch_data(dataset: str) -> list:
    """Fetch tabular data by dataset name.

    Available datasets:
    - "sales_2024": columns month, region, revenue, units
    - "employees": columns name, department, salary, years_exp, performance_rating
    - "website_traffic": columns date, page, visitors, bounce_rate, avg_duration
    - "inventory": columns product, category, stock, price, supplier
    """
    ...

async def create_chart(chart_type: str, title: str, labels: list, values: list) -> str:
    """Generate a self-contained Chart.js HTML snippet.

    Args:
        chart_type: One of "bar", "line", "pie", "doughnut".
        title: Chart title displayed above the canvas.
        labels: X-axis labels (or slice labels for pie/doughnut).
        values: Either a flat list of numbers, or a list of
                {"label": str, "data": list[number]} dicts for multi-series.
    """
    ...

async def calculate_statistics(data: list, column: str) -> dict:
    """Calculate descriptive statistics for a numeric column.

    Returns dict with keys: count, mean, median, min, max, std_dev.
    """
    ...

async def filter_data(data: list, column: str, operator: str, value: object) -> list:
    """Filter rows where column matches the condition.

    Operator: one of "==", "!=", ">", ">=", "<", "<=".
    """
    ...

ALL_TOOLS = {
    "fetch_data": fetch_data,
    "create_chart": create_chart,
    "calculate_statistics": calculate_statistics,
    "filter_data": filter_data,
}
```

The `ALL_TOOLS` dict is the single source of truth.
The agent introspects it to build the system prompt, and the sandbox uses it to resolve function calls.

## Example: programmatic tool-calling agent

The `CodeModeAgent` implements the generate-execute-retry loop:

```python
import flyte.sandbox
from _tools import ALL_TOOLS

class CodeModeAgent:
    def __init__(self, tools, *, model="claude-sonnet-4-6", max_retries=2):
        self._tools = tools
        self._model = model
        self._max_retries = max_retries
        # System prompt auto-generated from tool signatures + docstrings
        self.system_prompt = self._build_system_prompt()

    async def run(self, message: str, history: list[dict]) -> AgentResult:
        messages = [*history, {"role": "user", "content": message}]

        # Step 1: LLM generates Python code
        code = await generate_code(self._model, self.system_prompt, messages)

        # Step 2: Execute in Monty sandbox with registered tools
        for attempt in range(1 + self._max_retries):
            try:
                result = await flyte.sandbox.orchestrate_local(
                    code,
                    inputs={"_unused": 0},
                    tasks=list(self._tools.values()),
                )
                return AgentResult(code=code, charts=result.get("charts", []),
                                   summary=result.get("summary", ""))
            except Exception as exc:
                if attempt < self._max_retries:
                    # Step 3: Feed error back to LLM for retry
                    code = await generate_code(
                        self._model, self.system_prompt,
                        [*messages,
                         {"role": "assistant", "content": f"```python\n{code}\n```"},
                         {"role": "user", "content": f"Error: {exc}\nFix the code."}],
                    )
                    continue
                return AgentResult(code=code, error=str(exc))
```

The pattern:

1. **Generate**: The LLM receives tool signatures and the user's request, and outputs Python code.
2. **Execute**: The code runs in the Monty sandbox. Tool calls pause the sandbox, dispatch to real implementations, and resume with results.
3. **Retry**: If execution fails, the error message is fed back to the LLM, which generates a corrected version. This repeats up to `max_retries` times.

## Example: chat app

Wrap the agent in a FastAPI endpoint to create a conversational analytics assistant:

```python
from _agent import CodeModeAgent
from _tools import ALL_TOOLS
from fastapi import FastAPI

import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(title="Chat Data Analytics Agent")

env = FastAPIAppEnvironment(
    name="chat-analytics-agent",
    app=app,
    image=flyte.Image.from_debian_base().with_pip_packages(
        "fastapi", "uvicorn", "httpx", "pydantic-monty",
    ),
    secrets=flyte.Secret(key="anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
)

agent = CodeModeAgent(tools=ALL_TOOLS, max_retries=2)

@app.post("/api/chat")
async def chat(req: ChatRequest) -> ChatResponse:
    result = await agent.run(req.message, req.history)
    return ChatResponse(
        code=result.code,
        charts=result.charts,
        summary=result.summary,
        error=result.error,
    )
```

Users send natural language requests (`"Show me monthly revenue trends for 2024"`), the agent generates analysis code, the sandbox executes it with the registered tools, and the response includes charts and a text summary.

## Example: durable agent

For production workloads, wrap the tools as `@env.task` so the sandbox dispatches them as durable Flyte tasks through the controller.
This gives you execution history, retries, caching, and full observability.

```python
from _agent import CodeModeAgent
from _tools import ALL_TOOLS

import flyte
import flyte.report

env = flyte.TaskEnvironment(
    name="llm-code-mode",
    secrets=[flyte.Secret(key="anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "httpx", "pydantic-monty", "unionai-reuse",
    ),
)

# Wrap each tool as a durable task
@env.task
async def fetch_data(dataset: str) -> list:
    return await _tools.fetch_data(dataset)

@env.task
async def create_chart(chart_type: str, title: str, labels: list, values: list) -> str:
    return await _tools.create_chart(chart_type, title, labels, values)

# ... wrap remaining tools similarly ...

# Agent uses plain functions for prompt generation,
# @env.task versions for durable sandbox execution
durable_tools = {t.func.__name__: t for t in [fetch_data, create_chart, ...]}
agent = CodeModeAgent(tools=ALL_TOOLS, execution_tools=durable_tools)

@env.task(report=True)
async def analyze(request: str) -> str:
    """Run the code-mode agent and render an HTML report."""
    result = await agent.run(request, [])
    report_html = build_report(request, result)
    await flyte.report.replace.aio(report_html)
    await flyte.report.flush.aio()
    return result.summary
```

The key difference from the chat app: each tool call goes through the Flyte controller as a durable task.
If `fetch_data` fails, Flyte retries it automatically.
Every tool invocation is recorded and visible in the execution timeline.

Run it with:

```bash
flyte run durable_agent.py analyze \
    --request "Show me monthly revenue trends for 2024, broken down by region"
```

## References

- [Code execution with MCP](https://www.anthropic.com/engineering/code-execution-with-mcp): Anthropic engineering blog on the code execution pattern
- [Code Mode](https://blog.cloudflare.com/code-mode/): Cloudflare's introduction to code mode for LLM tool calling
- [Code Mode MCP](https://blog.cloudflare.com/code-mode-mcp/): Cloudflare's server-side code mode implementation
- [Code Mode Protocol](https://github.com/universal-tool-calling-protocol/code-mode): Open specification for the code mode pattern


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/agents/sandboxing/code-sandboxing ===

# Code sandboxing

`flyte.sandbox.create()` runs arbitrary Python code or shell commands inside an ephemeral, stateless Docker container.
The container is built on demand from declared dependencies, executed once, and discarded.
Each invocation starts from a clean slate: no filesystem state, environment variables, or side effects carry over between runs.

## Execution modes

`flyte.sandbox.create()` supports three mutually exclusive execution modes.

### Auto-IO mode

The default mode. Write only the business logic. Flyte generates the I/O boilerplate automatically.

How it works:

1. Flyte generates an `argparse` preamble that parses declared inputs from CLI arguments.
2. Declared inputs become local variables in scope.
3. After your code runs, Flyte writes declared scalar outputs to `/var/outputs/` automatically.

```python{hl_lines=[2, 4, 6, 11]}
import flyte
import flyte.sandbox

sandbox = flyte.sandbox.create(
    name="double",
    code="result = x * 2",
    inputs={"x": int},
    outputs={"result": int},
)

result = await sandbox.run.aio(x=21)  # returns 42
```

No imports, no argument parsing, no file writing. The variable `x` is available directly, and the variable `result` is captured automatically because it matches a declared output name.

A more involved example with third-party packages:

```python{hl_lines=["4-9", 12, 20, 24]}
import datetime

_stats_code = """\
import numpy as np
nums = np.array([float(v) for v in values.split(",")])
mean = float(np.mean(nums))
std  = float(np.std(nums))

window_end = dt + delta
"""

stats_sandbox = flyte.sandbox.create(
    name="numpy-stats",
    code=_stats_code,
    inputs={
        "values": str,
        "dt": datetime.datetime,
        "delta": datetime.timedelta,
    },
    outputs={"mean": float, "std": float, "window_end": datetime.datetime},
    packages=["numpy"],
)

mean, std, window_end = await stats_sandbox.run.aio(
    values="1,2,3,4,5",
    dt=datetime.datetime(2024, 1, 1),
    delta=datetime.timedelta(days=1),
)
```

When there are multiple outputs, `.run()` returns them as a tuple in declaration order.

### Verbatim mode

Set `auto_io=False` to run a complete Python script with full control over I/O.
Flyte runs the script exactly as written: no injected preamble, no automatic output collection.

Your script must:

- Read inputs from `/var/inputs/<name>` (files are bind-mounted at these paths)
- Write outputs to `/var/outputs/<name>`

```python{hl_lines=["4-9", 12, 17]}
from flyte.io import File

_etl_script = """\
import json, pathlib

payload = json.loads(pathlib.Path("/var/inputs/payload").read_text())
total = sum(payload["values"])

pathlib.Path("/var/outputs/total").write_text(str(total))
"""

etl_sandbox = flyte.sandbox.create(
    name="etl-script",
    code=_etl_script,
    inputs={"payload": File},
    outputs={"total": int},
    auto_io=False,
)

total = await etl_sandbox.run.aio(payload=payload_file)
```

Use verbatim mode when you need precise control over how inputs are read and outputs are written, or when your script has its own argument parsing.

### Command mode

Run any shell command, binary, or pipeline. Provide `command` instead of `code`.

```python{hl_lines=[5]}
from flyte.io import File

linecount_sandbox = flyte.sandbox.create(
    name="line-counter",
    command=[
        "/bin/bash",
        "-c",
        "grep -c . /var/inputs/data_file > /var/outputs/line_count || echo 0 > /var/outputs/line_count",
    ],
    inputs={"data_file": File},
    outputs={"line_count": str},
)

count = await linecount_sandbox.run.aio(data_file=data_file)
```

Command mode is useful for running test suites, compiled binaries, shell pipelines, or any non-Python workload.

Use `arguments` to pass positional arguments to the command.
File inputs are bind-mounted at `/var/inputs/<name>` and can be referenced in the arguments list:

```python{hl_lines=[4, 5]}
sandbox = flyte.sandbox.create(
    name="test-runner",
    command=["/bin/bash", "-c", pytest_cmd],
    arguments=["/var/inputs/solution.py", "/var/inputs/tests.py"],
    inputs={"solution.py": File, "tests.py": File},
    outputs={"exit_code": str},
)
```

## Executing a sandbox

Call `.run()` on the sandbox object to build the image and execute.

**Async execution**

```python
result = await sandbox.run.aio(x=21)
```

**Sync execution**

```python
result = sandbox.run(x=21)
```

Both forms build the container image (if not already built), start the container, execute the code or command, collect outputs, and discard the container.

`flyte.sandbox.create()` defines the sandbox configuration and can be called at module level or inside a task. The actual container execution happens when you call `.run()`, which must run inside a Flyte task (either locally or remotely on the cluster).

### Error handling

If the sandbox code fails (non-zero exit code, Python exception, or timeout), `.run()` raises an exception with the error details.
If `retries` is set, Flyte automatically retries the execution before surfacing the error.

If the image build fails due to an invalid package, an `InvalidPackageError` is raised with the package name and the underlying error message.

## Supported types

Inputs and outputs must use one of the following types:

| Category         | Types                                     |
| ---------------- | ----------------------------------------- |
| **Primitive**    | `int`, `float`, `str`, `bool`             |
| **Date/time**    | `datetime.datetime`, `datetime.timedelta` |
| **File handles** | `flyte.io.File`                           |

### How types are handled

**In auto-IO mode:**

- **Primitive and date/time inputs** are injected as local variables with the correct Python type. Flyte generates an `argparse` preamble behind the scenes. Your code just uses the variable names directly.
- **`File` inputs** are bind-mounted into the container. The input variable contains the file path as a string (e.g., `"/var/inputs/payload"`), so you can read it with `pathlib.Path(payload).read_text()`.
- **Primitive and date/time outputs** are written to `/var/outputs/<name>` automatically. Just assign the value to a variable matching the declared output name.
- **`File` outputs** are the exception: your code must write the file to `/var/outputs/<name>` manually.

**In verbatim mode:**

- All inputs (including primitives) are available at `/var/inputs/<name>`. Your script reads them directly from the filesystem.
- All outputs must be written to `/var/outputs/<name>` by your script.

**In command mode:**

- `File` inputs are bind-mounted at `/var/inputs/<name>`.
- All outputs must be written to `/var/outputs/<name>` by your command.

## Configuring the container image

### Python packages

Install PyPI packages with `packages`:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="data-analysis",
    code="...",
    inputs={"data": str},
    outputs={"result": str},
    packages=["numpy", "pandas>=2.0", "scikit-learn"],
)
```

### System packages

Install system-level (apt) packages with `system_packages`:

```python{hl_lines=[7]}
sandbox = flyte.sandbox.create(
    name="image-processor",
    code="...",
    inputs={"image": File},
    outputs={"result": File},
    packages=["Pillow"],
    system_packages=["libgl1-mesa-glx", "libglib2.0-0"],
)
```

> [!NOTE]
> `gcc`, `g++`, and `make` are included automatically in every sandbox image.

### Additional Dockerfile commands

For advanced image customization, use `additional_commands` to inject arbitrary `RUN` commands into the Dockerfile:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="custom-env",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    additional_commands=["curl -sSL https://example.com/setup.sh | bash"],
)
```

### Pre-built images

Skip the image build entirely by providing a pre-built image URI:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="prebuilt",
    code="result = x + 1",
    inputs={"x": int},
    outputs={"result": int},
    image="ghcr.io/my-org/my-sandbox-image:latest",
)
```

### Image configuration

Control the registry and Python version with `ImageConfig`:

```python{hl_lines=["8-12"]}
from flyte.sandbox import ImageConfig

sandbox = flyte.sandbox.create(
    name="custom-registry",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    image_config=ImageConfig(
        registry="ghcr.io/my-org",
        registry_secret="ghcr-credentials",
        python_version=(3, 12),
    ),
)
```

## Runtime configuration

### Resources

Set CPU and memory limits for the container:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="heavy-compute",
    code="...",
    inputs={"data": str},
    outputs={"result": str},
    resources=flyte.Resources(cpu=4, memory="8Gi"),
)
```

The default is 1 CPU and 1Gi memory.

### Retries

Automatically retry failed executions:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="flaky-task",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    retries=3,
)
```

### Timeout

Set a maximum execution time in seconds:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="bounded-task",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    timeout=300,  # 5 minutes
)
```

### Environment variables

Inject environment variables into the container:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="configured-task",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    env_vars={"LOG_LEVEL": "DEBUG", "FEATURE_FLAG": "true"},
)
```

### Secrets

Mount Flyte secrets into the container:

```python{hl_lines=[6]}
sandbox = flyte.sandbox.create(
    name="authenticated-task",
    code="...",
    inputs={"query": str},
    outputs={"result": str},
    secrets=[flyte.Secret(key="api-key", as_env_var="API_KEY")],
)
```

### Caching

Control output caching behavior:

```python{hl_lines=["6-8"]}
sandbox = flyte.sandbox.create(
    name="cached-task",
    code="...",
    inputs={"x": int},
    outputs={"y": int},
    cache="auto",      # default — Flyte decides based on inputs
    # cache="override"  # force re-execution and update cache
    # cache="disable"   # no caching
)
```

## Deploying a sandbox as a task

Use `.as_task()` to convert a sandbox into a deployable `ContainerTask`.
The returned task has the generated script pre-filled as a default input, so retriggers from the UI only require user-declared inputs.

This pattern is useful when you want to define a sandbox dynamically (for example, with LLM-generated code) and then deploy it as a standalone task that others can trigger from the UI.

```python{hl_lines=[4, 11, "33-38"]}
import flyte
import flyte.sandbox
from flyte.io import File
from flyte.sandbox import sandbox_environment

# sandbox_environment provides the base runtime image for code sandboxes.
# Include it in depends_on so Flyte builds the sandbox runtime before your task runs.
env = flyte.TaskEnvironment(
    name="sandbox-demo",
    image=flyte.Image.from_debian_base(name="sandbox-demo"),
    depends_on=[sandbox_environment],
)

@env.task
async def deploy_sandbox_task() -> str:
    # Initialize the Flyte client for in-cluster operations (image building, deployment)
    flyte.init_in_cluster()

    sandbox = flyte.sandbox.create(
        name="deployable-sandbox",
        # In auto-IO mode, File inputs become path strings — read with pathlib
        code="""\
import json, pathlib
data = json.loads(pathlib.Path(payload).read_text())
total = sum(data["values"])
""",
        inputs={"payload": File},
        outputs={"total": int},
        resources=flyte.Resources(cpu=1, memory="512Mi"),
    )

    # Build the image and get a ContainerTask with the script pre-filled
    task = await sandbox.as_task.aio()

    # Create a TaskEnvironment from the task and deploy it
    deploy_env = flyte.TaskEnvironment.from_task("deployable-sandbox", task)
    versions = flyte.deploy(deploy_env)

    return versions[0].summary_repr()
```

## End-to-end example

The following example defines sandboxes in all three modes, creates helper tasks, and runs everything in a single pipeline:

```
import datetime
from pathlib import Path

import flyte
import flyte.sandbox
from flyte.io import File
from flyte.sandbox import sandbox_environment

# {{docs-fragment create}}
# sandbox_environment provides the base runtime for code sandboxes.
# Include it in depends_on so the sandbox runtime is available when tasks execute.
env = flyte.TaskEnvironment(
    name="sandbox-demo",
    image=flyte.Image.from_debian_base(name="sandbox-demo"),
    depends_on=[sandbox_environment],
)

# Auto-IO mode: pure computation. The code string runs in an isolated sandbox;
# only the declared inputs go in and only the declared outputs come back.
sum_sandbox = flyte.sandbox.create(
    name="sum-to-n",
    code="total = sum(range(n + 1)) if conditional else 0",
    inputs={"n": int, "conditional": bool},
    outputs={"total": int},
)
# {{/docs-fragment create}}

# Auto-IO mode with packages
_stats_code = """\
import numpy as np
nums = np.array([float(v) for v in values.split(",")])
mean = float(np.mean(nums))
std  = float(np.std(nums))
window_end = dt + delta
"""

stats_sandbox = flyte.sandbox.create(
    name="numpy-stats",
    code=_stats_code,
    inputs={
        "values": str,
        "dt": datetime.datetime,
        "delta": datetime.timedelta,
    },
    outputs={"mean": float, "std": float, "window_end": datetime.datetime},
    packages=["numpy"],
)

# Verbatim mode: full script control
_etl_script = """\
import json, pathlib

payload = json.loads(pathlib.Path("/var/inputs/payload").read_text())
total = sum(payload["values"])

pathlib.Path("/var/outputs/total").write_text(str(total))
"""

etl_sandbox = flyte.sandbox.create(
    name="etl-script",
    code=_etl_script,
    inputs={"payload": File},
    outputs={"total": int},
    auto_io=False,
)

# Command mode: shell pipeline
linecount_sandbox = flyte.sandbox.create(
    name="line-counter",
    command=[
        "/bin/bash",
        "-c",
        "grep -c . /var/inputs/data_file > /var/outputs/line_count || echo 0 > /var/outputs/line_count",
    ],
    inputs={"data_file": File},
    outputs={"line_count": str},
)

@env.task
async def create_text_file() -> File:
    path = Path("/tmp/data.txt")
    path.write_text("line 1\n\nline 2\n")
    return await File.from_local(str(path))

@env.task
async def payload_generator() -> File:
    path = Path("/tmp/payload.json")
    path.write_text('{"values": [1, 2, 3, 4, 5]}')
    return await File.from_local(str(path))

@env.task
async def run_pipeline() -> dict:
    # Auto-IO: sum 1..10 = 55
    total = await sum_sandbox.run.aio(n=10, conditional=True)

    # Auto-IO with numpy
    mean, std, window_end = await stats_sandbox.run.aio(
        values="1,2,3,4,5",
        dt=datetime.datetime(2024, 1, 1),
        delta=datetime.timedelta(days=1),
    )

    # Verbatim ETL
    payload = await payload_generator()
    etl_total = await etl_sandbox.run.aio(payload=payload)

    # Command mode: line count
    data_file = await create_text_file()
    line_count = await linecount_sandbox.run.aio(data_file=data_file)

    return {
        "sum_1_to_10": total,
        "mean": round(mean, 4),
        "std": round(std, 4),
        "window_end": window_end.isoformat(),
        "etl_sum_1_to_10": etl_total,
        "line_count": line_count,
    }

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.run(run_pipeline)
    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/sandboxing/code_sandbox.py*

## API reference

### `flyte.sandbox.create()`

| Parameter             | Type              | Description                                                |
| --------------------- | ----------------- | ---------------------------------------------------------- |
| `name`                | `str`             | Sandbox name. Derives task and image names.                |
| `code`                | `str`             | Python source to run. Mutually exclusive with `command`.   |
| `inputs`              | `dict[str, type]` | Input type declarations.                                   |
| `outputs`             | `dict[str, type]` | Output type declarations.                                  |
| `command`             | `list[str]`       | Shell command to run. Mutually exclusive with `code`.      |
| `arguments`           | `list[str]`       | Arguments forwarded to `command`.                          |
| `packages`            | `list[str]`       | Python packages to install via pip.                        |
| `system_packages`     | `list[str]`       | System packages to install via apt.                        |
| `additional_commands` | `list[str]`       | Extra Dockerfile `RUN` commands.                           |
| `resources`           | `flyte.Resources` | CPU and memory limits. Default: 1 CPU, 1Gi memory.         |
| `image_config`        | `ImageConfig`     | Registry and Python version settings.                      |
| `image_name`          | `str`             | Explicit image name (overrides auto-generated).            |
| `image`               | `str`             | Pre-built image URI (skips build).                         |
| `auto_io`             | `bool`            | Auto-generate I/O wiring. Default: `True`.                 |
| `retries`             | `int`             | Number of retries on failure. Default: `0`.                |
| `timeout`             | `int`             | Timeout in seconds.                                        |
| `env_vars`            | `dict[str, str]`  | Environment variables for the container.                   |
| `secrets`             | `list[Secret]`    | Flyte secrets to mount.                                    |
| `cache`               | `str`             | `"auto"`, `"override"`, or `"disable"`. Default: `"auto"`. |

### Sandbox methods

| Method                            | Description                                                       |
| --------------------------------- | ----------------------------------------------------------------- |
| `sandbox.run(**kwargs)`           | Build the image and execute synchronously. Returns typed outputs. |
| `await sandbox.run.aio(**kwargs)` | Async version of `run()`.                                         |
| `sandbox.as_task()`               | Build the image and return a deployable `ContainerTask`.          |
| `await sandbox.as_task.aio()`     | Async version of `as_task()`.                                     |

Both `run()` and `as_task()` accept an optional `image` parameter to provide a pre-built image URI, skipping the build step.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/project-patterns ===

# Project patterns

The rest of the user guide explains what Flyte can do.
This section explains how we recommend you structure real projects.

These are opinionated guides. They represent patterns we've seen work well across many teams and production deployments. If you're starting a new project or scaling an existing one, start here.

### **Project patterns > Bring your own image**

Two patterns for teams that own their Docker images and want Flyte for orchestration without handing over their build pipeline.

### **Project patterns > Structuring Flyte projects with uv**

How to structure Flyte projects with uv, from single-package setups to multi-team monorepos with shared and independent lockfiles.

### **Project patterns > CI/CD deployments**

How to deploy a Flyte project from CI. Uses GitHub Actions as the reference, but the building blocks (API key, `flyte deploy`, commit-pinned versions) translate to any runner.

## Subpages
- **Project patterns > Bring your own image**
- **Project patterns > Structuring Flyte projects with uv**
- **Project patterns > CI/CD deployments**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/bring-your-own-image ===

# Bring your own image

This guide is for teams who own their Docker images and want Flyte for orchestration without handing over their build pipeline.

> [!NOTE]
> This guide does **not** cover `flyte.Image.from_debian_base()`, the Flyte-managed image builder. It assumes you already have images.

## The multi-team problem

Two teams. Two images. One workflow.

| | Team A (data-prep) | Team B (training) |
|---|---|---|
| Base | `python:3.11-slim` | `python:3.10-slim` (prod: CUDA) |
| Python | 3.11 | 3.10 |
| WORKDIR | `/app` | `/workspace` |
| Packages | pandas, pyarrow | torch, numpy |

The `prepare` task runs in Team A's container. It processes the input and calls `train`, which runs in Team B's container. One workflow, two images, different filesystem layouts.

Two patterns solve this. Pick based on who controls what:

| | Pattern 1: Pure BYOI | Pattern 2: Remote Builder |
|---|---|---|
| Who owns the image? | Each team owns everything | Each team owns the base |
| Flyte-aware? | Yes: code is baked in | No: Flyte adapts on top |
| Code change = image rebuild? | Yes | No |
| Use when | Teams can't let Flyte touch images | Teams can hand off a base |

## Pattern 1: Pure BYOI

Teams build complete, Flyte-aware images. Workflow code is COPYed into the Dockerfile. Flyte runs the container as a black box: it sends no code and modifies nothing.

### Dockerfiles

Both teams install `flyte` and COPY the shared `workflow_code/` into their image. The only difference is their base, Python version, and WORKDIR.

**Team A (data prep):**

```dockerfile
# Team A's image: data preparation
#
# This team owns this entire Dockerfile. They control Python version, WORKDIR,
# and PYTHONPATH. Flyte has no say here.
#
# Pure BYOI constraint: workflow code must be baked in because there is no
# code bundle. Every code change requires rebuilding and pushing this image.

FROM python:3.11-slim

# System deps this team needs
RUN apt-get update && apt-get install -y --no-install-recommends \
    libpq-dev \
    && rm -rf /var/lib/apt/lists/*

# Team A's WORKDIR. Python will find modules here because PYTHONPATH includes it.
WORKDIR /app

# Team A's Python packages. These are their own dependencies — Flyte doesn't
# install anything on top in pure BYOI mode.
RUN pip install --no-cache-dir \
    flyte \
    pandas==2.1.4 \
    pyarrow==14.0.1

# Bake the workflow code into the image.
# In pure BYOI, this is the ONLY way Flyte can find your task functions.
# Downside: every edit to tasks.py requires a new image tag + CI build.
COPY workflow_code/ /app/workflow_code/

# /app is on PYTHONPATH so `import workflow_code.tasks` resolves at runtime.
ENV PYTHONPATH=/app
CODE0dockerfile
# Team B's image: model training (GPU workload)
#
# Intentionally different from Team A:
#   - Different base: CUDA runtime instead of debian slim
#   - Different Python version: 3.10 (team B's standard)
#   - Different WORKDIR: /workspace (not /app)
#   - Different PYTHONPATH: /workspace
#   - Different system packages: CUDA tools
#
# Pure BYOI: Flyte injects nothing at runtime. Everything must be baked in.
# Each team owns their filesystem layout entirely.
#
# In practice this would be:
#   FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
# Using python:3.10-slim here so you can build/test without a GPU machine.

FROM python:3.10-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Team B uses /workspace, not /app. Each team controls their own layout.
WORKDIR /workspace

RUN pip install --no-cache-dir \
    flyte \
    torch==2.1.2 \
    numpy==1.26.4

# Bake in workflow code at /workspace/workflow_code/.
# Python finds it because PYTHONPATH includes /workspace.
COPY workflow_code/ /workspace/workflow_code/

# /workspace is on PYTHONPATH — same import path as Team A's image despite
# the different WORKDIR. Both images expose `import workflow_code.tasks`.
ENV PYTHONPATH=/workspace
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/pure_byoi/training/Dockerfile*

Both expose `import workflow_code.tasks` at runtime because each image's PYTHONPATH points to its own WORKDIR where the code was COPYed.

### Build and push

The build context is the `pure_byoi/` directory so that `workflow_code/` is available to both Dockerfiles:

CODE1

### Python code

**Environment definitions**. Image names are specified via `from_ref_name()`:

CODE2

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/pure_byoi/workflow_code/envs.py*

`from_ref_name()` is a placeholder resolved at runtime. The actual URIs are passed in the entry point via `init_from_config(images=...)`. This is necessary because the envs file is COPYed into both images. Hardcoding a URI would create a circular reference.

**Task definitions:**

CODE3

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/pure_byoi/workflow_code/tasks.py*

**Entry point**. This is where image URIs are wired in:

CODE4

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/pure_byoi/main.py*

### Run and deploy

CODE5

There is no separate deploy step. The image tag is the version. To ship a code change: edit tasks, rebuild both images, push new tags, update the tag constants in `main.py`, run again.

## Pattern 2: Remote builder

Teams hand you their base images. They built these images for their own purposes. Flyte was never a consideration. Your job is to adapt them.

### The base images

**Team A** uses `continuumio/miniconda3` as their base:

CODE6

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/remote_builder/data_prep/Dockerfile*

- Python at `/opt/conda/bin/python` (conda manages this)
- conda's Dockerfile already adds `/opt/conda/bin` to `PATH`
- No PYTHONPATH set
- WORKDIR `/app`

**Team B** uses `python:3.10-slim` with a pip venv at `/opt/venv`:

CODE7

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/remote_builder/training/Dockerfile*

- Python at `/opt/venv/bin/python`
- `PATH` does **not** include `/opt/venv/bin`: the venv was created but never activated
- No PYTHONPATH set
- WORKDIR `/workspace`

### Adapting with `flyte.Image`

`flyte.Image.from_base()` takes the base URI and lets you layer on top. This is where the adaptation happens:

CODE8

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/remote_builder/tasks/envs.py*

**Team A** only needs `flyte` installed and `PYTHONPATH` set. conda's PATH is already correct.

**Team B** needs three things: `flyte` installed in the venv, `PATH` updated so the venv's `python` is the default, and `PYTHONPATH` set. `$PATH` in an `ENV` instruction expands at Docker build time.

`.with_code_bundle()` tells Flyte to inject task source at runtime (dev) or bake it into the image at deploy time (prod).

### Task definitions

CODE9

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/remote_builder/tasks/tasks.py*

### Entry point

CODE10

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/bring-your-own-image/remote_builder/main.py*

### Run and deploy

CODE11

During development you only rebuild the base image when the Dockerfile changes. Code changes are free: they travel as a tarball at runtime.

## Decision matrix

| Scenario | Pattern |
|---|---|
| Teams own full images, can't let Flyte touch them | Pure BYOI |
| Teams hand off a base image (no Flyte knowledge required) | Remote Builder |
| Code change should not require image rebuild | Remote Builder + `with_code_bundle()` |
| Base has non-standard Python location | `.with_commands()` to fix PATH before Flyte uses it |
| Production deploy, self-contained containers | `copy_style="none"` in `flyte.deploy()` |


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/monorepo-with-uv ===

# Structuring Flyte projects with uv

## The two layers

Every Flyte + uv project involves two distinct layers. Understanding this distinction is the foundation for every decision that follows.

**The image** (slow-changing): the Python environment, installed packages, system dependencies, the interpreter. The SDK computes an MD5 hash of the image's layer stack and only rebuilds when a layer actually changes.

**The code bundle** (fast-changing): your task source code, packaged as a tarball and uploaded on every run. The container downloads and unpacks it at startup.

```
+-----------------------------------+
|       Docker Image (slow path)    |
|  Python interpreter               |
|  Installed packages (uv sync)     |  <- rebuilt only when deps change
|  System packages (apt)            |
|  Content-hashed, registry-cached  |
+-----------------------------------+
|       Code Bundle (fast path)     |
|  Your task source files           |  <- uploaded on every run
|  Local library code               |
|  Tarball extracted at startup     |
+-----------------------------------+
```

Keep these two layers separate. Your image definition should describe only the environment. Source code travels in the code bundle (in fast-deploy mode) or gets baked in at deploy time (in full-build mode).

Violating this principle (copying source into the image with `with_source_folder()` or using `install_project` mode for local code) means your image hash changes on every code edit, causing a full Docker build and push on every iteration.

## How the image gets built

`flyte.Image` is a frozen, content-addressed layer stack. Each `.with_*()` call appends an immutable layer. The final image tag is an MD5 hash of all layers.

The primary method for uv projects is `.with_uv_project()`:

```python
image = flyte.Image.from_debian_base().with_uv_project(
    pyproject_file=Path("my_app/pyproject.toml"),
)
```

**Two installation modes:**

- `dependencies_only` (default): Only `pyproject.toml` and `uv.lock` are included in the build context. The image hash covers only these two files. Your code does not affect the image hash.

- `install_project`: The entire project directory is copied into the build context. Any code change triggers a full image rebuild. Use this only when you need the project installed as a proper package (e.g., when you need package entry points or compiled extension modules).

## `with_code_bundle()`: one image for dev and prod

`with_code_bundle()` is how you write an image definition that works for both development and production without changing any code.

```python
image = (
    flyte.Image.from_debian_base()
    .with_uv_project(pyproject_file=Path("pyproject.toml"))
    .with_code_bundle()
)
```

Its behavior depends on `copy_style` at run time:

- **Fast deploy** (default): `with_code_bundle()` is a no-op. Source travels as a tarball. The image only rebuilds when `pyproject.toml` or `uv.lock` changes.
- **Full build** (`copy_style="none"`): `with_code_bundle()` resolves to a `COPY` instruction. Source is baked into the image. This is your production path.

```python
# Development
flyte.run(my_task)

# Production
flyte.deploy(my_env, copy_style="none", version="1.2.3")
```

## `root_dir`

`root_dir` tells Flyte where to look when building the code bundle and what path prefix to strip when packaging.

**The rule:** set `root_dir` to the directory you would `cd` into before running `python -c "import my_module"`.

For **src-layout** projects, set `root_dir` to `src/`:

```python
flyte.init_from_config(root_dir=Path(__file__).parent.parent)  # -> src/
```

For **flat layout** projects, set `root_dir` to the project root:

```python
flyte.init_from_config(root_dir=Path(__file__).parent)  # -> my_project/
```

## Monorepo patterns

Three patterns cover most cases:

| | Pattern A: Shared Lockfile | Pattern B: Independent Packages | Pattern C: uv workspace |
|---|---|---|---|
| Lockfile | One `uv.lock` for everything | Each package has its own | One `uv.lock` for the whole workspace |
| Package model | Single package; libraries are modules under one `src/` | Separate, independently-locked packages | Multiple installable members sharing the root lockfile |
| Sibling code reaches the container via | Code bundle (fast deploy) | `with_source_folder()` baked into the image | Code bundle, `root_dir` = workspace root (fast deploy) |
| Use when | Packages developed together, shared dep graph | Different release cadences, fully independent | Multiple versioned packages developed and locked together |

### Pattern A: Shared lockfile (recommended)

All packages live under one `src/` directory with a single `pyproject.toml` and `uv.lock`. Tasks install different subsets via dependency groups.

```
workspace_root/
├── pyproject.toml         <- defines dependency groups
├── uv.lock                <- one lockfile for everything
└── src/
    ├── workspace_app/
    │   ├── main.py
    │   └── tasks/
    │       ├── envs.py
    │       ├── etl_tasks.py
    │       └── ml_tasks.py
    ├── lib_transforms/
    │   └── ops.py
    └── lib_models/
        └── baseline.py
```

**`pyproject.toml`**: only external PyPI deps in dependency groups. Local libraries travel via the code bundle:

```toml
[project]
name = "workspace-app"
version = "0.1.0"
description = "uv workspace monorepo example for Flyte"
requires-python = ">=3.11"
dependencies = ["flyte>=2.0"]

[build-system]
requires = ["uv_build>=0.9,<0.10"]
build-backend = "uv_build"

[tool.uv]
package = true

[dependency-groups]
# Only external PyPI deps. lib_transforms and lib_models live under src/ alongside
# workspace_app and are included in the code bundle automatically.
etl = ["pandas"]
ml  = ["scikit-learn"]
dev = ["pytest", "ruff"]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/01_workspace_monorepo/pyproject.toml*

**Per-task images using dependency groups:**

```python
import pathlib

import flyte

WORKSPACE_ROOT = pathlib.Path(__file__).parent.parent.parent.parent  # -> 01_workspace_monorepo/

etl_env = flyte.TaskEnvironment(
    name="etl",
    resources=flyte.Resources(memory="512Mi", cpu="1"),
    image=flyte.Image.from_debian_base()
    .with_uv_project(
        pyproject_file=WORKSPACE_ROOT / "pyproject.toml",
        extra_args="--only-group etl",
    )
    .with_code_bundle(),
)

ml_env = flyte.TaskEnvironment(
    name="ml",
    resources=flyte.Resources(memory="1Gi", cpu="1"),
    image=flyte.Image.from_debian_base()
    .with_uv_project(
        pyproject_file=WORKSPACE_ROOT / "pyproject.toml",
        extra_args="--only-group ml",
    )
    .with_code_bundle(),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/01_workspace_monorepo/src/workspace_app/tasks/envs.py*

Both `etl_env` and `ml_env` point to the same `pyproject.toml` but install different dependency groups. The `extra_args` string is included in the image hash, so they produce separate images.

**ETL tasks** (use the shared `lib_transforms` library):

```python
from lib_transforms.ops import normalize

from workspace_app.tasks.envs import etl_env

@etl_env.task
async def load_data(n: int) -> list[float]:
    """Simulate loading raw data."""
    return [float(i * 1.5) for i in range(n)]

@etl_env.task
async def transform_data(raw: list[float]) -> list[float]:
    """Normalize raw data."""
    return normalize(raw)

@etl_env.task
async def etl_pipeline(n: int) -> list[float]:
    """Load and normalize data end-to-end."""
    raw = await load_data(n=n)
    return await transform_data(raw=raw)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/01_workspace_monorepo/src/workspace_app/tasks/etl_tasks.py*

**ML tasks** (use the shared `lib_models` library):

```python
from lib_models.baseline import predict, train_mean_predictor

from workspace_app.tasks.envs import ml_env

@ml_env.task
async def train(features: list[float], labels: list[float]) -> dict:
    """Train a simple model."""
    return train_mean_predictor(features, labels)

@ml_env.task
async def evaluate(model: dict, features: list[float]) -> float:
    """Evaluate the model on a set of features."""
    return predict(model, features)

@ml_env.task
async def ml_pipeline(features: list[float], labels: list[float]) -> float:
    """Train and evaluate a model end-to-end."""
    model = await train(features=features, labels=labels)
    return await evaluate(model=model, features=features)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/01_workspace_monorepo/src/workspace_app/tasks/ml_tasks.py*

**Entry point**: `root_dir` is set to `src/` so the code bundle covers all packages:

```python
import pathlib

import flyte
from workspace_app.tasks.etl_tasks import etl_pipeline
from workspace_app.tasks.ml_tasks import ml_pipeline

SRC_DIR = pathlib.Path(__file__).parent.parent  # -> 01_workspace_monorepo/src/

if __name__ == "__main__":
    flyte.init_from_config(root_dir=SRC_DIR)

    features = [1.5, 3.0, 4.5, 6.0, 7.5]
    labels = [0.0, 1.0, 2.0, 3.0, 4.0]

    # Development: fast deploy (code bundle delivers source at runtime)
    etl_run = flyte.run(etl_pipeline, n=10)
    print(f"ETL run: {etl_run.url}")

    ml_run = flyte.run(ml_pipeline, features=features, labels=labels)
    print(f"ML run: {ml_run.url}")

    # Production: bake source into the image (uncomment and set a version)
    # flyte.deploy(etl_env, copy_style="none", version="1.0.0")
    # flyte.deploy(ml_env, copy_style="none", version="1.0.0")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/01_workspace_monorepo/src/workspace_app/main.py*

### Pattern B: Independent packages

Each package has its own `pyproject.toml` and `uv.lock`. Fully independent image builds.

```
repo_root/
├── pyproject.toml         <- dev-only convenience (optional)
├── my_app/
│   ├── pyproject.toml     <- lists external deps + my-lib as editable path dep
│   ├── uv.lock            <- deployment lockfile
│   └── src/my_app/
│       ├── env.py
│       ├── main.py
│       └── tasks.py
└── my_lib/
    ├── pyproject.toml
    └── src/my_lib/
        └── stats.py
```

**Root `pyproject.toml`**: dev-only, installs both packages as editable for local development:

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "sibling-packages-dev"
version = "0.1.0"
description = "Dev-only root: installs both packages as editable for local development"
requires-python = ">=3.11"
dependencies = ["my-app", "my-lib"]

[tool.uv.sources]
my-app = { path = "my_app", editable = true }
my-lib = { path = "my_lib", editable = true }

[tool.hatch.build.targets.wheel]
packages = [
    "my_app/src/my_app",
    "my_lib/src/my_lib",
]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/02_sibling_packages/pyproject.toml*

**`my_app/pyproject.toml`**: declares `my-lib` as an editable path dep:

```toml
[project]
name = "my-app"
version = "0.1.0"
description = "Flyte app"
requires-python = ">=3.11"
dependencies = ["flyte>=2.0", "my-lib"]

[tool.uv]
package = true

[tool.uv.sources]
my-lib = { path = "../my_lib", editable = true }

[build-system]
requires = ["uv_build>=0.9,<0.10"]
build-backend = "uv_build"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/02_sibling_packages/my_app/pyproject.toml*

**Image definition**: sibling library baked into the image via `with_source_folder()`:

```python
import pathlib

import flyte

MY_APP_ROOT = pathlib.Path(__file__).parent.parent.parent  # -> my_app/
MY_LIB_PKG = MY_APP_ROOT.parent / "my_lib" / "src" / "my_lib"  # -> my_lib/src/my_lib/

env = flyte.TaskEnvironment(
    name="my_app",
    resources=flyte.Resources(memory="256Mi", cpu="1"),
    # my_lib is an editable path dep in pyproject.toml (so uv_build can find its source
    # during image build). Its package files are also baked into the image at /root/my_lib/
    # via with_source_folder, so they're importable at runtime without relying on the
    # editable install's .pth file (which points to a build-stage-only path).
    image=flyte.Image.from_debian_base()
    .with_uv_project(pyproject_file=MY_APP_ROOT / "pyproject.toml")
    .with_source_folder(MY_LIB_PKG)
    .with_code_bundle(),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/02_sibling_packages/my_app/src/my_app/env.py*

`with_source_folder(MY_LIB_PKG)` copies the `my_lib` package directory into the image at `/root/my_lib/`. This is necessary because the editable install's `.pth` file points to a path that only exists during the image build stage. The `my_lib` layer is part of the image hash, so the image rebuilds when `my_lib` changes: correct behavior for a dependency.

**Task definitions:**

```python
from my_app.env import env

@env.task
async def compute_stats(values: list[float]) -> dict:
    """Compute basic statistics using the my_lib utility library."""
    from my_lib.stats import mean, std

    return {
        "mean": mean(values),
        "std": std(values),
        "count": len(values),
    }

@env.task
async def summarize(stats: dict) -> str:
    return f"n={stats['count']}, mean={stats['mean']:.2f}, std={stats['std']:.2f}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/02_sibling_packages/my_app/src/my_app/tasks.py*

**Entry point**: `root_dir` covers only `my_app` source; `my_lib` is baked into the image:

```python
import pathlib

import flyte
from my_app.env import env
from my_app.tasks import compute_stats, summarize

MY_APP_ROOT = pathlib.Path(__file__).parent.parent.parent  # -> my_app/
SRC_DIR = MY_APP_ROOT / "src"  # -> my_app/src/

@env.task
async def stats_pipeline(values: list[float]) -> str:
    stats = await compute_stats(values=values)
    return await summarize(stats=stats)

if __name__ == "__main__":
    # my_lib is installed in the image; root_dir only needs to cover my_app source
    flyte.init_from_config(root_dir=SRC_DIR)

    # Development -- run a task directly, code bundle handles source delivery
    run = flyte.run(stats_pipeline, values=[1.0, 2.0, 3.0, 4.0, 5.0])
    print(f"Run URL: {run.url}")

    # Production -- deploy an environment with source baked into the image
    # flyte.deploy(env, copy_style="none", version="1.0.0")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/project-patterns/monorepo-with-uv/02_sibling_packages/my_app/src/my_app/main.py*

### Pattern C: uv workspace (`[tool.uv.workspace]`)

A uv *workspace* is uv's native mechanism for a multi-package repository: several packages share **one root `uv.lock`**, and any package can depend on its siblings through `[tool.uv.sources]` entries marked `{ workspace = true }`. Unlike Pattern A, each member is a real installable package rather than a plain module under a shared `src/`; unlike Pattern B, you manage a single lockfile for the whole tree instead of one per package.

Reach for a workspace when your packages are distinct, versioned distributions that are nonetheless developed and locked together.

```
albatross/
├── pyproject.toml         <- workspace root: [tool.uv.workspace] + shared deps
├── uv.lock                <- ONE lockfile for the whole workspace
├── src/
│   └── albatross/         <- the root package's source
│       ├── main.py
│       └── condor/
│           └── strategy.py
└── packages/              <- workspace members
    ├── bird_feeder/
    │   ├── pyproject.toml  <- member package
    │   └── src/bird_feeder/actions.py
    └── seeds/
        ├── pyproject.toml
        └── src/seeds/...
```

**Workspace root `pyproject.toml`** — declares the members and wires the sibling packages as workspace sources:

```toml
[project]
name = "albatross"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = ["bird-feeder"]

[tool.uv.sources]
bird-feeder = { workspace = true }
seeds = { workspace = true }

[tool.uv.workspace]
members = ["packages/*"]

[build-system]
requires = ["uv_build>=0.9.3,<0.10.0"]
build-backend = "uv_build"

[dependency-groups]
albatross = ["numpy", "bird-feeder"]
```

`members = ["packages/*"]` includes every package under `packages/`. The `{ workspace = true }` sources tell uv to resolve `bird-feeder` and `seeds` from the workspace instead of PyPI, so a single `uv.lock` covers the whole tree. A member package is an ordinary package that can itself depend on other members:

```toml
# packages/bird_feeder/pyproject.toml
[project]
name = "bird-feeder"
version = "0.1.0"
requires-python = ">=3.10"
dependencies = ["seeds"]

[tool.uv]
package = true
```

**Building the image** — point `.with_uv_project()` at the *workspace-root* `pyproject.toml`. Because `project_install_mode` defaults to `dependencies_only`, only the workspace's `pyproject.toml` and `uv.lock` enter the build context, so the image rebuilds only when dependencies change — fast registration is preserved. `extra_args` is forwarded to the `uv sync` that installs the dependencies (not to `pip install`), so `uv sync` flags apply — here `--only-group albatross` installs just that dependency group, keeping the image lean:

```python
from pathlib import Path

from bird_feeder.actions import bird_env, get_feeder
from seeds.actions import get_seed

import flyte
from albatross.condor.strategy import get_strategy

UV_WORKSPACE_ROOT = Path(__file__).parent.parent.parent  # -> albatross/

env = flyte.TaskEnvironment(
    name="uv_workspace",
    image=flyte.Image.from_debian_base().with_uv_project(
        pyproject_file=UV_WORKSPACE_ROOT / "pyproject.toml",
        extra_args="--only-group albatross",
    ),
    depends_on=[bird_env],
)
```

You do **not** need `with_source_folder()` to bake sibling code into the image (as Pattern B requires): in the default `dependencies_only` build `uv sync` resolves the members `bird-feeder` and `seeds` from the shared `uv.lock` (uv reads the member metadata during the sync), but their **code is not baked into the image** — only `pyproject.toml` and `uv.lock` enter the build context. The sibling source (`from bird_feeder.actions import ...`, `from seeds.actions import ...`) travels in the **code bundle** instead (the entry point below sets `root_dir` to the workspace root, so the bundle packages every member) and is on `sys.path` at runtime. The `[tool.uv.workspace]` config's job is to let uv resolve the siblings from the one lockfile; it does not put their code in the image.

**Entry point** — set `root_dir` to the *workspace root* (here `albatross/`), **not** to any single member's `src/`. A workspace spreads its members across several `src/` trees (`src/albatross/`, `packages/bird_feeder/src/`, …), so the workspace root is the one directory whose bundle captures them all — this is the `root_dir` rule from the start of this page applied to a multi-member tree, not an exception to it. With every member's source in the bundle, imports resolve the same way locally and at runtime:

```python
@env.task
async def albatross_task() -> str:
    get_feeder()
    get_strategy()
    seed = get_seed(seed_name="Sun Flower seed")
    return f"Get bird feeder and feed with {seed}"

if __name__ == "__main__":
    flyte.init_from_config(root_dir=UV_WORKSPACE_ROOT)
    run = flyte.run(albatross_task)
    print(run.url)
```

Each member can define its own `TaskEnvironment` pointing at the same workspace-root `pyproject.toml`; because they all resolve against the one lockfile, their images stay mutually consistent. For production, bake the bundle with `flyte.deploy(env, copy_style="none", version="1.2.3")` exactly as in the other patterns.

## The full build path (production)

For production deployments where you need immutable, self-contained images:

```python
flyte.deploy(my_env, copy_style="none", version="1.2.3")
```

`with_code_bundle()` on the image resolves to a `COPY` instruction. The image is fully self-contained.

Use a deterministic version string: a git commit SHA, a git tag, a CI build number. Avoid auto-generated strings so you can trace which code is in which image.

> [!IMPORTANT]
> Do not use `install_project` mode for production builds. `install_project` copies the entire project directory into the build context and hashes all of it. Every code change triggers a full image rebuild. `with_code_bundle()` + `copy_style="none"` is more surgical: only the files you select are in the image.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/cicd ===

# CI/CD deployments

This guide walks through deploying a Flyte project from CI. It uses GitHub Actions as the reference implementation, but the building blocks (a non-interactive credential, `flyte deploy`, and a commit-pinned version) translate to GitLab CI, Buildkite, CircleCI, or any runner that can run a Python script.

The examples below assume the project layout and image definitions from the [Monorepo with uv](./monorepo-with-uv) pattern; that guide covers how to structure `pyproject.toml`, `envs.py`, and task modules in a way that makes the `flyte deploy` commands shown here work cleanly.

## What CI needs to do

A deploy pipeline has three jobs:

1. **Install** the project and the `flyte` CLI.
2. **Authenticate** non-interactively against your instance.
3. **Run `flyte deploy`** for every `TaskEnvironment` in your project, pinned to the commit SHA.

Everything else (branch protections, approvals, notifications) is generic CI concerns and out of scope.

## Authentication: client credentials

Locally, `flyte deploy` typically authenticates via a browser login (PKCE). A CI runner has no browser and no human to click through a consent screen, so you need a credential the CLI can use without any prompts. For Flyte OSS this is an **OAuth2 client-credentials** application: a client ID and client secret for a machine ("service") identity that your instance's identity provider (IdP) trusts.

### Register a client-credentials application

Client-credentials applications are provisioned in your **identity provider**, not through the `flyte` CLI. (The `flyte create api-key` command referenced elsewhere is a Union feature from the `flyteplugins-union` package and isn't available in Flyte OSS.) The exact steps depend on the IdP your instance is configured against (Okta, Keycloak, Auth0, Google, Azure AD, and so on), but the result is always a **client ID** and a **client secret**.

Ask whoever administers your instance for:

- The **client ID** of a service application authorized to register tasks in the target project and domain.
- The matching **client secret**.
- The admin **endpoint**: the same host you pass to `flyte create config`.

> [!NOTE]
> The application must be granted whatever scopes your admin API requires. Provisioning and scoping the IdP application is an instance-administration task; see your instance's authentication setup for the specifics.

### Store the secret as a CI secret

Add the **client secret** to your CI system's secret store. However it's configured, the secret needs to:

- Be exposed to the deploy job as an environment variable (this guide uses `FLYTE_CLIENT_SECRET`).
- Be masked in logs (most CI systems do this automatically for secrets).
- Be scoped to the branches/environments that actually deploy: typically `main` or a release branch, not every feature branch or fork PR.

The client ID and endpoint aren't secret; they live in the `config.yaml` you check into the repo (see **Project patterns > CI/CD deployments > Project configuration**). Only the client secret goes in the secret store.

### Point the CLI at the credential

The `flyte` CLI reads client-credentials settings from `config.yaml` under `admin:`:

- **`authType: ClientSecret`** selects the OAuth2 client-credentials flow instead of the interactive PKCE default.
- **`clientId`** is the application's client ID.
- **`clientSecretEnvVar`** names the environment variable the CLI reads the secret from: `FLYTE_CLIENT_SECRET` here. (Alternatively, `clientSecretLocation` points at a file containing the secret, which suits runners that mount secrets as files rather than env vars.)

### Scope and rotation

Grant the service identity only the permissions CI needs: deploy rights on the target project/domain, nothing more. Rotate the client secret on a schedule (90 days is a reasonable default) in your IdP and update the CI secret to match.

## Project configuration

Two files drive `flyte deploy` behavior in CI: `pyproject.toml` (or `uv.lock`) for dependencies, and `config.yaml` for your endpoint and image-builder settings.

### `config.yaml`

Save this at `.flyte/config.yaml` (or `config.yaml`) in your repo and check it in. In CI the `flyte` CLI auto-discovers config from the repo checkout: repo-relative paths (`./config.yaml`, `./.flyte/config.yaml`, `<git-root>/.flyte/config.yaml`) take precedence over any home-directory config, so it's picked up automatically after checkout with no `--config` flag needed. See [the config discovery order](../../api-reference/flyte-sdk/flyte.config/_index#auto) for the full precedence; pass `--config <path>` only to point at a non-standard location.

It supplies the endpoint, the client-credentials auth settings, and the (local) image builder, everything `flyte deploy` needs beyond the client secret:

```yaml
admin:
  endpoint: dns:///<your-flyte-host>
  authType: ClientSecret
  clientId: <client-id>
  clientSecretEnvVar: FLYTE_CLIENT_SECRET
image:
  builder: local
task:
  project: <default-project>
  domain: development
```

The `clientId` and `endpoint` are safe to commit; only the client secret named by `clientSecretEnvVar` comes from the CI secret store. `builder: local` means images are built on the runner with Docker (Flyte OSS has no remote builder). See [Container images](../tasks/task-configuration/container-images#image-building).

## The GitHub Actions workflow

A minimal deploy workflow, one job, one step per `TaskEnvironment`:

The client ID and endpoint live in the checked-in `config.yaml` (auto-discovered from the repo, as above); each deploy step injects the client secret from the CI secret store.

```yaml
# .github/workflows/deploy.yml
name: Deploy to Flyte

on:
  push:
    branches: [main]
  workflow_dispatch:

env:
  FLYTE_PROJECT: my-project
  FLYTE_DOMAIN: development

jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Install uv
        uses: astral-sh/setup-uv@v5
        with:
          enable-cache: true

      - name: Sync dependencies
        run: uv sync --group etl --group ml

      - name: Deploy etl_env
        env:
          FLYTE_CLIENT_SECRET: ${{ secrets.FLYTE_CLIENT_SECRET }}
        run: |
          uv run flyte deploy \
            --copy-style none \
            --version ${{ github.sha }} \
            --project "$FLYTE_PROJECT" \
            --domain "$FLYTE_DOMAIN" \
            src/workspace_app/tasks/etl_tasks.py etl_env

      - name: Deploy ml_env
        env:
          FLYTE_CLIENT_SECRET: ${{ secrets.FLYTE_CLIENT_SECRET }}
        run: |
          uv run flyte deploy \
            --copy-style none \
            --version ${{ github.sha }} \
            --project "$FLYTE_PROJECT" \
            --domain "$FLYTE_DOMAIN" \
            src/workspace_app/tasks/ml_tasks.py ml_env
```

Because Flyte OSS builds images locally, the deploy steps need Docker available on the runner (the `ubuntu-latest` image includes it) and access to your container registry; add a `docker login` step for private registries before the first deploy.

### Key flag choices

- **`--copy-style none`**: bakes source into the image as part of the build layer. Combined with `.with_code_bundle()` on your `flyte.Image` (see [Monorepo with uv](./monorepo-with-uv)), this resolves to a `COPY` instruction so the image is fully self-contained. This is the production path: one immutable artifact per commit, no runtime code bundle download.
- **`--version ${{ github.sha }}`**: makes deploys idempotent and traceable. Re-running the same commit produces the same version identifier; tasks already registered at that version are no-ops.
- **Path argument points at the task file, not `envs.py`.** `flyte deploy` only imports the file you give it, so tasks decorated with `@env.task` in separate files won't register unless you point at (or transitively import) those files. Pointing at `etl_tasks.py` pulls in `envs.py` via its import chain and runs the `@etl_env.task` decorators. As an alternative, you can point at a directory and pass `--recursive` to load every task module under it in one command. For a `src/` layout project, also pass `--root-dir src` so shared modules like `envs.py` resolve to a single import path instead of being loaded twice:

  ```yaml
  - name: Deploy all envs
    env:
      FLYTE_CLIENT_SECRET: ${{ secrets.FLYTE_CLIENT_SECRET }}
    run: |
      uv run flyte deploy \
        --copy-style none \
        --version ${{ github.sha }} \
        --project "$FLYTE_PROJECT" \
        --domain "$FLYTE_DOMAIN" \
        --root-dir src --recursive src/workspace_app/tasks
  ```

### Splitting build from deploy

`flyte deploy` builds any missing images before it registers tasks. If you'd rather treat image builds as a separate CI concern (for clearer logs, independent retry, or parallel builds per env), run `flyte build` first and let deploy reuse the result:

```yaml
- name: Build etl image
  env:
    FLYTE_CLIENT_SECRET: ${{ secrets.FLYTE_CLIENT_SECRET }}
  run: |
    uv run flyte build \
      --copy-style none --root-dir src \
      src/workspace_app/tasks/etl_tasks.py etl_env

- name: Deploy etl_env
  env:
    FLYTE_CLIENT_SECRET: ${{ secrets.FLYTE_CLIENT_SECRET }}
  run: |
    uv run flyte deploy \
      --copy-style none \
      --version ${{ github.sha }} \
      --project "$FLYTE_PROJECT" --domain "$FLYTE_DOMAIN" \
      --root-dir src src/workspace_app/tasks/etl_tasks.py etl_env
```

Image tags are content hashes of the `flyte.Image` definition: `flyte build` pushes `<registry>:flyte-<hash>`, and `flyte deploy` computes the same hash, sees the image already in the registry, and skips rebuilding. `--copy-style` must match between the two commands; otherwise the hashes diverge and deploy will build again.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/run-scaling ===

# Scale your runs

This guide helps you understand and optimize the performance of your Flyte workflows. Whether you're building latency-sensitive applications or high-throughput data pipelines, these docs will help you make the right architectural choices.

## Understanding Flyte execution

Before optimizing performance, it's important to understand how Flyte executes your workflows:

- ****Scale your runs > Data flow****: Learn how data moves between tasks, including inline vs. raw (by-reference) data, caching mechanisms, and storage configuration.
- ****Scale your runs > Life of a run****: Understand what happens when you invoke `flyte.run()`, from code analysis and image building to task execution and state management.

## Performance optimization

Once you understand the fundamentals, dive into performance tuning:

- ****Scale your runs > Scale your workflows****: A comprehensive guide to optimizing workflow performance, covering latency vs. throughput, task overhead analysis, batching strategies, reusable containers, and more.

## Key concepts for scaling

When scaling your workflows, keep these principles in mind:

1. **Task overhead matters**: The overhead of creating a task (uploading data, enqueuing, creating containers) should be much smaller than the task runtime.
2. **Batch for throughput**: For large-scale data processing, batch multiple items into single tasks to reduce overhead.
3. **Reusable containers**: Eliminate container startup overhead and enable concurrent execution with reusable containers.
4. **Traces for lightweight ops**: Use traces instead of tasks for lightweight operations that need checkpointing.
5. **Limit fanout**: Keep the total number of actions per run below 50k (target 10k-20k for best performance).
6. **Choose the right data types**: Use reference types (files, directories, DataFrames) for large data and inline types for small data.

For detailed guidance on each of these topics, see **Scale your runs > Scale your workflows**.

## Subpages
- **Scale your runs > Data flow**
- **Scale your runs > Life of a run**
- **Scale your runs > Scale your workflows**
- **Scale your runs > Maximize GPU utilization for batch inference**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/data-flow ===

# Data flow

Understanding how data flows between tasks is critical for optimizing workflow performance in Flyte. Tasks take inputs and produce outputs, with data flowing through your workflow using an efficient transport layer.

> [!NOTE]
> This page focuses on **how** data moves at runtime. For the static map of **what** lives in the control plane database versus the data plane object store (including what *metadata*, *literals*, and *raw data* mean), see [Where your data lives](../get-started/core-concepts/where-data-lives).

## Overview

Flyte tasks are run to completion. Each task takes inputs and produces exactly one output. Even if multiple instances run concurrently (such as in retries), only one output will be accepted. This deterministic data flow model provides several key benefits:

1. **Reduced boilerplate**: Automatic handling of files, DataFrames, directories, custom types, data classes, Pydantic models, and primitive types without manual serialization.
2. **Type safety**: Optional type annotations enable deeper type understanding, automatic UI form generation, and runtime type validation.
3. **Efficient transport**: Data is passed by reference (files, directories, DataFrames) or by value (primitives) based on type.
4. **Durable storage**: All data is stored durably and accessible through APIs and the UI.
5. **Caching support**: Efficient caching using shallow immutable references for referenced data.

## Data types and transport

Flyte handles different data types with different transport mechanisms:

### Passed by reference

These types are not copied but passed as references to storage locations. That offloaded, by-reference content is called **raw data**:

- **Files**: `flyte.io.File`
- **Directories**: `flyte.io.Dir`
- **Dataframes**: `flyte.io.DataFrame`, `pd.DataFrame`, `pl.DataFrame`, etc.

Dataframes are automatically converted to Parquet format and read using Apache Arrow for zero-copy reads. Use `flyte.io.DataFrame` for lazy materialization to any supported type like pandas or polars. [Learn more about the Flyte Dataframe type](../tasks/task-programming/dataframes)

### Passed by value (inline I/O)

Primitive and structured types are serialized and passed inline:

| Type Category | Examples | Serialization |
|--------------|----------|---------------|
| **Primitives** | `int`, `float`, `str`, `bool`, `None` | MessagePack |
| **Time types** | `datetime.datetime`, `datetime.date`, `datetime.timedelta` | MessagePack |
| **Collections** | `list`, `dict`, `tuple` | MessagePack |
| **Data structures** | data classes, Pydantic `BaseModel` | MessagePack |
| **Enums** | `enum.Enum` subclasses | MessagePack |
| **Unions** | `Union[T1, T2]`, `Optional[T]` | MessagePack |
| **Protobuf** | `google.protobuf.Message` | Binary |

Flyte uses efficient MessagePack serialization for most types, providing compact binary representation with strong type safety.

> [!NOTE]
> If type annotations are not used, or if `typing.Any` or unrecognized types are used, data will be pickled. By default, pickled objects smaller than 10KB are passed inline, while larger pickled objects are automatically passed as a file. Pickling allows for progressive typing but should be used carefully.

## Task execution and data flow

### Input download

When a task starts:

1. **Inline inputs download**: The task downloads inline inputs from the configured Flyte object store.
2. **Size limits**: By default, inline inputs are limited to 10MB, but this can be adjusted using `flyte.TaskEnvironment`'s `max_inline_io` parameter.
3. **Memory consideration**: Inline data is materialized in memory, so adjust your task resources accordingly.
4. **Raw data materialization**: Raw data (files, directories) is passed using special types in `flyte.io`. Dataframes are automatically materialized if using `pd.DataFrame`. Use `flyte.io.DataFrame` to avoid automatic materialization.

### Output upload

When a task returns data:

1. **Inline data**: Uploaded to the Flyte object store configured at the organization, project, or domain level.
2. **Raw data**: Stored under the same bucket prefix by default, or routed to a different location using `flyte.with_runcontext(raw_data_path=...)`.
3. **Separate prefixes**: Each task creates one output per retry attempt in separate prefixes, making data incorruptible by design.

## Task-to-task data flow

When a task invokes downstream tasks:

1. **Input recording**: The input to the downstream task is recorded to the object store.
2. **Reference upload**: All referenced objects are uploaded (if not already present).
3. **Task invocation**: The downstream task is invoked on the remote server.
4. **Parallel execution**: When multiple tasks are invoked in parallel using `flyte.map` or `asyncio`, inputs are written in parallel.
5. **Storage layer**: Data writing uses the `flyte.storage` layer, backed by the Rust-based `object-store` crate and optionally `fsspec` plugins.
6. **Output download**: Once the downstream task completes, inline outputs are downloaded and returned to the calling task.

## Caching and data hashing

Understanding how Flyte caches data is essential for performance optimization.

### Cache key computation

A cache hit occurs when the following components match:

- **Task name**: The fully-qualified task name
- **Computed input hash**: Hash of all inputs (excluding `ignored_inputs`)
- **Task interface hash**: Hash of input and output types
- **Task config hash**: Hash of task configuration
- **Cache version**: User-specified or automatically computed

### Inline data caching

All inline data is cached using a consistent hashing system. The cache key is derived from the data content.

### Raw data hashing

Raw data (DataFrames, files, directories) is hashed shallowly by default using the hash of the storage location, so a downstream task does not cache-hit on identical content stored at a new path. To cache on content instead, attach a content hash at production time with `flyte.io.HashFunction`. See [Content-based caching for DataFrames, files, and directories](../tasks/task-configuration/caching#content-based-caching-for-dataframes-files-and-directories).

### Cache control

Control caching behavior using `flyte.with_runcontext`:

- **Scope**: Set `cache_lookup_scope` to `"global"` or `"project/domain"`.
- **Disable cache**: Set `overwrite_cache=True` to force re-execution.

For more details on caching configuration, see [Caching](../tasks/task-configuration/caching).

## Traces and data flow

When using [traces](../tasks/task-programming/traces), the data flow behavior is different:

1. **Full execution first**: The trace is fully executed before inputs and outputs are recorded.
2. **Checkpoint behavior**: Recording happens like a checkpoint at the end of trace execution.
3. **Streaming iterators**: The entire output is buffered and recorded after the stream completes. Buffering is pass-through, allowing caller functions to consume output while buffering.
4. **Chained traces**: All traces are recorded after the last one completes consumption.
5. **Same process with `asyncio`**: Traces run within the same Python process and support `asyncio` parallelism, so failures can be retried, effectively re-running the trace.
6. **Lightweight overhead**: Traces only have the overhead of data storage (no task orchestration overhead).

> [!NOTE]
> Traces are not a substitute for tasks if you need caching. Tasks provide full caching capabilities, while traces provide lightweight checkpointing with storage overhead. However, traces support concurrent execution using `asyncio` patterns within a single task.

## Object stores and latency considerations

By default, Flyte uses object stores like S3, GCS, Azure Storage, and R2 to persist task inputs, outputs, and offloaded raw data. These have high latency for smaller objects, so:

- **Minimum task duration**: Tasks should take at least a second to run to amortize storage overhead.
- **Future improvements**: High-performance key/value or relational stores like Redis and PostgreSQL may be supported in the future as alternative offload backends. Contact the Union team if you're interested.

## Configuring data storage

### Organization and project level

Object stores are configured at the organization level or per project/domain. Documentation for this configuration is coming soon.

### Per-run configuration

Configure raw data storage on a per-run basis using `flyte.with_runcontext`:

```python
run = flyte.with_runcontext(
    raw_data_path="s3://my-bucket/custom-path"
).run(my_task, input_data=data)
```

This allows you to control where raw data (files, directories, DataFrames) is stored for specific runs. See [Run context](../tasks/task-deployment/run-context) for the full set of options.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/life-of-a-run ===

# Life of a run

Understanding what happens when you invoke `flyte.run()` is crucial for optimizing workflow performance and debugging issues. This guide walks through each phase of task execution from submission to completion.

## Overview

When you execute `flyte.run()`, the system goes through several phases:

1. **Code analysis and preparation**: Discover environments and images
2. **Image building**: Build container images if changes are detected
3. **Code bundling**: Package your Python code
4. **Upload**: Transfer the code bundle to object storage
5. **Run creation**: Submit the run to the backend
6. **Task execution**: Execute the task in the data plane
7. **State management**: Track and persist execution state

> [!NOTE]
> This walkthrough follows a run's *control flow*. For what the inputs and outputs actually are (**literals**, stored inline or offloaded as **raw data**) and where each lives, see [Where your data lives](../get-started/core-concepts/where-data-lives).

## Phase 1: Code analysis and preparation

When `flyte.run()` is invoked:

1. **Environment discovery**: Flyte analyzes your code and finds all relevant `flyte.TaskEnvironment` instances by walking the `depends_on` hierarchy.
2. **Image identification**: Discovers unique `flyte.Image` instances used across all environments.
3. **Image building**: Starts the image building process. Images are only built if a change is detected.

> [!NOTE]
> If you invoke `flyte.run()` multiple times within the same Python process without changing code (such as in a notebook or script), the code bundling and image building steps are done only once. This can dramatically speed up iteration.

## Phase 2: Image building

Container images provide the runtime environment for your tasks:

- **Change detection**: Images are only rebuilt if changes are detected in dependencies or configuration.
- **Caching**: Previously built images are reused when possible.
- **Parallel builds**: Multiple images can be built concurrently.

For more details on container images, see [Container Images](../tasks/task-configuration/container-images).

## Phase 3: Code bundling

After images are built, your project files are bundled:

### Default: `copy_style="loaded_modules"`

By default, all Python modules referenced by the invoked tasks through module-level import statements are automatically copied. This provides a good balance between completeness and efficiency.

### Alternative: `copy_style="none"`

Skip bundling by setting `copy_style="none"` in `flyte.with_runcontext()` and adding all code into `flyte.Image`:

```python
# Add code to image
image = flyte.Image().with_source_code("/path/to/code")

# Or use Dockerfile
image = flyte.Image.from_dockerfile("Dockerfile")

# Skip bundling
run = flyte.with_runcontext(copy_style="none").run(my_task, input_data=data)
```

For more details on code packaging, see [Packaging](../tasks/task-deployment/packaging).

## Phase 4: Upload code bundle

Once the code bundle is created:

1. **Request signed URL**: The SDK sends the bundle checksum and target path to the control plane.
2. **Control plane obtains URL**: The control plane calls the data plane to obtain a signed URL for that checksum and path.
3. **Direct upload**: The signed URL is returned to the SDK, which uploads the code bundle directly to the object store.

## Phase 5: Run creation and queuing

1. **Upload inputs**: The SDK uploads the run's inputs to the data plane `dataproxy` service, which writes them to the object store. The input values never pass through the control plane.
2. **Invoke `CreateRun`**: The SDK calls the `CreateRun` API with a reference (URI) to the uploaded inputs, not the input values themselves.
3. **En-queue a run**: The run is queued into the Union control plane.
4. **Hand off to executor**: The Union control plane hands the task to the Executor Service in your data plane.
5. **Create action**: The parent task action (called `a0`) is created.

## Phase 6: Task execution in data plane

### Container startup

1. **Container starts**: The task container starts in your data plane.
2. **Download code bundle**: The Flyte runtime downloads the code bundle from object storage.
3. **Inflate task**: The task is inflated from the code bundle.
4. **Download inputs**: The task's inputs are downloaded from the object store.
5. **Execute task**: The task is executed with context and inputs.

### Invoking downstream tasks

If the task invokes other tasks:

1. **Controller thread**: A controller thread starts to communicate with the backend Queue Service.
2. **Monitor status**: The controller monitors the status of downstream actions.
3. **Crash recovery**: If the task crashes, the action identifier is deterministic, allowing the task to resurrect its state from Union control plane.
4. **Replay**: The controller efficiently replays state (even at large scale) to find missing completions and resume monitoring.

### Execution flow diagram

```mermaid
sequenceDiagram
    participant Client as SDK/Client
    participant Control as Control plane<br/>(Queue Service)
    participant Data as Data plane<br/>(Executor)
    participant ObjStore as Object Store
    participant Container as Task Container

    Client->>Client: Analyze code & discover environments
    Client->>Client: Build images (if changed)
    Client->>Client: Bundle code
    Client->>Control: Request signed URL (checksum, path)
    Control->>Data: Get signed URL for bundle
    Data-->>Control: Signed URL
    Control-->>Client: Signed URL
    Client->>ObjStore: Upload code bundle (signed URL)
    Client->>Data: Upload inputs (dataproxy)
    Data->>ObjStore: Write inputs
    Client->>Control: CreateRun API (input URI)
    Control->>Data: Queue task (create action a0)
    Data->>Container: Start container
    Container->>Data: Request code bundle
    Data->>ObjStore: Read code bundle
    ObjStore-->>Data: Code bundle
    Data-->>Container: Code bundle
    Container->>Container: Inflate task
    Container->>Data: Request inputs
    Data->>ObjStore: Read inputs
    ObjStore-->>Data: Inputs
    Data-->>Container: Inputs
    Container->>Container: Execute task

    alt Invokes downstream tasks
        Container->>Container: Start controller thread
        Container->>Control: Submit downstream tasks
        Control->>Data: Queue downstream actions
        Container->>Control: Monitor downstream status
        Control-->>Container: Status updates
    end

    Container->>Data: Upload outputs
    Data->>ObjStore: Write outputs
    Container->>Control: Complete
    Control-->>Client: Run complete
```

## Action identifiers and crash recovery

Flyte uses deterministic action identifiers to enable robust crash recovery:

- **Consistent identifiers**: Action identifiers are consistently computed based on task and invocation context.
- **Re-run identical**: In any re-run, the action identifier is identical for the same invocation.
- **Multiple invocations**: Multiple invocations of the same task receive unique identifiers.
- **Efficient resurrection**: On crash, the `a0` action resurrects its state from Union control plane efficiently, even at large scale.
- **Replay and resume**: The controller replays execution until it finds missing completions and starts watching them.

## Downstream task execution

When downstream tasks are invoked:

1. **Action creation**: Downstream actions are created with unique identifiers.
2. **Queue assignment**: Actions are handed to an executor, which can be selected using a queue or from the general pool.
3. **Parallel execution**: Multiple downstream tasks can execute in parallel.
4. **Result aggregation**: Results are aggregated and returned to the parent task.

## Reusable containers

When using [reusable containers](../tasks/task-configuration/reusable-containers), the execution model changes:

1. **Environment spin-up**: The container environment is first spun up with configured replicas.
2. **Task allocation**: Tasks are allocated to available replicas in the environment.
3. **Scaling**: If all replicas are busy, new replicas are spun up (up to the configured maximum), or tasks are backlogged in queues.
4. **Container reuse**: The same container handles multiple task executions, reducing startup overhead.
5. **Lifecycle management**: Containers are managed according to `ReusePolicy` settings (`idle_ttl`, `scaledown_ttl`, etc.).

### Reusable container execution flow

```mermaid
sequenceDiagram
    participant Control as Queue Service
    participant Executor as Executor Service
    participant Pool as Container Pool
    participant Replica as Container Replica

    Control->>Executor: Submit task

    alt Reusable containers enabled
        Executor->>Pool: Request available replica

        alt Replica available
            Pool->>Replica: Allocate task
            Replica->>Replica: Execute task
            Replica->>Pool: Task complete (ready for next)
        else No replica available
            alt Can scale up
                Executor->>Pool: Create new replica
                Pool->>Replica: Spin up new container
                Replica->>Replica: Execute task
                Replica->>Pool: Task complete
            else At max replicas
                Executor->>Pool: Queue task
                Pool-->>Executor: Wait for available replica
                Pool->>Replica: Allocate when available
                Replica->>Replica: Execute task
                Replica->>Pool: Task complete
            end
        end
    else No reusable containers
        Executor->>Replica: Create new container
        Replica->>Replica: Execute task
        Replica->>Executor: Complete & terminate
    end

    Replica-->>Control: Return results
```

## State replication and visualization

### Queue service to run service

1. **Reliable replication**: Queue Service reliably replicates execution state back to Run Service.
2. **Eventual consistency**: The Run Service may be slightly behind the actual execution state.
3. **Visualization**: Run Service paints the entire run onto the UI.

### UI limitations

- **Current limit**: The UI is currently limited to displaying 50k actions per run.
- **Future improvements**: This limit will be increased in future releases. Contact the Union team if you need higher limits.

## Optimization opportunities

Understanding the life of a run reveals several optimization opportunities:

1. **Reuse Python process**: Run `flyte.run()` multiple times in the same process to avoid re-bundling code.
2. **Skip bundling**: Use `copy_style="none"` and bake code into images for faster startup.
3. **Reusable containers**: Use reusable containers to eliminate container startup overhead.
4. **Parallel execution**: Invoke multiple downstream tasks concurrently using `flyte.map()` or `asyncio`.
5. **Efficient data flow**: Minimize data transfer by using reference types (files, directories) instead of inline data.
6. **Caching**: Enable task caching to avoid redundant computation.

For detailed performance tuning guidance, see [Scale your workflows](./scale-your-workflows).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/scale-your-workflows ===

# Scale your workflows

Performance optimization in Flyte involves understanding the interplay between task execution overhead, data transfer, and concurrency. This guide helps you identify bottlenecks and choose the right patterns for your workload.

## Understanding performance dimensions

Performance optimization focuses on two key dimensions:

### Latency

**Goal**: Minimize end-to-end execution time for individual workflows.

**Characteristics**:

- Fast individual actions (milliseconds to seconds)
- Total action count typically less than 1,000
- Critical for interactive applications and real-time processing
- Multi-step inference, with reusing model or data in memory (use reusable containers with [@alru.cache](https://pypi.org/project/async-lru/))

**Recommended approach**:

- Use tasks for orchestration and parallelism
- Use [traces](../tasks/task-programming/traces) for fine-grained checkpointing
- Model parallelism using `asyncio` and use things methods like `asyncio.as_completed` or `asyncio.gather` to join the parallelism
- Use [reusable containers](../tasks/task-configuration/reusable-containers) with concurrency to eliminate startup overhead and optimize resource utilization

### Throughput

**Goal**: Maximize the number of items processed per unit time.

**Characteristics**:

- Processing large datasets (millions of items)
- High total action count (10k to 50k actions)
- Batch processing, large-scale batch inference and ETL workflows

**Recommended approach**:

- Batch workloads to reduce overhead
- Limit fanout to manage system load
- Use reusable containers with concurrency for maximum utilization
- Balance task granularity with overhead

## Task execution overhead

Understanding task overhead is critical for performance optimization. When you invoke a task, several operations occur:

| Operation | Symbol | Description |
|-----------|--------|-------------|
| **Upload data** | `u` | Time to upload input data to object store |
| **Download data** | `d` | Time to download input data from object store |
| **Enqueue task** | `e` | Time to enqueue task in Queue Service |
| **Create instance** | `t` | Time to create task container instance |

**Total overhead per task**: `2u + 2d + e + t`

This overhead includes:

- Uploading inputs from the parent task (`u`)
- Downloading inputs in the child task (`d`)
- Uploading outputs from the child task (`u`)
- Downloading outputs in the parent task (`d`)
- Enqueuing the task (`e`)
- Creating the container instance (`t`)

### The overhead principle

For efficient execution, task overhead should be much smaller than task runtime:

```
Total overhead (2u + 2d + e + t) << Task runtime
```

If task runtime is comparable to or less than overhead, consider:

1. **Batching**: Combine multiple work items into a single task
2. **Traces**: Use traces instead of tasks for lightweight operations
3. **Reusable containers**: Eliminate container creation overhead (`t`)
4. **Local execution**: Run lightweight operations within the parent task

## System architecture and data flow

To optimize performance, understand how tasks flow through the system:

1. **Control plane to data plane**: Tasks flow from the control plane (Run Service, Queue Service) to the data plane (Executor Service).
2. **Data movement**: Data moves between tasks through object storage. See [Data flow](./data-flow) for details.
3. **State replication**: Queue Service reliably replicates state back to Run Service for visualization. The Run Service may be slightly behind actual execution.

For a detailed walkthrough of task execution, see [Life of a run](./life-of-a-run).

## Optimization strategies

### 1. Use reusable containers for concurrency

[Reusable containers](../tasks/task-configuration/reusable-containers) eliminate the container creation overhead (`t`) and enable concurrent task execution:

```python
import flyte
from datetime import timedelta

# Define reusable environment
env = flyte.TaskEnvironment(
    name="high-throughput",
    reuse_policy=flyte.ReusePolicy(
        replicas=(2, 10),           # Auto-scale from 2 to 10 replicas
        concurrency=5,              # 5 tasks per replica = 50 max concurrent
        scaledown_ttl=timedelta(minutes=10),
        idle_ttl=timedelta(hours=1)
    )
)

@env.task
async def process_item(item: dict) -> dict:
    # Process individual item
    return {"processed": item["id"]}
```

**Benefits**:

- Eliminates container startup overhead (`t ≈ 0`)
- Supports concurrent execution (multiple tasks per container)
- Auto-scales based on demand
- Reuses Python environment and loaded dependencies

**Limitations**:

- Concurrency is limited by CPU and I/O resources in the container
- Memory requirements scale with total working set size
- Best for I/O-bound tasks or async operations

### 2. Batch workloads to reduce overhead

For high-throughput processing, batch multiple items into a single task:

```python
@env.task
async def process_batch(items: list[dict]) -> list[dict]:
    """Process a batch of items in a single task."""
    results = []
    for item in items:
        result = await process_single_item(item)
        results.append(result)
    return results

@env.task
async def process_large_dataset(dataset: list[dict]) -> list[dict]:
    """Process 1M items with batching."""
    batch_size = 1000  # Adjust based on overhead calculation
    batches = [dataset[i:i + batch_size] for i in range(0, len(dataset), batch_size)]

    # Process batches in parallel (1000 tasks instead of 1M)
    results = await asyncio.gather(*[process_batch(batch) for batch in batches])

    # Flatten results
    return [item for batch_result in results for item in batch_result]
```

**Benefits**:

- Reduces total number of tasks (e.g., 1000 tasks instead of 1M)
- Amortizes overhead across multiple items
- Lower load on Queue Service and object storage

**Choosing batch size**:

1. Calculate overhead: `overhead = 2u + 2d + e + t`
2. Target task runtime: `runtime > 10 × overhead` (rule of thumb)
3. Adjust batch size to achieve target runtime
4. Consider memory constraints (larger batches require more memory)

### 3. Use traces for lightweight operations

[Traces](../tasks/task-programming/traces) provide fine-grained checkpointing with minimal overhead:

```python
@flyte.trace
async def fetch_data(url: str) -> dict:
    """Traced function for API call."""
    response = await http_client.get(url)
    return response.json()

@flyte.trace
async def transform_data(data: dict) -> dict:
    """Traced function for transformation."""
    return {"transformed": data}

@env.task
async def process_workflow(urls: list[str]) -> list[dict]:
    """Orchestrate using traces instead of tasks."""
    results = []
    for url in urls:
        data = await fetch_data(url)
        transformed = await transform_data(data)
        results.append(transformed)
    return results
```

**Benefits**:

- Only storage overhead (no task orchestration overhead)
- Runs in the same Python process with asyncio parallelism
- Provides checkpointing and resumption
- Visible in execution logs and UI

**Trade-offs**:

- No caching (use tasks for cacheable operations)
- Shares resources with the parent task (CPU, memory)
- Storage writes may still be slow due to object store latency

**When to use traces**:

- API calls and external service interactions
- Deterministic transformations that need checkpointing
- Operations taking more than 1 second (to amortize storage overhead)

### 4. Limit fanout for system stability

The UI and system have limits on the number of actions per run:

- **Current limit**: 50k actions per run
- **Future**: Higher limits will be supported (contact the Union team if needed)

**Example: Control fanout with batching**

```python
@env.task
async def process_million_items(items: list[dict]) -> list[dict]:
    """Process 1M items with controlled fanout."""
    # Target 10k tasks, each processing 100 items
    batch_size = 100
    max_fanout = 10000

    batches = [items[i:i + batch_size] for i in range(0, len(items), batch_size)]

    # Use flyte.map for parallel execution
    results = await flyte.map(process_batch, batches)

    return [item for batch in results for item in batch]
```

### 5. Optimize data transfer

Minimize data transfer overhead by choosing appropriate data types:

**Use reference types for large data**:

```python
from flyte.io import File, Directory, DataFrame

@env.task
async def process_large_file(input_file: File) -> File:
    """Files passed by reference, not copied."""
    # Download only when needed
    local_path = input_file.download()

    # Process file
    result_path = process(local_path)

    # Upload result
    return File.new_remote(result_path)
```

**Use inline types for small data**:

```python
@env.task
async def process_metadata(metadata: dict) -> dict:
    """Small dicts passed inline efficiently."""
    return {"processed": metadata}
```

**Guideline**:

- **< 10 MB**: Use inline types (primitives, small dicts, lists)
- **> 10 MB**: Use reference types (File, Directory, DataFrame)
- **Adjust**: Use `max_inline_io` in `TaskEnvironment` to change the threshold

See [Data flow](./data-flow) for details on data types and transport.

### 6. Use caching

Enable [caching](../tasks/task-configuration/caching) to avoid redundant computation:

```python
@env.task(cache="auto")
async def expensive_computation(input_data: dict) -> dict:
    """Automatically cached based on inputs."""
    # Expensive operation
    return result
```

**Benefits**:

- Skips re-execution for identical inputs
- Reduces overall workflow runtime
- Preserves resources for new computations

**When to use**:

- Deterministic tasks (same inputs → same outputs)
- Expensive computations (model training, large data processing)
- Stable intermediate results

### 7. Parallelize with `flyte.map`

Use [`flyte.map`](../tasks/task-programming/fanout) for data-parallel workloads:

```python
@env.task
async def process_item(item: dict) -> dict:
    return {"processed": item}

@env.task
async def parallel_processing(items: list[dict]) -> list[dict]:
    """Process items in parallel using map."""
    results = await flyte.map(process_item, items)
    return results
```

**Benefits**:

- Automatic parallelization
- Dynamic scaling based on available resources
- Built-in error handling and retries

**Best practices**:

- Combine with batching to control fanout
- Use with reusable containers for maximum throughput
- Consider memory and resource limits

## Performance tuning workflow

Follow this workflow to optimize your Flyte workflows:

1. **Profile**: Measure task execution times and identify bottlenecks.
2. **Calculate overhead**: Estimate `2u + 2d + e + t` for your tasks.
3. **Compare**: Check if `task runtime >> overhead`. If not, optimize.
4. **Batch**: Increase batch size to amortize overhead.
5. **Reusable containers**: Enable reusable containers to eliminate `t`.
6. **Traces**: Use traces for lightweight operations within tasks.
7. **Cache**: Enable caching for deterministic, expensive tasks.
8. **Limit fanout**: Keep total actions below 50k (target 10k-20k).
9. **Monitor**: Use the UI to monitor execution and identify issues.
10. **Iterate**: Continuously refine based on performance metrics.

## Real-world example: PyIceberg batch processing

For an example of efficient data processing with Flyte, see the [PyIceberg parallel batch aggregation example](https://github.com/flyteorg/flyte-sdk/blob/main/examples/data_processing/pyiceberg_example.py). This example demonstrates:

- **Zero-copy data passing**: Pass file paths instead of data between tasks
- **Reusable containers with concurrency**: Maximize CPU utilization across workers
- **Parallel file processing**: Use `asyncio.gather()` to process multiple files concurrently
- **Efficient batching**: Distribute parquet files across worker tasks

Key pattern from the example:

```python
# Instead of loading entire table, get file paths
file_paths = [task.file.file_path for task in table.scan().plan_files()]

# Distribute files across partitions (zero-copy!)
partition_files = distribute_files(file_paths, num_partitions)

# Process partitions in parallel
results = await asyncio.gather(*[
    aggregate_partition(files, partition_id)
    for partition_id, files in enumerate(partition_files)
])
```

This approach achieves true parallel file processing without loading the entire dataset into memory.

## Example: Optimizing a data pipeline

### Before optimization

```python
@env.task
async def process_item(item: dict) -> dict:
    # Very fast operation (~100ms)
    return {"processed": item["id"]}

@env.task
async def process_dataset(items: list[dict]) -> list[dict]:
    # Create 1M tasks
    results = await asyncio.gather(*[process_item(item) for item in items])
    return results
```

**Issues**:

- 1M tasks created (exceeds UI limit)
- Task overhead >> task runtime (100ms task, seconds of overhead)
- High load on Queue Service and object storage

### After optimization

```python
# Use reusable containers
env = flyte.TaskEnvironment(
    name="optimized-pipeline",
    reuse_policy=flyte.ReusePolicy(
        replicas=(5, 20),
        concurrency=10,
        scaledown_ttl=timedelta(minutes=10),
        idle_ttl=timedelta(hours=1)
    )
)

@env.task
async def process_batch(items: list[dict]) -> list[dict]:
    # Process batch of items
    return [{"processed": item["id"]} for item in items]

@env.task
async def process_dataset(items: list[dict]) -> list[dict]:
    # Create 1000 tasks (batch size 1000)
    batch_size = 1000
    batches = [items[i:i + batch_size] for i in range(0, len(items), batch_size)]

    results = await flyte.map(process_batch, batches)
    return [item for batch in results for item in batch]
```

**Improvements**:

- 1000 tasks instead of 1M (within limits)
- Batch runtime ~100 seconds (100ms × 1000 items)
- Reusable containers eliminate startup overhead
- Concurrency enables high throughput (200 concurrent tasks max)

## When to contact the Union team

Reach out to the Union team if you:

- Need more than 50k actions per run
- Want to use high-performance metastores (Redis, PostgreSQL) instead of object stores
- Have specific performance requirements or constraints
- Need help profiling and optimizing your workflows


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/batch-inference ===

# Maximize GPU utilization for batch inference

GPUs are expensive. When running batch inference, the single biggest cost driver is **idle GPU time**: cycles where the GPU sits waiting with nothing to do. Understanding why this happens and how to fix it is the key to cost-effective batch inference.

## Why GPU utilization drops

A typical inference task does three things:

1. **Load data**: read from storage, deserialize, preprocess (CPU/IO-bound)
2. **Run inference**: forward pass through the model (GPU-bound)
3. **Post-process**: format results, write outputs (CPU/IO-bound)

When these steps run sequentially, the GPU is idle during steps 1 and 3. For many workloads, data loading and preprocessing dominate wall-clock time, leaving the GPU busy for only a fraction of the total:

```mermaid
gantt
    title Sequential execution — GPU idle during CPU/IO work
    dateFormat X
    axisFormat %s

    section Task 1
    Load data (CPU/IO)     :a1, 0, 3
    Inference (GPU)        :a2, after a1, 2
    Post-process (CPU/IO)  :a3, after a2, 1

    section Task 2
    Load data (CPU/IO)     :b1, after a3, 3
    Inference (GPU)        :b2, after b1, 2
    Post-process (CPU/IO)  :b3, after b2, 1

    section GPU
    Idle                   :crit, g1, 0, 3
    Busy                   :active, g2, 3, 5
    Idle                   :crit, g3, 5, 9
    Busy                   :active, g4, 9, 11
    Idle                   :crit, g5, 11, 12
```

In this example, the GPU is busy for only 4 out of 12 time units: **33% utilization**. The rest is wasted waiting for CPU and IO operations.

## Serving vs in-process batch inference

There are two common approaches to batch inference: sending requests to a **hosted model server** (serving), or running the model **in-process** alongside data loading. Each has distinct trade-offs:

| | Hosted serving | In-process (Flyte) |
|---|---|---|
| **Architecture** | Separate inference server (e.g. Triton, vLLM server, TGI) accessed over the network | Model loaded directly in the task process, inference via `DynamicBatcher` |
| **Data transfer** | Every request serialized over the network; large payloads add latency | Zero-copy: data stays in-process, no serialization overhead |
| **Backpressure** | Hard to implement; push-based architecture can overwhelm the server or drop requests | Two levels: `DynamicBatcher` queue blocks producers when full, and Flyte's task scheduling automatically queues new inference tasks when replicas are busy; backpressure propagates end-to-end without any extra code |
| **Utilization** | Servers are often over-provisioned to maintain availability, leading to low average utilization | Batcher continuously fills the GPU with work from concurrent producers |
| **Multi-model** | Each model needs its own serving deployment, load balancer, and scaling config | Multiple models can time-share the same GPU: when one model finishes, the next is loaded automatically via reusable containers, no container orchestration required |
| **Scaling** | Requires separate infrastructure for the serving layer (load balancers, autoscalers, health checks) | Scales with Flyte: replicas auto-scale based on demand |
| **Cost** | Pay for always-on serving infrastructure even during low-traffic periods | Pay only for the duration of the batch job |
| **Fault tolerance** | Need retries, circuit breakers, and timeout handling for network failures | Failures are local; Flyte handles retries and recovery at the task level |
| **Best for** | Real-time / low-latency serving with unpredictable request patterns | Large-scale batch processing with known datasets |

For batch workloads, in-process inference eliminates the network overhead and infrastructure complexity of a serving layer while achieving higher GPU utilization through intelligent batching.

## Solution: `DynamicBatcher`

`DynamicBatcher` from `flyte.extras` solves the utilization problem by **separating data loading from inference** and running them concurrently. Multiple async producers load and preprocess data while a single consumer feeds the GPU in optimally-sized batches:

```mermaid
flowchart LR
    subgraph producers ["Concurrent producers (CPU/IO)"]
        P1["Stream 1: load + preprocess"]
        P2["Stream 2: load + preprocess"]
        P3["Stream N: load + preprocess"]
    end

    subgraph batcher ["DynamicBatcher"]
        Q["Queue with backpressure"]
        A["Aggregation loop<br/>(assembles cost-budgeted batches)"]
        Q --> A
    end

    subgraph consumer ["Processing loop (GPU)"]
        G["process_fn / inference_fn<br/>(batched forward pass)"]
    end

    P1 --> Q
    P2 --> Q
    P3 --> Q
    A --> G
```

The batcher runs two internal loops:

1. **Aggregation loop**: drains the submission queue and assembles batches that respect a cost budget (`target_batch_cost`), a maximum size (`max_batch_size`), and a timeout (`batch_timeout_s`). This ensures the GPU always receives optimally-sized batches.
2. **Processing loop**: pulls assembled batches and calls your processing function, resolving each record's future with its result.

This pipelining means the GPU is processing batch N while data for batch N+1 is being loaded and assembled, **eliminating idle time**.

### Basic usage

```python
from flyte.extras import DynamicBatcher

async def process(batch: list[dict]) -> list[str]:
    """Your batch processing function. Must return results in the same order as the input."""
    return [heavy_computation(item) for item in batch]

async with DynamicBatcher(
    process_fn=process,
    target_batch_cost=1000,   # cost budget per batch
    max_batch_size=64,        # hard cap on records per batch
    batch_timeout_s=0.05,     # max wait time before dispatching a partial batch
    max_queue_size=5_000,     # queue size for backpressure
) as batcher:
    futures = []
    for record in my_records:
        future = await batcher.submit(record, estimated_cost=10)
        futures.append(future)
    results = await asyncio.gather(*futures)
```

Each call to `submit()` is non-blocking; it enqueues the record and immediately returns a `Future`. When the queue is full, `submit()` awaits until space is available, providing natural backpressure to prevent producers from overwhelming the GPU.

### Cost estimation

The batcher uses cost estimates to decide how many records to group into each batch. You can provide costs in several ways (checked in order of precedence):

1. **Explicit**: pass `estimated_cost` to `submit()`
2. **Estimator function**: pass `cost_estimator` to the constructor
3. **Protocol**: implement `estimate_cost()` on your record type
4. **Default**: falls back to `default_cost` (default: 1)

## `TokenBatcher` for LLM inference

For LLM workloads, `TokenBatcher` is a convenience subclass that uses token-aware parameter names:

```python
from dataclasses import dataclass
from flyte.extras import TokenBatcher

@dataclass
class Prompt:
    text: str

    def estimate_tokens(self) -> int:
        """Rough token estimate (~4 chars per token)."""
        return len(self.text) // 4 + 1

async def inference(batch: list[Prompt]) -> list[str]:
    """Run batched inference through your model."""
    texts = [p.text for p in batch]
    outputs = model.generate(texts, sampling_params)
    return [o.outputs[0].text for o in outputs]

async with TokenBatcher(
    inference_fn=inference,
    target_batch_tokens=32_000,  # token budget per batch
    max_batch_size=256,
) as batcher:
    future = await batcher.submit(Prompt(text="What is 2+2?"))
    result = await future
```

`TokenBatcher` checks the `TokenEstimator` protocol (`estimate_tokens()`) in addition to `CostEstimator` (`estimate_cost()`), making it natural to work with prompt types.

## Combining with app environments

[`DynamicBatcher`](../../api-reference/flyte-sdk/flyte.extras/dynamicbatcher) on its own improves utilization within a single task, but the model has to be loaded from scratch on every invocation. To amortize that cost across many task runs, host the model inside a long-lived [`AppEnvironment`](../../api-reference/flyte-sdk/flyte.app/appenvironment) and have driver tasks call it over HTTP:

- **Amortized model loading**: the model is loaded once when the app starts and stays in memory for the lifetime of the replica
- **Cross-task batching**: every concurrent HTTP request submits to the **same shared [`TokenBatcher`](../../api-reference/flyte-sdk/flyte.extras/tokenbatcher)**, so the GPU always has a full queue of work
- **Automatic scaling**: the app autoscales between min and max replicas based on a concurrency target, and each replica maintains its own model and batcher

```mermaid
flowchart LR
    D["Driver task<br/>fans out chunks<br/>(concurrency cap)"]

    subgraph calls ["infer_batch tasks (HTTP clients)"]
        T1["call 1"]
        T2["call 2"]
        T3["call N"]
    end

    D --> T1
    D --> T2
    D --> T3

    subgraph app ["FastAPI app environment (GPU)"]
        FA["POST /generate"]
        B["Shared TokenBatcher"]
        M["vLLM model<br/>(loaded in lifespan)"]
        FA --> B --> M
    end

    T1 --> FA
    T2 --> FA
    T3 --> FA
```

The two key techniques are:

1. **Use FastAPI's `lifespan`** to load the model and start the `TokenBatcher` exactly once per replica, then attach the batcher to `app.state` so request handlers can reach it.
2. **Cap driver concurrency** with `flyte.map.aio(..., concurrency=N)` so the orchestrator doesn't overload the app with more in-flight requests than its scaling target can serve.

### Example: batch LLM inference with vLLM behind a FastAPI app

This example loads math problems from HuggingFace's gsm8k dataset and solves them by calling a FastAPI app that runs vLLM with a shared `TokenBatcher`.

#### 1. Load the model and batcher once via FastAPI lifespan

The FastAPI `lifespan` runs on startup and shutdown. Use it to load the vLLM model and start the `TokenBatcher` exactly once per replica, then attach the batcher to `app.state` so request handlers can reach it:

```python
import asyncio
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel

import flyte
import flyte.app
from flyte.app.extras import FastAPIAppEnvironment
from flyte.extras import TokenBatcher

logger = logging.getLogger(__name__)

@dataclass
class Prompt:
    task_id: str
    index: int
    text: str

@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Load the vLLM model and start the TokenBatcher once at startup."""
    from vllm import LLM, SamplingParams

    llm = LLM(
        model="Qwen/Qwen2.5-7B-Instruct",
        gpu_memory_utilization=0.9,
        max_model_len=4096,
    )
    params = SamplingParams(temperature=0.7, max_tokens=512)
    logger.info("vLLM model loaded")

    async def inference(batch: list[Prompt]) -> list[str]:
        texts = [p.text for p in batch]
        outputs = llm.generate(texts, params)
        return [o.outputs[0].text for o in outputs]

    batcher = TokenBatcher[Prompt, str](
        inference_fn=inference,
        target_batch_tokens=32_000,
        max_batch_size=256,
        batch_timeout_s=0.05,
        max_queue_size=5_000,
    )
    await batcher.start()
    logger.info("TokenBatcher started")

    app.state.batcher = batcher
    yield
    await batcher.stop()

app = FastAPI(title="Batched Inference Service", lifespan=lifespan)
```

Stashing the batcher on `app.state` means every request handler can grab the same shared instance via `request.app.state.batcher`, so concurrent requests all feed into one queue.

#### 2. Add an endpoint that submits to the shared batcher

Each request just enqueues records; the batcher aggregates records across concurrent requests into token-budgeted batches before hitting the GPU:

```python
class GenerateRequest(BaseModel):
    prompts: list[str]
    task_id: str

@app.post("/generate")
async def generate(request_body: GenerateRequest, request: Request):
    if not request_body.prompts:
        raise HTTPException(status_code=400, detail="No prompts provided")

    batcher: TokenBatcher[Prompt, str] = request.app.state.batcher

    futures: list[asyncio.Future[str]] = []
    for idx, text in enumerate(request_body.prompts):
        record = Prompt(task_id=request_body.task_id, index=idx, text=text)
        future = await batcher.submit(record)
        futures.append(future)

    results = await asyncio.gather(*futures)
    return {"results": results}
```

#### 3. Define the app environment and driver task environment

The app uses a [`FastAPIAppEnvironment`](../../api-reference/flyte-sdk/flyte.app.extras/fastapiappenvironment) on a GPU and autoscales via [`Scaling`](../../api-reference/flyte-sdk/flyte.app/scaling). The driver runs in a CPU-only [`TaskEnvironment`](../../api-reference/flyte-sdk/flyte/taskenvironment) that `depends_on` the app so the app is deployed before the driver runs:

```python
image = (
    flyte.Image.from_debian_base()
    .with_pip_packages("vllm", "hf-transfer", "fastapi", "uvicorn")
    .with_env_vars({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

app_env = FastAPIAppEnvironment(
    name="batch-inference-saturate-app",
    app=app,
    image=image,
    resources=flyte.Resources(cpu=6, memory="24Gi", gpu="L4:1", disk="64Gi"),
    scaling=flyte.app.Scaling(
        replicas=(0, 2),
        metric=flyte.app.Scaling.Concurrency(val=10),
        scaledown_after=300,
    ),
    requires_auth=False,
)

driver_env = flyte.TaskEnvironment(
    name="batch_inference_saturate_app_driver",
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    image=image,
    depends_on=[app_env],
)
```

With `replicas=(0, 2)` and a concurrency target of `10`, the app scales between 0 and 2 GPU replicas and aims for ~10 concurrent in-flight requests per replica, so up to ~20 requests can be served in parallel.

#### 4. Define a driver task that calls the app

The driver task POSTs prompt chunks to the app's endpoint. Use generous timeouts and retries to absorb cold starts and transient failures during scaling events:

```python
import httpx

@driver_env.task(retries=20)
async def infer_batch(
    endpoint: str,
    prompts: list[str],
    task_id: str,
) -> list[str]:
    url = f"{endpoint}/generate"
    async with httpx.AsyncClient(
        timeout=httpx.Timeout(connect=60.0, read=600.0, write=30.0, pool=10.0),
    ) as client:
        response = await client.post(
            url,
            json={"prompts": prompts, "task_id": task_id},
        )
        response.raise_for_status()
        return response.json()["results"]
```

#### 5. Fan out chunks with a concurrency cap

The orchestrator chunks the dataset and submits each chunk as a separate `infer_batch` call. Use `flyte.map.aio(..., concurrency=max_concurrency)` to cap the number of in-flight HTTP calls so the task doesn't overload the app with more requests than its scaling target can serve:

```python
@driver_env.task
async def main(
    num_questions: int = 500,
    chunk_size: int = 50,
    max_concurrency: int = 10,
) -> dict[str, list[str]]:
    questions = await fetch_gsm8k_questions(num_questions)
    endpoint = app_env.endpoint

    chunks = [
        questions[i : i + chunk_size]
        for i in range(0, len(questions), chunk_size)
    ]
    task_ids = [f"gsm8k_{i:03d}" for i in range(len(chunks))]

    all_results = [
        result
        async for result in flyte.map.aio(
            infer_batch,
            [endpoint] * len(chunks),
            chunks,
            task_ids,
            concurrency=max_concurrency,
        )
    ]
    return dict(zip(task_ids, all_results))
```

> [!IMPORTANT]
> Match `max_concurrency` to the app's scaling configuration. In this example, the app autoscales up to 2 replicas with a concurrency target of 10, so ~20 requests can be in flight at once. Setting `max_concurrency=10` keeps the driver from queueing requests far beyond what the app can absorb, which would otherwise stack up behind the batcher's `max_queue_size`, exhaust HTTP timeouts, and waste retry budget.

## Monitoring utilization

`DynamicBatcher` exposes a `stats` property with real-time metrics:

```python
stats = batcher.stats
print(f"Utilization: {stats.utilization:.1%}")         # fraction of time spent processing
print(f"Records processed: {stats.total_completed}")
print(f"Batches dispatched: {stats.total_batches}")
print(f"Avg batch size: {stats.avg_batch_size:.1f}")
print(f"Busy time: {stats.busy_time_s:.1f}s")
print(f"Idle time: {stats.idle_time_s:.1f}s")
```

| Metric | Description |
|---|---|
| `utilization` | Fraction of wall-clock time spent inside `process_fn` (0.0 to 1.0). Target: > 0.9. |
| `total_submitted` | Total records submitted via `submit()` |
| `total_completed` | Total records whose futures have been resolved |
| `total_batches` | Number of batches dispatched to `process_fn` |
| `avg_batch_size` | Running average records per batch |
| `avg_batch_cost` | Running average cost per batch |
| `busy_time_s` | Cumulative seconds spent inside `process_fn` |
| `idle_time_s` | Cumulative seconds the processing loop waited for batches |

If utilization is low, consider:

- **Increasing concurrency**: more concurrent producers means the batcher has more records to assemble into batches
- **Reducing `batch_timeout_s`**: dispatch partial batches faster instead of waiting
- **Increasing `max_queue_size`**: allow more records to be buffered ahead of the GPU
- **Adding more data streams**: ensure the GPU always has work queued up


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/advanced-project ===

# Advanced project: LLM reporting agent

This example demonstrates a resilient agentic report generator that showcases
Flyte 2.0's advanced features for building production-grade AI workflows.

## What you'll build

A batch report generator that:

1. Processes multiple topics in parallel
2. Iteratively critiques and refines each report until it meets a quality threshold
3. Produces multiple output formats (Markdown, HTML, summary) for each report
4. Serves results through an interactive UI

## Concepts covered

| Feature | Description |
|---------|-------------|
| `ReusePolicy` | Keep containers warm for high-throughput batch processing |
| `@flyte.trace` | Checkpoint LLM calls for recovery and observability |
| `RetryStrategy` | Handle transient API failures gracefully |
| `flyte.group` | Organize parallel batches and iterations in the UI |
| `asyncio.gather` | Fan out to process multiple topics concurrently |
| Pydantic models | Structured LLM outputs |
| `AppEnvironment` | Deploy interactive Streamlit apps |
| `RunOutput` | Connect apps to pipeline outputs |

## Architecture

```mermaid
flowchart TD
    A[Topics List] --> B

    B["report_batch_pipeline<br/><i>driver_env</i>"]

    subgraph B1 ["refine_all (parallel)"]
        direction LR
        R1["refine_report<br/>topic 1"]
        R2["refine_report<br/>topic 2"]
        R3["refine_report<br/>topic N"]
    end
    B --> B1

    subgraph B2 ["format_all (parallel)"]
        direction LR
        F1["format_outputs<br/>report 1"]
        F2["format_outputs<br/>report 2"]
        F3["format_outputs<br/>report N"]
    end
    B1 --> B2

    B2 --> C["Output: List of Dirs"]
```

Each `refine_report` task runs in a reusable container (`llm_env`) and performs
multiple LLM calls through traced functions:

```mermaid
flowchart TD
    A[Topic] --> B["generate_initial_draft<br/><i>@flyte.trace</i>"]
    B --> C

    subgraph C ["refinement_loop"]
        direction TB
        D["critique_content<br/><i>@flyte.trace</i>"] -->|score >= threshold| E[exit loop]
        D -->|score < threshold| F["revise_content<br/><i>@flyte.trace</i>"]
        F --> D
    end
    C --> G[Refined Report]
```

## Prerequisites

- A Flyte account with an active project
- An OpenAI API key stored as a secret named `openai-api-key`

To create the secret:

```bash
flyte secret create openai-api-key
```

## Parts

1. ****Advanced project: LLM reporting agent > Resilient generation****: Set up reusable environments, traced LLM calls, and retry strategies
2. ****Advanced project: LLM reporting agent > Agentic refinement****: Build the iterative critique-and-revise loop
3. ****Advanced project: LLM reporting agent > Parallel outputs****: Generate multiple formats concurrently
4. ****Advanced project: LLM reporting agent > Serving app****: Deploy an interactive UI for report generation

## Key takeaways

1. **Reusable environments for batch processing**: `ReusePolicy` keeps containers warm,
   enabling efficient processing of multiple topics without cold start overhead. With
   5 topics × ~7 LLM calls each, the reusable pool handles ~35 calls efficiently.

2. **Checkpointed LLM calls**: `@flyte.trace` provides automatic checkpointing at the
   function level, enabling recovery without re-running expensive API calls.

3. **Agentic patterns**: The generate-critique-revise loop demonstrates how to build
   self-improving AI workflows with clear observability through `flyte.group`.

4. **Parallel fan-out**: `asyncio.gather` processes multiple topics concurrently,
   maximizing throughput by running refinement tasks in parallel across the batch.

## Subpages
- **Advanced project: LLM reporting agent > Resilient generation**
- **Advanced project: LLM reporting agent > Agentic refinement**
- **Advanced project: LLM reporting agent > Parallel outputs**
- **Advanced project: LLM reporting agent > Serving app**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/resilient-generation ===

# Resilient generation

This section covers the foundational patterns for building resilient LLM-powered
tasks: reusable environments, traced function calls, and retry strategies.

## Two environments

This example uses two task environments with different characteristics:

1. **`llm_env`** (reusable): For tasks that make many LLM calls in a loop or
   process batches in parallel. Container reuse avoids cold starts.
2. **`driver_env`** (standard): For orchestration tasks that fan out work to
   other tasks but don't make LLM calls themselves.

### Reusable environment for LLM work

When processing a batch of topics, each topic goes through multiple LLM calls
(generate, critique, revise, repeat). With 5 topics × ~7 calls each, that's ~35
LLM calls. `ReusePolicy` keeps containers warm to handle this efficiently:

```python
# Reusable environment for tasks that make many LLM calls in a loop.
# The ReusePolicy keeps containers warm, reducing cold start latency for iterative work.
llm_env = flyte.TaskEnvironment(
    name="llm-worker",
    secrets=[] if MOCK_MODE else [flyte.Secret(key="openai-api-key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "unionai-reuse>=0.1.10",
        "openai>=1.0.0",
        "pydantic>=2.0.0",
    ),
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    reusable=flyte.ReusePolicy(
        replicas=2,              # Keep 2 container instances ready
        concurrency=4,           # Allow 4 concurrent tasks per container
        scaledown_ttl=timedelta(minutes=5),   # Wait 5 min before scaling down
        idle_ttl=timedelta(minutes=30),       # Shut down after 30 min idle
    ),
    cache="auto",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### ReusePolicy parameters

| Parameter | Description |
|-----------|-------------|
| `replicas` | Number of container instances to keep ready (or `(min, max)` tuple) |
| `concurrency` | Maximum tasks per container at once |
| `scaledown_ttl` | Minimum wait before scaling down a replica |
| `idle_ttl` | Time after which idle containers shut down completely |

The configuration above keeps 2 containers ready, allows 4 concurrent tasks per
container, waits 5 minutes before scaling down, and shuts down after 30 minutes
of inactivity.

> **📝 Note**
>
> Both `scaledown_ttl` and `idle_ttl` must be at least 30 seconds.

### Standard environment for orchestration

The driver environment doesn't need container reuse; it just coordinates work.
The `depends_on` parameter declares that tasks in this environment call tasks
in `llm_env`, ensuring both environments are deployed together:

```python
# Standard environment for orchestration tasks that don't need container reuse.
# depends_on declares that this environment's tasks call tasks in llm_env.
driver_env = flyte.TaskEnvironment(
    name="driver",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "pydantic>=2.0.0",
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    depends_on=[llm_env],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

## Traced LLM calls

The `@flyte.trace` decorator provides automatic checkpointing at the function level.
When a traced function completes successfully, its result is cached. If the task
fails and restarts, previously completed traced calls return their cached results
instead of re-executing.

```python
@flyte.trace
async def call_llm(prompt: str, system: str, json_mode: bool = False) -> str:
    """
    Make an LLM call with automatic checkpointing.

    The @flyte.trace decorator provides:
    - Automatic caching of results for identical inputs
    - Recovery from failures without re-running successful calls
    - Full observability in the Flyte UI

    Args:
        prompt: The user prompt to send
        system: The system prompt defining the LLM's role
        json_mode: Whether to request JSON output

    Returns:
        The LLM's response text
    """
    # Use mock responses for testing without API keys
    if MOCK_MODE:
        import asyncio
        await asyncio.sleep(0.5)  # Simulate API latency

        if "critique" in prompt.lower() or "critic" in system.lower():
            # Return good score if draft has been revised (contains revision marker)
            if "[REVISED]" in prompt:
                return MOCK_CRITIQUE_GOOD
            return MOCK_CRITIQUE_NEEDS_WORK
        elif "summary" in system.lower():
            return MOCK_SUMMARY
        elif "revis" in system.lower():
            # Return revised version with marker
            return MOCK_REPORT.replace("## Introduction", "[REVISED]\n\n## Introduction")
        else:
            return MOCK_REPORT

    from openai import AsyncOpenAI

    client = AsyncOpenAI()

    kwargs = {
        "model": "gpt-4o-mini",
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt},
        ],
        "max_tokens": 2000,
    }

    if json_mode:
        kwargs["response_format"] = {"type": "json_object"}

    response = await client.chat.completions.create(**kwargs)
    return response.choices[0].message.content
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### Benefits of tracing

1. **Cost savings**: Failed tasks don't re-run expensive API calls that already succeeded
2. **Faster recovery**: Resuming from checkpoints skips completed work
3. **Observability**: Each traced call appears in the Flyte UI with timing data

### When to use @flyte.trace

Use `@flyte.trace` for:

- LLM API calls (OpenAI, Anthropic, etc.)
- External API requests
- Any expensive operation you don't want to repeat on retry

Don't use `@flyte.trace` for:

- Simple computations (overhead outweighs benefit)
- Operations with side effects that shouldn't be skipped

## Traced helper functions

The LLM-calling functions are decorated with `@flyte.trace` rather than being
separate tasks. This keeps the architecture simple while still providing
checkpointing:

```python
@flyte.trace
async def generate_initial_draft(topic: str) -> str:
    """
    Generate the initial report draft.

    The @flyte.trace decorator provides checkpointing - if the task fails
    after this completes, it won't re-run on retry.

    Args:
        topic: The topic to write about

    Returns:
        The initial draft in markdown format
    """
    print(f"Generating initial draft for topic: {topic}")

    prompt = f"Write a comprehensive report on the following topic:\n\n{topic}"
    draft = await call_llm(prompt, GENERATOR_SYSTEM_PROMPT)

    print(f"Generated initial draft ({len(draft)} characters)")
    return draft
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

These traced functions run inside the `refine_report` task. If the task fails
and retries, completed traced calls return cached results instead of re-executing.

## Retry strategies

The task that orchestrates the LLM calls uses `retries` to handle transient failures:

```python
@llm_env.task(retries=3)
async def refine_report(topic: str, ...) -> str:
    # Traced functions are called here
    draft = await generate_initial_draft(topic)
    ...
```

### Configuring retries

You can specify retries as a simple integer:

```python
@llm_env.task(retries=3)
async def my_task():
    ...
```

Or use `RetryStrategy` for more control:

```python
@llm_env.task(retries=flyte.RetryStrategy(count=3))
async def my_task():
    ...
```

### Combining tracing with retries

When you combine `@flyte.trace` with task-level retries, you get the best of both:

1. Task fails after completing some traced calls
2. Flyte retries the task
3. Previously completed traced calls return cached results
4. Only the failed operation (and subsequent ones) re-execute

This pattern is essential for multi-step LLM workflows where you don't want to
re-run the entire chain when a single call fails.

## Structured prompts

The example uses a separate `prompts.py` module for system prompts and Pydantic models:

```python
GENERATOR_SYSTEM_PROMPT = """You are an expert report writer. Generate a well-structured,
informative report on the given topic. The report should include:

1. An engaging introduction that sets context
2. Clear sections with descriptive headings
3. Specific facts, examples, or data points where relevant
4. A conclusion that summarizes key takeaways

Write in a professional but accessible tone. Use markdown formatting for structure.
Aim for approximately 500-800 words."""

CRITIC_SYSTEM_PROMPT = """You are a demanding but fair editor reviewing a report draft.
Evaluate the report on these criteria:

- Clarity: Is the writing clear and easy to follow?
- Structure: Is it well-organized with logical flow?
- Depth: Does it provide sufficient detail and insight?
- Accuracy: Are claims supported and reasonable?
- Engagement: Is it interesting to read?

Provide your response as JSON matching this schema:
{
    "score": <1-10 integer>,
    "strengths": ["strength 1", "strength 2", ...],
    "improvements": ["improvement 1", "improvement 2", ...],
    "summary": "brief overall assessment"
}

Be specific in your feedback. A score of 8+ means the report is ready for publication."""

REVISER_SYSTEM_PROMPT = """You are an expert editor revising a report based on feedback.
Your task is to improve the report by addressing the specific improvements requested
while preserving its strengths.

Guidelines:
- Address each improvement point specifically
- Maintain the original voice and style
- Keep the same overall structure unless restructuring is requested
- Preserve any content that was praised as a strength
- Ensure the revised version is cohesive and flows well

Return only the revised report in markdown format, no preamble or explanation."""

SUMMARY_SYSTEM_PROMPT = """Create a concise executive summary (2-3 paragraphs) of the
following report. Capture the key points and main takeaways. Write in a professional
tone suitable for busy executives who need the essential information quickly."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/prompts.py*

### Pydantic models for structured output

LLM responses can be unpredictable. Using Pydantic models with JSON mode ensures
you get structured, validated data:

```python
class Critique(BaseModel):
    """Structured critique response from the LLM."""

    score: int = Field(
        ge=1,
        le=10,
        description="Quality score from 1-10, where 10 is publication-ready",
    )
    strengths: list[str] = Field(
        description="List of strengths in the current draft",
    )
    improvements: list[str] = Field(
        description="Specific improvements needed",
    )
    summary: str = Field(
        description="Brief summary of the critique",
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/prompts.py*

The `Critique` model validates that:

- `score` is an integer between 1 and 10
- `strengths` and `improvements` are lists of strings
- All required fields are present

If the LLM returns malformed JSON, Pydantic raises a validation error, which
triggers a retry (if configured).

## Next steps

With resilient generation in place, you're ready to build the
[agentic refinement loop](./agentic-refinement).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/agentic-refinement ===

# Agentic refinement

The core of this example is an agentic refinement loop: generate content, critique
it, revise based on feedback, and repeat until quality meets a threshold. This
pattern is fundamental to building self-improving AI systems.

## The agentic pattern

Traditional pipelines are linear: input → process → output. Agentic workflows
are iterative: they evaluate their own output and improve it through multiple
cycles.

```mermaid
flowchart TD
    A[Generate] --> B[Critique]
    B -->|score >= threshold| C[Done]
    B -->|score < threshold| D[Revise]
    D --> B
```

## Critique function

The critique function evaluates the current draft and returns structured feedback.
It's a traced function (not a separate task) that runs inside `refine_report`:

```python
@flyte.trace
async def critique_content(draft: str) -> Critique:
    """
    Critique the current draft and return structured feedback.

    Uses Pydantic models to parse the LLM's JSON response into
    a typed object for reliable downstream processing.

    Args:
        draft: The current draft to critique

    Returns:
        Structured critique with score, strengths, and improvements
    """
    print("Critiquing current draft...")

    response = await call_llm(
        f"Please critique the following report:\n\n{draft}",
        CRITIC_SYSTEM_PROMPT,
        json_mode=True,
    )

    # Parse the JSON response into our Pydantic model
    critique_data = json.loads(response)
    critique = Critique(**critique_data)

    print(f"Critique score: {critique.score}/10")
    print(f"Strengths: {len(critique.strengths)}, Improvements: {len(critique.improvements)}")

    return critique
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

Key points:

- Uses `json_mode=True` to ensure the LLM returns valid JSON
- Parses the response into a Pydantic `Critique` model
- Returns a typed object for reliable downstream processing
- `@flyte.trace` provides checkpointing: if the task retries, completed critiques aren't re-run

## Revise function

The revise function takes the current draft and specific improvements to address:

```python
@flyte.trace
async def revise_content(draft: str, improvements: list[str]) -> str:
    """
    Revise the draft based on critique feedback.

    Args:
        draft: The current draft to revise
        improvements: List of specific improvements to address

    Returns:
        The revised draft
    """
    print(f"Revising draft to address {len(improvements)} improvements...")

    improvements_text = "\n".join(f"- {imp}" for imp in improvements)
    prompt = f"""Please revise the following report to address these improvements:

IMPROVEMENTS NEEDED:
{improvements_text}

CURRENT DRAFT:
{draft}"""

    revised = await call_llm(prompt, REVISER_SYSTEM_PROMPT)

    print(f"Revision complete ({len(revised)} characters)")
    return revised
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

The prompt includes:

1. The list of improvements from the critique
2. The current draft to revise

This focused approach helps the LLM make targeted changes rather than rewriting
from scratch.

## The refinement loop

The `refine_report` task orchestrates the iterative refinement. It runs in the
reusable `llm_env` because it makes multiple LLM calls through traced functions:

```python
@llm_env.task(retries=3)
async def refine_report(
    topic: str,
    max_iterations: int = 3,
    quality_threshold: int = 8,
) -> str:
    """
    Iteratively refine a report until it meets the quality threshold.

    This task runs in a reusable container because it makes multiple LLM calls
    in a loop. The traced helper functions provide checkpointing, so if the
    task fails mid-loop, completed LLM calls won't be re-run on retry.

    Args:
        topic: The topic to write about
        max_iterations: Maximum refinement cycles (default: 3)
        quality_threshold: Minimum score to accept (default: 8)

    Returns:
        The final refined report
    """
    # Generate initial draft
    draft = await generate_initial_draft(topic)

    # Iterative refinement loop
    for i in range(max_iterations):
        with flyte.group(f"refinement_{i + 1}"):
            # Get critique
            critique = await critique_content(draft)

            # Check if we've met the quality threshold
            if critique.score >= quality_threshold:
                print(f"Quality threshold met at iteration {i + 1}!")
                print(f"Final score: {critique.score}/10")
                break

            # Revise based on feedback
            print(f"Score {critique.score} < {quality_threshold}, revising...")
            draft = await revise_content(draft, critique.improvements)
    else:
        print(f"Reached max iterations ({max_iterations})")

    return draft
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### How it works

1. **Generate initial draft**: Creates the first version of the report
2. **Enter refinement loop**: Iterates up to `max_iterations` times
3. **Critique**: Evaluates the current draft and assigns a score
4. **Check threshold**: If score meets `quality_threshold`, exit early
5. **Revise**: If below threshold, revise based on improvements
6. **Repeat**: Continue until threshold met or iterations exhausted

All the LLM calls (generate, critique, revise) are traced functions inside this
single task. This keeps the task graph simple while the reusable container handles
the actual LLM work efficiently.

### Early exit

The `if critique.score >= quality_threshold: break` pattern enables early exit
when quality is sufficient. This saves compute costs and time: no need to run
all iterations if the first draft is already good.

## Grouping iterations with flyte.group

Each refinement iteration is wrapped in `flyte.group`:

```python
for i in range(max_iterations):
    with flyte.group(f"refinement_{i + 1}"):
        critique = await critique_content(draft)
        # ...
```

### Why use flyte.group?

Groups provide hierarchical organization in the Flyte UI. Since critique and
revise are traced functions (not separate tasks), groups help organize them:

```
refine_report
├── generate_initial_draft (traced)
├── refinement_1
│   ├── critique_content (traced)
│   └── revise_content (traced)
├── refinement_2
│   ├── critique_content (traced)
│   └── revise_content (traced)
└── [returns refined report]
```

Benefits:

- **Clarity**: See exactly how many iterations occurred
- **Debugging**: Quickly find which iteration had issues
- **Observability**: Track time spent in each refinement cycle

### Group context

Groups are implemented as context managers. All traced calls and nested groups
within the `with flyte.group(...)` block are associated with that group.

## Configuring the loop

The refinement loop accepts parameters to tune its behavior:

| Parameter | Default | Description |
|-----------|---------|-------------|
| `max_iterations` | 3 | Upper bound on refinement cycles |
| `quality_threshold` | 8 | Minimum score (1-10) to accept |

### Choosing thresholds

- **Higher threshold** (9-10): More refinement cycles, higher quality, more API costs
- **Lower threshold** (6-7): Faster completion, may accept lower quality
- **More iterations**: Safety net for difficult topics
- **Fewer iterations**: Cost control, faster turnaround

A good starting point is `quality_threshold=8` with `max_iterations=3`. Adjust
based on your quality requirements and budget.

## Best practices for agentic loops

1. **Always set max iterations**: Prevent infinite loops if the quality threshold
   is never reached.

2. **Use structured critiques**: Pydantic models ensure you can reliably extract
   the score and improvements from LLM responses.

3. **Log iteration progress**: Print statements help debug when reviewing logs:

   ```python
   print(f"Iteration {i + 1}: score={critique.score}")
   ```

4. **Consider diminishing returns**: After 3-4 iterations, improvements often
   become marginal. Set `max_iterations` accordingly.

5. **Use groups for observability**: `flyte.group` makes the iterative nature
   visible in the UI, essential for debugging and monitoring.

## Next steps

With the agentic refinement loop complete, learn how to
[generate multiple outputs in parallel](./parallel-outputs).


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/parallel-outputs ===

# Parallel outputs

After refining the report, the pipeline generates multiple output formats in
parallel. This demonstrates how to use `asyncio.gather` for concurrent execution
within a task.

## The formatting functions

The pipeline generates three outputs: markdown, HTML, and an executive summary.
Only `generate_summary` uses `@flyte.trace` because it makes an LLM call.
The markdown and HTML functions are simple, deterministic transformations that
don't benefit from checkpointing:

```python
async def format_as_markdown(content: str) -> str:
    """Format the report as clean markdown."""
    # Content is already markdown, but we could add TOC, metadata, etc.
    return f"""---
title: Generated Report
date: {__import__('datetime').datetime.now().isoformat()}
---

{content}
"""

async def format_as_html(content: str) -> str:
    """Convert the report to HTML."""
    # Simple markdown to HTML conversion
    import re

    html = content
    # Convert headers
    html = re.sub(r"^### (.+)$", r"<h3>\1</h3>", html, flags=re.MULTILINE)
    html = re.sub(r"^## (.+)$", r"<h2>\1</h2>", html, flags=re.MULTILINE)
    html = re.sub(r"^# (.+)$", r"<h1>\1</h1>", html, flags=re.MULTILINE)
    # Convert bold/italic
    html = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", html)
    html = re.sub(r"\*(.+?)\*", r"<em>\1</em>", html)
    # Convert paragraphs
    html = re.sub(r"\n\n", r"</p><p>", html)

    return f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>Generated Report</title>
    <style>
        body {{ font-family: system-ui, sans-serif; max-width: 800px; margin: 2rem auto; padding: 0 1rem; }}
        h1, h2, h3 {{ color: #333; }}
        p {{ line-height: 1.6; }}
    </style>
</head>
<body>
<p>{html}</p>
</body>
</html>
"""

@flyte.trace
async def generate_summary(content: str) -> str:
    """Generate an executive summary of the report."""
    return await call_llm(content, SUMMARY_SYSTEM_PROMPT)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### When to trace and when not to

Use `@flyte.trace` for operations that are expensive, non-deterministic, or
call external APIs (like `generate_summary`). Skip it for cheap, deterministic
transformations (like `format_as_markdown` and `format_as_html`) where
re-running on retry is trivial.

## Parallel execution with asyncio.gather

The `format_outputs` task runs all formatters concurrently:

```python
@llm_env.task
async def format_outputs(content: str) -> Dir:
    """
    Generate multiple output formats in parallel.

    Uses asyncio.gather to run all formatting operations concurrently,
    maximizing efficiency when each operation is I/O-bound.

    Args:
        content: The final report content

    Returns:
        Directory containing all formatted outputs
    """
    print("Generating output formats in parallel...")

    with flyte.group("formatting"):
        # Run all formatting operations in parallel
        markdown, html, summary = await asyncio.gather(
            format_as_markdown(content),
            format_as_html(content),
            generate_summary(content),
        )

    # Write outputs to a directory
    output_dir = tempfile.mkdtemp()

    with open(os.path.join(output_dir, "report.md"), "w") as f:
        f.write(markdown)

    with open(os.path.join(output_dir, "report.html"), "w") as f:
        f.write(html)

    with open(os.path.join(output_dir, "summary.txt"), "w") as f:
        f.write(summary)

    print(f"Created outputs in {output_dir}")
    return await Dir.from_local(output_dir)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### How asyncio.gather works

`asyncio.gather` takes multiple coroutines and runs them concurrently:

```python
markdown, html, summary = await asyncio.gather(
    format_as_markdown(content),  # Starts immediately
    format_as_html(content),      # Starts immediately
    generate_summary(content),    # Starts immediately
)
# All three run concurrently, results returned in order
```

Without `gather`, these would run sequentially:

```python
# Sequential (slower)
markdown = await format_as_markdown(content)  # Wait for completion
html = await format_as_html(content)          # Then start this
summary = await generate_summary(content)     # Then start this
```

### When to use asyncio.gather

Use `asyncio.gather` when:

- Operations are independent (don't depend on each other's results)
- Operations are I/O-bound (API calls, file operations)
- You want to minimize total execution time

Don't use `asyncio.gather` when:

- Operations depend on each other
- Operations are CPU-bound (use process pools instead)
- Order of execution matters for side effects

## Grouping parallel operations

The parallel formatting is wrapped in a group for UI clarity:

```python
with flyte.group("formatting"):
    markdown, html, summary = await asyncio.gather(...)
```

In the Flyte UI, the traced call within the group is visible:

```
format_outputs
└── formatting
    ├── format_as_markdown
    ├── format_as_html
    └── generate_summary (traced)
```

## Collecting outputs in a directory

The formatted outputs are written to a temporary directory and returned as a
`Dir` artifact:

```python
output_dir = tempfile.mkdtemp()

with open(os.path.join(output_dir, "report.md"), "w") as f:
    f.write(markdown)

with open(os.path.join(output_dir, "report.html"), "w") as f:
    f.write(html)

with open(os.path.join(output_dir, "summary.txt"), "w") as f:
    f.write(summary)

return await Dir.from_local(output_dir)
```

The `Dir.from_local()` call uploads the directory to Flyte's
artifact storage, making it available to downstream tasks or applications.

## The batch pipeline

The batch pipeline processes multiple topics in parallel, demonstrating where
`ReusePolicy` truly shines:

```python
@driver_env.task
async def report_batch_pipeline(
    topics: list[str],
    max_iterations: int = 3,
    quality_threshold: int = 8,
) -> list[Dir]:
    """
    Generate reports for multiple topics in parallel.

    This is where ReusePolicy shines: with N topics, each going through
    up to max_iterations refinement cycles, the reusable container pool
    handles potentially N × 7 LLM calls efficiently without cold starts.

    Args:
        topics: List of topics to write about
        max_iterations: Maximum refinement cycles per topic
        quality_threshold: Minimum quality score to accept

    Returns:
        List of directories, each containing a report's formatted outputs
    """
    print(f"Starting batch pipeline for {len(topics)} topics...")

    # Fan out: refine all reports in parallel
    # Each refine_report makes 2-7 LLM calls, all hitting the reusable pool
    with flyte.group("refine_all"):
        reports = await asyncio.gather(*[
            refine_report(topic, max_iterations, quality_threshold)
            for topic in topics
        ])

    print(f"All {len(reports)} reports refined, formatting outputs...")

    # Fan out: format all reports in parallel
    with flyte.group("format_all"):
        outputs = await asyncio.gather(*[
            format_outputs(report)
            for report in reports
        ])

    print(f"Batch pipeline complete! Generated {len(outputs)} reports.")
    return outputs
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

### Pipeline flow

1. **Fan out refine_all**: Process all topics in parallel using `asyncio.gather`
2. **Fan out format_all**: Format all reports in parallel
3. **Return list of Dirs**: Each directory contains one report's outputs

With 5 topics, each making ~7 LLM calls, the reusable container pool handles
~35 LLM calls efficiently without cold starts.

## Running the pipeline

To run the batch pipeline:

```python
if __name__ == "__main__":
    flyte.init_from_config()

    # Multiple topics to generate reports for
    topics = [
        "The Impact of Large Language Models on Software Development",
        "Edge Computing: Bringing AI to IoT Devices",
        "Quantum Computing: Current State and Near-Term Applications",
        "The Rise of Rust in Systems Programming",
        "WebAssembly: The Future of Browser-Based Applications",
    ]

    print(f"Submitting batch run for {len(topics)} topics...")
    import sys
    sys.stdout.flush()

    # Run the batch pipeline - this will generate all reports in parallel,
    # with the reusable container pool handling 5 topics × ~7 LLM calls each
    run = flyte.run(
        report_batch_pipeline,
        topics=topics,
        max_iterations=3,
        quality_threshold=8,
    )
    print(f"Batch report generation run URL: {run.url}")
    sys.stdout.flush()
    print("Waiting for pipeline to complete (Ctrl+C to skip)...")
    try:
        run.wait()
        print(f"Pipeline complete! Outputs: {run.outputs()}")
    except KeyboardInterrupt:
        print(f"\nSkipped waiting. Check status at: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/generate.py*

```bash
uv run generate.py
```

The pipeline will:

1. Process all topics in parallel (each with iterative refinement)
2. Format all reports in parallel
3. Return a list of directories, each containing a report's outputs

## Cost optimization tips

### 1. Choose the right model

The example uses `gpt-4o-mini` for cost efficiency. For higher quality (at higher
cost), you could use `gpt-4o` or `gpt-4-turbo`:

```python
response = await client.chat.completions.create(
    model="gpt-4o",  # More capable, more expensive
    ...
)
```

### 2. Tune iteration parameters

Fewer iterations mean fewer API calls:

```python
run = flyte.run(
    report_batch_pipeline,
    topics=["Topic A", "Topic B"],
    max_iterations=2,      # Limit iterations
    quality_threshold=7,   # Accept slightly lower quality
)
```

### 3. Use caching effectively

The `cache="auto"` setting on the environment caches task outputs. Running the
same pipeline with the same inputs returns cached results instantly:

```python
llm_env = flyte.TaskEnvironment(
    ...
    cache="auto",  # Cache task outputs
)
```

### 4. Scale the batch

The batch pipeline already processes topics in parallel. To handle larger batches,
adjust the `ReusePolicy`:

```python
reusable=flyte.ReusePolicy(
    replicas=4,           # More containers for larger batches
    concurrency=4,        # Tasks per container
    ...
)
```

With 4 replicas × 4 concurrency = 16 slots, you can process 16 topics' refinement
tasks concurrently.

## Next steps

Learn how to [deploy a serving app](./serving-app) that connects to the pipeline
outputs and provides an interactive UI for report generation.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/serving-app ===

# Serving app

The final piece is a serving application that displays generated reports and
provides an interactive interface. This demonstrates how to connect apps to
pipeline outputs using `RunOutput`.

## App environment configuration

The `AppEnvironment` defines how the Streamlit application runs and connects to
the batch report pipeline:

```python
# Define the app environment
env = AppEnvironment(
    name="report-generator-app",
    description="Interactive report generator with AI-powered refinement",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "streamlit>=1.41.0",
    ),
    args=["streamlit", "run", "app.py", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    parameters=[
        # Connect to the batch pipeline output (list of report directories)
        Parameter(
            name="reports",
            value=RunOutput(
                task_name="driver.report_batch_pipeline",
                type="directory",
            ),
            download=True,
            env_var="REPORTS_PATH",
        ),
    ],
    include=["app.py"],
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/serve.py*

### Key configuration

| Setting | Purpose |
|---------|---------|
| `args` | Command to run the Streamlit app |
| `port` | Port the app listens on |
| `parameters` | Inputs to the app, including pipeline connections |
| `include` | Additional files to bundle with the app |

### Connecting to pipeline output with RunOutput

The `RunOutput` parameter connects the app to the batch pipeline's output:

```python
Parameter(
    name="reports",
    value=RunOutput(
        task_name="driver.report_batch_pipeline",
        type="directory",
    ),
    download=True,
    env_var="REPORTS_PATH",
)
```

This configuration:

1. **Finds the latest run** of `report_batch_pipeline` in the `driver` environment
2. **Downloads the output** to local storage (`download=True`)
3. **Sets an environment variable** with the path (`REPORTS_PATH`)

The app can then scan this directory for all generated reports.

## The Streamlit application

The app loads and displays all generated reports from the batch pipeline:

```python
def load_report_from_dir(report_dir: str) -> dict | None:
    """Load a single report from a directory."""
    if not os.path.isdir(report_dir):
        return None

    report = {"path": report_dir, "name": os.path.basename(report_dir)}

    md_path = os.path.join(report_dir, "report.md")
    if os.path.exists(md_path):
        with open(md_path) as f:
            report["markdown"] = f.read()

    html_path = os.path.join(report_dir, "report.html")
    if os.path.exists(html_path):
        with open(html_path) as f:
            report["html"] = f.read()

    summary_path = os.path.join(report_dir, "summary.txt")
    if os.path.exists(summary_path):
        with open(summary_path) as f:
            report["summary"] = f.read()

    # Only return if we found at least markdown content
    return report if "markdown" in report else None

def load_all_reports() -> list[dict]:
    """Load all reports from the batch pipeline output."""
    reports_path = os.environ.get("REPORTS_PATH")
    if not reports_path or not os.path.exists(reports_path):
        return []

    reports = []

    # Check if this is a single report directory (has report.md directly)
    if os.path.exists(os.path.join(reports_path, "report.md")):
        report = load_report_from_dir(reports_path)
        if report:
            report["name"] = "Report"
            reports.append(report)
    else:
        # Batch output: scan subdirectories for reports
        for entry in sorted(os.listdir(reports_path)):
            entry_path = os.path.join(reports_path, entry)
            report = load_report_from_dir(entry_path)
            if report:
                reports.append(report)

    return reports
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/app.py*

### Displaying multiple reports

The app provides a sidebar for selecting between reports when multiple are available:

```python
reports = load_all_reports()

if reports:
    # Sidebar for report selection if multiple reports
    if len(reports) > 1:
        st.sidebar.header("Select Report")
        report_names = [f"Report {i+1}: {r['name']}" for i, r in enumerate(reports)]
        selected_idx = st.sidebar.selectbox(
            "Choose a report to view:",
            range(len(reports)),
            format_func=lambda i: report_names[i],
        )
        selected_report = reports[selected_idx]
        st.sidebar.markdown(f"**Viewing {len(reports)} reports**")
    else:
        selected_report = reports[0]

    st.header(f"Generated Report: {selected_report['name']}")

    # Summary section
    if "summary" in selected_report:
        with st.expander("Executive Summary", expanded=True):
            st.write(selected_report["summary"])

    # Tabbed view for different formats
    tab_md, tab_html = st.tabs(["Markdown", "HTML Preview"])

    with tab_md:
        st.markdown(selected_report.get("markdown", ""))

    with tab_html:
        if "html" in selected_report:
            st.components.v1.html(selected_report["html"], height=600, scrolling=True)

    # Download options
    st.subheader("Download")
    col1, col2, col3 = st.columns(3)

    with col1:
        if "markdown" in selected_report:
            st.download_button(
                label="Download Markdown",
                data=selected_report["markdown"],
                file_name="report.md",
                mime="text/markdown",
            )

    with col2:
        if "html" in selected_report:
            st.download_button(
                label="Download HTML",
                data=selected_report["html"],
                file_name="report.html",
                mime="text/html",
            )

    with col3:
        if "summary" in selected_report:
            st.download_button(
                label="Download Summary",
                data=selected_report["summary"],
                file_name="summary.txt",
                mime="text/plain",
            )

else:
    st.info("No reports generated yet. Run the batch pipeline to create reports.")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/advanced-project/app.py*

Features:

- **Report selector**: Sidebar navigation when multiple reports exist
- **Executive summary**: Expandable section with key takeaways
- **Tabbed views**: Switch between Markdown and HTML preview
- **Download buttons**: Export in any format

### Generation instructions

The app includes instructions for generating new reports:

```python
st.divider()
st.header("Generate New Reports")
st.write("""
To generate reports, run the batch pipeline:

```bash
uv run generate.py
```

This generates reports for multiple topics in parallel, demonstrating
how ReusePolicy efficiently handles many concurrent LLM calls.
""")

# Show pipeline parameters info
with st.expander("Pipeline Parameters"):
    st.markdown("""
    **Available parameters:**

    | Parameter | Default | Description |
    |-----------|---------|-------------|
    | `topics` | (required) | List of topics to write about |
    | `max_iterations` | 3 | Maximum refinement cycles per topic |
    | `quality_threshold` | 8 | Minimum score (1-10) to accept |

    **Example:**
    CODE5
    """)
CODE6python
if __name__ == "__main__":
    flyte.init_from_config()

    # Deploy the report generator app
    print("Deploying report generator app...")
    deployment = flyte.serve(env)
    print(f"App deployed at: {deployment.url}")
CODE7bash
uv run serve.py
CODE8bash
   uv run generate.py
   CODE9bash
   uv run serve.py
   ```

3. **Access the app** at the provided URL and browse all generated reports

The app automatically picks up the latest pipeline run, so you can generate
new batches and always see the most recent results.

## Automatic updates with RunOutput

The `RunOutput` connection is evaluated at app startup. Each time the app
restarts or redeploys, it fetches the latest batch pipeline output.

For real-time updates without redeployment, you could:

1. Poll for new runs using the Flyte API
2. Implement a webhook that triggers app refresh
3. Use a database to track run status

## Complete example structure

Here's the full project structure:

CODE10

## Running the complete example

1. **Set up the secret**:

   CODE11

2. **Run the pipeline**:

   CODE12

3. **Deploy the app**:

   CODE13

4. **Open the app URL** and view your generated report

## Summary

This example demonstrated:

| Feature | What it does |
|---------|--------------|
| `ReusePolicy` | Keeps containers warm for high-throughput batch processing |
| `@flyte.trace` | Checkpoints LLM calls for recovery and observability |
| `RetryStrategy` | Handles transient API failures gracefully |
| `flyte.group` | Organizes parallel batches and iterations in the UI |
| `asyncio.gather` | Fans out to process multiple topics concurrently |
| Pydantic models | Structured LLM outputs |
| `AppEnvironment` | Deploys interactive Streamlit apps |
| `RunOutput` | Connects apps to pipeline outputs |

These patterns form the foundation for building production-grade AI workflows
that are resilient, observable, and cost-efficient at scale.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration ===

# Migration

Guides for migrating to Flyte 2 from other systems.

### **Migration > From Flyte 1 to 2**

What's new in Flyte 2 (pure Python execution, simplified API, fine-grained reproducibility) and how to port a Flyte 1 codebase.

### **Migration > From Airflow to Flyte**

Mapping from Airflow concepts (DAGs, operators, schedules, XCom, trigger rules) to their Flyte 2 equivalents.

## Subpages
- **Migration > From Flyte 1 to 2**
- **Migration > From Airflow to Flyte**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2 ===

# From Flyte 1 to 2

Flyte 2 represents a fundamental shift in how Flyte workflows are written and executed.

## Pure Python execution

Write workflows in pure Python, enabling a more natural development experience and removing the constraints of a
domain-specific language (DSL).

### Sync Python

```
import flyte

env = flyte.TaskEnvironment("sync_example_env")

@env.task
def hello_world(name: str) -> str:
    return f"Hello, {name}!"

@env.task
def main(name: str) -> str:
    for i in range(10):
        hello_world(name)
    return "Done"

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, name="World")
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/sync_example.py*

### Async Python

```
import asyncio
import flyte

env = flyte.TaskEnvironment("async_example_env")

@env.task
async def hello_world(name: str) -> str:
    return f"Hello, {name}!"

@env.task
async def main(name: str) -> str:
    results = []
    for i in range(10):
        results.append(hello_world(name))
    await asyncio.gather(*results)
    return "Done"

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, name="World")
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async_example.py*

As you can see in the hello world example, workflows can be constructed at runtime, allowing for more flexible and
adaptive behavior. Flyte 2 supports:

- Python's asynchronous programming model to express parallelism.
- Python's native error handling with `try-except` to overridden configurations, like resource requests.
- Python's native conditional `if`/`elif`/`else` syntax and `for`/`while` loops.

## Simplified API

The new API is more intuitive, with fewer abstractions to learn and a focus on simplicity.

| Use case                      | Flyte 1                     | Flyte 2                                 |
| ----------------------------- | --------------------------- | --------------------------------------- |
| Environment management        | `N/A`                       | `TaskEnvironment`                       |
| Perform basic computation     | `@task`                     | `@env.task`                             |
| Combine tasks into a workflow | `@workflow`                 | `@env.task`                             |
| Create dynamic workflows      | `@dynamic`                  | `@env.task`                             |
| Fanout parallelism            | `flytekit.map`              | Python `for` loop with `asyncio.gather` |
| Conditional execution         | `flytekit.conditional`      | Python `if-elif-else`                   |
| Catching workflow failures    | `@workflow(on_failure=...)` | Python `try-except`                     |

There is no `@workflow` decorator. Instead, "workflows" are authored through a pattern of tasks calling tasks.
Tasks are defined within environments, which encapsulate the context and resources needed for execution.

## Fine-grained reproducibility and recoverability

As in Flyte 1, Flyte 2 supports caching at the task level (via `@env.task(cache=...)`), but it further enables recovery at the finer-grained, sub-task level through a feature called tracing (via `@flyte.trace`).

```
import flyte

env = flyte.TaskEnvironment(name="trace_example_env")

@flyte.trace
async def call_llm(prompt: str) -> str:
    return "Initial response from LLM"

@env.task
async def finalize_output(output: str) -> str:
    return "Finalized output"

@env.task(cache=flyte.Cache(behavior="auto"))
async def main(prompt: str) -> str:
    output = await call_llm(prompt)
    output = await finalize_output(output)
    return output

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, prompt="Prompt to LLM")
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/trace.py*

Here `call_llm` runs in the same container as `main` and acts as an automated checkpoint with full observability in the UI.
If the task fails due to a system error (e.g., node preemption or infrastructure failure), Flyte can recover and replay from the
last successful trace rather than restarting from the beginning.

Note that tracing is distinct from caching: traces are recovered only if there is a system failure
whereas with cached outputs are persisted for reuse across separate runs.

## Improved remote functionality

Flyte 2 provides full management of the workflow lifecycle through a standardized API through the CLI and the Python SDK.

| Use case      | CLI                | Python SDK          |
| ------------- | ------------------ | ------------------- |
| Run a task    | `flyte run ...`    | `flyte.run(...)`    |
| Deploy a task | `flyte deploy ...` | `flyte.deploy(...)` |

You can also fetch and run remote (previously deployed) tasks within the course of a running workflow.

```
import flyte
from flyte import remote

env_1 = flyte.TaskEnvironment(name="env_1")
env_2 = flyte.TaskEnvironment(name="env_2")
env_1.add_dependency(env_2)

@env_2.task
async def remote_task(x: str) -> str:
    return "Remote task processed: " + x

@env_1.task
async def main() -> str:
    remote_task_ref = remote.Task.get("env_2.remote_task", auto_version="latest")
    r = await remote_task_ref(x="Hello")
    return "main called remote and recieved: " + r

if __name__ == "__main__":
    flyte.init_from_config()
    d = flyte.deploy(env_1)
    print(d[0].summary_repr())
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/remote.py*

## Native notebook support

Author and run workflows and fetch workflow metadata (I/O and logs) directly from Jupyter notebooks.

![Native Notebook](../../../_static/images/user-guide/notebook.png)

## Subpages
- **Migration > From Flyte 1 to 2 > Migration overview**
- **Migration > From Flyte 1 to 2 > Tasks and workflows**
- **Migration > From Flyte 1 to 2 > Task configuration**
- **Migration > From Flyte 1 to 2 > CLI and configuration**
- **Migration > From Flyte 1 to 2 > Control flow**
- **Migration > From Flyte 1 to 2 > Parallelism and fan-out**
- **Migration > From Flyte 1 to 2 > Data types and I/O**
- **Migration > From Flyte 1 to 2 > ML workloads**
- **Migration > From Flyte 1 to 2 > New in Flyte 2: patterns that weren't possible in Flyte 1**
- **Migration > From Flyte 1 to 2 > Hybrid v1 and v2 pipelines**
- **Migration > From Flyte 1 to 2 > Gotchas and caveats**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/overview ===

# Migration overview

This section walks through the Flyte 1 workload patterns you already know — data ETL, model training, hyperparameter sweeps, batch inference — and their Flyte 2 equivalents. Every pattern is a complete, runnable v1↔v2 example pair in the [`unionai-examples`](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/migration/flyte-2) repository.

Two conceptual shifts motivate almost every change — **pure Python execution** and the **asynchronous model** — after which most migrations come down to a couple of mechanical moves. This page covers both shifts, the terminology mapping, and a quick-reference module.

## Pure Python execution

In Flyte 1, `@workflow` functions were constrained to a DSL subset of Python that compiled to a static DAG. In Flyte 2 there is **no `@workflow` decorator**: everything is a `@env.task`, and a "workflow" is simply a task that calls other tasks. Orchestration runs as real Python at runtime, so loops, conditionals, and `try`/`except` work anywhere.

| Flyte 1 | Flyte 2 |
| --- | --- |
| `@workflow` functions are constrained to a subset of Python defining a static DAG. | **No `@workflow` decorator**: your top-level "workflow" is just a task that calls other tasks. |
| `@task` functions had the full power of Python, but only within a single container execution. | `@env.task`s call other tasks and build dynamic structures with any Python construct, anywhere. |
| Workflows compiled to static DAGs at registration time. | Workflows are tasks calling tasks; compile-time safety is coming via `compiled_task`. |

### Flyte 1

```python
import flytekit

image = flytekit.ImageSpec(
    name="hello-world-image",
    packages=["requests"],
)

@flytekit.task(container_image=image)
def mean(data: list[float]) -> float:
    return sum(list) / len(list)

@flytekit.workflow
def main(data: list[float]) -> float:
    output = mean(data)

    # ❌ performing trivial operations in a workflow is not allowed
    # output = output / 100

    # ❌ if/else is not allowed
    # if output < 0:
    #     raise ValueError("Output cannot be negative")

    return output
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/pure-python/flyte_1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(
    "hello_world",
    image=flyte.Image.from_debian_base().with_pip_packages("requests"),
)

@env.task
def mean(data: list[float]) -> float:
    return sum(data) / len(data)

@env.task
def main(data: list[float]) -> float:
    output = mean(data)

    # ✅ performing trivial operations in a workflow is allowed
    output = output / 100

    # ✅ if/else is allowed
    if output < 0:
        raise ValueError("Output cannot be negative")

    return output
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/pure-python/flyte_2.py*

This unlocks workflows that adapt to runtime conditions, native `try`/`except` error handling, and the intuitive composition of ordinary Python functions.

## Asynchronous model

Flyte 2 is built on Python's `asyncio`, with a crucial twist: **the Flyte orchestrator acts as the event loop**, scheduling awaited tasks across distributed infrastructure. This makes `async`/`await` the natural way to express parallelism.

| | Flyte 1 | Flyte 2 |
| --- | --- | --- |
| Parallelism | The DSL auto-parallelized independent tasks; the `map` operator ran a task over many inputs. | Python's `asyncio` expresses parallelism, with the Flyte orchestrator acting as the event loop across distributed infrastructure. |

The core async keywords carry Flyte-specific meaning:

- **`async def`** declares a coroutine.
- **`await`** signals where a task can be scheduled in parallel — not just an I/O yield point.
- **`asyncio.gather`** tells the orchestrator that a set of tasks are independent and can be distributed across separate compute resources.

Consider this pattern for parallel data processing:

```
import asyncio
import flyte

env = flyte.TaskEnvironment("data_pipeline")

@env.task
async def process_chunk(chunk_id: int, data: str) -> str:
    # This could be any computational work - CPU or I/O bound
    await asyncio.sleep(1)  # Simulating work
    return f"Processed chunk {chunk_id}: {data}"

@env.task
async def parallel_pipeline(data_chunks: list[str]) -> list[str]:
    # Create coroutines for all chunks
    tasks = []
    for i, chunk in enumerate(data_chunks):
        tasks.append(process_chunk(i, chunk))

    # Execute all chunks in parallel
    results = await asyncio.gather(*tasks)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

In standard Python this mainly benefits I/O-bound work; in Flyte 2 the orchestrator schedules each `process_chunk` task on its own pod.

### True parallelism for all workloads

Async syntax in Flyte 2 is **not just for I/O-bound operations**. When the orchestrator encounters `await asyncio.gather(...)`, it runs those independent tasks simultaneously across compute resources — achieving true parallelism for CPU-bound work (model training, heavy math), I/O-bound work (queries, API calls), and mixed workloads alike.

### Calling sync tasks from async tasks

You don't need to rewrite existing synchronous code. Flyte automatically "asyncifies" sync functions; just call them from an async context with `.aio()`:

```
@env.task
def legacy_computation(x: int) -> int:
    # Existing synchronous function works unchanged
    return x * x + 2 * x + 1

@env.task
async def modern_workflow(numbers: list[int]) -> list[int]:
    # Call sync tasks from async context using .aio()
    tasks = []
    for num in numbers:
        tasks.append(legacy_computation.aio(num))

    results = await asyncio.gather(*tasks)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

### The `flyte.map` function: Familiar patterns

For code that used Flyte 1's `map`, `flyte.map` is a direct replacement that works in both sync and async contexts:

### Sync Map

```
@env.task
def sync_map_example(n: int) -> list[str]:
    # Synchronous version for easier migration
    results = []
    for result in flyte.map(process_item, range(n)):
        if isinstance(result, Exception):
            raise result
        results.append(result)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

### Async Map

```
@env.task
async def async_map_example(n: int) -> list[str]:
    # Async version using flyte.map - exact pattern from SDK examples
    results = []
    async for result in flyte.map.aio(process_item, range(n), return_exceptions=True):
        if isinstance(result, Exception):
            raise result
        results.append(result)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

It provides dual interfaces (`flyte.map.aio()` and `flyte.map()`), `return_exceptions` for graceful failure handling (a more flexible replacement for Flyte 1's `min_success_ratio`), automatic UI grouping, and optional concurrency limits.

## Terminology and concept mapping

Several Flyte 1 concepts were renamed or reshaped in Flyte 2. The table below maps the ones you'll meet most often.

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `flytekit` (package) | `flyte` (package) | The Python SDK was renamed; imports change from `import flytekit` to `import flyte`. |
| `pyflyte` (CLI) | `flyte` (CLI) | The command-line tool was renamed. |
| `@task` / `@workflow` / `@dynamic` | `@env.task` | A single task decorator off a `flyte.TaskEnvironment`. Workflows and dynamic tasks are no longer distinct constructs: everything is a task, and orchestration is plain Python. |
| `map_task()` | `flyte.map()` | Plus `asyncio.gather()` for async fan-out. |
| `conditional()` | native `if` / `elif` / `else` | Branching is now ordinary Python control flow, not a DSL. |
| `ImageSpec` | `flyte.Image` | Container image definition. |
| `current_context()` | `flyte.ctx()` | Runtime context access. |
| `FlyteFile` / `FlyteDirectory` | `flyte.io.File` / `flyte.io.Dir` | Offloaded file and directory references. |
| `StructuredDataset` | `flyte.io.DataFrame` | Offloaded tabular data. |
| `LaunchPlan` | `flyte.Trigger` | Scheduling and parameterized entry points. |
| `CronSchedule` | `flyte.Cron` | Cron-based scheduling, used with a `flyte.Trigger`. |
| Decks (`enable_deck=True`) | Reports (`report=True`) | Custom HTML rendered in the UI during/after a run. See [Reports](../../tasks/task-programming/reports). |

## Package imports

The package is renamed from `flytekit` to `flyte`, and the workflow/dynamic/map_task imports disappear:

### Flyte 1

```python
import flytekit
from flytekit import task, workflow, dynamic, map_task
from flytekit import ImageSpec, Resources, Secret
from flytekit import current_context, LaunchPlan, CronSchedule
```

### Flyte 2

```python
import flyte
from flyte import TaskEnvironment, Resources, Secret
from flyte import Image, Trigger, Cron
```

## The two mechanical changes behind (almost) every migration

Most of a migration comes down to two moves:

### 1. Move task configuration into a `TaskEnvironment`

Instead of configuring the image, resources, and caching on each task decorator, configure them once on a `flyte.TaskEnvironment` and share it across tasks:

```python
env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    cache="auto",
)
```

### 2. Replace `@task` / `@workflow` / `@dynamic` with `@env.task`

Every decorated function becomes an `@env.task`. There is no separate workflow or dynamic construct: a "workflow" is simply a task that calls other tasks, and orchestration is plain Python.

> **📝 Note**
>
> The `env` in `@env.task` is just the variable you assigned your `TaskEnvironment` to. Name it whatever you like.

## In this section

The migration patterns are grouped by theme. Start with **Tasks and workflows**, then jump to whatever your workload needs:

- **[Tasks and workflows](./tasks-and-workflows)** — the structural shift: `@task`/`@workflow` → `@env.task`, sequential ordering, nested "subworkflows", and the `@task` → `TaskEnvironment` parameter mapping.
- **[Task configuration](./configuration)** — moving image/resources/cache to the `TaskEnvironment`, GPUs, secrets, caching, and scheduling with triggers.
- **[CLI and configuration](./cli-and-configuration)** — `pyflyte` → `flyte` command mapping and config-file changes.
- **[Control flow](./control-flow)** — `conditional()` and `@dynamic` become plain Python `if`/loops, and `on_failure` becomes `try`/`except`.
- **[Parallelism and fan-out](./parallelism)** — `map_task` → `flyte.map` / `asyncio.gather`, plus a data-backfill example.
- **[Data types and I/O](./data-io)** — `FlyteFile`/`FlyteDirectory` → `flyte.io.File`/`Dir`, `StructuredDataset` → `flyte.io.DataFrame`, dataclasses, and an ETL example.
- **[ML workloads](./ml-workloads)** — small-model training, hyperparameter optimization, deep learning, batch inference, and an end-to-end pipeline.
- **[New in Flyte 2](./new-in-flyte-2)** — patterns that weren't possible in Flyte 1 at all: real-time model serving, batch inference, apps, and sandboxed code execution.

- **[Hybrid v1 and v2 pipelines](./hybrid-pipelines)** — calling between v1 and v2 in both directions during the transition.
- **[Gotchas and caveats](./gotchas-and-caveats)** — common gotchas plus the deeper caveats of the new execution model, including non-deterministic behavior and keeping orchestration lightweight.

## Quick reference

A minimal Flyte 2 module, end to end:

```python
import asyncio
import flyte

# 1. Define an image
image = (
    flyte.Image.from_debian_base(python_version=(3, 11))
    .with_pip_packages("pandas", "numpy")
)

# 2. Create a TaskEnvironment
env = flyte.TaskEnvironment(
    name="my_env",
    image=image,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

# 3. Define tasks
@env.task
async def process(x: int) -> int:
    return x * 2

# 4. Define the entrypoint task
@env.task
async def main(items: list[int]) -> list[int]:
    results = await asyncio.gather(*[process(x) for x in items])
    return list(results)

# 5. Run it
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, items=[1, 2, 3, 4, 5])
    print(run.url)
    run.wait()
```

```bash
# CLI
flyte run my_module.py main --items '[1,2,3,4,5]'   # remote (default)
flyte run --local my_module.py main --items '[1,2,3,4,5]'
flyte deploy my_module.py my_env
```


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/tasks-and-workflows ===

# Tasks and workflows

The biggest structural change in Flyte 2 is that **everything is a task**. The `@task`, `@workflow`, and `@dynamic` decorators all collapse into a single `@env.task` on a `flyte.TaskEnvironment`, and a "workflow" is just a task that calls other tasks. This page covers the basic structure; see [Task configuration](./configuration) for the environment settings and [Migration](./overview) for the big picture.

## Hello world: tasks and workflows

A `@task` plus `@workflow` becomes two `@env.task`s, where the entrypoint task calls the others. Sequential calls are naturally ordered — no `>>` operator required.

### Flyte 1

```python
import flytekit

@flytekit.task
def say_hello(name: str) -> str:
    return f"Hello, {name}!"

@flytekit.task
def to_upper(greeting: str) -> str:
    return greeting.upper()

@flytekit.workflow
def main(name: str) -> str:
    greeting = say_hello(name=name)
    return to_upper(greeting=greeting)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hello_world_v1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(name="hello_world")

@env.task
def say_hello(name: str) -> str:
    return f"Hello, {name}!"

@env.task
def to_upper(greeting: str) -> str:
    return greeting.upper()

# The "workflow" is now just a task that calls other tasks.
@env.task
def main(name: str) -> str:
    greeting = say_hello(name)
    return to_upper(greeting)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hello_world_v2.py*

## Chaining and ordering

In Flyte 1 you sometimes used `>>` to force ordering between tasks with no data dependency. In Flyte 2, sequential (synchronous) calls run in the order they're written, and `await`ing async tasks in sequence does the same. The `>>` operator is gone.

### Flyte 1

```python
from flytekit import task, workflow

@task
def clear_staging_table() -> None:
    # Side effect only: truncate the staging table.
    print("cleared staging table")

@task
def load_into_staging() -> None:
    # Side effect only: load fresh rows into staging.
    print("loaded staging table")

@task
def publish_to_prod() -> None:
    # Side effect only: swap staging into the production table.
    print("published to prod")

@workflow
def main() -> None:
    clear = clear_staging_table()
    load = load_into_staging()
    publish = publish_to_prod()

    # These tasks pass no data between them, so use the >> operator to force
    # ordering: clear must finish before load, which must finish before publish.
    clear >> load >> publish
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/chained_tasks_v1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(name="staging_publish")

@env.task
def clear_staging_table() -> None:
    print("cleared staging table")

@env.task
def load_into_staging() -> None:
    print("loaded staging table")

@env.task
def publish_to_prod() -> None:
    print("published to prod")

# Sequential (synchronous) calls run in the order they're written, even when no
# data flows between them. The Flyte 1 `>>` ordering operator is gone.
@env.task
def main() -> None:
    clear_staging_table()
    load_into_staging()
    publish_to_prod()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/chained_tasks_v2.py*

## Subworkflows

A `@workflow` invoked by another `@workflow` (for example, a reusable preprocessing pipeline) becomes a task that calls other tasks — nest them as deeply as you like.

### Flyte 1

CODE3

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/subworkflow_v1.py*

### Flyte 2

CODE4

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/subworkflow_v2.py*

## TaskEnvironment configuration

The `TaskEnvironment` holds the configuration that Flyte 1 spread across the `@task` decorator. The task decorator can still override a few settings per-task.

CODE5

## Parameter mapping: `@task` → `TaskEnvironment` + `@env.task`

| Flyte 1 `@task` parameter | Flyte 2 location | Notes |
|---|---|---|
| `container_image` | `TaskEnvironment(image=...)` | Env-level only |
| `requests` | `TaskEnvironment(resources=...)` | Env-level only |
| `limits` | `TaskEnvironment(resources=...)` | Combined with requests (single value) |
| `environment` | `TaskEnvironment(env_vars=...)` | Env-level only |
| `secret_requests` | `TaskEnvironment(secrets=...)` | Env-level only |
| `cache` | Both | Can override at task level |
| `cache_version` | `flyte.Cache(version_override=...)` | In a `Cache` object |
| `retries` | `@env.task(retries=...)` | Task-level only |
| `timeout` | `@env.task(timeout=...)` | Task-level only |
| `interruptible` | Both | Can override at task level |
| `pod_template` | Both | Can override at task level |
| `deprecated` | N/A | Not in Flyte 2 |
| `docs` | `@env.task(docs=...)` | Task-level only |

For image, resource, secret, and caching detail, see [Task configuration](./configuration).

## Next

- [Task configuration](./configuration) — image, resources, caching, secrets, and scheduling
- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Parallelism and fan-out](./parallelism) — running tasks in parallel


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/configuration ===

# Task configuration

In Flyte 1, image, resources, caching, secrets, and scheduling were configured per-task on the `@task` decorator or per-workflow on a `LaunchPlan`. In Flyte 2 most of this moves to the `flyte.TaskEnvironment`, so it's declared once and shared. See [Migration](./overview) for the overall approach.

## Image, resources, and caching

Image, resources, and caching move from the `@task` decorator to the `TaskEnvironment`. Per-task settings like `retries` and `timeout` stay on `@env.task`. Note that `mem` is renamed to `memory`, and there are no separate `requests`/`limits` — a single `Resources` value serves as both.

### Flyte 1

```python
from datetime import timedelta

import flytekit
from flytekit import Resources

image = flytekit.ImageSpec(
    name="training-image",
    packages=["scikit-learn", "pandas"],
)

@flytekit.task(
    container_image=image,
    requests=Resources(cpu="2", mem="4Gi"),
    limits=Resources(cpu="4", mem="8Gi"),
    cache=True,
    cache_version="1.0",
    retries=3,
    timeout=timedelta(minutes=30),
)
def train_epoch(step: int) -> float:
    # A stand-in for a training step that returns the current loss.
    return 1.0 / (step + 1)

@flytekit.workflow
def main(step: int) -> float:
    return train_epoch(step=step)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/task_config_v1.py*

### Flyte 2

```
from datetime import timedelta

import flyte

# Image, resources, and caching move to the TaskEnvironment, so they are declared
# once and shared by every task in the environment.
env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),  # "memory", not "mem"
    cache="auto",
)

# retries and timeout stay on the task decorator.
@env.task(retries=3, timeout=timedelta(minutes=30))
def train_epoch(step: int) -> float:
    # A stand-in for a training step that returns the current loss.
    return 1.0 / (step + 1)

@env.task
def main(step: int) -> float:
    return train_epoch(step)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/task_config_v2.py*

For the full `@task` → `TaskEnvironment` + `@env.task` parameter mapping, see [Tasks and workflows](./tasks-and-workflows).

## Container images

Flyte 1's `ImageSpec` is replaced by Flyte 2's `flyte.Image` with a fluent builder API.

```python
from flyte import Image

image = (
    Image.from_debian_base(name="my-image", registry="ghcr.io/myorg", python_version=(3, 11))
    .with_pip_packages("pandas", "numpy")
    .with_apt_packages("curl", "git")
    .with_env_vars({"MY_VAR": "value"})
)
```

Instead of one constructor with many arguments, you start from a base and chain builder methods:

| Constructor | Use case |
|---|---|
| `Image.from_debian_base()` | Most common; includes the Flyte SDK |
| `Image.from_base(image_uri)` | Start from any existing image |
| `Image.from_dockerfile(path)` | Complex custom builds |
| `Image.from_uv_script(path)` | UV-based projects |

Common chainable builder methods: `.with_pip_packages(...)`, `.with_requirements(path)`, `.with_uv_project(path)`, `.with_apt_packages(...)`, `.with_commands([...])`, `.with_source_file(path, dst=...)`, `.with_source_folder(path, dst=...)`, `.with_env_vars({...})`, and `.with_workdir(...)`.

| Flyte 1 `ImageSpec` | Flyte 2 `Image` | Notes |
|---|---|---|
| `name` | `name` (constructor) | Same |
| `registry` | `registry` (constructor) | Same |
| `python_version` | `python_version` (tuple) | `"3.11"` becomes `(3, 11)` |
| `packages` | `.with_pip_packages()` | Method instead of param |
| `apt_packages` | `.with_apt_packages()` | Method instead of param |
| `requirements` | `.with_requirements()` | Supports txt, poetry.lock, uv.lock |
| `env` | `.with_env_vars()` | Method instead of param |
| `commands` | `.with_commands()` | Method instead of param |
| `copy` / `source_root` | `.with_source_file()` / `.with_source_folder()` | More explicit methods |
| `base_image` | `Image.from_base()` | Different constructor |
| `builder` | Config file or `flyte.init()` | Global setting |
| `platform` | `platform` (constructor) | Tuple: `("linux/amd64", "linux/arm64")` |

For a private registry, create an image-pull secret and reference it:

```shell
flyte create secret --type image_pull my-registry-secret --from-file ~/.docker/config.json
```

```python
image = Image.from_debian_base(
    registry="private.registry.com",
    name="my-image",
    registry_secret="my-registry-secret",
)
```

See [Container images](../../tasks/task-configuration/container-images) for more.

## Resources

A single `flyte.Resources` value serves as both request and limit — there are no separate `requests`/`limits`. Several parameters were renamed.

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `cpu="1"` | `cpu="1"` | Same |
| `mem="2Gi"` | `memory="2Gi"` | Renamed |
| `gpu="1"` | `gpu="A100:1"` | `Type:count` format |
| `ephemeral_storage="10Gi"` | `disk="10Gi"` | Renamed |
| N/A | `shm="auto"` | New: shared memory |

GPU type and count are combined into one string, replacing the separate Flyte 1 `accelerator=` argument:

```python
env = flyte.TaskEnvironment(
    name="gpu_env",
    resources=flyte.Resources(
        cpu="4",
        memory="32Gi",
        gpu="A100:2",              # Type:count
        # gpu="A100 80G:1"         # 80GB variant
        # gpu=flyte.GPU("A100", count=1, partition="1g.5gb")   # MIG partition
    ),
)
```

Supported GPU types include A10, A10G, A100, A100 80G, B200, H100, H200, L4, L40s, T4, V100, RTX PRO 6000, and GB10. See [Resources](../../tasks/task-configuration/resources) for more.

## Caching

Caching is enabled at the env level with `cache="auto"` (or per-task on `@env.task`). The explicit `cache_version` string moves into a `flyte.Cache` object.

| Behavior | Description |
|---|---|
| `"auto"` | Cache results and reuse if available |
| `"override"` | Always execute and overwrite the cache |
| `"disable"` | No caching (default for a `TaskEnvironment`) |

```python
# Flyte 1: @task(cache=True, cache_version="1.0")
# Flyte 2:
@env.task(cache="auto")
def cached_task(x: int) -> int:
    return x * 2

# Advanced control (replaces cache_version, serialize, ignored_inputs, ...)
@env.task(cache=flyte.Cache(
    behavior="auto",
    version_override="v1.0",
    serialize=True,
    ignored_inputs=("debug",),
))
def advanced(x: int, debug: bool = False) -> int:
    return x * 2
```

See [Caching](../../tasks/task-configuration/caching) for more.

## Secrets

Secrets move from `secret_requests` on the task to `secrets` on the `TaskEnvironment`, and you read them from environment variables instead of `current_context().secrets` — for example, an API key for a model registry or hosted LLM.

### Flyte 1

```python
from flytekit import task, workflow, Secret, current_context

@task(secret_requests=[Secret(group="openai", key="api_key")])
def call_api() -> str:
    token = current_context().secrets.get(group="openai", key="api_key")
    return f"token has {len(token)} chars"

@workflow
def main() -> str:
    return call_api()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/secrets_v1.py*

### Flyte 2

```
import os

import flyte

# Secrets are declared on the TaskEnvironment and injected as environment
# variables (instead of read through current_context().secrets).
env = flyte.TaskEnvironment(
    name="secrets",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
)

@env.task
def call_api() -> str:
    token = os.getenv("OPENAI_API_KEY", "")
    return f"token has {len(token)} chars"

@env.task
def main() -> str:
    return call_api()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/secrets_v2.py*

A `flyte.Secret` can be mounted as an environment variable or as a file, and the access convention changes:

```python
flyte.Secret(key="openai-key", as_env_var="OPENAI_API_KEY")   # mount as env var
flyte.Secret(key="access-key", group="aws")                    # env var: AWS_ACCESS_KEY
flyte.Secret(key="ssl-cert", mount="/etc/flyte/secrets")       # mount as a file
```

| Flyte 1 pattern | Flyte 2 pattern |
|---|---|
| `ctx.secrets.get(key="mykey", group="mygroup")` | `os.environ["MYGROUP_MYKEY"]` (auto-named) |
| `ctx.secrets.get(key="mykey")` | `os.environ["MY_SECRET"]` (with `as_env_var="MY_SECRET"`) |

Create and manage secrets from the CLI:

```bash
flyte create secret MY_SECRET_KEY --value my_secret_value
flyte create secret MY_SECRET_KEY --from-file /path/to/secret
flyte get secret
flyte delete secret MY_SECRET_KEY
```

See [Secrets](../../tasks/task-configuration/secrets) for more.

## Scheduling

A `LaunchPlan` with a `CronSchedule` (say, a nightly retraining job) becomes a `flyte.Trigger` attached directly to the task. Use `flyte.TriggerTime` to bind the scheduled fire time to an input, and deploy the trigger with `flyte deploy`.

### Flyte 1

```python
from flytekit import task, workflow, LaunchPlan, CronSchedule

@task
def retrain(kickoff_time: str) -> str:
    return f"retrained model at {kickoff_time}"

@workflow
def main(kickoff_time: str) -> str:
    return retrain(kickoff_time=kickoff_time)

# A LaunchPlan attaches a schedule (and default inputs) to a workflow.
nightly_retrain = LaunchPlan.get_or_create(
    workflow=main,
    name="nightly_retrain",
    schedule=CronSchedule(
        schedule="0 2 * * *",  # 2 AM daily
        kickoff_time_input_arg="kickoff_time",
    ),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/scheduling_v1.py*

### Flyte 2

```
from datetime import datetime

import flyte

env = flyte.TaskEnvironment(name="scheduling")

# A Trigger replaces LaunchPlan + CronSchedule. It is attached directly to the
# task and deployed with it (flyte deploy). flyte.TriggerTime binds the
# scheduled fire time to a task input.
nightly_retrain = flyte.Trigger(
    name="nightly_retrain",
    automation=flyte.Cron("0 2 * * *"),  # 2 AM daily
    inputs={"trigger_time": flyte.TriggerTime},
    auto_activate=True,
)

@env.task(triggers=nightly_retrain)
def main(trigger_time: datetime = datetime(2024, 1, 1, 2, 0)) -> str:
    return f"retrained model at {trigger_time.isoformat()}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/scheduling_v2.py*

Triggers support `flyte.Cron("0 9 * * *", timezone="America/New_York")` and `flyte.FixedRate(timedelta(hours=1))` as automations, plus convenience constructors like `flyte.Trigger.hourly()` and `flyte.Trigger.daily()`. See [Triggers](../../tasks/task-configuration/triggers) for more.

## Next

- [CLI and configuration](./cli-and-configuration) — `pyflyte` → `flyte` command and config-file mapping
- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Data types and I/O](./data-io) — files, DataFrames, and dataclasses


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/cli-and-configuration ===

# CLI and configuration

The command-line tool is renamed from `pyflyte` to `flyte`, and the config file is trimmed down. See [Migration](./overview) for the overall approach.

## CLI command mapping

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `pyflyte run` | `flyte run` | Similar, different flags |
| `pyflyte run --remote` | `flyte run` | Remote is the default in Flyte 2 |
| `pyflyte run` (local) | `flyte run --local` | Local execution is now explicit |
| `pyflyte register` | `flyte deploy` | Different concept |
| `pyflyte package` | N/A | Not needed in Flyte 2 |
| `pyflyte serialize` | N/A | Not needed in Flyte 2 |

### Running tasks

### Flyte 1

```shell
# Local
pyflyte run my_module.py my_workflow --arg1 value1

# Remote
pyflyte --config config.yaml run --remote my_module.py my_workflow --arg1 value1
```

### Flyte 2

```shell
# Remote (default)
flyte run my_module.py my_task --arg1 value1

# Local
flyte run --local my_module.py my_task --arg1 value1

# With an explicit config file
flyte --config config.yaml run my_module.py my_task --arg1 value1
```

### Deploying

In Flyte 1 you registered a module; in Flyte 2 you deploy task environments.

### Flyte 1

```shell
pyflyte register my_module.py -p my-project -d development
```

### Flyte 2

```shell
# Deploy a task environment
flyte deploy my_module.py my_env --project my-project --domain development

# Deploy all environments in a file
flyte deploy --all my_module.py

# Deploy with an explicit version, or recursively
flyte deploy --version v1.0.0 my_module.py my_env
flyte deploy --recursive --all ./src
```

### Key flag differences

| Flyte 1 flag | Flyte 2 flag | Notes |
|---|---|---|
| `--remote` | (default) | Remote is the default |
| `--copy-all` | `--copy-style all` | File copying |
| N/A | `--copy-style loaded_modules` | Default: only imported modules |
| `-p, --project` | `--project` | Same |
| `-d, --domain` | `--domain` | Same |
| `-i, --image` | `--image` | Same format |
| N/A | `--follow, -f` | Follow execution logs |

## Configuration files

The config file lives in the same place (`~/.flyte/config.yaml`), but the environment variable changes from `FLYTECTL_CONFIG` to `FLYTE_CONFIG`, and the format is simpler.

### Flyte 1

```yaml
admin:
  endpoint: dns:///your-cluster.hosted.unionai.cloud
  insecure: false
  authType: Pkce
```

### Flyte 2

```yaml
admin:
  endpoint: dns:///your-cluster.hosted.unionai.cloud

image:
  builder: remote  # or "local"

task:
  domain: development
  org: your-org
  project: your-project
```

| Setting | Flyte 1 | Flyte 2 |
|---|---|---|
| Endpoint | `admin.endpoint` | `admin.endpoint` |
| Auth type | `admin.authType` | Auto-detected (PKCE default) |
| Project | CLI flag `-p` | `task.project` (default) |
| Domain | CLI flag `-d` | `task.domain` (default) |
| Organization | CLI flag `--org` | `task.org` (default) |
| Image builder | N/A | `image.builder` (`local` or `remote`) |

### Configuring in code

```python
import flyte

# From a config file (auto-discovers, or pass a path)
flyte.init_from_config()
flyte.init_from_config("path/to/config.yaml")

# Programmatically
flyte.init(
    endpoint="flyte.example.com",
    project="my-project",
    domain="development",
)
```

For API-key authentication in non-interactive environments, use `flyte.init_from_api_key()` — see [Run on a remote cluster](../../get-started/run-modes/running-remote).

## Next

- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Hybrid v1 and v2 pipelines](./hybrid-pipelines) — calling between v1 and v2 during the transition


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/control-flow ===

# Control flow

Flyte 1 expressed branching, dynamic fan-out, and failure handling through DSL constructs (`conditional()`, `@dynamic`, `@workflow(on_failure=...)`). In Flyte 2 these are all ordinary Python, because orchestration runs as real Python at runtime. See [Migration](./overview) for the overall approach.

## Conditional execution

The `conditional()` DSL becomes ordinary Python `if` / `elif` / `else` — for example, choosing a model based on dataset size.

### Flyte 1

```python
from flytekit import task, workflow, conditional

@task
def train_gradient_boosting(n_rows: int) -> str:
    return f"trained gradient boosting on {n_rows} rows"

@task
def train_logistic_regression(n_rows: int) -> str:
    return f"trained logistic regression on {n_rows} rows"

@workflow
def main(n_rows: int) -> str:
    # Pick the model based on dataset size.
    return (
        conditional("model_choice")
        .if_(n_rows > 10_000)
        .then(train_gradient_boosting(n_rows=n_rows))
        .else_()
        .then(train_logistic_regression(n_rows=n_rows))
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/conditional_v1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(name="conditional")

@env.task
def train_gradient_boosting(n_rows: int) -> str:
    return f"trained gradient boosting on {n_rows} rows"

@env.task
def train_logistic_regression(n_rows: int) -> str:
    return f"trained logistic regression on {n_rows} rows"

# Branching is now ordinary Python control flow -- no conditional() DSL.
@env.task
def main(n_rows: int) -> str:
    if n_rows > 10_000:
        return train_gradient_boosting(n_rows)
    return train_logistic_regression(n_rows)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/conditional_v2.py*

## Dynamic workflows

`@dynamic` existed so a task could generate a variable number of subtask calls at runtime (e.g. one per data partition discovered at runtime). In Flyte 2 every task can do this natively, so `@dynamic` simply disappears — loop over runtime data in an ordinary `@env.task`.

### Flyte 1

```python
from flytekit import task, workflow, dynamic

@task
def list_partitions(n: int) -> list[int]:
    return list(range(n))

@task
def process_partition(partition_id: int) -> int:
    # Aggregate one data partition.
    return partition_id * 2

@dynamic
def process_all(partitions: list[int]) -> list[int]:
    results = []
    for partition_id in partitions:
        results.append(process_partition(partition_id=partition_id))
    return results

@workflow
def main(n: int) -> list[int]:
    partitions = list_partitions(n=n)
    return process_all(partitions=partitions)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dynamic_v1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(name="dynamic")

@env.task
def process_partition(partition_id: int) -> int:
    # Aggregate one data partition.
    return partition_id * 2

# No @dynamic decorator needed: a plain task can loop over runtime data (e.g. a
# variable number of partitions discovered at runtime) and call other tasks.
@env.task
def main(n: int) -> list[int]:
    return [process_partition(partition_id) for partition_id in range(n)]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dynamic_v2.py*

## Error handling

Flyte 1's `@workflow(on_failure=...)` handler becomes ordinary Python `try` / `except` — catch a failed training run, run cleanup, and recover or re-raise.

### Flyte 1

```python
from flytekit import task, workflow

@task
def train_fold(max_depth: int) -> float:
    if max_depth <= 0:
        raise ValueError("max_depth must be positive")
    # Return validation accuracy for this hyperparameter.
    return 0.90 + 0.001 * max_depth

@task
def notify_failure() -> None:
    print("training run failed -- sending alert")

# The on_failure handler runs if any node in the workflow fails. There is no
# try/except inside a Flyte 1 workflow.
@workflow(on_failure=notify_failure)
def main(max_depth: int) -> float:
    return train_fold(max_depth=max_depth)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/error_handling_v1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(name="error_handling")

@env.task
async def train_fold(max_depth: int) -> float:
    if max_depth <= 0:
        raise ValueError("max_depth must be positive")
    return 0.90 + 0.001 * max_depth

# Failure handling is ordinary Python try/except -- no on_failure handler.
@env.task
async def main(max_depth: int) -> float:
    try:
        return await train_fold(max_depth)
    except ValueError as e:
        print(f"invalid hyperparameter ({e}); falling back to a safe default")
        # Recover with a safe default instead of failing the whole run.
        return await train_fold(max_depth=6)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/error_handling_v2.py*

Flyte 2 also exposes typed errors, so you can catch a specific failure and retry with more resources — a common need for memory-hungry training jobs:

```python
try:
    return await train_fold(sample_size)
except flyte.errors.OOMError:
    # Retry the same task with a larger memory request.
    return await train_fold.override(
        resources=flyte.Resources(memory="16Gi")
    )(sample_size)
```

See [Error handling](../../tasks/task-programming/error-handling) for more.

## Next

- [Parallelism and fan-out](./parallelism) — running many tasks in parallel
- [ML workloads](./ml-workloads) — training, HPO, and inference


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/parallelism ===

# Parallelism and fan-out

Flyte 1's `map_task` becomes `flyte.map`, and the idiomatic Flyte 2 approach to fan-out is Python `async`/`await` with `asyncio.gather`. See the [Asynchronous model](./overview#asynchronous-model) guide for the concepts, and [Migration](./overview) for the overall approach.

## Fan-out: `map_task`

`map_task()` becomes `flyte.map()`, a near drop-in replacement. The one catch: `flyte.map` returns a generator, so wrap it in `list()`. For new code, the idiomatic approach is Python `async`/`await` with `asyncio.gather()`, which gives you finer control over concurrency and error handling.

### Flyte 1

```python
from functools import partial

from flytekit import task, workflow, map_task

@task
def get_shards(n: int) -> list[int]:
    return list(range(n))

@task
def score_shard(shard_id: int, model_version: int) -> int:
    # Score one shard of records with the given model version.
    return shard_id * model_version

@workflow
def main(n: int, model_version: int) -> list[int]:
    shards = get_shards(n=n)
    return map_task(
        partial(score_shard, model_version=model_version),
        concurrency=10,
    )(shard_id=shards)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v1.py*

### Flyte 2 (flyte.map)

```
@env.task
def score_shard(shard_id: int, model_version: int) -> int:
    # Score one shard of records with the given model version.
    return shard_id * model_version

@env.task
def main(n: int, model_version: int) -> list[int]:
    bound = partial(score_shard, model_version=model_version)
    # flyte.map is a drop-in for map_task, but it returns a generator, so wrap
    # it in list() to materialize the results.
    return list(flyte.map(bound, range(n), concurrency=10))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v2.py*

### Flyte 2 (asyncio.gather)

```
@env.task
async def score_shard_async(shard_id: int, model_version: int) -> int:
    return shard_id * model_version

@env.task
async def main_async(n: int, model_version: int) -> list[int]:
    # asyncio.gather is the idiomatic Flyte 2 way to fan out.
    coros = [score_shard_async(i, model_version) for i in range(n)]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v2.py*

See the [Asynchronous model](./overview#asynchronous-model) guide for the concepts behind async execution.

### Choosing `flyte.map` vs `asyncio.gather`

| Feature | `flyte.map` (sync) | `asyncio.gather` (async) |
|---|---|---|
| Syntax | `list(flyte.map(fn, items))` | `await asyncio.gather(*tasks)` |
| Concurrency limit | Built-in `concurrency=N` | Use `asyncio.Semaphore` |
| Streaming / as-completed | No | Yes, via `asyncio.as_completed()` |
| Error handling | `return_exceptions=True` | Check return type |

Use **`flyte.map`** for the smallest change from Flyte 1 `map_task`, or when you're stuck in synchronous code. Use **`asyncio.gather`** for new code where you want streaming results or fine-grained concurrency control.

### Concurrency control and error handling

`map_task`'s `concurrency` and `min_success_ratio` become an `asyncio.Semaphore` and `return_exceptions=True`:

```python
import asyncio

@env.task
async def main(items: list[int], max_concurrent: int = 5) -> list[str]:
    sem = asyncio.Semaphore(max_concurrent)

    async def process_with_limit(item: int) -> str:
        async with sem:
            return await process_item(item)

    tasks = [process_with_limit(i) for i in items]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    return [r for r in results if not isinstance(r, Exception)]
```

## Data backfills

Reprocessing a range of dates is a textbook `@dynamic` use case in Flyte 1, because the number of days is only known at runtime. In Flyte 2 it's a plain task that builds the date range and fans the days out with `asyncio.gather`.

### Flyte 1

```python
from datetime import date, timedelta

from flytekit import task, workflow, dynamic

@task
def process_day(day: str) -> int:
    # Reprocess a single day's partition; return the row count.
    return len(day)

# @dynamic is needed because the number of days is only known at runtime.
@dynamic
def backfill(start: str, days: int) -> list[int]:
    base = date.fromisoformat(start)
    results = []
    for i in range(days):
        day = (base + timedelta(days=i)).isoformat()
        results.append(process_day(day=day))
    return results

@workflow
def main(start: str, days: int) -> list[int]:
    return backfill(start=start, days=days)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_backfill_v1.py*

### Flyte 2

```
import asyncio
from datetime import date, timedelta

import flyte

env = flyte.TaskEnvironment(name="data_backfill")

@env.task
async def process_day(day: str) -> int:
    # Reprocess a single day's partition; return the row count.
    return len(day)

# A plain task builds the date range at runtime and fans the days out in
# parallel with asyncio.gather -- no @dynamic and no map_task needed.
@env.task
async def main(start: str, days: int) -> list[int]:
    base = date.fromisoformat(start)
    coros = [
        process_day((base + timedelta(days=i)).isoformat())
        for i in range(days)
    ]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_backfill_v2.py*

For fine-grained concurrency control (semaphores, `as_completed`, error handling), see [Fanout](../../tasks/task-programming/fanout).

## Next

- [Data types and I/O](./data-io) — files, DataFrames, and dataclasses
- [ML workloads](./ml-workloads) — training, HPO, and batch inference


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/data-io ===

# Data types and I/O

Flyte 2 renames the offloaded-data types and makes their I/O `async`, but the mental model is the same: pass lightweight references to large data between tasks. See [Migration](./overview) for the overall approach.

## Files and directories

`FlyteFile` and `FlyteDirectory` become `flyte.io.File` and `flyte.io.Dir` — the way you pass model artifacts and datasets between tasks. The I/O is now `async`: use `await File.from_local(...)` to upload and `async with file.open(...)` to read. Like their Flyte 1 counterparts, these are lightweight references to offloaded data, not the materialized bytes.

### Flyte 1

```python
import os

from flytekit import task, workflow, current_context
from flytekit.types.file import FlyteFile

@task
def write_file(content: str) -> FlyteFile:
    path = os.path.join(current_context().working_directory, "out.txt")
    with open(path, "w") as f:
        f.write(content)
    return FlyteFile(path=path)

@task
def read_file(f: FlyteFile) -> str:
    with open(f.download()) as fh:
        return fh.read()

@workflow
def main(content: str) -> str:
    f = write_file(content=content)
    return read_file(f=f)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/files_v1.py*

### Flyte 2

```
import flyte
from flyte.io import File

env = flyte.TaskEnvironment(name="files")

@env.task
async def write_file(content: str) -> File:
    with open("out.txt", "w") as f:
        f.write(content)
    # File.from_local uploads the file to blob storage and returns a reference
    # (a lightweight pointer, not the materialized bytes).
    return await File.from_local("out.txt")

@env.task
async def read_file(f: File) -> str:
    async with f.open("rb") as fh:
        return (await fh.read()).decode("utf-8")

@env.task
async def main(content: str) -> str:
    f = await write_file(content)
    return await read_file(f)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/files_v2.py*

See [Files and directories](../../tasks/task-programming/files-and-directories) for more.

## DataFrames

`StructuredDataset` becomes `flyte.io.DataFrame`. Construct one with `flyte.io.DataFrame.from_df(df)` and read it back with `await df.open(pandas.DataFrame).all()`.

### Flyte 1

```python
import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def make_df() -> StructuredDataset:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    return StructuredDataset(dataframe=df)

@task
def total_payroll(sd: StructuredDataset) -> float:
    df = sd.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@workflow
def main() -> float:
    return total_payroll(sd=make_df())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataframe_v1.py*

### Flyte 2

```
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="dataframe",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def make_df() -> flyte.io.DataFrame:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(df)

@env.task
async def total_payroll(fdf: flyte.io.DataFrame) -> float:
    df = await fdf.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@env.task
async def main() -> float:
    return await total_payroll(await make_df())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataframe_v2.py*

See [DataFrames](../../tasks/task-programming/dataframes) for more.

## Dataclasses and structured types

Flyte 1 required a `@dataclass_json` mixin for dataclass I/O. In Flyte 2, plain dataclasses (and Pydantic `BaseModel`s) work directly as task inputs and outputs — handy for passing around a training config.

### Flyte 1

```python
from dataclasses import dataclass

from dataclasses_json import dataclass_json
from flytekit import task, workflow

@dataclass_json
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@workflow
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate=learning_rate, n_estimators=n_estimators)
    return train(config=config)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataclasses_v1.py*

### Flyte 2

```
from dataclasses import dataclass

import flyte

env = flyte.TaskEnvironment(name="dataclasses")

# Plain dataclasses work directly as task I/O -- no @dataclass_json mixin needed.
# Pydantic BaseModels work the same way.
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@env.task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@env.task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@env.task
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate, n_estimators)
    return train(config)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataclasses_v2.py*

## Data ETL

Putting the data types together: extract, clean, aggregate, and write out a feature table. `StructuredDataset` becomes `flyte.io.DataFrame`, and the tasks become `async`.

### Flyte 1

```python
import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@task
def transform(df: pd.DataFrame) -> StructuredDataset:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    return StructuredDataset(dataframe=agg)

@task
def load(sd: StructuredDataset) -> int:
    df = sd.open(pd.DataFrame).all()
    return len(df)

@workflow
def main() -> int:
    raw = extract()
    features = transform(df=raw)
    return load(sd=features)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_etl_v1.py*

### Flyte 2

```
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="data_etl",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@env.task
async def transform(df: pd.DataFrame) -> flyte.io.DataFrame:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(agg)

@env.task
async def load(sd: flyte.io.DataFrame) -> int:
    df = await sd.open(pd.DataFrame).all()
    return len(df)

@env.task
async def main() -> int:
    raw = await extract()
    features = await transform(raw)
    return await load(features)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_etl_v2.py*

## Next

- [ML workloads](./ml-workloads) — training, HPO, and batch inference
- [Parallelism and fan-out](./parallelism) — processing partitions in parallel


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/ml-workloads ===

# ML workloads

The core patterns compose into the workloads data scientists and ML engineers run every day. Each of these is a complete v1→v2 pair in the [examples repo](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/migration/flyte-2). See [Migration](./overview) for the overall approach.

## Small model training (scikit-learn / XGBoost)

Train a model, persist it as a `File`, and evaluate it. Image, resources, and caching move to the `TaskEnvironment`; `FlyteFile` becomes `flyte.io.File`.

### Flyte 1

```python
import os

import joblib
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

image = ImageSpec(
    name="xgb-image",
    packages=["xgboost", "scikit-learn", "joblib"],
)

@task(container_image=image, requests=Resources(cpu="2", mem="4Gi"))
def train_model(n_estimators: int, max_depth: int) -> FlyteFile:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(current_context().working_directory, "model.json")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile) -> float:
    model = joblib.load(model_file.download())
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@workflow
def main(n_estimators: int, max_depth: int) -> float:
    model = train_model(n_estimators=n_estimators, max_depth=max_depth)
    return evaluate(model_file=model)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_xgboost_v1.py*

### Flyte 2

```
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

env = flyte.TaskEnvironment(
    name="train_xgboost",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "xgboost", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
)

@env.task
async def train_model(n_estimators: int, max_depth: int) -> File:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(os.getcwd(), "model.json")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@env.task
async def main(n_estimators: int, max_depth: int) -> float:
    model = await train_model(n_estimators, max_depth)
    return await evaluate(model)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_xgboost_v2.py*

## Hyperparameter optimization

Fan out one training run per hyperparameter, then pick the best. In Flyte 1 the grid search runs through `map_task` and the "pick the best" step must itself be a task. In Flyte 2 you `gather` the runs and select the winner in plain Python.

### Flyte 1

```python
from flytekit import task, workflow, map_task
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

@task
def get_grid() -> list[int]:
    return [2, 4, 8, 16]

@task
def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@task
def best_score(scores: list[float]) -> float:
    return max(scores)

@workflow
def main() -> float:
    grid = get_grid()
    # Fan out one training run per hyperparameter value.
    scores = map_task(train_eval)(max_depth=grid)
    return best_score(scores=scores)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hpo_v1.py*

### Flyte 2

```
import asyncio

import flyte
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

env = flyte.TaskEnvironment(
    name="hpo",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn"),
)

@env.task
async def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@env.task
async def main() -> dict:
    grid = [2, 4, 8, 16]
    # Fan out one training run per hyperparameter value...
    scores = await asyncio.gather(*[train_eval(d) for d in grid])
    # ...then pick the best in plain Python (impossible in a Flyte 1 workflow).
    best_idx = max(range(len(scores)), key=lambda i: scores[i])
    return {"best_max_depth": grid[best_idx], "best_score": scores[best_idx]}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hpo_v2.py*

## Large model training (deep learning)

GPU configuration moves to the `TaskEnvironment`: the Flyte 1 `Resources(gpu="1")` plus a separate `accelerator=T4` become a single `gpu="T4:1"` string on `flyte.Resources`.

### Flyte 1

```python
from flytekit import task, workflow, ImageSpec, Resources
from flytekit.extras.accelerators import T4
import torch
import torch.nn as nn

image = ImageSpec(
    name="dl-image",
    packages=["torch"],
)

@task(
    container_image=image,
    requests=Resources(cpu="4", mem="16Gi", gpu="1"),
    accelerator=T4,
)
def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@workflow
def main(epochs: int) -> float:
    return train(epochs=epochs)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_deep_learning_v1.py*

### Flyte 2

```
import flyte
import torch
import torch.nn as nn

# GPU type and count go in a single "T4:1"-style string. For multi-node
# distributed training, wrap the training task with the torch elastic plugin.
env = flyte.TaskEnvironment(
    name="train_deep_learning",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="T4:1"),
)

@env.task
async def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@env.task
async def main(epochs: int) -> float:
    return await train(epochs)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_deep_learning_v2.py*

For multi-node distributed training (PyTorch elastic, etc.), see [Resources](../../tasks/task-configuration/resources) and the plugin integrations.

## Batch inference

Load a trained model once and score many batches in parallel. `map_task` with a `partial`-bound model becomes `asyncio.gather` over the batches, reusing the same model reference.

### Flyte 1

```python
import os
from functools import partial

import joblib
from flytekit import task, workflow, map_task, ImageSpec, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(name="inference-image", packages=["scikit-learn", "joblib"])

@task(container_image=image)
def train_model() -> FlyteFile:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def get_batches() -> list[list[list[float]]]:
    data = load_iris()
    rows = data.data.tolist()
    # Split the rows into batches of 30.
    return [rows[i : i + 30] for i in range(0, len(rows), 30)]

@task(container_image=image)
def score_batch(model_file: FlyteFile, batch: list[list[float]]) -> list[int]:
    model = joblib.load(model_file.download())
    return [int(p) for p in model.predict(batch)]

@workflow
def main() -> list[list[int]]:
    model = train_model()
    batches = get_batches()
    return map_task(partial(score_batch, model_file=model))(batch=batches)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/batch_inference_v1.py*

### Flyte 2

```
import asyncio
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

env = flyte.TaskEnvironment(
    name="batch_inference",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "joblib"),
)

@env.task
async def train_model() -> File:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def score_batch(model_file: File, batch: list[list[float]]) -> list[int]:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    return [int(p) for p in model.predict(batch)]

@env.task
async def main() -> list[list[int]]:
    model = await train_model()
    rows = load_iris().data.tolist()
    batches = [rows[i : i + 30] for i in range(0, len(rows), 30)]
    # Score every batch in parallel, reusing the same model reference.
    coros = [score_batch(model, batch) for batch in batches]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/batch_inference_v2.py*

## A complete example: end-to-end ML pipeline

Putting it together — a load / train / evaluate pipeline shows the image, resources, caching, file I/O, and orchestration changes in one place.

### Flyte 1

```python
import os

import joblib
import pandas as pd
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(
    name="ml-image",
    packages=["pandas", "scikit-learn", "joblib"],
)

@task(
    container_image=image,
    requests=Resources(cpu="2", mem="4Gi"),
    cache=True,
    cache_version="1.0",
)
def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@task(container_image=image)
def train_model(data: pd.DataFrame) -> FlyteFile:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile, data: pd.DataFrame) -> float:
    model = joblib.load(model_file.download())
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

@workflow
def main() -> float:
    data = load_data()
    model = train_model(data=data)
    return evaluate(model_file=model, data=data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/ml_pipeline_v1.py*

### Flyte 2

```
import os

import joblib
import pandas as pd
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

# Image, resources, and cache are set once on the TaskEnvironment.
env = flyte.TaskEnvironment(
    name="ml_pipeline",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "pandas", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    cache="auto",
)

@env.task
async def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@env.task
async def train_model(data: pd.DataFrame) -> File:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File, data: pd.DataFrame) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

# The "workflow" is just an orchestrating task.
@env.task
async def main() -> float:
    data = await load_data()
    model = await train_model(data)
    return await evaluate(model, data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/ml_pipeline_v2.py*

## Next

- [New in Flyte 2](./new-in-flyte-2) — real-time serving, apps, and sandboxing
- [Gotchas and caveats](./gotchas-and-caveats) — caveats of the new execution model


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/new-in-flyte-2 ===

# New in Flyte 2: patterns that weren't possible in Flyte 1

Flyte 1 was a batch orchestration system: everything ran as a finite DAG that started, did work, and finished. Flyte 2 keeps all of that and adds **long-running services**, **high-throughput batch inference**, and **sandboxed code execution** — so the same project that trains your model can also serve it, host a dashboard, saturate a GPU, or safely run LLM-generated code. There is no v1 counterpart to migrate here; these are net-new capabilities.

## Real-time inference and model serving

Instead of scoring a batch and exiting, you can stand up an always-on REST endpoint from a `FastAPIAppEnvironment` and deploy it with `flyte.deploy`. The app can load a model artifact produced by one of your training tasks.

```
app = FastAPI(title="ML Model API")

# Define request/response models
class PredictionRequest(BaseModel):
    feature1: float
    feature2: float
    feature3: float

class PredictionResponse(BaseModel):
    prediction: float
    probability: float

# Load model (you would typically load this from storage)
model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model_path = os.getenv("MODEL_PATH", "/app/models/model.joblib")
    # In production, load from your storage
    if os.path.exists(model_path):
        with open(model_path, "rb") as f:
            model = joblib.load(f)
    yield

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    # Make prediction
    # prediction = model.predict([[request.feature1, request.feature2, request.feature3]])

    # Dummy prediction for demo
    prediction = 0.85
    probability = 0.92

    return PredictionResponse(
        prediction=prediction,
        probability=probability,
    )

env = FastAPIAppEnvironment(
    name="ml-model-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "scikit-learn",
        "pydantic",
        "joblib",
    ),
    parameters=[
        flyte.app.Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.joblib"),
            mount="/app/models",
            env_var="MODEL_PATH",
        ),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/ml_model_serving.py*

See [Build apps](../../apps/build-apps/_index) and [Serve and deploy apps](../../apps/serve-and-deploy-apps/_index).

## LLM serving

For serving large language models, the `flyteplugins-vllm` integration gives you a production-grade vLLM server (with autoscaling to zero) in a few lines.

```
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by vLLM
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/basic_vllm.py*

See [vLLM](../../apps/native-app-integrations/vllm-app) and the other [Native app integrations](../../apps/native-app-integrations/_index) (SGLang, Streamlit, FastAPI).

## App serving (dashboards and APIs)

Any web app — a Streamlit dashboard for exploring results, a Gradio demo, a Flask backend — runs as an `AppEnvironment`. Configure the image, resources, port, autoscaling, and a custom subdomain, then `flyte.serve` it.

```
app_env = flyte.app.AppEnvironment(
    name="hello-world-app",
    image=image,
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
    domain=flyte.app.Domain(subdomain="hello"),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app.py*

See [Introducing apps](../../get-started/core-concepts/introducing-apps) and [Configure apps](../../apps/configure-apps/_index).

## Dynamic batching for GPU inference

For in-process batch inference, `DynamicBatcher` from `flyte.extras` keeps an expensive GPU saturated: async producers load and preprocess data concurrently while a single consumer feeds the model in optimally-sized batches, with built-in backpressure. This replaces the Flyte 1 pattern of standing up a separate inference server just to get request batching.

```python
import asyncio
from flyte.extras import DynamicBatcher

async with DynamicBatcher(
    process_fn=run_inference,   # takes a batch, returns results in the same order
    target_batch_cost=1000,     # cost budget per batch
    max_batch_size=64,          # hard cap on records per batch
    batch_timeout_s=0.05,       # max wait before dispatching a partial batch
) as batcher:
    futures = [await batcher.submit(record) for record in records]
    results = await asyncio.gather(*futures)
```

`submit()` is non-blocking and returns a `Future`; when the queue is full it applies backpressure automatically. See [Batch inference](../../run-scaling/batch-inference), which also covers `TokenBatcher` for token-aware LLM batching.

## Sandboxed code execution

`flyte.sandbox.create()` runs arbitrary Python code or shell commands inside an ephemeral, single-use Docker container — built on demand from declared dependencies, executed once, then discarded. Only declared inputs go in and only declared outputs come back, which makes it the safe way to run **untrusted code, most importantly code generated by an LLM**.

```
# sandbox_environment provides the base runtime for code sandboxes.
# Include it in depends_on so the sandbox runtime is available when tasks execute.
env = flyte.TaskEnvironment(
    name="sandbox-demo",
    image=flyte.Image.from_debian_base(name="sandbox-demo"),
    depends_on=[sandbox_environment],
)

# Auto-IO mode: pure computation. The code string runs in an isolated sandbox;
# only the declared inputs go in and only the declared outputs come back.
sum_sandbox = flyte.sandbox.create(
    name="sum-to-n",
    code="total = sum(range(n + 1)) if conditional else 0",
    inputs={"n": int, "conditional": bool},
    outputs={"total": int},
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/sandboxing/code_sandbox.py*

Call it from a task with `await sum_sandbox.run.aio(n=10, conditional=True)`. See [Code sandboxing](../../agents/sandboxing/code-sandboxing). This also powers **code mode** (programmatic tool calling), where an agent writes a whole program instead of emitting one tool call at a time — see [Programmatic tool calling for agents](../../agents/sandboxing/code-mode).

## Next

- [Gotchas and caveats](./gotchas-and-caveats) — caveats of the new execution model
- [Hybrid v1 and v2 pipelines](./hybrid-pipelines) — calling between v1 and v2 during the transition


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/hybrid-pipelines ===

# Hybrid v1 and v2 pipelines

Migrations rarely happen all at once. For a while you'll have Flyte 1 and Flyte 2 workloads running side by side, and you'll want them to call each other: a Flyte 1 workflow that kicks off a newly ported Flyte 2 task, or a Flyte 2 task that triggers a workflow that hasn't been migrated yet.

You can bridge the two in both directions. The idea is the same each way: one task installs **both** SDKs, authenticates to the **other** control plane, fetches the entity it wants to run, and launches it.

> **📝 Note**
>
> The bridging task acts as a driver: it authenticates to a remote control plane and launches work there. Keep it lightweight and focused on orchestration — see [Gotchas and caveats](./gotchas-and-caveats).

## Running a Flyte 2 task from a Flyte 1 workflow

The bridge is a single Flyte 1 task — call it `launch_v2_from_v1` — that runs the Flyte 2 client.

**High-level steps:**

1. Give the `launch_v2_from_v1` task an image with **both** `flytekit` (Flyte 1) and `flyte` (Flyte 2) installed.
2. Give it a Flyte 2 **API key** (see **Migration > From Flyte 1 to 2 > Hybrid v1 and v2 pipelines > Running a Flyte 2 task from a Flyte 1 workflow > 1. Create the API key and store it** below — how you create one differs between Union and open-source Flyte). The key is the `export FLYTE_API_KEY="..."` value it produces.
3. Make the key available to the task, either by storing it as a secret it can read, or by injecting it as the `FLYTE_API_KEY` environment variable.
4. Authenticate inside the task with `flyte.init_from_api_key()`.
5. Fetch the deployed Flyte 2 task with `flyte.remote.Task.get(...)` and run it with `flyte.run(...)`.

### 1. Create the API key and store it

Obtain a Flyte 2 API key from your control plane's authentication setup and store it as a secret the bridging task can read — see [Run on a remote cluster](../../get-started/run-modes/running-remote) for the authentication options.

You can also make the same value available to your Flyte 1 task as a secret through your existing Flyte 1 secret workflow, or set it directly as the `FLYTE_API_KEY` environment variable on the task.

### 2. Write the bridging task

```python
import flytekit
from flytekit import task, workflow, ImageSpec, Secret, current_context

# The bridge image needs BOTH the v1 (flytekit) and v2 (flyte) SDKs.
bridge_image = ImageSpec(
    name="v1-to-v2-bridge",
    packages=["flytekit", "flyte"],
)

@task(
    container_image=bridge_image,
    secret_requests=[Secret(group="flyte", key="flyte_api_key")],
)
def launch_v2_from_v1(x: int) -> str:
    import flyte
    import flyte.remote

    # Authenticate to the Flyte 2 control plane with the API key.
    # Option A: read the mounted secret and pass it explicitly.
    api_key = current_context().secrets.get(group="flyte", key="flyte_api_key")
    flyte.init_from_api_key(api_key=api_key)

    # Option B: if FLYTE_API_KEY is set as an env var, no argument is needed:
    #     flyte.init_from_api_key()

    # Fetch the deployed Flyte 2 task and run it.
    remote_v2_task = flyte.remote.Task.get(
        "my_v2_env.process",
        auto_version="latest",
    )
    run = flyte.run(remote_v2_task, x=x)
    run.wait()  # optional: block until the v2 run finishes
    return run.url

@workflow
def main(x: int) -> str:
    return launch_v2_from_v1(x=x)
```

The referenced Flyte 2 task (`my_v2_env.process` above) must be **deployed** before the bridge runs — `flyte.remote.Task.get()` looks it up by `env_name.task_name`. See [Remote tasks](../../tasks/task-programming/remote-tasks) for versioning options (`auto_version="latest"`, `version="v1.2.3"`) and `flyte.run` details.

> **📝 Note**
>
> `flyte.init_from_api_key()` is required here — do **not** use `flyte.init_from_config()`, which reads a `config.yaml` that has no API-key field. See [Run on a remote cluster](../../get-started/run-modes/running-remote) for the authentication methods.

## Running a Flyte 1 workflow from a Flyte 2 task

The reverse works the same way: a Flyte 2 task installs the Flyte 1 client and uses `FlyteRemote` to launch a Flyte 1 workflow.

**High-level steps:**

1. Give the `launch_v1_from_v2` `TaskEnvironment` an image with the Flyte 1 client (`flytekit`) installed.
2. Provide the task with credentials for the Flyte 1 control plane (a config file or API key, supplied as a secret).
3. Instantiate a `FlyteRemote` client inside the task.
4. Fetch the Flyte 1 workflow with `fetch_workflow(...)`.
5. Launch it with `execute(...)`.

```python
import flyte

env = flyte.TaskEnvironment(
    name="v2_to_v1_bridge",
    # The image needs the Flyte 1 client installed.
    image=flyte.Image.from_debian_base().with_pip_packages("flytekit"),
    # Supply credentials for the Flyte 1 control plane (config or API key).
    secrets=[flyte.Secret(key="v1_client_secret", as_env_var="V1_CLIENT_SECRET")],
)

@env.task
async def launch_v1_from_v2(x: int) -> str:
    from flytekit.remote import FlyteRemote
    from flytekit.configuration import Config

    # Point the client at your Flyte 1 cluster.
    remote = FlyteRemote(
        config=Config.for_endpoint(endpoint="my-v1-cluster.example.com"),
        default_project="flytesnacks",
        default_domain="development",
    )

    # Fetch the deployed Flyte 1 workflow and execute it.
    wf = remote.fetch_workflow(name="my_v1_module.main", version="v1.2.3")
    execution = remote.execute(wf, inputs={"x": x}, wait=True)
    return execution.id.name
```

## Considerations

- **Both SDKs in one image.** The bridging task installs `flytekit` and `flyte` together. Pin versions and watch for dependency conflicts; keep the bridge image minimal.
- **Deploy the callee first.** For the v1→v2 direction, the Flyte 2 task must be deployed (`flyte deploy`) before `flyte.remote.Task.get()` can resolve it. For the v2→v1 direction, the Flyte 1 workflow must be registered on its cluster.
- **Wait vs. fire-and-forget.** Both `run.wait()` (v2) and `execute(..., wait=True)` (v1) block until the launched run finishes. Omit them to launch and return immediately, then poll or hand off the execution URL.
- **Credentials cross a boundary.** The bridge authenticates to a *different* control plane than the one it runs on. Store the API key or client credentials as a secret — never hard-code them. See [Secrets](../../tasks/task-configuration/secrets) and [Run on a remote cluster](../../get-started/run-modes/running-remote).
- **Keep the bridge lightweight.** Like any orchestrating task, it should mostly launch and assemble results rather than do heavy compute — see [Gotchas and caveats](./gotchas-and-caveats).

## See also

- [Remote tasks](../../tasks/task-programming/remote-tasks) — fetching and running deployed Flyte 2 tasks
- [Run on a remote cluster](../../get-started/run-modes/running-remote) — authentication methods, including `flyte.init_from_api_key()`

- [Migration](./overview) — mapping Flyte 1 workload patterns to Flyte 2


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/gotchas-and-caveats ===

# Gotchas and caveats

Flyte 2 is a substantial change from Flyte 1: each Python task can act as its own engine, launching sub-tasks and assembling their outputs. That flexibility is powerful, but it warrants some caveats when authoring your tasks.

This page starts with a quick list of common gotchas you'll hit during a migration, then covers the deeper execution-model considerations in detail.

## Common gotchas

- **`flyte.map` returns a generator.** Wrap it in `list()` to materialize results, unlike `map_task` which returned a list directly.
- **`memory`, not `mem`.** The `Resources` parameter was renamed, and there are no separate `requests`/`limits` — a single value serves as both.
- **GPUs use a `"T4:1"` string.** Type and count are combined; the separate `accelerator=` argument is gone.
- **Image, resources, and cache live on the `TaskEnvironment`.** Set them once at the env level instead of repeating them on every task decorator.
- **`current_context()` is gone.** Read secrets from environment variables and use `flyte.ctx()` for runtime context.
- **The `>>` ordering operator is gone.** Sequential (sync) calls and sequential `await`s are naturally ordered.
- **Retries no longer have a platform cap.** In Flyte 1 the control plane capped attempts at 3; in Flyte 2 total attempts equal `retries + 1`. Audit any large `retries` values before deploying.
- **You can only `await` async tasks.** Call a sync task from an async context with `.aio()`; see the [Asynchronous model](./overview#asynchronous-model).
- **Pick an entrypoint task name.** There's no `@workflow`, so the top-level task is just a task (commonly `main`); run it with `flyte run module.py main`.
- **Type annotations are more lenient.** Flyte 2 will pickle untyped I/O rather than rejecting it at registration.
- **Keep orchestration lightweight.** A task that calls other tasks acts as a driver pod. Avoid heavy CPU work in it — see **Migration > From Flyte 1 to 2 > Gotchas and caveats > Driver pod requirements** below.

## Non-deterministic behavior

When a task launches another task, a new Action ID is determined.
This ID is a hash of the inputs to the task, the task definition itself, along with some other information.
The fact that this ID is consistently hashed is important when it comes to things like recovery and replay.

For example, assume you have the following tasks

```python
@env.task
async def t1():
    val = get_int_input()
    await t2(int=val)

@env.task
async def t2(val: int): ...
```

If you run `t1`, and it launches the downstream `t2` task, and then the pod executing `t1` fails, when Flyte restarts `t1` it will automatically detect that `t2` is still running and will just use that.
If `t2` ends up finishing in the interim, those results would just be used.

However, if you introduce non-determinism into the picture, then that guarantee is no longer there.
To give a contrived example:

```python
@env.task
async def t1():
    val = get_int_input()
    now = datetime.now()

    if now.second % 2 == 0:
        await t2(int=val)
    else:
        await t3(int=val)
```

Here, depending on what time it is, either `t2` or `t3` may end up running.
In the earlier scenario, if `t1` crashes unexpectedly, and Flyte retries the execution, a different downstream task may get kicked off instead.

### Dealing with non-determinism

As a developer, the best way to manage non-deterministic behavior (if it is unavoidable) is to be able to observe it and see exactly what is happening in your code. Flyte 2 provides precisely the tool needed to enable this: Traces.

With this feature you decorate the sub-task functions in your code with `@trace`, enabling checkpointing, reproducibility and recovery at a fine-grained level. See [Traces](../../tasks/task-programming/traces) for more details.

## Type safety

In Flyte 1, the top-level workflow was defined by a Python-like DSL that was compiled into a static DAG composed of tasks, each of which was, internally, defined in real Python.
The system was able to guarantee type safety across task boundaries because the task definitions were static and the inputs and outputs were defined in a way that Flytekit could validate them.

In Flyte 2, the top-level workflow is defined by Python code that runs at runtime (unless using a compiled task).
This means that the system can no longer guarantee type safety at the workflow level.

Happily, the Python ecosystem has evolved considerably since Flyte 1, and Python type hints are now a standard way to define types.

Consequently, in Flyte 2, developers should use Python type hints and type checkers like `mypy` to ensure type safety at all levels, including the top-most task (i.e., the "workflow" level).

## No global state

A core principle of Flyte 2 (that is also shared with Flyte 1) is that you should not try to maintain global state across your workflow.
It will not be translated across tasks containers,

In a single process Python program, global variables are available across functions.
In the distributed execution model of Flyte, each task runs in its own container, and each container is isolated from the others.

If there is some state that needs to be preserved, it must be reconstructable through repeated deterministic execution.

## Driver pod requirements

Tasks don't have to kick off downstream tasks of course and may themselves represent a leaf level atomic unit of compute.
However, when tasks do run other tasks, and more so if they assemble the outputs of those other tasks, then that parent task becomes a driver
pod of sorts.
In Flyte 1, this assembling of intermediate outputs was done by Flyte Propeller.
In 2, it's done by the parent task.

This means that the pod running your parent task must be appropriately sized, and should ideally not be CPU-bound, otherwise it slow down downstream evaluation and kickoff of tasks.

For example, if you had this also scenario,

```python
@env.task
async def t_main():
    await t1()
    local_cpu_intensive_function()
    await t2()
```

The pod running `t_main` will hang in between tasks `t1` and `t2`. Your parent tasks should ideally focus only on orchestration.

## OOM risk from materialized I/O

Something maybe more nuanced is that if you're not using the soon-to-be-released ref mode, outputs are actually
materialized. That is, if you have the following scenario,

```python
@env.task
async def produce_1gb_list() -> List[float]: ...

@env.task
async def t1():
    list_floats = produce_1gb_list()
    t2(floats=list_floats)
```

The pod running `t1` needs to have memory to handle that 1 GB of floats. Those numbers will be materialized in that pod's memory.
This can lead to out of memory issues.

Note that `flyte.io.File`, `flyte.io.Dir` and `flyte.io.DataFrame` will not suffer from this because while those are materialized, they're only materialized as pointers to offloaded data, so their memory footprint is much lower.


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/from-airflow ===

# From Airflow to Flyte

A staged guide to migrating Airflow DAGs to Flyte 2.

The migration is split by operator family. Each section maps an Airflow construct to its Flyte 2 equivalent.

### **Migration > From Airflow to Flyte > Part 1: vanilla operators**

PythonOperator, TaskFlow, BashOperator, KubernetesPodOperator, plus DAG schedules, the driver task model, and orchestration patterns (parallelism, conditionals, error handling).

> [!NOTE]
> **Part 2** (later) covers provider operators: Beam, Dataproc, BigQuery, Databricks, Spark, and sensors.

## Subpages
- **Migration > From Airflow to Flyte > Part 1: vanilla operators**


=== PAGE: https://www.union.ai/docs/v2/flyte/user-guide/migration/from-airflow/part-1-vanilla-operators ===

# Part 1: vanilla operators

This is the first part of the [Airflow → Flyte migration guide](./_index). It covers:

1. Where dependencies are specified
2. The driver task (in place of a DAG definition)
3. Triggers (in place of DAG schedules)
4. PythonOperator → `@env.task`
5. TaskFlow → `@env.task`
6. BashOperator → ContainerTask
7. KubernetesPodOperator → TaskEnvironment + PodTemplate
8. Orchestration: parallelism, conditionals, error handling

**Part 2** (later) covers provider operators: Beam, Dataproc, BigQuery, Databricks, Spark, sensors.

---

## 1. Where dependencies are specified

In Airflow, dependencies are specified at the platform level. A single Airflow deployment has a base image with a fixed Python environment; every DAG author writes against the same set of installed libraries. Adding a new library means modifying the deployment (Helm `extraPipPackages`, a custom base image, or a redeploy), or working around the shared env with `PythonVirtualenvOperator`, `DockerOperator`, or `KubernetesPodOperator`.

In Flyte, dependencies are specified at the code level. Each task declares its `TaskEnvironment`, which includes the image and its dependencies. The image is the unit of isolation, and a single workflow can fan out across tasks running in different images.

Two ways to declare the image:

```python
# (a) Build with flyte.Image — from a base, add pip/apt packages, env vars, etc.
etl_env = flyte.TaskEnvironment(
    name="etl",
    image=flyte.Image.from_debian_base()
        .with_pip_packages("pandas", "pyarrow"),
)

# (b) Pass a string reference to an existing image — for example the same one
# you already deploy with in your Airflow KubernetesExecutor / KPO setup.
gpu_env = flyte.TaskEnvironment(
    name="gpu",
    image="registry.example.com/my-org/gpu-training:2026.04.01",
)
```

Docs: [TaskEnvironment](../../get-started/core-concepts/task-environment) · [Container Images](../../tasks/task-configuration/container-images)

---

## 2. The driver task (in place of DAGs)

Airflow DAGs are static graphs. The `with DAG(...)` block runs at parse time; the scheduler compiles the node/edge structure and then traverses it.

Flyte has no parse-time graph. The driver task is Python code that runs at execution time; the graph is built dynamically as the driver calls other tasks. There is no compilation step.

Flyte tasks are async-native: `@env.task` functions are typically declared `async def` and tasks are invoked with `await`. Plain `def` tasks are also supported when you don't need concurrency.

```python
@env.task
async def driver(start: date, regions: list[str]) -> list[Summary]:
    data = await fetch(start)
    summaries = await asyncio.gather(
        *[summarize(data, region) for region in regions]
    )
    return summaries
```

The driver is just a task that calls other tasks; there's no separate workflow object.

---

## 3. Triggers (in place of schedules)

Airflow's `schedule=` on a DAG maps to a Flyte `Trigger` attached to a task.

| Airflow | Flyte |
|---|---|
| `schedule="@hourly"` | `flyte.Trigger.hourly()` |
| `schedule="@daily"` | `flyte.Trigger.daily()` |
| `schedule="0 5 * * *"` | `flyte.Trigger("nightly", flyte.Cron("0 5 * * *"))` |
| `schedule="30 9 * * 1-5"` + `timezone=...` | `flyte.Trigger("biz_hours", flyte.Cron("30 9 * * 1-5", timezone="America/New_York"))` |

```python
@env.task(
    triggers=flyte.Trigger(
        "daily_report",
        flyte.Cron("0 6 * * *"),
        inputs={"trigger_time": flyte.TriggerTime},
    )
)
def generate_report(trigger_time: datetime) -> str:
    ...
```

Multiple triggers per task and parameterized trigger inputs are supported; see the [Triggers docs](../../tasks/task-configuration/triggers).

---

## 4. PythonOperator to `@env.task`

Airflow's `PythonOperator` runs a Python callable in the worker's environment. The callable's return value becomes XCom, and inputs arrive through three channels: `op_args`/`op_kwargs` passed at operator construction, the Airflow context injected as `**kwargs` (`ti`, `ds`, `dag_run`, `logical_date`, …), and `ti.xcom_pull(...)` for data from upstream tasks.

Flyte's equivalent is `@env.task` on a plain function. The function's parameters and return type are the interface; there is no separate context channel and no XCom step.

```python
# Airflow
def fetch_events(**context):
    ds = context["ds"]
    return _fetch(ds)  # returned value is serialized to XCom

def summarize(**context):
    ti = context["ti"]
    records = ti.xcom_pull(task_ids="fetch_events")
    return f"{len(records)} events on {context['ds']}"

with DAG("events", schedule="@daily", ...) as dag:
    t1 = PythonOperator(task_id="fetch_events", python_callable=fetch_events)
    t2 = PythonOperator(task_id="summarize", python_callable=summarize)
    t1 >> t2
```

```python
# Flyte
env = flyte.TaskEnvironment(name="events", image=...)

@env.task
async def fetch_events(ds: str) -> list[dict]:
    return _fetch(ds)

@env.task
async def summarize(ds: str, records: list[dict]) -> str:
    return f"{len(records)} events on {ds}"

@env.task
async def driver(ds: str) -> str:
    records = await fetch_events(ds)
    return await summarize(ds, records)
```

A few things change in the move:

- **Inputs are the function parameters.** No `**context`. If the task needs the run's date, declare it as a parameter (`ds: str`) and the driver passes it in. The driver itself can receive trigger time when a [Trigger](../../tasks/task-configuration/triggers) fires it.
- **Data flows through `await`, not XCom.** The value returned by `fetch_events` is the value `summarize` receives: the function call graph IS the dependency graph. No `xcom_pull` and no `t1 >> t2` to maintain separately from the data flow.
- **Types are part of the signature.** Flyte uses the hints to serialize between tasks, but keep expectations calibrated: the runtime is more like typed JSON than a fully enforced contract. It is useful as documentation and for tooling, not as a strict static check.
- **Async-native, sync-also-works.** Tasks are typically `async def` and invoked with `await`. Plain `def` tasks are fully supported if you'd rather stay in a sync codebase; you just give up some of the flexibility async offers.

The driver above has nothing in it but task calls, for readability. It doesn't have to. A driver is just a `@env.task`, and any code that belongs in a Python function belongs in a driver: plain expressions, loops, `if`/`try`, helpers. Turn something into a `@env.task` when you want it to have its own resources, image, retries, caching, or parallelism. Otherwise leave it as regular Python and call it inline.

Docs: [Tasks](../../get-started/core-concepts/tasks)

### File and Dir: for data that doesn't fit in a return value

Primitive and JSON-serializable values (`int`, `str`, `list`, `dict`, dataclasses, Pydantic models) flow directly as return values: the SDK serializes them. Same shape as XCom, but typed. Most tasks will use these and nothing else.

Flyte adds `File` and `Dir` for the cases where the payload is too big or too binary to inline. In Airflow this is where pipelines step outside the framework: XCom is a Postgres row with a soft ~48KB limit, so larger payloads are written to a shared filesystem or object storage and a path is passed as a string; the upload, the lifecycle, and the cleanup are the author's responsibility, outside Airflow's model.

Flyte covers both cases with the same interface. A task that returns `File` or `Dir` is declaring that its output is an offloaded blob, and the SDK handles the upload on write and the download on read.

```python
import flyte
from flyte.io import File, Dir

@env.task
async def extract(ds: str) -> File:
    # Stream straight to remote storage — no local temp needed.
    file = File.new_remote()
    async with file.open("wb") as f:
        await f.write(b"col1,col2\nfoo,bar\n")
    return file

@env.task
async def count_rows(csv: File) -> int:
    async with csv.open("rb") as f:
        data = await f.read()
    return data.count(b"\n") - 1
```

The `File` object travels between tasks the same way an `int` does: as a typed argument. Underneath, it carries a remote path. Common methods:

- `File.new_remote()`: new reference in the run's scratch area, for streaming writes.
- `File.from_local(path)` / `from_local_sync(path)`: upload a local file, get a `File` back.
- `File.from_existing_remote(uri)`: wrap an existing remote URI (for example, a path produced by an upstream system).
- `async with file.open("rb")` / `async with file.open("wb")` / `with file.open_sync(...)`: stream read/write.
- `await file.download()` / `file.download_sync()`: materialize to a local path and return it.

`Dir` has the same surface for directories, plus `walk()` and `list_files()` to iterate entries.

Docs: [Files and directories](../../tasks/task-programming/files-and-directories)

---

## 5. TaskFlow to `@env.task`

If the DAG you're porting uses Airflow's TaskFlow API (`@task`, `@dag`), the surface move is small: `@task` becomes `@env.task`, the function's return value is the data (no `ti.xcom_pull`), and function calls ARE the dependencies (no `>>`). A lot of TaskFlow code compiles to Flyte with little more than a find-and-replace on the decorator.

```python
# Airflow TaskFlow
from airflow.decorators import dag, task

@dag(schedule="@daily", start_date=datetime(2026, 1, 1), catchup=False)
def events():
    @task
    def fetch_events() -> list[dict]:
        return _fetch()

    @task
    def summarize(records: list[dict]) -> str:
        return f"{len(records)} events"

    summarize(fetch_events())

events()
```

```python
# Flyte
env = flyte.TaskEnvironment(name="events", image=...)

@env.task
async def fetch_events() -> list[dict]:
    return _fetch()

@env.task
async def summarize(records: list[dict]) -> str:
    return f"{len(records)} events"

@env.task(triggers=flyte.Trigger.daily())
async def driver() -> str:
    return await summarize(await fetch_events())
```

The thing worth internalizing, and the main place this stops being a find-and-replace, is what the outer function is doing.

An Airflow `@dag` function runs at **parse time**. Calling `fetch_events()` inside it doesn't run `fetch_events`; it registers a task and an edge in the static graph. The scheduler later traverses that graph. By the time the tasks actually execute, the `@dag` function is long gone.

A Flyte driver is a `@env.task` that runs at **execution time**. There is no parse-time graph-building step. Calling `await fetch_events()` actually calls `fetch_events`. That means the driver (and any task) is just Python: `if`/`else`, `try`/`except`, loops, recursion, calling other tasks from inside other tasks, nested drivers, reading a value from one task and deciding what to do next. All of it works because there is no static graph to fit into.

To make the point concrete, a task can call itself:

```python
@env.task
async def countdown(n: int) -> int:
    if n == 0:
        return 0
    return 1 + await countdown(n - 1)
```

Each `await countdown(...)` call is a real task invocation; the graph grows as the computation runs. This is impossible to express in Airflow's `@dag` model, where the graph has to be known before execution.

The practical effect: patterns that Airflow encodes with its own primitives (`BranchPythonOperator`, `ShortCircuitOperator`, `trigger_rule`, `.expand()` for dynamic mapping, custom `XComArg` gymnastics) are just Python constructs in Flyte. Branching is `if`. Short-circuit is `return`. Dynamic mapping is `asyncio.gather` or `flyte.map`. Error handling is `try`/`except`/`finally`. Section 8 covers these with runnable examples.

### TaskFlow decorator variants

TaskFlow ships several decorators beyond `@task`. Rough mapping:

| TaskFlow | Flyte equivalent |
|---|---|
| `@task` | `@env.task` |
| `@task.bash` | **Migration > From Airflow to Flyte > Part 1: vanilla operators > 6. BashOperator to ContainerTask** |
| `@task.virtualenv` | `@env.task` on a `TaskEnvironment` with its own image |
| `@task.docker` | `@env.task` on a `TaskEnvironment` with `image=...` |
| `@task.kubernetes` | **Migration > From Airflow to Flyte > Part 1: vanilla operators > 7. KubernetesPodOperator to TaskEnvironment + PodTemplate** |
| `@task.branch` | plain `if` in the driver |
| `@task.short_circuit` | plain `return` in the driver |

---

## 6. BashOperator to ContainerTask

Airflow's `BashOperator` runs a shell command in the Airflow worker's image, with inputs rendered into the command via Jinja templating and output captured as the last line of stdout.

```python
BashOperator(
    task_id="extract",
    bash_command="gsutil cp gs://bucket/data-{{ ds }}.csv /tmp/data.csv "
                 "&& wc -l /tmp/data.csv | awk '{print $1}'",
    do_xcom_push=True,
)
```

Flyte's equivalent is a `ContainerTask`: specify an image, a command, typed inputs, and typed outputs. Inputs are substituted via `{{.inputs.<name>}}`; outputs are files the container writes to `output_data_dir`, which Flyte reads back with the declared types.

```python
import flyte
from flyte.extras import ContainerTask

extract = ContainerTask(
    name="extract",
    image=flyte.Image.from_base("google/cloud-sdk:slim"),
    inputs={"date": str},
    outputs={"row_count": int},
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    command=[
        "/bin/sh", "-c",
        "gsutil cp gs://bucket/data-{{.inputs.date}}.csv /tmp/data.csv && "
        "wc -l /tmp/data.csv | awk '{print $1}' > /var/outputs/row_count",
    ],
)
```

A `ContainerTask` is invoked the same way as any other task, by calling it from a driver with `await`:

```python
container_env = flyte.TaskEnvironment.from_task("extract_env", extract)
env = flyte.TaskEnvironment(
    name="pipeline",
    image=flyte.Image.from_debian_base().with_uv_project(pyproject_file="pyproject.toml"),
    depends_on=[container_env],
)

@env.task
async def driver(date: str) -> int:
    return await extract(date=date)
```

Two things about the invocation:

- `TaskEnvironment.from_task(...)` wraps the container task in an environment so it can be registered alongside the driver.
- The driver's env `depends_on=[container_env]` so Flyte registers both together. The driver's own image needs Flyte installed (that's what `from_uv_project` does: builds an image from your `pyproject.toml`, which includes `flyte`). The container task's image does *not* need Flyte; it just needs the tools its command invokes.

### When to use `ContainerTask`

`ContainerTask` is the right choice when the container shouldn't or can't have Flyte installed, for example:

- The tool is not Python (a Go/C CLI, a bioinformatics binary, an ML framework container)
- You want to reuse an existing production image without modifying it
- You want to stay out of Python entirely for the task body

If you already have Python in the loop and just need to shell out for one step, a regular `@env.task` with `subprocess` is simpler:

```python
@env.task
async def extract(date: str) -> int:
    import subprocess
    subprocess.run(["gsutil", "cp", f"gs://bucket/data-{date}.csv", "/tmp/data.csv"], check=True)
    out = subprocess.check_output(["wc", "-l", "/tmp/data.csv"])
    return int(out.split()[0])
```

### How the arguments map

| BashOperator | ContainerTask |
|---|---|
| `bash_command` (string) | `command=[...]` (list; you choose the shell) |
| (implicit worker image) | `image=` (explicit, per task) |
| `env` / `append_env` | `flyte.Image.from_...().with_env_vars(...)` or the `TaskEnvironment` |
| `{{ ds }}`, `{{ ti.xcom_pull(...) }}` | `{{.inputs.<name>}}` |
| `do_xcom_push=True` (last stdout line) | `outputs={...}`, written to files in `output_data_dir` |
| `cwd` | `cd ... && ...` inside the command |

Docs: [Container Tasks](../../tasks/task-programming/container-tasks)

---

## 7. KubernetesPodOperator to TaskEnvironment + PodTemplate

`KubernetesPodOperator` (KPO) gives you the full pod spec: image, commands, env, secrets, resources, volumes, node selectors, tolerations, service accounts, affinity, plus XCom via a sidecar writing to `/airflow/xcom/return.json`.

In Flyte, the same knobs live in three places:

1. **`TaskEnvironment(...)`**: the common knobs. Image, resources, env vars, secrets, interruptible/spot, and an option to add a `pod_template` for every task in the env.
2. **`@env.task(...)`**: per-task overrides on top of the env: retries, timeout, cache, triggers, and a task-level `pod_template` if this one task needs to differ.
3. **`flyte.PodTemplate(...)`**: raw Kubernetes escape hatch. Wraps `kubernetes.client.V1PodSpec`, so anything in the pod spec (volumes, node selectors, tolerations, affinity, service accounts, sidecars, init containers, image pull secrets, security contexts, lifecycle hooks) is available.

XCom has no equivalent; the task's typed return value is the output, and large payloads use `File`/`Dir` (Section 4). The sidecar-writing-to-`/airflow/xcom/return.json` contract doesn't exist.

### Where every KPO knob lands

| KPO argument | Flyte location |
|---|---|
| `image` | `TaskEnvironment(image=...)` |
| `cmds`, `arguments` | function body of `@env.task` |
| `env_vars` (dict) | `TaskEnvironment(env_vars={...})` |
| `secrets=[Secret(...)]` | `TaskEnvironment(secrets=[flyte.Secret(...)])` |
| `container_resources` (requests/limits) | `TaskEnvironment(resources=flyte.Resources(cpu=(1,4), memory="2Gi", gpu="T4:1"))` |
| `node_selector`, `tolerations`, `affinity` | `flyte.PodTemplate(pod_spec=V1PodSpec(node_selector=..., tolerations=..., affinity=...))` |
| `service_account_name` | `PodTemplate(pod_spec=V1PodSpec(service_account_name=...))` |
| `volumes`, `volume_mounts` | `PodTemplate(pod_spec=V1PodSpec(volumes=[...], containers=[V1Container(volume_mounts=[...])]))` |
| `image_pull_secrets` | `PodTemplate(pod_spec=V1PodSpec(image_pull_secrets=[V1LocalObjectReference(name=...)]))` |
| `security_context` | `PodTemplate(pod_spec=V1PodSpec(security_context=...))` |
| `labels`, `annotations` | `PodTemplate(labels={...}, annotations={...})` |
| `init_containers`, sidecars | `PodTemplate(pod_spec=V1PodSpec(init_containers=[...], containers=[primary, ...]))` |
| `retries`, `retry_delay` | `@env.task(retries=...)` |
| `execution_timeout` | `@env.task(timeout=timedelta(...))` |
| `do_xcom_push` + sidecar contract | function return type: primitives/dataclasses inline, large payloads via `File`/`Dir` |
| `on_finish_action` / pod cleanup | handled by Flyte: pods are cleaned up per run lifecycle |

### What a fully-specified task looks like

```python
from datetime import timedelta
from kubernetes.client import V1Container, V1PodSpec

import flyte

pod_template = flyte.PodTemplate(
    primary_container_name="primary",
    labels={"team": "etl"},
    pod_spec=V1PodSpec(
        service_account_name="etl-runner",
        init_containers=[
            V1Container(
                name="warm-cache",
                image="busybox:1.36",
                command=["sh", "-c", "echo warming cache && sleep 1"],
            ),
        ],
    ),
)

etl_env = flyte.TaskEnvironment(
    name="etl",
    image="registry.example.com/etl:2026.04.01",
    resources=flyte.Resources(cpu=(1, 4), memory="2Gi", gpu="T4:1"),
    env_vars={"LOG_LEVEL": "INFO"},
    secrets=[flyte.Secret(key="db-password", as_env_var="DB_PASSWORD")],
    pod_template=pod_template,
    interruptible=True,
)

@etl_env.task(retries=3, timeout=timedelta(minutes=30))
async def load_warehouse(ds: str) -> int:
    ...
```

You don't have to list the primary container in the pod_spec; Flyte fills it in from the env's image, the function's command, and the decorator's resources. Add a `V1Container(name="primary", ...)` entry only when you need to put fields on it directly (volume mounts, extra env, security context).

Docs: [TaskEnvironment](../../get-started/core-concepts/task-environment) · [Secrets](../../tasks/task-configuration/secrets) · [PodTemplate / advanced k8s config](../../tasks/task-configuration/pod-templates)

---

## 8. Orchestration: parallelism, conditionals, error handling

Airflow encodes orchestration in first-class primitives: `[t1, t2, t3] >> merge` for fan-out, `.expand()` for dynamic mapping, `BranchPythonOperator` / `@task.branch` for branching, `ShortCircuitOperator` for early exit, `trigger_rule` for post-branch merges, and `on_failure_callback` / `trigger_rule=ALL_DONE` for failure paths.

In Flyte these are plain Python inside a driver task, because the driver runs at execution time. There is no static graph to encode into.

### Parallelism

Static fan-out in Airflow:

```python
fetch >> [summarize_us, summarize_eu, summarize_apac] >> merge
```

Flyte, concurrent awaits with `asyncio.gather`:

```python
@env.task
async def driver(ds: str) -> Summary:
    raw = await fetch(ds)
    us, eu, apac = await asyncio.gather(
        summarize(raw, "us"),
        summarize(raw, "eu"),
        summarize(raw, "apac"),
    )
    return await merge(us, eu, apac)
```

Each `summarize(...)` returns a coroutine; `asyncio.gather` runs them concurrently and awaits all of them. Tasks called concurrently run in their own pods; the concurrency is real, not just asyncio on one worker.

### Dynamic mapping

Airflow uses `.expand()` to fan out over values known only at runtime:

```python
process.partial(batch_size=100).expand(shard_id=list_shards())
```

Flyte, regular comprehension over the runtime list:

```python
shards = await list_shards()
results = await asyncio.gather(*(process(shard) for shard in shards))
```

To bound concurrency (for example, when the downstream is rate-limited), wrap the call in an `asyncio.Semaphore`:

```python
sem = asyncio.Semaphore(20)

async def process_one(shard):
    async with sem:
        return await process(shard)

results = await asyncio.gather(
    *(process_one(s) for s in shards),
    return_exceptions=True,
)
```

`return_exceptions=True` collects per-item failures instead of failing the batch. The semaphore is also the pattern when different tasks in the same fan-out need different concurrency limits.

If your codebase is sync, `list(flyte.map(process, shards, concurrency=20))` is the sync equivalent of the pattern above.

Docs: [Controlling parallelism](../../tasks/task-programming/controlling-parallelism) · [Fanout](../../tasks/task-programming/fanout)

### Conditionals

Airflow: `BranchPythonOperator` (or `@task.branch`) returns the task_id(s) to run next; `ShortCircuitOperator` skips the rest of the branch; a `trigger_rule=NONE_FAILED_MIN_ONE_SUCCESS` on the merge task reconciles skipped upstreams.

Flyte:

```python
@env.task
async def driver(ds: str) -> Summary:
    size = await inspect(ds)
    if size == 0:
        return Summary(status="empty")
    if size < 1_000_000:
        return await fast_path(ds)
    return await slow_path(ds)
```

Plain `if` / `elif` / `else`, plain `return`. There is no `trigger_rule` to set because there are no skipped tasks to reconcile; the code below the branch simply doesn't run.

### Error handling

Per-task retries and timeouts live on the `@env.task` decorator:

```python
@env.task(retries=3, timeout=timedelta(minutes=15))
async def flaky(ds: str) -> int:
    ...
```

Orchestration-level error handling (Airflow's `trigger_rule=ALL_DONE` cleanup and `on_failure_callback` alerts) is `try` / `except` / `finally` in the driver:

```python
@env.task
async def driver(ds: str) -> Summary:
    try:
        result = await heavy_step(ds)
        return await publish(result)
    except Exception as e:
        await alert(f"{ds} failed: {e}")
        raise
    finally:
        await cleanup(ds)
```

`finally` runs on both success and failure. `except` catches task failures at the `await` site after the task's own retries are exhausted. Specific failure modes live in `flyte.errors` (`OOMError`, `TaskTimeoutError`, `RetriesExhaustedError`, `ActionAbortedError`) and can be caught by type. A common pattern is retrying an OOM with larger resources via `.override(...)`:

```python
import flyte.errors

env = flyte.TaskEnvironment(
    name="transforms",
    image=...,
    resources=flyte.Resources(cpu=1, memory="250Mi"),
)

@env.task
async def transform(ds: str) -> int:
    ...

@env.task
async def driver(ds: str) -> int:
    try:
        return await transform(ds)
    except flyte.errors.OOMError:
        return await transform.override(
            resources=flyte.Resources(cpu=1, memory="2Gi"),
        )(ds)
```

Docs: [Retries and timeouts](../../tasks/task-configuration/retries-and-timeouts) · [Error handling](../../tasks/task-programming/error-handling)

---

## What's next

Once the port is in place, a few Flyte features don't have direct Airflow counterparts and are worth knowing about.

### Caching

A task can be marked cacheable; subsequent calls with the same inputs short-circuit to the previous output instead of re-running. The cache key is derived from inputs and a task version, so bumping the version invalidates.

```python
@env.task(cache="auto")
async def expensive(ds: str) -> Result:
    ...
```

Airflow has no equivalent; XCom stores outputs but doesn't short-circuit on re-execution.

Docs: [Caching](../../tasks/task-configuration/caching)

### Reusable containers

By default each task call gets a fresh pod. A reuse policy keeps the container warm across calls so follow-up invocations skip pod startup and image pull.

```python
warm_env = flyte.TaskEnvironment(
    name="warm",
    image=...,
    reusable=flyte.ReusePolicy(replicas=(1, 3), concurrency=2),
)
```

Useful when a fan-out issues many short tasks against a heavy image.

Docs: [Reusable containers](../../tasks/task-configuration/reusable-containers)

### Reports

A task can emit an HTML report (tables, plots, logs) attached to the run and viewable in the UI. Written from inside the task with `flyte.report`.

```python
@env.task(report=True)
async def summarize(ds: str) -> Summary:
    tab = flyte.report.get_tab("main")
    tab.log(f"<h2>{ds}</h2>")
    tab.log(dataframe.to_html())
    await flyte.report.flush.aio()
    ...
```

Docs: [Reports](../../tasks/task-programming/reports)

### Apps

A long-running HTTP server (FastAPI, Panel, Streamlit, a webhook endpoint) can be deployed alongside your tasks. The app has a URL and can call tasks via the Flyte API. This is the path for webhook-triggered runs, a UI on top of a pipeline, or a custom inference endpoint.

Docs: [Serve and deploy apps](../../apps/serve-and-deploy-apps/_index) · [Build apps](../../apps/build-apps/_index)


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials ===

# Tutorials

This section contains tutorials that showcase relevant use cases and provide step-by-step instructions on how to implement various features using Flyte and Union. Tutorials are organized by **industry vertical** and by **technical topic**.

## Industry verticals

### **Biotech & healthcare**

Bioinformatics, medical imaging, and other life-sciences workloads.

### **Geospatial**

Satellite imagery, remote sensing, and earth and atmospheric modeling workloads.

### **Financial services & fintech**

Financial research, trading, and other fintech workloads.

### **Frontier AI**

Frontier-model pretraining, automated experimentation, and large-scale AI workloads.

## Technical topics

### **Computer vision**

Image and vision-language model workloads.

### **Agents**

Agentic workflows and autonomous LLM-powered systems.

### **Context engineering**

Prompt engineering, prompt optimization, and context construction.

### **Model training**

Training, fine-tuning, and hyperparameter optimization of models at scale.

## Subpages
- **Biotech & healthcare**
- **Geospatial**
- **Financial services & fintech**
- **Frontier AI**
- **Computer vision**
- **Agents**
- **Context engineering**
- **Model training**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare ===

# Biotech & healthcare

Tutorials for bioinformatics, medical imaging, and other life-sciences workloads.

### **Biotech & healthcare > Genomic alignment**

Align sequencing reads to a reference genome with a cached, parallel Bowtie 2 pipeline.

### **Biotech & healthcare > Cross-species gene comparison**

Compare homologous genes across species with Carbon scoring, sequence alignment, and ESMFold 3D structures.

### **Biotech & healthcare > Genomic variant effect prediction**

Zero-shot pathogenicity scoring with HuggingFace Carbon and interactive VEP reports.

### **Biotech & healthcare > Brain tumor MRI classification**

Classify brain MRI scans with a two-phase EfficientNet-B4 pipeline featuring resumable GPU checkpointing and in-UI reports.

### **Biotech & healthcare > Drug molecule screening agent**

Agentic virtual screening with RDKit stage tools, Lipinski filters, and ranked drug-likeness reports.

## Subpages
- **Biotech & healthcare > Genomic alignment**
- **Biotech & healthcare > Brain tumor MRI classification**
- **Biotech & healthcare > Cross-species gene comparison**
- **Biotech & healthcare > Genomic variant effect prediction**
- **Biotech & healthcare > Drug molecule screening agent**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/genomic-alignment ===

# Genomic alignment

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic_alignment).

This tutorial builds a bioinformatics pipeline that aligns raw sequencing reads to a reference genome. The workflow downloads a reference genome and paired-end sequencing data, performs quality filtering, builds a reference index, and aligns the filtered reads with the [Bowtie 2](https://bowtie-bio.sourceforge.net/bowtie2/index.shtml) aligner, running each sample in parallel.

It's a good showcase of how Flyte handles real bioinformatics workloads:

- **Per-task resources** so quality filtering, indexing, and alignment each request exactly the CPU and memory they need.
- **`cache="auto"`** on the download and indexing steps, so re-runs skip work that hasn't changed.
- **Fan-out parallelism** across samples with `asyncio.gather`.
- **System dependencies** (`fastp`, `bowtie2`) installed into the container image with `apt`.

## Define the container image

Because the pipeline shells out to bioinformatics tools, we build a custom image with `flyte.Image.from_uv_script` and install `fastp` (quality filtering) and `bowtie2` (alignment) via `apt`.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

The Python dependencies are declared at the top of the file using the `uv` script style:

CODE0

## Define the task environments

Each stage runs in its own `TaskEnvironment` with tailored resources. The top-level `base_env` declares the others as `depends_on` so the tasks it calls are available at run time.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

## Define the data classes

We model the reference genome, sequencing reads, and alignment results as dataclasses. `flyte.io.File` and `flyte.io.Dir` reference offloaded data in blob storage, so large genomic files are passed between tasks by reference rather than copied through the orchestrator.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

## Fetch assets

The first task downloads the reference genome and paired-end reads from remote URLs and materializes them as `File`/`Dir` objects. It's cached, so repeat runs skip the download.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

## Quality filtering with fastp

`pyfastp` removes duplicate and low-quality reads. It requests extra memory so it can process larger read files efficiently.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
CODE1
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
CODE2
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

## Orchestrate the workflow

The top-level task fetches the assets, filters every sample in parallel, builds the index, and aligns all samples. Parallelism across samples is achieved with `asyncio.gather` rather than a separate `@dynamic` decorator.

```
# # Genomic Alignment
#
# This tutorial demonstrates how to use Flyte to build a workflow that
# performs genomic alignment on sequencing data. The workflow takes as input
# a reference genome and raw sequencing data, performs quality filtering and
# preprocessing on the raw data, generates an index for the reference genome,
# and aligns the filtered data to the reference genome using the Bowtie 2 aligner.

# {{run-on-union}}

# The tutorial is divided into the following sections:
# 1. Define the container image
# 2. Define the data classes
# 3. Define the tasks
# 4. Define the workflow

# /// script
# requires-python = "3.12"
# dependencies = [
#    "flyte",
#    "requests",
# ]
# main = "alignment_wf"
# params = ""
# ///

import asyncio
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List

import requests
import flyte
from flyte.io import Dir, File

# ## Defining a Container Image
#
# We define a custom container image using `flyte.Image`. Since we need bioinformatics
# tools — `fastp` for quality filtering and `bowtie2` for alignment — we install them
# via apt. This approach replaces the v1 `ImageSpec` with conda channels.

# {{docs-fragment image}}
main_img = (
    flyte.Image.from_uv_script(
        __file__,
        name="alignment-tutorial",
    )
    .with_apt_packages("fastp", "bowtie2")
)
# {{/docs-fragment image}}

# We define per-task environments with different resource requirements, then a
# top-level `base_env` that declares all of them as dependencies (required because
# `alignment_wf` and `bowtie2_align_samples` call tasks that live in those environments).

# {{docs-fragment envs}}
fetch_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fetch",
    image=main_img,
    cache="auto",
)

fastp_env = flyte.TaskEnvironment(
    name="alignment-tutorial-fastp",
    image=main_img,
    resources=flyte.Resources(memory="2Gi"),
)

index_env = flyte.TaskEnvironment(
    name="alignment-tutorial-index",
    image=main_img,
    resources=flyte.Resources(memory="10Gi"),
    cache="auto",
)

align_env = flyte.TaskEnvironment(
    name="alignment-tutorial-align",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="10Gi"),
)

base_env = flyte.TaskEnvironment(
    name="alignment-tutorial",
    image=main_img,
    depends_on=[fetch_env, fastp_env, index_env, align_env],
)
# {{/docs-fragment envs}}

# ## Defining Data Classes
#
# We define three data classes to represent the reference genome, sequencing reads,
# and alignment results. We'll first define a convenience function to download files,
# which we'll use within the fetch task to materialize assets from their remote locations.

def fetch_file(url: str, local_dir: str) -> Path:
    """
    Downloads a file from the specified URL.

    Args:
        url (str): The URL of the file to download.
        local_dir (str): The directory where you would like this file saved.

    Returns:
        Path: The local path to the file.

    Raises:
        requests.HTTPError: If an HTTP error occurs while downloading the file.
    """
    url_parts = url.split("/")
    fname = url_parts[-1]
    local_path = Path(local_dir) / fname

    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)

    return local_path

# Reference genomes are used extensively throughout bioinformatics workflows. We define a
# `Reference` data class to represent a reference genome and its associated index files.

# {{docs-fragment dataclasses}}
@dataclass
class Reference:
    """
    Represents a reference FASTA and associated index files.

    Attributes:
        ref_name (str): Name or identifier of the reference file.
        ref_dir (Dir): Directory containing the reference and any index files.
        index_name (str): Index string to pass to tools requiring it.
        indexed_with (str): Name of tool used to create the index.
    """

    ref_name: str
    ref_dir: Dir
    index_name: str | None = None
    indexed_with: str | None = None

# Sequencing reads are the raw data generated from a sequencing experiment.

@dataclass
class Reads:
    """
    Represents a sequencing reads sample via its associated FastQ files.

    Attributes:
        sample (str): The name or identifier of the raw sequencing sample.
        read1 (File): A File object representing the path to the raw R1 read file.
        read2 (File): A File object representing the path to the raw R2 read file.
    """

    sample: str
    read1: File | None = None
    read2: File | None = None

    def get_read_fnames(self):
        return (
            f"{self.sample}_1.fastq.gz",
            f"{self.sample}_2.fastq.gz",
        )

# Finally, we define an `Alignment` data class to represent an alignment file.

@dataclass
class Alignment:
    """
    Represents an alignment file and its associated sample.

    Attributes:
        sample (str): The name or identifier of the sample.
        aligner (str): The name of the aligner used to generate the alignment file.
        format (str): The format of the alignment file (e.g., SAM, BAM).
        alignment (File): A File object representing the path to the alignment file.
    """

    sample: str
    aligner: str
    format: str | None = None
    alignment: File | None = None

    def get_alignment_fname(self):
        return f"{self.sample}_{self.aligner}_aligned.{self.format}"
# {{/docs-fragment dataclasses}}

# ## Tasks
#
# We define a series of tasks to perform the following operations:
# 1. Fetch assets from remote URLs
# 2. Perform quality filtering and preprocessing using FastP
# 3. Generate Bowtie2 index files from a reference genome
# 4. Perform alignment using Bowtie2 on a filtered sample
#
# The first task fetches the reference genome and sequencing reads. It is cached
# so that re-runs skip the download step.

# {{docs-fragment fetch_assets}}
@fetch_env.task
async def fetch_assets(
    ref_url: str, read_urls: List[str]
) -> tuple[Reference, List[Reads]]:
    """
    Fetch assets from remote URLs.
    """
    # Download reference genome
    ref_dir = Path("/tmp/reference_genome")
    ref_dir.mkdir(exist_ok=True, parents=True)
    ref = fetch_file(ref_url, str(ref_dir))
    ref_obj = Reference(
        ref_name=ref.name,
        ref_dir=await Dir.from_local(str(ref_dir)),
    )

    # Download sequencing reads
    dl_loc = Path("/tmp/reads")
    dl_loc.mkdir(exist_ok=True, parents=True)

    samples: dict[str, Reads] = {}
    for url in read_urls:
        fp = fetch_file(url, str(dl_loc))
        sample = fp.stem.split("_")[0]

        if sample not in samples:
            samples[sample] = Reads(sample=sample)

        if ".fastq.gz" in fp.name or "fasta" in fp.name:
            mate = fp.name.strip(".fastq.gz").strip(".filt").split("_")[-1]
            if "1" in mate:
                samples[sample].read1 = await File.from_local(str(fp))
            elif "2" in mate:
                samples[sample].read2 = await File.from_local(str(fp))

    return ref_obj, list(samples.values())
# {{/docs-fragment fetch_assets}}

# The second task performs quality filtering and preprocessing using FastP on a Reads object.
# FastP is a performant tool for removing duplicate or low-quality reads. We increase
# the memory request for this task so FastP can efficiently process reads from larger files.

# {{docs-fragment pyfastp}}
@fastp_env.task
async def pyfastp(rs: Reads) -> Reads:
    """
    Perform quality filtering and preprocessing using Fastp on a Reads object.

    Args:
        rs (Reads): A Reads object containing raw sequencing data to be processed.

    Returns:
        Reads: A Reads object representing the filtered and preprocessed data.
    """
    ldir = Path(tempfile.mkdtemp())
    samp = Reads(rs.sample)
    o1, o2 = samp.get_read_fnames()
    o1p = ldir / o1
    o2p = ldir / o2

    assert rs.read1 is not None and rs.read2 is not None
    r1 = await rs.read1.download()
    r2 = await rs.read2.download()

    cmd = [
        "fastp",
        "-i", str(r1),
        "-I", str(r2),
        "-o", str(o1p),
        "-O", str(o2p),
    ]
    subprocess.run(cmd, check=True)

    samp.read1 = await File.from_local(str(o1p))
    samp.read2 = await File.from_local(str(o2p))

    return samp
# {{/docs-fragment pyfastp}}

# Next, we define a task to generate Bowtie2 index files from a reference genome. As the index
# for a given tool and reference seldom changes, we cache this task.

# {{docs-fragment bowtie2_index}}
@index_env.task
async def bowtie2_index(ref: Reference) -> Reference:
    """
    Generate Bowtie2 index files from a reference genome.

    Args:
        ref (Reference): A Reference object representing the reference genome.

    Returns:
        Reference: The same reference object with the index_name and indexed_with attributes set.
    """
    ref_dir = await ref.ref_dir.download()
    idx_name = "bt2_idx"
    cmd = [
        "bowtie2-build",
        str(Path(str(ref_dir)) / ref.ref_name),
        str(Path(str(ref_dir)) / idx_name),
    ]
    subprocess.run(cmd, check=True)
    return Reference(
        ref.ref_name,
        await Dir.from_local(str(ref_dir)),
        idx_name,
        "bowtie2",
    )
# {{/docs-fragment bowtie2_index}}

# The next task performs paired-end alignment using Bowtie 2 on a single sample.

# {{docs-fragment bowtie2_align}}
@align_env.task
async def bowtie2_align_paired_reads(idx: Reference, fs: Reads) -> Alignment:
    """
    Perform paired-end alignment using Bowtie 2 on a filtered sample.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        fs (Reads): A filtered Reads object containing sample data to be aligned.

    Returns:
        Alignment: An Alignment object representing the alignment result.
    """
    assert idx.indexed_with == "bowtie2", "Reference index must be generated with bowtie2"
    assert idx.index_name is not None
    assert fs.read1 is not None and fs.read2 is not None

    ref_dir = await idx.ref_dir.download()
    r1 = await fs.read1.download()
    r2 = await fs.read2.download()

    ldir = Path(tempfile.mkdtemp())
    alignment = Alignment(fs.sample, "bowtie2", "sam")
    al = ldir / alignment.get_alignment_fname()

    cmd = [
        "bowtie2",
        "-x", str(Path(str(ref_dir)) / idx.index_name),
        "-1", str(r1),
        "-2", str(r2),
        "-S", str(al),
    ]
    subprocess.run(cmd, check=True)

    alignment.alignment = await File.from_local(str(al))
    return alignment
# {{/docs-fragment bowtie2_align}}

# In place of the v1 `@dynamic` workflow, we use a plain async task with `asyncio.gather`
# to run alignments for all samples in parallel.

@base_env.task
async def bowtie2_align_samples(
    idx: Reference, samples: List[Reads]
) -> List[Alignment]:
    """
    Process samples through bowtie2 in parallel.

    Args:
        idx (Reference): A Reference object containing the Bowtie 2 index.
        samples (List[Reads]): A list of Reads objects to be aligned.

    Returns:
        List[Alignment]: A list of Alignment objects representing the alignment results.
    """
    tasks = [bowtie2_align_paired_reads(idx=idx, fs=sample) for sample in samples]
    return list(await asyncio.gather(*tasks))

# ## End-to-End Workflow
#
# We tie everything together in a final task that fetches assets, filters them, generates
# an index, and aligns the samples. In place of the v1 `@workflow`, we use a top-level
# `@base_env.task`. Parallelism across samples is achieved with `asyncio.gather`.

# {{docs-fragment workflow}}
@base_env.task
async def alignment_wf() -> List[Alignment]:
    # Prepare raw samples from remote URLs
    ref, samples = await fetch_assets(
        ref_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38_short.fasta",
        read_urls=[
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_1.fastq.gz",
            "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny_2.fastq.gz",
        ],
    )

    # Filter all samples in parallel
    filtered_samples = list(
        await asyncio.gather(*[pyfastp(rs=s) for s in samples])
    )

    # Generate a bowtie2 index or load it from cache
    bowtie2_idx = await bowtie2_index(ref=ref)

    # Generate alignments using bowtie2
    sams = await bowtie2_align_samples(idx=bowtie2_idx, samples=filtered_samples)

    return sams
# {{/docs-fragment workflow}}

# You can now run the workflow using the command in the dropdown at the top of the page!

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(alignment_wf)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_alignment/genomic_alignment.py*

## Run the workflow

This example has no secrets or external API keys; it pulls public test data from GitHub.

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic_alignment), run it as a `uv` script:

CODE3

Or submit it with the Flyte CLI:

CODE4

When the run completes, each returned `Alignment` points to a SAM file in blob storage that you can download from the run's outputs in the UI.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/tumor-detection ===

# Brain tumor MRI classification

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/tumor_detection).

This tutorial builds a medical-imaging pipeline that classifies brain MRI scans into four categories (Glioma, Meningioma, No Tumor, and Pituitary) using a two-phase EfficientNet-B4 transfer-learning strategy. The pipeline downloads the dataset, trains on a GPU with fault-tolerant checkpointing, and renders training curves and a confusion matrix directly in the Flyte UI.

The example is split into focused modules:

- `config.py`: container image, task environments, and the `TrainingConfig` hyperparameters.
- `dataset.py`: downloads the Hugging Face dataset, builds class-balanced data loaders.
- `model.py` / `training.py`: the Lightning module and the two-phase training loop.
- `utils.py`: plotting helpers for the report.
- `run.py`: the three Flyte tasks and the pipeline driver.

Flyte handles the production concerns:

- **Per-task resources**: CPU for download/reporting, a GPU for training.
- **`cache="auto"`** on dataset download and training, so reruns with the same data and config are free.
- **`retries=3`** plus **Flyte checkpointing** on the training task so a preempted GPU job resumes from the last epoch.
- **Built-in reports** to visualize metrics without separate dashboard infrastructure.

## Define the container image

A single GPU-ready image is shared by all tasks. `with_source_folder` copies the local modules (`dataset.py`, `model.py`, etc.) into the image.

```
"""
Configuration for brain tumor MRI classification pipeline.

Defines task environments, resource requirements, and training hyperparameters.
"""

import pathlib

import flyte

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(
    name="tumor_detection_gpu"
).with_pip_packages(
    "torch",
    "lightning",
    "torchvision",
    "timm",
    "pillow",
    "scikit-learn",
    "plotly",
    "numpy",
    "pandas",
    "torchmetrics",
    "datasets",
    "typing_extensions",
).with_source_folder(
    pathlib.Path(__file__).parent,
    copy_contents_only=True,
)
# {{/docs-fragment image}}

# {{docs-fragment envs}}
# Downloads raw MRI JPEG files — CPU only, no auth needed, result is cached
dataset_env = flyte.TaskEnvironment(
    name="tumor_dataset",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi", disk="8Gi"),
    cache="auto",
)

# GPU training — result is cached so re-running with the same data + config is free
training_env = flyte.TaskEnvironment(
    name="tumor_gpu_training",
    image=image,
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    env_vars={
        "CUDA_VISIBLE_DEVICES": "0",
        "CUDA_LAUNCH_BLOCKING": "1",
        "TORCH_CUDA_MEMORY_FRACTION": "1.0",
        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    },
    cache="auto",
)

# Report generation — CPU only, reads training results and renders Union UI panels
report_env = flyte.TaskEnvironment(
    name="tumor_report",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)

# Pipeline driver — lightweight orchestrator that calls the three tasks above
pipeline_env = flyte.TaskEnvironment(
    name="tumor_pipeline",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[dataset_env, training_env, report_env],
)
# {{/docs-fragment envs}}

class TrainingConfig:
    """Unified training configuration for brain tumor MRI classification."""

    def __init__(
        self,
        image_size: int = 380,
        num_classes: int = 4,
        model_name: str = "efficientnet_b4",
        pretrained: bool = True,
        phase1_epochs: int = 8,
        phase1_lr: float = 1e-3,
        phase1_freeze_backbone: bool = True,
        phase2_epochs: int = 25,
        phase2_lr: float = 5e-5,
        batch_size: int = 16,
        num_workers: int = 0,
        val_split: float = 0.2,
        weight_decay: float = 1e-4,
        warmup_steps: int = 200,
        focal_gamma: float = 2.0,
        mixup_alpha: float = 0.0,
        log_interval: int = 50,
    ):
        self.image_size = image_size
        self.num_classes = num_classes

        self.model_name = model_name
        self.pretrained = pretrained

        self.phase1_epochs = phase1_epochs
        self.phase1_lr = phase1_lr
        self.phase1_freeze_backbone = phase1_freeze_backbone

        self.phase2_epochs = phase2_epochs
        self.phase2_lr = phase2_lr

        self.batch_size = batch_size
        self.num_workers = num_workers
        self.val_split = val_split

        self.weight_decay = weight_decay
        self.warmup_steps = warmup_steps

        self.focal_gamma = focal_gamma
        self.mixup_alpha = mixup_alpha

        self.log_interval = log_interval

    def to_dict(self) -> dict:
        return self.__dict__
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/config.py*

## Define the task environments

Each stage declares the resources it needs. The lightweight `pipeline_env` orchestrates the others via `depends_on`.

```
"""
Configuration for brain tumor MRI classification pipeline.

Defines task environments, resource requirements, and training hyperparameters.
"""

import pathlib

import flyte

# {{docs-fragment image}}
image = flyte.Image.from_debian_base(
    name="tumor_detection_gpu"
).with_pip_packages(
    "torch",
    "lightning",
    "torchvision",
    "timm",
    "pillow",
    "scikit-learn",
    "plotly",
    "numpy",
    "pandas",
    "torchmetrics",
    "datasets",
    "typing_extensions",
).with_source_folder(
    pathlib.Path(__file__).parent,
    copy_contents_only=True,
)
# {{/docs-fragment image}}

# {{docs-fragment envs}}
# Downloads raw MRI JPEG files — CPU only, no auth needed, result is cached
dataset_env = flyte.TaskEnvironment(
    name="tumor_dataset",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi", disk="8Gi"),
    cache="auto",
)

# GPU training — result is cached so re-running with the same data + config is free
training_env = flyte.TaskEnvironment(
    name="tumor_gpu_training",
    image=image,
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    env_vars={
        "CUDA_VISIBLE_DEVICES": "0",
        "CUDA_LAUNCH_BLOCKING": "1",
        "TORCH_CUDA_MEMORY_FRACTION": "1.0",
        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    },
    cache="auto",
)

# Report generation — CPU only, reads training results and renders Union UI panels
report_env = flyte.TaskEnvironment(
    name="tumor_report",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)

# Pipeline driver — lightweight orchestrator that calls the three tasks above
pipeline_env = flyte.TaskEnvironment(
    name="tumor_pipeline",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[dataset_env, training_env, report_env],
)
# {{/docs-fragment envs}}

class TrainingConfig:
    """Unified training configuration for brain tumor MRI classification."""

    def __init__(
        self,
        image_size: int = 380,
        num_classes: int = 4,
        model_name: str = "efficientnet_b4",
        pretrained: bool = True,
        phase1_epochs: int = 8,
        phase1_lr: float = 1e-3,
        phase1_freeze_backbone: bool = True,
        phase2_epochs: int = 25,
        phase2_lr: float = 5e-5,
        batch_size: int = 16,
        num_workers: int = 0,
        val_split: float = 0.2,
        weight_decay: float = 1e-4,
        warmup_steps: int = 200,
        focal_gamma: float = 2.0,
        mixup_alpha: float = 0.0,
        log_interval: int = 50,
    ):
        self.image_size = image_size
        self.num_classes = num_classes

        self.model_name = model_name
        self.pretrained = pretrained

        self.phase1_epochs = phase1_epochs
        self.phase1_lr = phase1_lr
        self.phase1_freeze_backbone = phase1_freeze_backbone

        self.phase2_epochs = phase2_epochs
        self.phase2_lr = phase2_lr

        self.batch_size = batch_size
        self.num_workers = num_workers
        self.val_split = val_split

        self.weight_decay = weight_decay
        self.warmup_steps = warmup_steps

        self.focal_gamma = focal_gamma
        self.mixup_alpha = mixup_alpha

        self.log_interval = log_interval

    def to_dict(self) -> dict:
        return self.__dict__
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/config.py*

## Configure training

Hyperparameters are gathered in a single `TrainingConfig`, serialized to JSON, and passed into the training task so the exact configuration is captured alongside the run.

```
"""
Flyte/Union pipeline for brain tumor MRI classification.

Three-task pipeline:
1. load_dataset  — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU)
2. train_model   — two-phase EfficientNet-B4 training with focal loss (GPU)
3. create_report — render training curves and confusion matrix in the Union UI (CPU)
"""

import json

import flyte
from flyte.io import Dir

from config import TrainingConfig, dataset_env, pipeline_env, report_env, training_env
from dataset import download_tumor_dataset

# {{docs-fragment config}}
TRAINING_CONFIG = TrainingConfig(
    phase1_epochs=8,
    phase2_epochs=25,
    phase1_lr=1e-3,
    phase2_lr=5e-5,
    batch_size=16,
    num_workers=0,
    log_interval=50,
    mixup_alpha=0.0,
    image_size=380,
    focal_gamma=3.0,
)
# {{/docs-fragment config}}

# {{docs-fragment load_dataset}}
@dataset_env.task
async def load_dataset() -> Dir:
    """
    Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir.
    Runs once — result is reused on subsequent pipeline runs (cache="auto").
    """
    return await download_tumor_dataset()
# {{/docs-fragment load_dataset}}

# {{docs-fragment train_model}}
@training_env.task(retries=3)
async def train_model(dataset_dir: Dir, config_json: str) -> Dir:
    """
    Download the raw dataset Dir, run two-phase training,
    and return training metrics and final predictions as a Dir for the report task.
    """
    from pathlib import Path

    local_dir = Path("/tmp/tumor_local")
    local_dir.mkdir(parents=True, exist_ok=True)
    await dataset_dir.download(local_path=str(local_dir))

    from training import train_tumor_classifier
    config = TrainingConfig(**json.loads(config_json))
    result = train_tumor_classifier(config=config, dataset_path=str(local_dir))

    output_dir = Path("/tmp/training_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(result["metrics"]))
    (output_dir / "predictions.json").write_text(json.dumps({
        "preds": result["final_preds"],
        "targets": result["final_targets"],
    }))

    return await Dir.from_local(str(output_dir))
# {{/docs-fragment train_model}}

# {{docs-fragment create_report}}
@report_env.task(report=True)
async def create_report(results_dir: Dir) -> None:
    """
    Download training metrics and render loss/accuracy curves, confusion matrix,
    and per-class F1 chart in the Union UI report panel.
    """
    import numpy as np
    from pathlib import Path

    from utils import create_confusion_matrix_plot, create_metrics_plots, create_per_class_f1_plot

    local_dir = Path("/tmp/tumor_report")
    local_dir.mkdir(parents=True, exist_ok=True)
    await results_dir.download(local_path=str(local_dir))

    matches = list(local_dir.glob("**/metrics.json"))
    if not matches:
        raise RuntimeError(f"metrics.json not found under {local_dir}")
    local_path = matches[0].parent

    history = json.loads((local_path / "metrics.json").read_text())
    predictions = json.loads((local_path / "predictions.json").read_text())

    preds = np.array(predictions["preds"])
    targets = np.array(predictions["targets"])

    loss_fig, acc_fig = create_metrics_plots(history)
    cm_fig = create_confusion_matrix_plot(preds, targets)
    f1_fig = create_per_class_f1_plot(preds, targets)

    combined_html = (
        acc_fig.to_html(include_plotlyjs=True, full_html=False)
        + loss_fig.to_html(include_plotlyjs=False, full_html=False)
        + cm_fig.to_html(include_plotlyjs=False, full_html=False)
        + f1_fig.to_html(include_plotlyjs=False, full_html=False)
    )
    flyte.report.log(combined_html, do_flush=True)
# {{/docs-fragment create_report}}

# {{docs-fragment pipeline}}
@pipeline_env.task
async def tumor_detection_pipeline() -> None:
    """Orchestrate dataset loading, GPU training, and report generation."""
    dataset_dir = await load_dataset()
    results_dir = await train_model(
        dataset_dir=dataset_dir,
        config_json=json.dumps(TRAINING_CONFIG.to_dict()),
    )
    await create_report(results_dir=results_dir)
# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    import pathlib
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext().run(tumor_detection_pipeline)
    print(f"\n✓ Pipeline submitted!")
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/run.py*

## Load the dataset

The first task downloads the public [Brain Tumor MRI dataset](https://huggingface.co/datasets/AIOmarRehan/Brain_Tumor_MRI_Dataset) from Hugging Face (no auth required) and stores it as a `flyte.io.Dir`. It's cached, so subsequent runs reuse it.

```
"""
Flyte/Union pipeline for brain tumor MRI classification.

Three-task pipeline:
1. load_dataset  — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU)
2. train_model   — two-phase EfficientNet-B4 training with focal loss (GPU)
3. create_report — render training curves and confusion matrix in the Union UI (CPU)
"""

import json

import flyte
from flyte.io import Dir

from config import TrainingConfig, dataset_env, pipeline_env, report_env, training_env
from dataset import download_tumor_dataset

# {{docs-fragment config}}
TRAINING_CONFIG = TrainingConfig(
    phase1_epochs=8,
    phase2_epochs=25,
    phase1_lr=1e-3,
    phase2_lr=5e-5,
    batch_size=16,
    num_workers=0,
    log_interval=50,
    mixup_alpha=0.0,
    image_size=380,
    focal_gamma=3.0,
)
# {{/docs-fragment config}}

# {{docs-fragment load_dataset}}
@dataset_env.task
async def load_dataset() -> Dir:
    """
    Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir.
    Runs once — result is reused on subsequent pipeline runs (cache="auto").
    """
    return await download_tumor_dataset()
# {{/docs-fragment load_dataset}}

# {{docs-fragment train_model}}
@training_env.task(retries=3)
async def train_model(dataset_dir: Dir, config_json: str) -> Dir:
    """
    Download the raw dataset Dir, run two-phase training,
    and return training metrics and final predictions as a Dir for the report task.
    """
    from pathlib import Path

    local_dir = Path("/tmp/tumor_local")
    local_dir.mkdir(parents=True, exist_ok=True)
    await dataset_dir.download(local_path=str(local_dir))

    from training import train_tumor_classifier
    config = TrainingConfig(**json.loads(config_json))
    result = train_tumor_classifier(config=config, dataset_path=str(local_dir))

    output_dir = Path("/tmp/training_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(result["metrics"]))
    (output_dir / "predictions.json").write_text(json.dumps({
        "preds": result["final_preds"],
        "targets": result["final_targets"],
    }))

    return await Dir.from_local(str(output_dir))
# {{/docs-fragment train_model}}

# {{docs-fragment create_report}}
@report_env.task(report=True)
async def create_report(results_dir: Dir) -> None:
    """
    Download training metrics and render loss/accuracy curves, confusion matrix,
    and per-class F1 chart in the Union UI report panel.
    """
    import numpy as np
    from pathlib import Path

    from utils import create_confusion_matrix_plot, create_metrics_plots, create_per_class_f1_plot

    local_dir = Path("/tmp/tumor_report")
    local_dir.mkdir(parents=True, exist_ok=True)
    await results_dir.download(local_path=str(local_dir))

    matches = list(local_dir.glob("**/metrics.json"))
    if not matches:
        raise RuntimeError(f"metrics.json not found under {local_dir}")
    local_path = matches[0].parent

    history = json.loads((local_path / "metrics.json").read_text())
    predictions = json.loads((local_path / "predictions.json").read_text())

    preds = np.array(predictions["preds"])
    targets = np.array(predictions["targets"])

    loss_fig, acc_fig = create_metrics_plots(history)
    cm_fig = create_confusion_matrix_plot(preds, targets)
    f1_fig = create_per_class_f1_plot(preds, targets)

    combined_html = (
        acc_fig.to_html(include_plotlyjs=True, full_html=False)
        + loss_fig.to_html(include_plotlyjs=False, full_html=False)
        + cm_fig.to_html(include_plotlyjs=False, full_html=False)
        + f1_fig.to_html(include_plotlyjs=False, full_html=False)
    )
    flyte.report.log(combined_html, do_flush=True)
# {{/docs-fragment create_report}}

# {{docs-fragment pipeline}}
@pipeline_env.task
async def tumor_detection_pipeline() -> None:
    """Orchestrate dataset loading, GPU training, and report generation."""
    dataset_dir = await load_dataset()
    results_dir = await train_model(
        dataset_dir=dataset_dir,
        config_json=json.dumps(TRAINING_CONFIG.to_dict()),
    )
    await create_report(results_dir=results_dir)
# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    import pathlib
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext().run(tumor_detection_pipeline)
    print(f"\n✓ Pipeline submitted!")
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/run.py*

## Train the model

The training task downloads the dataset `Dir`, runs two-phase training (frozen backbone, then full fine-tuning), and writes metrics and predictions to an output `Dir`. It sets `retries=3` so a preempted GPU node restarts the task.

```
"""
Flyte/Union pipeline for brain tumor MRI classification.

Three-task pipeline:
1. load_dataset  — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU)
2. train_model   — two-phase EfficientNet-B4 training with focal loss (GPU)
3. create_report — render training curves and confusion matrix in the Union UI (CPU)
"""

import json

import flyte
from flyte.io import Dir

from config import TrainingConfig, dataset_env, pipeline_env, report_env, training_env
from dataset import download_tumor_dataset

# {{docs-fragment config}}
TRAINING_CONFIG = TrainingConfig(
    phase1_epochs=8,
    phase2_epochs=25,
    phase1_lr=1e-3,
    phase2_lr=5e-5,
    batch_size=16,
    num_workers=0,
    log_interval=50,
    mixup_alpha=0.0,
    image_size=380,
    focal_gamma=3.0,
)
# {{/docs-fragment config}}

# {{docs-fragment load_dataset}}
@dataset_env.task
async def load_dataset() -> Dir:
    """
    Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir.
    Runs once — result is reused on subsequent pipeline runs (cache="auto").
    """
    return await download_tumor_dataset()
# {{/docs-fragment load_dataset}}

# {{docs-fragment train_model}}
@training_env.task(retries=3)
async def train_model(dataset_dir: Dir, config_json: str) -> Dir:
    """
    Download the raw dataset Dir, run two-phase training,
    and return training metrics and final predictions as a Dir for the report task.
    """
    from pathlib import Path

    local_dir = Path("/tmp/tumor_local")
    local_dir.mkdir(parents=True, exist_ok=True)
    await dataset_dir.download(local_path=str(local_dir))

    from training import train_tumor_classifier
    config = TrainingConfig(**json.loads(config_json))
    result = train_tumor_classifier(config=config, dataset_path=str(local_dir))

    output_dir = Path("/tmp/training_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(result["metrics"]))
    (output_dir / "predictions.json").write_text(json.dumps({
        "preds": result["final_preds"],
        "targets": result["final_targets"],
    }))

    return await Dir.from_local(str(output_dir))
# {{/docs-fragment train_model}}

# {{docs-fragment create_report}}
@report_env.task(report=True)
async def create_report(results_dir: Dir) -> None:
    """
    Download training metrics and render loss/accuracy curves, confusion matrix,
    and per-class F1 chart in the Union UI report panel.
    """
    import numpy as np
    from pathlib import Path

    from utils import create_confusion_matrix_plot, create_metrics_plots, create_per_class_f1_plot

    local_dir = Path("/tmp/tumor_report")
    local_dir.mkdir(parents=True, exist_ok=True)
    await results_dir.download(local_path=str(local_dir))

    matches = list(local_dir.glob("**/metrics.json"))
    if not matches:
        raise RuntimeError(f"metrics.json not found under {local_dir}")
    local_path = matches[0].parent

    history = json.loads((local_path / "metrics.json").read_text())
    predictions = json.loads((local_path / "predictions.json").read_text())

    preds = np.array(predictions["preds"])
    targets = np.array(predictions["targets"])

    loss_fig, acc_fig = create_metrics_plots(history)
    cm_fig = create_confusion_matrix_plot(preds, targets)
    f1_fig = create_per_class_f1_plot(preds, targets)

    combined_html = (
        acc_fig.to_html(include_plotlyjs=True, full_html=False)
        + loss_fig.to_html(include_plotlyjs=False, full_html=False)
        + cm_fig.to_html(include_plotlyjs=False, full_html=False)
        + f1_fig.to_html(include_plotlyjs=False, full_html=False)
    )
    flyte.report.log(combined_html, do_flush=True)
# {{/docs-fragment create_report}}

# {{docs-fragment pipeline}}
@pipeline_env.task
async def tumor_detection_pipeline() -> None:
    """Orchestrate dataset loading, GPU training, and report generation."""
    dataset_dir = await load_dataset()
    results_dir = await train_model(
        dataset_dir=dataset_dir,
        config_json=json.dumps(TRAINING_CONFIG.to_dict()),
    )
    await create_report(results_dir=results_dir)
# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    import pathlib
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext().run(tumor_detection_pipeline)
    print(f"\n✓ Pipeline submitted!")
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/run.py*

### Resumable checkpointing

To make retries cheap, training mirrors its Lightning checkpoint directory to a `flyte.Checkpoint` after every epoch, and resumes from the latest checkpoint when the task restarts.

```
"""
Training pipeline for brain tumor MRI classification.

Implements two-phase training:
- Phase 1: Frozen backbone (feature extractor), train classification head
- Phase 2: Fine-tune full model with differential LRs + cosine annealing
"""

from config import TrainingConfig
from dataset import compute_class_weights, create_data_loaders
from model import TumorClassifierLightningModule
from utils import get_model_size, get_trainable_params

def train_tumor_classifier(
    config: TrainingConfig,
    dataset_path: str,
) -> dict:
    """
    Run two-phase training on the preprocessed dataset and return metrics + final predictions.

    dataset_path: local directory where the flyte.io.Dir was downloaded by the training task.
    """
    import pathlib

    import flyte
    import lightning as L
    import torch
    from lightning.pytorch.callbacks import ModelCheckpoint
    from typing_extensions import override

    # {{docs-fragment flyte_checkpoint}}
    class FlyteLightningCheckpointCallback(ModelCheckpoint):
        """Mirrors the checkpoint directory to Flyte after every epoch so retries can resume."""

        def __init__(self, flyte_checkpoint: flyte.Checkpoint, *, dirpath: str, **kwargs):
            super().__init__(dirpath=dirpath, **kwargs)
            self._flyte_checkpoint = flyte_checkpoint

        @override
        def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
            super().on_train_epoch_end(trainer, pl_module)
            if self.dirpath:
                self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))
    # {{/docs-fragment flyte_checkpoint}}

    class MetricsLoggerCallback(L.Callback):
        def __init__(self, phase1_epochs: int):
            super().__init__()
            self.phase1_epochs = phase1_epochs
            self.history = []

        def on_validation_epoch_end(self, trainer, _pl_module):
            epoch = trainer.current_epoch
            metrics = trainer.callback_metrics
            self.history.append({
                "epoch": epoch,
                "phase": 1 if epoch < self.phase1_epochs else 2,
                "train_loss": float(metrics.get("train/loss_epoch", 0)),
                "val_loss": float(metrics.get("val/loss", 0)),
                "val_acc": float(metrics.get("val/acc", 0)),
                "macro_f1": float(metrics.get("val/macro_f1", 0)),
            })

    class PhaseChangeCallback(L.Callback):
        def __init__(self, phase1_epochs: int, phase2_lr: float):
            super().__init__()
            self.phase1_epochs = phase1_epochs
            self.phase2_lr = phase2_lr
            self.phase_changed = False

        def on_train_epoch_end(self, trainer, pl_module):
            if not self.phase_changed and (trainer.current_epoch + 1) == self.phase1_epochs:
                print("\n" + "=" * 80)
                print("TRANSITIONING TO PHASE 2: UNFREEZING BACKBONE AND ADJUSTING LR")
                print("=" * 80 + "\n")

                pl_module.model.unfreeze_backbone()

                for param_group in trainer.optimizers[0].param_groups:
                    param_group["lr"] = self.phase2_lr

                # Add backbone params to optimizer with 10x lower LR.
                # Backbone was excluded at init because it was frozen.
                backbone_lr = self.phase2_lr * 0.1
                backbone_decay, backbone_no_decay = [], []
                for param in pl_module.model.backbone.parameters():
                    if param.ndim >= 2:
                        backbone_decay.append(param)
                    else:
                        backbone_no_decay.append(param)
                optimizer = trainer.optimizers[0]
                optimizer.add_param_group({"params": backbone_decay, "lr": backbone_lr, "weight_decay": pl_module.weight_decay})
                optimizer.add_param_group({"params": backbone_no_decay, "lr": backbone_lr, "weight_decay": 0.0})

                # Fresh cosine schedule over remaining Phase 2 steps to avoid
                # the Phase 1 schedule arriving near-zero before Phase 2 begins.
                steps_remaining = trainer.estimated_stepping_batches - trainer.global_step
                new_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                    trainer.optimizers[0],
                    T_max=max(1, steps_remaining),
                    eta_min=1e-6,
                )
                for lr_scheduler_config in trainer.lr_scheduler_configs:
                    lr_scheduler_config.scheduler = new_scheduler

                print(f"Phase 2 started: lr={self.phase2_lr}")
                print(f"Total parameters: {get_model_size(pl_module.model):,}")
                print(f"Trainable parameters: {get_trainable_params(pl_module.model):,}")
                self.phase_changed = True

    print("\n" + "=" * 80)
    print("BRAIN TUMOR MRI CLASSIFICATION WITH EFFICIENTNET-B4")
    print("=" * 80)
    print(f"Config: {config.to_dict()}\n")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

    print("\nLoading MRI images...")
    train_loader, val_loader = create_data_loaders(
        dataset_path=dataset_path,
        image_size=config.image_size,
        batch_size=config.batch_size,
        num_workers=config.num_workers,
        val_split=config.val_split,
    )
    print(f"Data loaders created: {len(train_loader)} train batches, {len(val_loader)} val batches")

    print("\nComputing class weights for focal loss...")
    class_weights = compute_class_weights(dataset_path)
    print(f"Class weights: {class_weights.tolist()}")

    # Per-class gamma: Meningioma gets 7.0, all others 3.0.
    # CLASS_NAMES alphabetical order: Glioma=0, Meningioma=1, No Tumor=2, Pituitary=3
    gamma_per_class = torch.tensor([3.0, 7.0, 3.0, 3.0])

    print("\nInitializing model...")
    model = TumorClassifierLightningModule(
        num_classes=config.num_classes,
        model_name=config.model_name,
        pretrained=config.pretrained,
        learning_rate=config.phase1_lr,
        freeze_backbone=config.phase1_freeze_backbone,
        weight_decay=config.weight_decay,
        warmup_steps=config.warmup_steps,
        max_epochs=config.phase1_epochs + config.phase2_epochs,
        focal_gamma=config.focal_gamma,
        mixup_alpha=config.mixup_alpha,
        class_weights=class_weights,
        gamma_per_class=gamma_per_class,
    )

    print(f"Model: {config.model_name}")
    print(f"Total parameters: {get_model_size(model.model):,}")
    print(f"Trainable parameters: {get_trainable_params(model.model):,}")

    from pathlib import Path
    checkpoint_dir = Path("/tmp/tumor_checkpoints")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    # {{docs-fragment resume}}
    # --- Flyte checkpoint: resume from previous attempt if one exists ---
    resume_ckpt: str | None = None
    ctx = flyte.ctx()
    flyte_checkpoint = getattr(ctx, "checkpoint", None) if ctx else None

    if flyte_checkpoint:
        prev_path = flyte_checkpoint.load_sync()
        if prev_path:
            last = flyte.latest_checkpoint(prev_path)
            if last:
                ck = torch.load(str(last), map_location="cpu", weights_only=False)
                epoch_start = int(ck.get("epoch", 0))
                resume_ckpt = str(last)
                print(f"Resuming from epoch {epoch_start}, checkpoint: {last}")
    # --------------------------------------------------------------------
    # {{/docs-fragment resume}}

    metrics_cb = MetricsLoggerCallback(phase1_epochs=config.phase1_epochs)

    resume_callback = (
        FlyteLightningCheckpointCallback(
            flyte_checkpoint,
            dirpath=str(checkpoint_dir),
            filename="last",
            save_last=True,
            save_top_k=1,
        )
        if flyte_checkpoint else
        ModelCheckpoint(
            dirpath=str(checkpoint_dir),
            filename="best-{epoch:03d}-{val_acc:.3f}",
            monitor="val/acc",
            mode="max",
            save_top_k=3,
            verbose=True,
            auto_insert_metric_name=False,
        )
    )

    callbacks = [
        resume_callback,
        metrics_cb,
        PhaseChangeCallback(
            phase1_epochs=config.phase1_epochs,
            phase2_lr=config.phase2_lr,
        ),
    ]

    trainer = L.Trainer(
        max_epochs=config.phase1_epochs + config.phase2_epochs,
        accelerator="gpu" if torch.cuda.is_available() else "cpu",
        devices=1,
        precision="16-mixed",
        callbacks=callbacks,
        enable_progress_bar=True,
        enable_model_summary=True,
        log_every_n_steps=config.log_interval,
        gradient_clip_val=1.0,
    )

    trainer.fit(model, train_loader, val_loader, ckpt_path=resume_ckpt)

    best_checkpoint = trainer.checkpoint_callback.best_model_path
    print(f"\n✓ Training complete!")
    print(f"Best checkpoint: {best_checkpoint}")

    # Final inference with TTA (test-time augmentation): average logits over
    # original + h-flip + v-flip + 90° rotations for a free accuracy boost.
    print("\nRunning final inference with TTA for confusion matrix...")
    import numpy as np
    import torchvision.transforms.functional as TF
    model.eval()
    model.to(device)
    all_preds, all_targets = [], []
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            aug_logits = [
                model.model(images),
                model.model(TF.hflip(images)),
                model.model(TF.vflip(images)),
                model.model(torch.rot90(images, k=1, dims=[2, 3])),
                model.model(torch.rot90(images, k=3, dims=[2, 3])),
            ]
            avg_logits = torch.stack(aug_logits).mean(dim=0)
            all_preds.append(avg_logits.argmax(dim=1).cpu())
            all_targets.append(labels.cpu())
    final_preds = torch.cat(all_preds).numpy()
    final_targets = torch.cat(all_targets).numpy()

    return {
        "best_checkpoint": best_checkpoint,
        "metrics": metrics_cb.history,
        "final_preds": final_preds.tolist(),
        "final_targets": final_targets.tolist(),
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/training.py*

On startup, the training loop looks for a checkpoint from a previous attempt and resumes from it if present:

```
"""
Training pipeline for brain tumor MRI classification.

Implements two-phase training:
- Phase 1: Frozen backbone (feature extractor), train classification head
- Phase 2: Fine-tune full model with differential LRs + cosine annealing
"""

from config import TrainingConfig
from dataset import compute_class_weights, create_data_loaders
from model import TumorClassifierLightningModule
from utils import get_model_size, get_trainable_params

def train_tumor_classifier(
    config: TrainingConfig,
    dataset_path: str,
) -> dict:
    """
    Run two-phase training on the preprocessed dataset and return metrics + final predictions.

    dataset_path: local directory where the flyte.io.Dir was downloaded by the training task.
    """
    import pathlib

    import flyte
    import lightning as L
    import torch
    from lightning.pytorch.callbacks import ModelCheckpoint
    from typing_extensions import override

    # {{docs-fragment flyte_checkpoint}}
    class FlyteLightningCheckpointCallback(ModelCheckpoint):
        """Mirrors the checkpoint directory to Flyte after every epoch so retries can resume."""

        def __init__(self, flyte_checkpoint: flyte.Checkpoint, *, dirpath: str, **kwargs):
            super().__init__(dirpath=dirpath, **kwargs)
            self._flyte_checkpoint = flyte_checkpoint

        @override
        def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
            super().on_train_epoch_end(trainer, pl_module)
            if self.dirpath:
                self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))
    # {{/docs-fragment flyte_checkpoint}}

    class MetricsLoggerCallback(L.Callback):
        def __init__(self, phase1_epochs: int):
            super().__init__()
            self.phase1_epochs = phase1_epochs
            self.history = []

        def on_validation_epoch_end(self, trainer, _pl_module):
            epoch = trainer.current_epoch
            metrics = trainer.callback_metrics
            self.history.append({
                "epoch": epoch,
                "phase": 1 if epoch < self.phase1_epochs else 2,
                "train_loss": float(metrics.get("train/loss_epoch", 0)),
                "val_loss": float(metrics.get("val/loss", 0)),
                "val_acc": float(metrics.get("val/acc", 0)),
                "macro_f1": float(metrics.get("val/macro_f1", 0)),
            })

    class PhaseChangeCallback(L.Callback):
        def __init__(self, phase1_epochs: int, phase2_lr: float):
            super().__init__()
            self.phase1_epochs = phase1_epochs
            self.phase2_lr = phase2_lr
            self.phase_changed = False

        def on_train_epoch_end(self, trainer, pl_module):
            if not self.phase_changed and (trainer.current_epoch + 1) == self.phase1_epochs:
                print("\n" + "=" * 80)
                print("TRANSITIONING TO PHASE 2: UNFREEZING BACKBONE AND ADJUSTING LR")
                print("=" * 80 + "\n")

                pl_module.model.unfreeze_backbone()

                for param_group in trainer.optimizers[0].param_groups:
                    param_group["lr"] = self.phase2_lr

                # Add backbone params to optimizer with 10x lower LR.
                # Backbone was excluded at init because it was frozen.
                backbone_lr = self.phase2_lr * 0.1
                backbone_decay, backbone_no_decay = [], []
                for param in pl_module.model.backbone.parameters():
                    if param.ndim >= 2:
                        backbone_decay.append(param)
                    else:
                        backbone_no_decay.append(param)
                optimizer = trainer.optimizers[0]
                optimizer.add_param_group({"params": backbone_decay, "lr": backbone_lr, "weight_decay": pl_module.weight_decay})
                optimizer.add_param_group({"params": backbone_no_decay, "lr": backbone_lr, "weight_decay": 0.0})

                # Fresh cosine schedule over remaining Phase 2 steps to avoid
                # the Phase 1 schedule arriving near-zero before Phase 2 begins.
                steps_remaining = trainer.estimated_stepping_batches - trainer.global_step
                new_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                    trainer.optimizers[0],
                    T_max=max(1, steps_remaining),
                    eta_min=1e-6,
                )
                for lr_scheduler_config in trainer.lr_scheduler_configs:
                    lr_scheduler_config.scheduler = new_scheduler

                print(f"Phase 2 started: lr={self.phase2_lr}")
                print(f"Total parameters: {get_model_size(pl_module.model):,}")
                print(f"Trainable parameters: {get_trainable_params(pl_module.model):,}")
                self.phase_changed = True

    print("\n" + "=" * 80)
    print("BRAIN TUMOR MRI CLASSIFICATION WITH EFFICIENTNET-B4")
    print("=" * 80)
    print(f"Config: {config.to_dict()}\n")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

    print("\nLoading MRI images...")
    train_loader, val_loader = create_data_loaders(
        dataset_path=dataset_path,
        image_size=config.image_size,
        batch_size=config.batch_size,
        num_workers=config.num_workers,
        val_split=config.val_split,
    )
    print(f"Data loaders created: {len(train_loader)} train batches, {len(val_loader)} val batches")

    print("\nComputing class weights for focal loss...")
    class_weights = compute_class_weights(dataset_path)
    print(f"Class weights: {class_weights.tolist()}")

    # Per-class gamma: Meningioma gets 7.0, all others 3.0.
    # CLASS_NAMES alphabetical order: Glioma=0, Meningioma=1, No Tumor=2, Pituitary=3
    gamma_per_class = torch.tensor([3.0, 7.0, 3.0, 3.0])

    print("\nInitializing model...")
    model = TumorClassifierLightningModule(
        num_classes=config.num_classes,
        model_name=config.model_name,
        pretrained=config.pretrained,
        learning_rate=config.phase1_lr,
        freeze_backbone=config.phase1_freeze_backbone,
        weight_decay=config.weight_decay,
        warmup_steps=config.warmup_steps,
        max_epochs=config.phase1_epochs + config.phase2_epochs,
        focal_gamma=config.focal_gamma,
        mixup_alpha=config.mixup_alpha,
        class_weights=class_weights,
        gamma_per_class=gamma_per_class,
    )

    print(f"Model: {config.model_name}")
    print(f"Total parameters: {get_model_size(model.model):,}")
    print(f"Trainable parameters: {get_trainable_params(model.model):,}")

    from pathlib import Path
    checkpoint_dir = Path("/tmp/tumor_checkpoints")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    # {{docs-fragment resume}}
    # --- Flyte checkpoint: resume from previous attempt if one exists ---
    resume_ckpt: str | None = None
    ctx = flyte.ctx()
    flyte_checkpoint = getattr(ctx, "checkpoint", None) if ctx else None

    if flyte_checkpoint:
        prev_path = flyte_checkpoint.load_sync()
        if prev_path:
            last = flyte.latest_checkpoint(prev_path)
            if last:
                ck = torch.load(str(last), map_location="cpu", weights_only=False)
                epoch_start = int(ck.get("epoch", 0))
                resume_ckpt = str(last)
                print(f"Resuming from epoch {epoch_start}, checkpoint: {last}")
    # --------------------------------------------------------------------
    # {{/docs-fragment resume}}

    metrics_cb = MetricsLoggerCallback(phase1_epochs=config.phase1_epochs)

    resume_callback = (
        FlyteLightningCheckpointCallback(
            flyte_checkpoint,
            dirpath=str(checkpoint_dir),
            filename="last",
            save_last=True,
            save_top_k=1,
        )
        if flyte_checkpoint else
        ModelCheckpoint(
            dirpath=str(checkpoint_dir),
            filename="best-{epoch:03d}-{val_acc:.3f}",
            monitor="val/acc",
            mode="max",
            save_top_k=3,
            verbose=True,
            auto_insert_metric_name=False,
        )
    )

    callbacks = [
        resume_callback,
        metrics_cb,
        PhaseChangeCallback(
            phase1_epochs=config.phase1_epochs,
            phase2_lr=config.phase2_lr,
        ),
    ]

    trainer = L.Trainer(
        max_epochs=config.phase1_epochs + config.phase2_epochs,
        accelerator="gpu" if torch.cuda.is_available() else "cpu",
        devices=1,
        precision="16-mixed",
        callbacks=callbacks,
        enable_progress_bar=True,
        enable_model_summary=True,
        log_every_n_steps=config.log_interval,
        gradient_clip_val=1.0,
    )

    trainer.fit(model, train_loader, val_loader, ckpt_path=resume_ckpt)

    best_checkpoint = trainer.checkpoint_callback.best_model_path
    print(f"\n✓ Training complete!")
    print(f"Best checkpoint: {best_checkpoint}")

    # Final inference with TTA (test-time augmentation): average logits over
    # original + h-flip + v-flip + 90° rotations for a free accuracy boost.
    print("\nRunning final inference with TTA for confusion matrix...")
    import numpy as np
    import torchvision.transforms.functional as TF
    model.eval()
    model.to(device)
    all_preds, all_targets = [], []
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            aug_logits = [
                model.model(images),
                model.model(TF.hflip(images)),
                model.model(TF.vflip(images)),
                model.model(torch.rot90(images, k=1, dims=[2, 3])),
                model.model(torch.rot90(images, k=3, dims=[2, 3])),
            ]
            avg_logits = torch.stack(aug_logits).mean(dim=0)
            all_preds.append(avg_logits.argmax(dim=1).cpu())
            all_targets.append(labels.cpu())
    final_preds = torch.cat(all_preds).numpy()
    final_targets = torch.cat(all_targets).numpy()

    return {
        "best_checkpoint": best_checkpoint,
        "metrics": metrics_cb.history,
        "final_preds": final_preds.tolist(),
        "final_targets": final_targets.tolist(),
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/training.py*

## Generate the report

The reporting task reads the metrics and predictions, then renders accuracy/loss curves, a confusion matrix, and a per-class F1 chart with Plotly. `report=True` surfaces the HTML directly in the run's report panel.

```
"""
Flyte/Union pipeline for brain tumor MRI classification.

Three-task pipeline:
1. load_dataset  — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU)
2. train_model   — two-phase EfficientNet-B4 training with focal loss (GPU)
3. create_report — render training curves and confusion matrix in the Union UI (CPU)
"""

import json

import flyte
from flyte.io import Dir

from config import TrainingConfig, dataset_env, pipeline_env, report_env, training_env
from dataset import download_tumor_dataset

# {{docs-fragment config}}
TRAINING_CONFIG = TrainingConfig(
    phase1_epochs=8,
    phase2_epochs=25,
    phase1_lr=1e-3,
    phase2_lr=5e-5,
    batch_size=16,
    num_workers=0,
    log_interval=50,
    mixup_alpha=0.0,
    image_size=380,
    focal_gamma=3.0,
)
# {{/docs-fragment config}}

# {{docs-fragment load_dataset}}
@dataset_env.task
async def load_dataset() -> Dir:
    """
    Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir.
    Runs once — result is reused on subsequent pipeline runs (cache="auto").
    """
    return await download_tumor_dataset()
# {{/docs-fragment load_dataset}}

# {{docs-fragment train_model}}
@training_env.task(retries=3)
async def train_model(dataset_dir: Dir, config_json: str) -> Dir:
    """
    Download the raw dataset Dir, run two-phase training,
    and return training metrics and final predictions as a Dir for the report task.
    """
    from pathlib import Path

    local_dir = Path("/tmp/tumor_local")
    local_dir.mkdir(parents=True, exist_ok=True)
    await dataset_dir.download(local_path=str(local_dir))

    from training import train_tumor_classifier
    config = TrainingConfig(**json.loads(config_json))
    result = train_tumor_classifier(config=config, dataset_path=str(local_dir))

    output_dir = Path("/tmp/training_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(result["metrics"]))
    (output_dir / "predictions.json").write_text(json.dumps({
        "preds": result["final_preds"],
        "targets": result["final_targets"],
    }))

    return await Dir.from_local(str(output_dir))
# {{/docs-fragment train_model}}

# {{docs-fragment create_report}}
@report_env.task(report=True)
async def create_report(results_dir: Dir) -> None:
    """
    Download training metrics and render loss/accuracy curves, confusion matrix,
    and per-class F1 chart in the Union UI report panel.
    """
    import numpy as np
    from pathlib import Path

    from utils import create_confusion_matrix_plot, create_metrics_plots, create_per_class_f1_plot

    local_dir = Path("/tmp/tumor_report")
    local_dir.mkdir(parents=True, exist_ok=True)
    await results_dir.download(local_path=str(local_dir))

    matches = list(local_dir.glob("**/metrics.json"))
    if not matches:
        raise RuntimeError(f"metrics.json not found under {local_dir}")
    local_path = matches[0].parent

    history = json.loads((local_path / "metrics.json").read_text())
    predictions = json.loads((local_path / "predictions.json").read_text())

    preds = np.array(predictions["preds"])
    targets = np.array(predictions["targets"])

    loss_fig, acc_fig = create_metrics_plots(history)
    cm_fig = create_confusion_matrix_plot(preds, targets)
    f1_fig = create_per_class_f1_plot(preds, targets)

    combined_html = (
        acc_fig.to_html(include_plotlyjs=True, full_html=False)
        + loss_fig.to_html(include_plotlyjs=False, full_html=False)
        + cm_fig.to_html(include_plotlyjs=False, full_html=False)
        + f1_fig.to_html(include_plotlyjs=False, full_html=False)
    )
    flyte.report.log(combined_html, do_flush=True)
# {{/docs-fragment create_report}}

# {{docs-fragment pipeline}}
@pipeline_env.task
async def tumor_detection_pipeline() -> None:
    """Orchestrate dataset loading, GPU training, and report generation."""
    dataset_dir = await load_dataset()
    results_dir = await train_model(
        dataset_dir=dataset_dir,
        config_json=json.dumps(TRAINING_CONFIG.to_dict()),
    )
    await create_report(results_dir=results_dir)
# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    import pathlib
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext().run(tumor_detection_pipeline)
    print(f"\n✓ Pipeline submitted!")
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/run.py*

## Orchestrate the pipeline

The driver task wires the three steps together.

```
"""
Flyte/Union pipeline for brain tumor MRI classification.

Three-task pipeline:
1. load_dataset  — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU)
2. train_model   — two-phase EfficientNet-B4 training with focal loss (GPU)
3. create_report — render training curves and confusion matrix in the Union UI (CPU)
"""

import json

import flyte
from flyte.io import Dir

from config import TrainingConfig, dataset_env, pipeline_env, report_env, training_env
from dataset import download_tumor_dataset

# {{docs-fragment config}}
TRAINING_CONFIG = TrainingConfig(
    phase1_epochs=8,
    phase2_epochs=25,
    phase1_lr=1e-3,
    phase2_lr=5e-5,
    batch_size=16,
    num_workers=0,
    log_interval=50,
    mixup_alpha=0.0,
    image_size=380,
    focal_gamma=3.0,
)
# {{/docs-fragment config}}

# {{docs-fragment load_dataset}}
@dataset_env.task
async def load_dataset() -> Dir:
    """
    Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir.
    Runs once — result is reused on subsequent pipeline runs (cache="auto").
    """
    return await download_tumor_dataset()
# {{/docs-fragment load_dataset}}

# {{docs-fragment train_model}}
@training_env.task(retries=3)
async def train_model(dataset_dir: Dir, config_json: str) -> Dir:
    """
    Download the raw dataset Dir, run two-phase training,
    and return training metrics and final predictions as a Dir for the report task.
    """
    from pathlib import Path

    local_dir = Path("/tmp/tumor_local")
    local_dir.mkdir(parents=True, exist_ok=True)
    await dataset_dir.download(local_path=str(local_dir))

    from training import train_tumor_classifier
    config = TrainingConfig(**json.loads(config_json))
    result = train_tumor_classifier(config=config, dataset_path=str(local_dir))

    output_dir = Path("/tmp/training_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(result["metrics"]))
    (output_dir / "predictions.json").write_text(json.dumps({
        "preds": result["final_preds"],
        "targets": result["final_targets"],
    }))

    return await Dir.from_local(str(output_dir))
# {{/docs-fragment train_model}}

# {{docs-fragment create_report}}
@report_env.task(report=True)
async def create_report(results_dir: Dir) -> None:
    """
    Download training metrics and render loss/accuracy curves, confusion matrix,
    and per-class F1 chart in the Union UI report panel.
    """
    import numpy as np
    from pathlib import Path

    from utils import create_confusion_matrix_plot, create_metrics_plots, create_per_class_f1_plot

    local_dir = Path("/tmp/tumor_report")
    local_dir.mkdir(parents=True, exist_ok=True)
    await results_dir.download(local_path=str(local_dir))

    matches = list(local_dir.glob("**/metrics.json"))
    if not matches:
        raise RuntimeError(f"metrics.json not found under {local_dir}")
    local_path = matches[0].parent

    history = json.loads((local_path / "metrics.json").read_text())
    predictions = json.loads((local_path / "predictions.json").read_text())

    preds = np.array(predictions["preds"])
    targets = np.array(predictions["targets"])

    loss_fig, acc_fig = create_metrics_plots(history)
    cm_fig = create_confusion_matrix_plot(preds, targets)
    f1_fig = create_per_class_f1_plot(preds, targets)

    combined_html = (
        acc_fig.to_html(include_plotlyjs=True, full_html=False)
        + loss_fig.to_html(include_plotlyjs=False, full_html=False)
        + cm_fig.to_html(include_plotlyjs=False, full_html=False)
        + f1_fig.to_html(include_plotlyjs=False, full_html=False)
    )
    flyte.report.log(combined_html, do_flush=True)
# {{/docs-fragment create_report}}

# {{docs-fragment pipeline}}
@pipeline_env.task
async def tumor_detection_pipeline() -> None:
    """Orchestrate dataset loading, GPU training, and report generation."""
    dataset_dir = await load_dataset()
    results_dir = await train_model(
        dataset_dir=dataset_dir,
        config_json=json.dumps(TRAINING_CONFIG.to_dict()),
    )
    await create_report(results_dir=results_dir)
# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    import pathlib
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    run = flyte.with_runcontext().run(tumor_detection_pipeline)
    print(f"\n✓ Pipeline submitted!")
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor_detection/run.py*

## Run the pipeline

This example has no secrets: the dataset is public. Because the pipeline imports sibling modules and uses `with_source_folder`, run it from inside the example directory so the local files are picked up.

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/tumor_detection):

```
cd v2/tutorials/tumor_detection
python run.py
```

Or submit it with the Flyte CLI from the same directory:

```
flyte run run.py tumor_detection_pipeline
```

When the run completes, open the `create_report` task in the UI to view the training curves, confusion matrix, and per-class F1 scores.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/genomic-gene-comparison ===

# Cross-species gene comparison

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic_gene_comparison).

This tutorial builds a bioinformatics pipeline that compares homologous genes across species. The workflow loads curated gene sequences (insulin, hemoglobin, or p53 by default), scores each sequence with the [Carbon](https://huggingface.co/HuggingFaceBio/Carbon-3B) genomic language model, aligns DNA and translated protein sequences, folds proteins with [ESMFold](https://github.com/facebookresearch/esm), and renders interactive HTML reports with identity heatmaps, phylogenetic trees, and 3D structure viewers.

Flyte makes the multi-stage GPU/CPU pipeline reliable:

- **Separate CPU and GPU `TaskEnvironment`s** so alignment runs on modest CPU boxes while Carbon scoring and ESMFold run on GPUs.
- **`report=True`** on every stage for live HTML progress and final summaries in the Flyte UI.
- **Cached data loading** and orchestrated fan-out across pipeline stages.

## Define the task environments

GPU tasks handle Carbon log-likelihood scoring and ESMFold structure prediction; CPU tasks load gene sets, run Needleman-Wunsch alignments, and generate the final summary.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="genomic-gene-comparison", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="genomic-gene-comparison-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="32Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="genomic-gene-comparison-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Homologous gene sets - same gene across species
# ------------------------------------------------------------------
# Full-length coding sequences from NCBI RefSeq (stop codon excluded).

GENE_SETS = {
    "insulin": {
        "gene_name": "Insulin",
        "description": "Insulin regulates blood sugar in all vertebrates. Highly conserved across 500M+ years of evolution - even fish insulin can lower blood sugar in humans. Comparing across species reveals which regions are functionally essential (conserved) vs free to drift.",
        "sequences": {
            "Human": {
                "dna": "ATGGCCCTGTGGATGCGCCTCCTGCCCCTGCTGGCGCTGCTGGCCCTCTGGGGACCTGACCCAGCCGCAGCCTTTGTGAACCAACACCTGTGCGGCTCACACCTGGTGGAAGCTCTCTACCTAGTGTGCGGGGAACGAGGCTTCTTCTACACACCCAAGACCCGCCGGGAGGCAGAGGACCTGCAGGTGGGGCAGGTGGAGCTGGGCGGGGGCCCTGGTGCAGGCAGCCTGCAGCCCTTGGCCCTGGAGGGGTCCCTGCAGAAGCGTGGCATTGTGGAACAATGCTGTACCAGCATCTGCTCCCTCTACCAGCTGGAGAACTACTGCAAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGGCCCTGTGGATGCGCTTCCTGCCCCTGCTGGCCCTGCTCTTCCTCTGGGAGTCCCACCCCACCCAGGCTTTTGTCAAGCAGCACCTTTGTGGTTCCCACCTGGTGGAGGCTCTCTACCTGGTGTGTGGGGAGCGTGGCTTCTTCTACACACCCATGTCCCGCCGTGAAGTGGAGGACCCACAAGTGGCACAACTGGAGCTGGGTGGAGGCCCGGGAGCAGGTGACCTTCAGACCTTGGCACTGGAGGTGGCCCAGCAGAAGCGTGGCATTGTAGATCAGTGCTGCACCAGCATCTGCTCCCTCTACCAGCTGGAGAACTACTGCAAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGCTCTCTGGATCCGATCACTGCCTCTTCTGGCTCTCCTTGTCTTTTCTGGCCCTGGAACCAGCTATGCAGCTGCCAACCAGCACCTCTGTGGCTCCCACTTGGTGGAGGCTCTCTACCTGGTGTGTGGAGAGCGTGGCTTCTTCTACTCCCCCAAAGCCCGACGGGATGTCGAGCAGCCCCTAGTGAGCAGTCCCTTGCGTGGCGAGGCAGGAGTGCTGCCTTTCCAGCAGGAGGAATACGAGAAAGTCAAGCGAGGGATTGTTGAGCAATGCTGCCATAACACGTGTTCCCTCTACCAACTGGAGAACTACTGCAAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGCAGTGTGGCTTCAGGCTGGTGCTCTGTTGGTCCTGTTGGTCGTGTCCAGTGTAAGCACTAACCCAGGCACACCGCAGCACCTGTGTGGATCTCATCTGGTCGATGCCCTTTATCTGGTCTGTGGCCCAACAGGCTTCTTCTACAACCCCAAGAGAGACGTTGAGCCCCTTCTGGGTTTCCTTCCTCCTAAATCTGCCCAGGAAACTGAGGTGGCTGACTTTGCATTTAAAGATCATGCCGAGCTGATAAGGAAGAGAGGCATTGTAGAGCAGTGCTGCCACAAACCCTGCAGCATCTTTGAGCTGCAGAACTACTGTAAC",
                "common_name": "Danio rerio",
            },
            "Frog": {
                "dna": "ATGGCTCTATGGATGCAGTGTCTGCCCCTGGTTCTTGTCCTCTTTTTCTCTACACCCAACACCGAAGCTCTAGTTAACCAGCACTTGTGTGGGTCTCACCTGGTAGAAGCCCTGTACTTAGTATGTGGGGATCGAGGCTTCTTCTACTACCCTAAGGTCAAACGGGACATGGAACAAGCACTTGTCAGTGGACCCCAGGATAATGAGTTGGATGGAATGCAGCTCCAGCCTCAGGAGTATCAGAAAATGAAGAGGGGGATTGTGGAGCAATGTTGCCACAGCACATGTTCTCTCTTCCAGCTGGAGAGTTACTGCAAC",
                "common_name": "Xenopus laevis",
            },
            "Cow": {
                "dna": "ATGGCCCTGTGGACACGCCTGGCGCCCCTGCTGGCCCTGCTGGCGCTCTGGGCCCCCGCCCCGGCCCGCGCCTTCGTCAACCAGCATCTGTGTGGCTCCCACCTGGTGGAGGCGCTGTACCTGGTGTGCGGAGAGCGCGGCTTCTTCTACACGCCCAAGGCCCGCCGGGAGGTGGAGGGCCCCCAGGTGGGGGCGCTGGAGCTGGCCGGAGGCCCGGGCGCGGGCGGCCTGGAGGGGCCCCCGCAGAAGCGTGGCATCGTGGAGCAGTGCTGTGCCAGCGTCTGCTCGCTCTACCAGCTGGAGAACTACTGTAAC",
                "common_name": "Bos taurus",
            },
        },
    },
    "hemoglobin": {
        "gene_name": "Hemoglobin Beta",
        "description": "Beta-globin carries oxygen from lungs to tissues. The most studied gene in molecular evolution - sequence differences power the 'molecular clock' hypothesis. Sickle cell mutation (E6V) in humans shows how a single base change creates devastating disease.",
        "sequences": {
            "Human": {
                "dna": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTGGTCTACCCTTGGACCCAGAGGTTCTTTGAGTCCTTTGGGGATCTGTCCACTCCTGATGCTGTTATGGGCAACCCTAAGGTGAAGGCTCATGGCAAGAAAGTGCTCGGTGCCTTTAGTGATGGCCTGGCTCACCTGGACAACCTCAAGGGCACCTTTGCCACACTGAGTGAGCTGCACTGTGACAAGCTGCACGTGGATCCTGAGAACTTCAGGCTCCTGGGCAACGTGCTGGTCTGTGTGCTGGCCCATCACTTTGGCAAAGAATTCACCCCACCAGTGCAGGCTGCCTATCAGAAAGTGGTGGCTGGTGTGGCTAATGCCCTGGCCCACAAGTATCAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGGTGCACCTGACTGATGCTGAGAAGGCTGCTGTCTCTGGCCTGTGGGGAAAGGTGAACGCCGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTTGTCTACCCTTGGACCCAGCGGTACTTTGATAGCTTTGGAGACCTATCCTCTGCCTCTGCTATCATGGGTAATGCCAAAGTGAAGGCCCATGGCAAGAAAGTGATAACTGCCTTTAACGATGGCCTGAATCACTTGGACAGCCTCAAGGGCACCTTTGCCAGCCTCAGTGAGCTCCACTGTGACAAGCTGCATGTGGATCCTGAGAACTTCAGGCTCCTGGGCAATATGATCGTGATTGTGCTGGGCCACCACCTGGGCAAGGATTTCACCCCCGCTGCACAGGCTGCCTTCCAGAAGGTGGTGGCTGGAGTGGCTGCTGCCCTGGCTCACAAGTACCAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGTGCACTGGACTGCTGAGGAGAAGCAGCTCATCACCGGCCTCTGGGGCAAGGTCAATGTGGCCGAATGTGGGGCTGAAGCCCTGGCCAGGCTGCTGATCGTCTACCCCTGGACCCAGAGGTTCTTTGCGTCCTTTGGGAACCTCTCCAGCCCCACTGCCATCCTTGGCAACCCCATGGTCCGCGCCCATGGCAAGAAAGTGCTCACCTCCTTTGGGGATGCTGTGAAGAACCTGGACAACATCAAGAACACCTTCTCCCAACTGTCCGAACTGCATTGTGACAAGCTGCATGTGGACCCCGAGAACTTCAGGCTCCTGGGTGACATCCTCATCATTGTCCTGGCCGCCCACTTCAGCAAGGACTTCACTCCTGAATGCCAGGCTGCCTGGCAGAAGCTGGTCCGCGTGGTGGCCCATGCCCTGGCTCGCAAGTACCAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGTTGAGTGGACAGATGCCGAGCGCACAGCCATCCTTGGCCTGTGGGGAAAGCTCAATATCGATGAAATCGGACCTCAGGCCCTATCCAGATGTCTGATCGTGTATCCCTGGACTCAGAGATATTTCGCCACATTCGGCAACCTGTCAAGCCCCGCTGCGATCATGGGTAACCCCAAAGTGGCAGCTCATGGGAGGACTGTGATGGGAGGTCTTGAGAGAGCCATCAAGAACATGGACAACGTCAAGAACACCTATGCCGCCCTCAGTGTGATGCACTCTGAGAAACTGCATGTGGATCCCGACAACTTCAGGCTTCTCGCTGATTGCATCACCGTTTGCGCTGCCATGAAGTTCGGCCAAGCTGGTTTCAATGCTGATGTCCAGGAGGCCTGGCAGAAGTTTCTGGCTGTGGTCGTTTCTGCTCTGTGCAGACAGTACCAC",
                "common_name": "Danio rerio",
            },
            "Frog": {
                "dna": "ATGGTTCATTGGACAGCTGAAGAGAAGGCCGCCATCACCTCTGTGTGGCAGGAGGTCAACCAGGAGCAAGATGGCCATGATGCACTCACAAGGCTGCTGGTTGTGTACCCCTGGACCCAGAGATACTTCAGCAGTTTTGGAAATCTCGGTAATGCCACAGCTATTGCTGGAAATGTCAAGGTGCGTGCCCATGGCAAGAAGGTTCTTTCAGCTGTTGGTGATGCCATCGCCCATCTTGACAACGTGAAGGGAACTCTCCATGACCTCAGTGTGGTCCACGCCTTCAAGCTCTATGTGGATCCTGAGAACTTCAAGCGTCTTGGTGAAGTTCTGGTCATTGTCTTGGCTTCCAAACTGGGATCAGCCTTTACTCCTCAAGTCCAGGGAGCCTGGGAGAAATTTGTTGCTGTTCTGGTTGATGCCCTCAGCCAAGGATACAAC",
                "common_name": "Xenopus laevis",
            },
            "Cow": {
                "dna": "ATGCTGACTGCTGAGGAGAAGGCTGCCGTCACCGCCTTTTGGGGCAAGGTGAAAGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTTGTCTACCCCTGGACTCAGAGGTTCTTTGAGTCCTTTGGGGACTTGTCCACTGCTGATGCTGTTATGAACAACCCTAAGGTGAAGGCCCATGGCAAGAAGGTGCTAGATTCCTTTAGTAATGGCATGAAGCATCTCGATGACCTCAAGGGCACCTTTGCTGCGCTGAGTGAGCTGCACTGTGATAAGCTGCATGTGGATCCTGAGAACTTCAAGCTCCTGGGCAACGTGCTAGTGGTTGTGCTGGCTCGCAATTTTGGCAAGGAATTCACCCCGGTGCTGCAGGCTGACTTTCAGAAGGTGGTGGCTGGTGTGGCCAATGCCCTGGCCCACAGATATCAT",
                "common_name": "Bos taurus",
            },
        },
    },
    "p53": {
        "gene_name": "p53 (TP53)",
        "description": "The 'guardian of the genome' - p53 detects DNA damage and triggers repair or cell death. Mutated in >50% of human cancers. Elephants have 20 copies of p53 (humans have 1), which may explain their extremely low cancer rates despite their size (Peto's paradox).",
        "sequences": {
            "Human": {
                "dna": "ATGGAGGAGCCGCAGTCAGATCCTAGCGTCGAGCCCCCTCTGAGTCAGGAAACATTTTCAGACCTATGGAAACTACTTCCTGAAAACAACGTTCTGTCCCCCTTGCCGTCCCAAGCAATGGATGATTTGATGCTGTCCCCGGACGATATTGAACAATGGTTCACTGAAGACCCAGGTCCAGATGAAGCTCCCAGAATGCCAGAGGCTGCTCCCCCCGTGGCCCCTGCACCAGCAGCTCCTACACCGGCGGCCCCTGCACCAGCCCCCTCCTGGCCCCTGTCATCTTCTGTCCCTTCCCAGAAAACCTACCAGGGCAGCTACGGTTTCCGTCTGGGCTTCTTGCATTCTGGGACAGCCAAGTCTGTGACTTGCACGTACTCCCCTGCCCTCAACAAGATGTTTTGCCAACTGGCCAAGACCTGCCCTGTGCAGCTGTGGGTTGATTCCACACCCCCGCCCGGCACCCGCGTCCGCGCCATGGCCATCTACAAGCAGTCACAGCACATGACGGAGGTTGTGAGGCGCTGCCCCCACCATGAGCGCTGCTCAGATAGCGATGGTCTGGCCCCTCCTCAGCATCTTATCCGAGTGGAAGGAAATTTGCGTGTGGAGTATTTGGATGACAGAAACACTTTTCGACATAGTGTGGTGGTGCCCTATGAGCCGCCTGAGGTTGGCTCTGACTGTACCACCATCCACTACAACTACATGTGTAACAGTTCCTGCATGGGCGGCATGAACCGGAGGCCCATCCTCACCATCATCACACTGGAAGACTCCAGTGGTAATCTACTGGGACGGAACAGCTTTGAGGTGCGTGTTTGTGCCTGTCCTGGGAGAGACCGGCGCACAGAGGAAGAGAATCTCCGCAAGAAAGGGGAGCCTCACCACGAGCTGCCCCCAGGGAGCACTAAGCGAGCACTGCCCAACAACACCAGCTCCTCTCCCCAGCCAAAGAAGAAACCACTGGATGGAGAATATTTCACCCTTCAGATCCGTGGGCGTGAGCGCTTCGAGATGTTCCGAGAGCTGAATGAGGCCTTGGAACTCAAGGATGCCCAGGCTGGGAAGGAGCCAGGGGGGAGCAGGGCTCACTCCAGCCACCTGAAGTCCAAAAAGGGTCAGTCTACCTCCCGCCATAAAAAACTCATGTTCAAGACAGAAGGGCCTGACTCAGAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGACTGCCATGGAGGAGTCACAGTCGGATATCAGCCTCGAGCTCCCTCTGAGCCAGGAGACATTTTCAGGCTTATGGAAACTACTTCCTCCAGAAGATATCCTGCCATCACCTCACTGCATGGACGATCTGTTGCTGCCCCAGGATGTTGAGGAGTTTTTTGAAGGCCCAAGTGAAGCCCTCCGAGTGTCAGGAGCTCCTGCAGCACAGGACCCTGTCACCGAGACCCCTGGGCCAGTGGCCCCTGCCCCAGCCACTCCATGGCCCCTGTCATCTTTTGTCCCTTCTCAAAAAACTTACCAGGGCAACTATGGCTTCCACCTGGGCTTCCTGCAGTCTGGGACAGCCAAGTCTGTTATGTGCACGTACTCTCCTCCCCTCAATAAGCTATTCTGCCAGCTGGCGAAGACGTGCCCTGTGCAGTTGTGGGTCAGCGCCACACCTCCAGCTGGGAGCCGTGTCCGCGCCATGGCCATCTACAAGAAGTCACAGCACATGACGGAGGTCGTGAGACGCTGCCCCCACCATGAGCGCTGCTCCGATGGTGATGGCCTGGCTCCTCCCCAGCATCTTATCCGGGTGGAAGGAAATTTGTATCCCGAGTATCTGGAAGACAGGCAGACTTTTCGCCACAGCGTGGTGGTACCTTATGAGCCACCCGAGGCCGGCTCTGAGTATACCACCATCCACTACAAGTACATGTGTAATAGCTCCTGCATGGGGGGCATGAACCGCCGACCTATCCTTACCATCATCACACTGGAAGACTCCAGTGGGAACCTTCTGGGACGGGACAGCTTTGAGGTTCGTGTTTGTGCCTGCCCTGGGAGAGACCGCCGTACAGAAGAAGAAAATTTCCGCAAAAAGGAAGTCCTTTGCCCTGAACTGCCCCCAGGGAGCGCAAAGAGAGCGCTGCCCACCTGCACAAGCGCCTCTCCCCCGCAAAAGAAAAAACCACTTGATGGAGAGTATTTCACCCTCAAGATCCGCGGGCGTAAACGCTTCGAGATGTTCCGGGAGCTGAATGAGGCCTTAGAGTTAAAGGATGCCCATGCTACAGAGGAGTCTGGAGACAGCAGGGCTCACTCCAGCTACCTGAAGACCAAGAAGGGCCAGTCTACTTCCCGCCATAAAAAAACAATGGTCAAGAAAGTGGGGCCTGACTCAGAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGCGGAGGAGATGGAACCATTGCTGGAACCCACTGAGGTCTTCATGGACCTCTGGAGCATGCTCCCCTATAGCATGCAACAGCTGCCCCTCCCTGAGGATCACAGCAACTGGCAGGAGCTGAGCCCCCTGGAACCCAGCGACCCCCCCCCACCACCGCCACCACCACCTCTGCCATTGGCCGCCGCCGCCCCCCCCCCATTAAACCCCCCCACCCCCCCCCGCGCTGCCCCCTCCCCGGTGGTCCCATCCACGGAGGATTATGGGGGGGACTTCGACTTCCGGGTGGGGTTCGTGGAGGCGGGCACAGCCAAATCGGTCACCTGCACTTACTCCCCGGTGCTGAATAAGGTCTATTGCCGCCTGGCCAAGCCGTGCCCGGTGCAGGTGAGGGTGGGGGTGGCGCCCCCCCCCGGTTCCTCCCTCCGCGCCGTGGCCGTCTATAAGAAATCAGAGCACGTGGCCGAAGTGGTGCGGCGCTGCCCCCACCACGAGCGCTGCGGGGGGGGCACCGACGGCCTGGCCCCCGCACAGCACCTCATCCGGGTGGAGGGGAACCCCCAGGCGCGTTACCACGACGACGAGACCACCAAACGGCACAGCGTCGTCGTCCCCTATGAGCCCCCCGAGGTGGGCTCTGACTGTACCACGGTGCTGTACAACTTCATGTGCAACAGTTCCTGCATGGGGGGGATGAACCGCCGCCCCATCCTCACCATCCTTACACTGGAGGGGCCGGGGGGGCAGCTGTTGGGGCGGCGCTGCTTCGAGGTGCGCGTGTGCGCATGTCCGGGGAGGGACCGCAAGATCGAGGAGGAGAACTTCCGCAAGAGGGGCGGGGCCGGGGGCGTGGCTAAGCGAGCCATGTCGCCCCCAACCGAAGCCCCCGAGCCCCCCAAGAAGCGCGTGCTGAACCCCGACAATGAGATATTCTACCTGCAGGTGCGCGGGCGCCGCCGCTATGAGATGCTGAAGGAGATCAATGAGGCGCTGCAGCTCGCCGAGGGGGGGTCCGCACCGCGGCCTTCCAAAGGCCGCCGTGTGAAGGTGGAGGGACCCCAACCCAGCTGCGGGAAGAAACTGCTGCAAAAAGGCTCGGAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGCGCAAAACGACAGCCAAGAGTTCGCGGAGCTCTGGGAGAAGAATTTGATTATTCAGCCCCCAGGTGGTGGCTCTTGCTGGGACATCATTAATGATGAGGAGTACTTGCCGGGATCGTTTGACCCCAATTTTTTTGAAAATGTGCTTGAAGAACAGCCTCAGCCATCCACTCTCCCACCAACATCCACTGTTCCGGAGACAAGCGACTATCCCGGCGATCATGGATTTAGGCTCAGGTTCCCGCAGTCTGGCACAGCAAAATCTGTAACTTGCACTTATTCACCGGACCTGAATAAACTCTTCTGTCAGCTGGCAAAAACTTGCCCCGTTCAAATGGTGGTGGACGTTGCCCCTCCACAGGGCTCCGTGGTTCGAGCCACTGCCATCTATAAGAAGTCCGAGCATGTGGCTGAAGTGGTCCGCAGATGCCCCCATCATGAGCGAACCCCGGATGGAGATAACTTGGCGCCTGCTGGTCATTTGATAAGAGTGGAGGGCAATCAGCGAGCAAATTACAGGGAAGATAACATCACTTTAAGGCATAGTGTTTTTGTCCCATATGAAGCACCACAGCTTGGTGCTGAATGGACAACTGTGCTACTAAACTACATGTGCAATAGCAGCTGCATGGGGGGGATGAACCGCAGGCCCATCCTCACAATCATCACTCTGGAGACTCAGGAAGGTCAGTTGCTGGGCCGGAGGTCTTTTGAGGTGCGTGTGTGTGCATGTCCAGGCAGAGACAGGAAAACTGAGGAGAGCAACTTCAAGAAAGACCAAGAGACCAAAACCATGGCCAAAACCACCACTGGGACCAAACGTAGTTTGGTGAAAGAATCTTCTTCAGCTACATTACGACCTGAGGGGAGCAAAAAGGCCAAGGGCTCCAGCAGCGATGAGGAGATCTTTACCCTGCAGGTGAGGGGCAGGGAGCGTTATGAAATTTTAAAGAAATTGAACGACAGTCTGGAGTTAAGTGATGTGGTGCCTGCCTCAGATGCTGAAAAGTATCGTCAGAAATTCATGACAAAAAACAAAAAAGAGAATCGTGAATCATCTGAGCCCAAACAGGGAAAGAAGCTGATGGTGAAGGACGAAGGAAGAAGCGACTCTGAT",
                "common_name": "Danio rerio",
            },
            "Elephant": {
                "dna": "ATGGAGGAGCCCCAGTCAGATCTCAGCACCGAGCTCCCTCTGAGTCAAGAGACGTTTTCATACTTATGGGAACTCCTTCCTGAGAATCCGGTTCTGTCCCCCACACTACCCCCGGCAGTGGAGGTCATGGACGATCTGCTACTCTCAGAAGACACTGCAAACTGGCTAGAAAGCCAAGTTGAGGCTCAGGGAATGTCCACAACCCCTGCACCAGCCACCCCTACACCGGTGGCCCCCGCACCAGCCACCTCCTGGACCCTGTCATCTTCCGTCCCTTCCCAAAAGACCTACCCTGGCACCTATGGTTTCCGTCTGGGCTTCCTACATTCTGGGACAGCCAAGTCCGTCACCTGCACGTACTCCCCTGACCTTAACAAGCTGTTTTGCCAGCTGGCAAAAACCTGCCCAGTGCAGCTGTGGGTCGCCTCACCACCCCCGCCCGGCACCCGTGTTCGCACCATGGCCATCTACAAGAAGTCAGAGCATATGACGGAGGTCGTCAAGCGCTGCCCCCACCATGAGCGCTGCTCTGACTCTAGCGATGGCCTGGCCCCTCCTCAGCACCTCATCCGGGTGGAAGGAAACCTGCGTGCTGAGTATCTGGAGGACAGCATCACTCTCCGACACAGTGTGGTGGTGCCCTACGAGCCGCCCGAGGTTGGGTCTGACTGTACCACCATCCACTTCAACTTCATGTGTAACAGCTCCTGCATGGGGGGCATGAACCGGCGGCCCATCCTCACCATCATCACACTGGAAGACTCCAGTGGTAATCTGCTGGGACGTAACAGCTTTGAGGTGCGCATTTGTGCCTGTCCTGGAAGAGACAGACGTACAGAAGAAGAAAATTTCCACAAGAAGGGAGAGCCTTGCCCAGAGCCGCCACCCCCTGGGAGGAGCACTAAGCGAGCACTGCCCACCAACACCAGCTCCTCTACCCAGCCAAAGAAGAAGCCACTGGATGAAGAATATTTCACCCTTCAGATCCGTGGGCGTGAACGCTTCAAGATGTTCCTAGAGCTAAATGAGGCCTTGGAGCTGAAGGATGCCCAGGCTGGGAAGGAGCCAGAGGGGAGCCGGGCTCACTCCAGCCCTTCGAAGTCTAAGAAGGGACAGTCTACCTCCCGCCATAAAAAACCAATGTTCAAGAGAGAGGGACCTGACTCAGAC",
                "common_name": "Loxodonta africana",
            },
            "Dog": {
                "dna": "ATGGAGGAGTCGCAGTCAGAGCTCAATATCGACCCCCCTCTGAGCCAGGAGACATTTTCAGAATTGTGGAACCTGCTTCCTGAAAACAATGTTCTGTCTTCGGAGCTGTGCCCAGCAGTGGATGAGCTGCTGCTCCCAGAGAGCGTCGTGAACTGGCTAGACGAAGACTCAGATGATGCTCCCAGGATGCCAGCCACTTCTGCCCCCACAGCCCCTGGACCGGCCCCCTCGTGGCCCCTATCATCCTCTGTCCCTTCCCCGAAGACCTACCCTGGCACCTATGGGTTCCGTTTGGGGTTCCTGCATTCCGGGACAGCCAAGTCTGTTACTTGGACGTACTCCCCTCTCCTCAACAAGTTGTTTTGCCAGCTGGCGAAGACCTGCCCCGTGCAGCTGTGGGTCAGCTCCCCACCCCCACCCAATACCTGCGTCCGCGCTATGGCCATCTATAAGAAGTCGGAGTTCGTGACCGAGGTTGTGCGGCGCTGCCCCCACCATGAACGCTGCTCTGACAGTAGTGACGGTCTTGCCCCTCCTCAGCATCTCATCCGAGTGGAAGGAAATTTGCGGGCCAAGTACCTGGACGACAGAAACACTTTTCGACACAGTGTGGTGGTGCCTTATGAGCCACCCGAGGTTGGCTCTGACTATACCACCATCCACTACAACTACATGTGTAACAGTTCCTGCATGGGAGGCATGAACCGGCGGCCCATCCTCACTATCATCACCCTGGAAGACTCCAGTGGAAACGTGCTGGGACGCAACAGCTTTGAGGTACGCGTTTGTGCCTGTCCCGGGAGAGACCGCCGGACTGAGGAGGAGAATTTCCACAAGAAGGGGGAGCCTTGTCCTGAGCCACCCCCCGGGAGTACCAAGCGAGCACTGCCTCCCAGCACCAGCTCCTCTCCCCCGCAAAAGAAGAAGCCACTAGATGGAGAATATTTCACCCTTCAGATCCGTGGGCGTGAACGCTATGAGATGTTCAGGAATCTGAATGAAGCCTTGGAGCTGAAGGATGCCCAGAGTGGAAAGGAGCCAGGGGGAAGCAGGGCTCACTCCAGCCACCTGAAGGCAAAGAAGGGGCAATCTACCTCTCGCCATAAAAAACTGATGTTCAAGAGAGAAGGGCTTGACTCAGAC",
                "common_name": "Canis lupus familiaris",
            },
        },
    },
}

# Standard genetic code
CODON_TABLE = {
    "TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L", "CTT": "L", "CTC": "L",
    "CTA": "L", "CTG": "L", "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M",
    "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V", "TCT": "S", "TCC": "S",
    "TCA": "S", "TCG": "S", "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
    "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T", "GCT": "A", "GCC": "A",
    "GCA": "A", "GCG": "A", "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*",
    "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q", "AAT": "N", "AAC": "N",
    "AAA": "K", "AAG": "K", "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E",
    "TGT": "C", "TGC": "C", "TGA": "*", "TGG": "W", "CGT": "R", "CGC": "R",
    "CGA": "R", "CGG": "R", "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
    "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G",
}

BASE_COLORS = {"A": "#2ecc71", "T": "#e74c3c", "G": "#f39c12", "C": "#3498db"}

def _translate(dna: str) -> str:
    """Translate DNA to protein in reading frame 0."""
    dna = dna.upper()
    protein = []
    for i in range(0, len(dna) - 2, 3):
        codon = dna[i:i + 3]
        aa = CODON_TABLE.get(codon, "X")
        if aa == "*":
            break
        protein.append(aa)
    return "".join(protein)

def _gc_content(seq: str) -> float:
    if not seq:
        return 0.0
    return sum(1 for b in seq.upper() if b in "GC") / len(seq)

def _sequence_identity(seq1: str, seq2: str, match: int = 2, mismatch: int = -1, gap: int = -2) -> float:
    """Percent identity via Needleman-Wunsch global alignment."""
    if not seq1 or not seq2:
        return 0.0
    n, m = len(seq1), len(seq2)

    # Build score matrix
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        dp[i][0] = dp[i - 1][0] + gap
    for j in range(1, m + 1):
        dp[0][j] = dp[0][j - 1] + gap
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            s = match if seq1[i - 1] == seq2[j - 1] else mismatch
            dp[i][j] = max(dp[i - 1][j - 1] + s, dp[i - 1][j] + gap, dp[i][j - 1] + gap)

    # Traceback to count matches and alignment length
    i, j = n, m
    matches = 0
    aligned = 0
    while i > 0 or j > 0:
        if i > 0 and j > 0:
            s = match if seq1[i - 1] == seq2[j - 1] else mismatch
            if dp[i][j] == dp[i - 1][j - 1] + s:
                if seq1[i - 1] == seq2[j - 1]:
                    matches += 1
                aligned += 1
                i -= 1
                j -= 1
                continue
        if i > 0 and dp[i][j] == dp[i - 1][j] + gap:
            aligned += 1
            i -= 1
        else:
            aligned += 1
            j -= 1

    return matches / aligned if aligned else 0.0

# ------------------------------------------------------------------
# Report styling
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #1e3a5f; border-bottom: 2px solid #2563eb; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #1e40af; margin-top: 20px; }
  .report .card { background: #eff6ff; border: 1px solid #bfdbfe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #dbeafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #1e3a5f; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #1e3a5f; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dbeafe; }
  .report tr:nth-child(even) { background: #eff6ff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-warning { background: #fef3c7; color: #92400e; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #dbeafe; color: #1e40af; }
  .report .chart-container { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #eff6ff; border-left: 4px solid #2563eb; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .structure-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(340px, 1fr)); gap: 16px; margin: 12px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    width: int = 600,
    height: int = 500,
    value_format: str = ".1f",
    color_scale: str = "blue",
) -> str:
    """Generate an SVG heatmap."""
    n_rows = len(matrix)
    n_cols = len(matrix[0]) if matrix else 0
    if not n_rows or not n_cols:
        return ""

    show_values = n_rows <= 10 and n_cols <= 10
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    if color_scale == "blue":
        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 30))
            g = int(255 - t * (255 - 58))
            b = int(255 - t * (255 - 95))
            return f"rgb({r},{g},{b})"
    else:  # green
        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 6))
            g = int(255 - t * (255 - 95))
            b = int(255 - t * (255 - 70))
            return f"rgb({r},{g},{b})"

    ml = max(80, max(len(l) for l in row_labels) * 7 + 10) if row_labels else 80
    mr = 20
    mt = 80
    mb = 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>')

    for j, label in enumerate(col_labels):
        cx = ml + j * cell_w + cell_w / 2
        svg.append(f'<text x="{cx:.1f}" y="{mt - 8}" text-anchor="start" font-size="10" fill="#374151" transform="rotate(-45, {cx:.1f}, {mt - 8})">{label}</text>')

    for i, row_label in enumerate(row_labels):
        ry = mt + i * cell_h + cell_h / 2
        svg.append(f'<text x="{ml - 8}" y="{ry + 4:.1f}" text-anchor="end" font-size="10" fill="#374151">{row_label}</text>')
        for j in range(n_cols):
            val = matrix[i][j]
            color = get_color(val)
            cx = ml + j * cell_w
            cy = mt + i * cell_h
            svg.append(f'<rect x="{cx:.1f}" y="{cy:.1f}" width="{cell_w:.1f}" height="{cell_h:.1f}" fill="{color}" stroke="#fff" stroke-width="1"/>')
            if show_values:
                t = (val - v_min) / v_range
                text_color = "#fff" if t > 0.55 else "#1a1a2e"
                fs = min(10, int(cell_w / 4), int(cell_h / 2.5))
                fs = max(7, fs)
                svg.append(f'<text x="{cx + cell_w / 2:.1f}" y="{cy + cell_h / 2 + 3:.1f}" text-anchor="middle" font-size="{fs}" fill="{text_color}">{val:{value_format}}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_dendrogram(
    names: list[str],
    matrix: list[list[float]],
    title: str = "",
    width: int = 700,
    height: int = 350,
    color: str = "#2563eb",
) -> str:
    """Generate an SVG dendrogram from a similarity matrix using UPGMA."""
    n = len(names)
    if n < 2:
        return ""

    dist = [[1.0 - matrix[i][j] for j in range(n)] for i in range(n)]

    clusters = [{"members": [i], "height": 0.0, "left": None, "right": None} for i in range(n)]
    active = list(range(n))

    while len(active) > 1:
        best_d = float("inf")
        bi, bj = 0, 1
        for ii in range(len(active)):
            for jj in range(ii + 1, len(active)):
                ci, cj = active[ii], active[jj]
                d = 0
                count = 0
                for mi in clusters[ci]["members"]:
                    for mj in clusters[cj]["members"]:
                        d += dist[mi][mj]
                        count += 1
                avg_d = d / count if count else 0
                if avg_d < best_d:
                    best_d = avg_d
                    bi, bj = ii, jj

        ci, cj = active[bi], active[bj]
        new_cluster = {
            "members": clusters[ci]["members"] + clusters[cj]["members"],
            "height": best_d,
            "left": clusters[ci],
            "right": clusters[cj],
        }
        clusters.append(new_cluster)
        new_idx = len(clusters) - 1
        active.pop(bj)
        active.pop(bi)
        active.append(new_idx)

    root = clusters[active[0]]

    max_label_len = max((len(n) for n in names), default=0)
    ml, mr, mt, mb = max(50, max_label_len * 5 + 10), 30, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb
    max_h = root["height"] or 1

    leaf_positions = {}
    leaf_counter = [0]

    def assign_leaves(node):
        if node["left"] is None and node["right"] is None:
            leaf_positions[node["members"][0]] = leaf_counter[0]
            leaf_counter[0] += 1
        else:
            if node["left"]:
                assign_leaves(node["left"])
            if node["right"]:
                assign_leaves(node["right"])

    assign_leaves(root)
    n_leaves = len(leaf_positions)
    leaf_spacing = cw / max(n_leaves - 1, 1)

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="13" font-weight="600" fill="#1a1a2e">{title}</text>')

    def get_x(node):
        if node["left"] is None and node["right"] is None:
            return ml + leaf_positions[node["members"][0]] * leaf_spacing
        return (get_x(node["left"]) + get_x(node["right"])) / 2

    def get_y(h):
        return mt + ch - (h / max_h) * ch

    def draw_node(node):
        if node["left"] is None and node["right"] is None:
            return
        lx = get_x(node["left"])
        rx = get_x(node["right"])
        ly = get_y(node["left"]["height"])
        ry = get_y(node["right"]["height"])
        my = get_y(node["height"])

        svg.append(f'<line x1="{lx:.1f}" y1="{ly:.1f}" x2="{lx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')
        svg.append(f'<line x1="{rx:.1f}" y1="{ry:.1f}" x2="{rx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')
        svg.append(f'<line x1="{lx:.1f}" y1="{my:.1f}" x2="{rx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')

        if node["left"]:
            draw_node(node["left"])
        if node["right"]:
            draw_node(node["right"])

    draw_node(root)

    for idx, pos in leaf_positions.items():
        x = ml + pos * leaf_spacing
        svg.append(
            f'<text x="{x:.1f}" y="{mt + ch + 14}" text-anchor="start" font-size="10" fill="#374151" '
            f'transform="rotate(40, {x:.1f}, {mt + ch + 14})">{names[idx]}</text>'
        )

    for i in range(5):
        d = max_h * i / 4
        y = get_y(d)
        svg.append(f'<text x="{ml - 4}" y="{y + 3:.1f}" text-anchor="end" font-size="9" fill="#9ca3af">{d:.3f}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    value_format: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart."""
    if not labels:
        return ""

    default_colors = ["#2563eb", "#059669", "#f59e0b", "#dc2626", "#7c3aed"]
    colors = colors or default_colors

    ml, mr, mt, mb = 60, 20, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_min = min(all_vals) if all_vals else 0
    y_max = max(all_vals) if all_vals else 1
    if y_min >= 0:
        y_min_plot = 0
        y_max_plot = y_max * 1.15 or 1
    else:
        y_range = y_max - y_min or 1
        y_min_plot = y_min - y_range * 0.05
        y_max_plot = y_max + y_range * 0.15

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - ((v - y_min_plot) / (y_max_plot - y_min_plot)) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    for i in range(6):
        y_tick = y_min_plot + (y_max_plot - y_min_plot) * i / 5
        py = sy(y_tick)
        svg.append(f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" stroke="#e9ecef" stroke-width="1"/>')
        svg.append(f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" font-size="11" fill="#6c757d">{y_tick:{value_format}}</text>')

    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" height="{max(0, bh):.1f}" fill="{color}" rx="2"/>')
            svg.append(f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" text-anchor="middle" font-size="9" fill="#1a1a2e">{val:{value_format}}</text>')
        lx = gx + n_series * bar_width / 2
        svg.append(f'<text x="{lx:.1f}" y="{mt + ch + 14}" text-anchor="start" font-size="10" fill="#6c757d" transform="rotate(35, {lx:.1f}, {mt + ch + 14})">{label}</text>')

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>')

    if n_series > 1:
        lx = ml + cw - len(series) * 110
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(f'<rect x="{lx + si * 110}" y="{mt + ch + 55}" width="12" height="12" rx="2" fill="{color}"/>')
            svg.append(f'<text x="{lx + si * 110 + 16}" y="{mt + ch + 66}" font-size="11" fill="#1a1a2e">{name}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_plddt_sparkline(values: list[float], width: int = 400, height: int = 50) -> str:
    """pLDDT sparkline with AlphaFold-style coloring."""
    if not values or len(values) < 2:
        return ""

    pad = 4
    cw = width - 2 * pad
    ch = height - 2 * pad
    seg_w = cw / len(values)

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
    ]

    for i, v in enumerate(values):
        x = pad + i * seg_w
        bar_h = (v / 100) * ch
        y = pad + ch - bar_h

        if v >= 90:
            color = "#0053d6"
        elif v >= 70:
            color = "#65cbf3"
        elif v >= 50:
            color = "#ffdb13"
        else:
            color = "#ff7d45"

        svg.append(f'<rect x="{x:.1f}" y="{y:.1f}" width="{max(seg_w, 1):.1f}" height="{bar_h:.1f}" fill="{color}"/>')

    ref_y = pad + ch - (70 / 100) * ch
    svg.append(f'<line x1="{pad}" y1="{ref_y:.1f}" x2="{pad + cw}" y2="{ref_y:.1f}" stroke="#adb5bd" stroke-width="0.5" stroke-dasharray="3,2"/>')

    svg.append("</svg>")
    return "\n".join(svg)

def _outputs_to_pdb(outputs, sequence: str) -> str:
    """Convert ESMFold outputs to PDB format string."""
    import numpy as np

    pos = outputs.positions[0]
    if pos.dim() == 4:
        pos = pos[-1]
    positions = pos.cpu().numpy()
    atom_names = ["N", "CA", "C", "O"]
    aa_3letter = {
        "A": "ALA", "R": "ARG", "N": "ASN", "D": "ASP", "C": "CYS",
        "Q": "GLN", "E": "GLU", "G": "GLY", "H": "HIS", "I": "ILE",
        "L": "LEU", "K": "LYS", "M": "MET", "F": "PHE", "P": "PRO",
        "S": "SER", "T": "THR", "W": "TRP", "Y": "TYR", "V": "VAL",
    }

    pdb_lines = []
    atom_idx = 1
    for res_idx, aa in enumerate(sequence):
        res_name = aa_3letter.get(aa, "UNK")
        for atom_i, atom_name in enumerate(atom_names):
            if atom_i >= positions.shape[1]:
                break
            x, y, z = positions[res_idx, atom_i]
            if any(math.isnan(c) for c in (x, y, z)):
                continue
            pdb_lines.append(
                f"ATOM  {atom_idx:5d}  {atom_name:<3s} {res_name} A{res_idx + 1:4d}    "
                f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {atom_name[0]:>2s}"
            )
            atom_idx += 1
    pdb_lines.append("END")
    return "\n".join(pdb_lines)

# ------------------------------------------------------------------
# Task 1: Load gene set
# ------------------------------------------------------------------

@cpu_env.task()
async def load_genes(
    gene_set: str = "insulin",
    custom_json: str = "",
) -> flyte.io.Dir:
    """Load a set of homologous genes from different species."""
    if custom_json:
        data = json.loads(custom_json)
    elif gene_set in GENE_SETS:
        data = GENE_SETS[gene_set]
    else:
        available = ", ".join(GENE_SETS.keys())
        raise ValueError(f"Unknown gene set '{gene_set}'. Available: {available}")

    log.info(f"Loaded gene set: {data['gene_name']} - {len(data['sequences'])} species")

    out_dir = tempfile.mkdtemp(prefix="gene_compare_")
    with open(os.path.join(out_dir, "genes.json"), "w") as f:
        json.dump(data, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Score sequences with Carbon
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def score_sequences(
    genes_dir: flyte.io.Dir,
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> str:
    """Score each species' gene with Carbon-3B genomic language model.

    Returns per-species log-likelihood scores and sequence metadata.
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info(f"Loading Carbon model: {model_name}")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, trust_remote_code=True,
        dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    ).to(device)
    model.eval()

    genes_path = await genes_dir.download()
    with open(os.path.join(genes_path, "genes.json")) as f:
        data = json.load(f)

    species_names = list(data["sequences"].keys())
    n = len(species_names)

    scores = {}
    for i, species in enumerate(species_names):
        await flyte.report.replace.aio(_wrap_report(
            f"<h2>Carbon Scoring</h2>"
            f"<p>Scoring {species} ({i + 1}/{n})...</p>"
        ), do_flush=True)

        dna = data["sequences"][species]["dna"]
        prompt = f"<dna>{dna}"
        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            output = model(**inputs, labels=inputs["input_ids"])
            loss = output.loss.item()
            ll = -loss * inputs["input_ids"].shape[1]

        protein = _translate(dna)
        scores[species] = {
            "log_likelihood": round(ll, 4),
            "loss": round(loss, 4),
            "gc_content": round(_gc_content(dna), 4),
            "length": len(dna),
            "protein": protein,
            "protein_length": len(protein),
            "common_name": data["sequences"][species]["common_name"],
        }
        log.info(f"  {species} ({data['sequences'][species]['common_name']}): LL={ll:.2f}, GC={_gc_content(dna):.1%}")

    # Report
    html_parts = [
        f"<h2>{data['gene_name']} - Carbon Scoring</h2>",
        f'<div class="note">{data["description"]}</div>',
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{n}</div><div class="label">Species</div></div>',
        f'<div class="stat"><div class="value">{data["gene_name"]}</div><div class="label">Gene</div></div>',
        f'<div class="stat"><div class="value">{model_name.split("/")[-1]}</div><div class="label">Model</div></div>',
        "</div>",
    ]

    html_parts.append(
        "<table><tr><th>Species</th><th>Scientific Name</th><th>DNA Length</th>"
        "<th>GC%</th><th>Protein Length</th><th>Carbon LL</th></tr>"
    )
    for species in species_names:
        s = scores[species]
        html_parts.append(
            f'<tr><td><b>{species}</b></td><td><i>{s["common_name"]}</i></td>'
            f'<td>{s["length"]}bp</td><td>{s["gc_content"]:.1%}</td>'
            f'<td>{s["protein_length"]}aa</td>'
            f'<td>{s["log_likelihood"]:.2f}</td></tr>'
        )
    html_parts.append("</table>")

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species_names,
        {"Log-Likelihood": [scores[s]["log_likelihood"] for s in species_names]},
        title="Carbon Log-Likelihood per Species",
        value_format=".1f",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    result = {
        "gene_name": data["gene_name"],
        "description": data["description"],
        "species": species_names,
        "scores": scores,
    }
    return json.dumps(result)

# ------------------------------------------------------------------
# Task 3: Align sequences and compute similarity
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def align_and_compare(
    scores_json: str,
    genes_dir: flyte.io.Dir,
) -> str:
    """Align sequences with Needleman-Wunsch and compute pairwise identity.

    Translates DNA to protein, builds DNA and protein identity matrices,
    and generates phylogenetic trees from sequence divergence.
    """
    scores_data = json.loads(scores_json)
    species_names = scores_data["species"]
    scores = scores_data["scores"]
    gene_name = scores_data["gene_name"]
    n = len(species_names)

    genes_path = await genes_dir.download()
    with open(os.path.join(genes_path, "genes.json")) as f:
        data = json.load(f)

    await flyte.report.replace.aio(_wrap_report(
        f"<h2>{gene_name} - Sequence Alignment</h2>"
        f"<p>Aligning {n} species with Needleman-Wunsch...</p>"
    ), do_flush=True)

    # Pairwise DNA identity matrix
    identity_matrix = []
    for sp1 in species_names:
        row = []
        for sp2 in species_names:
            dna1 = data["sequences"][sp1]["dna"]
            dna2 = data["sequences"][sp2]["dna"]
            identity = _sequence_identity(dna1, dna2)
            row.append(round(identity, 4))
        identity_matrix.append(row)

    # Pairwise protein identity matrix
    protein_matrix = []
    for sp1 in species_names:
        row = []
        for sp2 in species_names:
            identity = _sequence_identity(scores[sp1]["protein"], scores[sp2]["protein"])
            row.append(round(identity, 4))
        protein_matrix.append(row)

    # Average pairwise identities (exclude diagonal)
    dna_pairs = [identity_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    prot_pairs = [protein_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    avg_dna = sum(dna_pairs) / len(dna_pairs) if dna_pairs else 0
    avg_prot = sum(prot_pairs) / len(prot_pairs) if prot_pairs else 0

    # Most/least similar pair
    best_pair = max(range(len(dna_pairs)), key=lambda k: dna_pairs[k])
    worst_pair = min(range(len(dna_pairs)), key=lambda k: dna_pairs[k])
    pair_indices = [(i, j) for i in range(n) for j in range(i + 1, n)]
    best_sp = f"{species_names[pair_indices[best_pair][0]]}-{species_names[pair_indices[best_pair][1]]}"
    worst_sp = f"{species_names[pair_indices[worst_pair][0]]}-{species_names[pair_indices[worst_pair][1]]}"

    # Report
    html_parts = [
        f"<h2>{gene_name} - Sequence Alignment</h2>",
        f'<div class="note">Pairwise alignment using Needleman-Wunsch (match=2, mismatch=-1, gap=-2). '
        f"Identity is computed as matches / aligned length from the optimal global alignment.</div>",
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{n}</div><div class="label">Species Aligned</div></div>',
        f'<div class="stat"><div class="value">{n * (n - 1) // 2}</div><div class="label">Pairwise Alignments</div></div>',
        f'<div class="stat"><div class="value">{avg_dna:.0%}</div><div class="label">Avg DNA Identity</div></div>',
        f'<div class="stat"><div class="value">{avg_prot:.0%}</div><div class="label">Avg Protein Identity</div></div>',
        f'<div class="stat"><div class="value">{best_sp}</div><div class="label">Most Similar</div></div>',
        f'<div class="stat"><div class="value">{worst_sp}</div><div class="label">Most Divergent</div></div>',
        "</div>",
    ]

    # DNA identity heatmap
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        identity_matrix, species_names, species_names,
        title="Pairwise DNA Sequence Identity (%)",
        value_format=".0%",
    ))
    html_parts.append("</div>")

    # Protein identity heatmap
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        protein_matrix, species_names, species_names,
        title="Pairwise Protein Sequence Identity (%)",
        value_format=".0%",
        color_scale="green",
    ))
    html_parts.append("</div>")

    # DNA phylogenetic tree
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species_names, identity_matrix,
        title=f"{gene_name} - Phylogenetic Tree (DNA Identity)",
    ))
    html_parts.append("</div>")

    # Protein phylogenetic tree
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species_names, protein_matrix,
        title=f"{gene_name} - Phylogenetic Tree (Protein Identity)",
        color="#059669",
    ))
    html_parts.append("</div>")

    # DNA vs Protein conservation comparison
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species_names,
        {
            "DNA vs Human": [identity_matrix[0][j] for j in range(n)],
            "Protein vs Human": [protein_matrix[0][j] for j in range(n)],
        },
        title=f"Conservation vs {species_names[0]} (DNA and Protein)",
        value_format=".0%",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    result = {
        "gene_name": gene_name,
        "description": scores_data["description"],
        "species": species_names,
        "scores": scores,
        "dna_identity_matrix": identity_matrix,
        "protein_identity_matrix": protein_matrix,
    }
    return json.dumps(result)

# ------------------------------------------------------------------
# Task 4: Fold proteins with ESMFold
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def fold_proteins(
    comparison_json: str,
    max_length: int = 400,
) -> str:
    """Fold each species' translated protein with ESMFold for 3D comparison.

    Returns PDB strings and pLDDT confidence scores for each species.
    """
    import torch
    import numpy as np
    from transformers import AutoTokenizer, EsmForProteinFolding

    comparison = json.loads(comparison_json)
    species_names = comparison["species"]
    scores = comparison["scores"]

    log.info("Loading ESMFold model...")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
    model = model.to(device)
    model.eval()

    structure_data = {}
    n = len(species_names)

    for idx, species in enumerate(species_names):
        protein = scores[species]["protein"]

        if len(protein) > max_length:
            log.info(f"Skipping {species} ({len(protein)} aa > {max_length} max)")
            continue

        log.info(f"ESMFold [{idx + 1}/{n}]: {species} ({len(protein)} aa)")
        await flyte.report.replace.aio(_wrap_report(
            f"<h2>ESMFold - 3D Structure Prediction</h2>"
            f"<p>Folding {species} ({idx + 1}/{n}): {len(protein)} residues...</p>"
        ), do_flush=True)

        inputs = tokenizer(protein, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            outputs = model(**inputs)

        pdb_str = _outputs_to_pdb(outputs, protein)

        plddt_raw = outputs.plddt[0].cpu().numpy()
        if plddt_raw.ndim == 2:
            plddt_raw = plddt_raw[-1]
        plddt = plddt_raw.flatten()[:len(protein)]
        if plddt.max() <= 1.0:
            plddt = plddt * 100
        plddt_mean = float(np.mean(plddt))

        structure_data[species] = {
            "pdb_str": pdb_str,
            "plddt_mean": round(plddt_mean, 1),
            "plddt_per_residue": [round(float(v), 1) for v in plddt[:len(protein)]],
            "protein_length": len(protein),
        }
        log.info(f"  → mean pLDDT: {plddt_mean:.1f}")

    # Report with 3D viewers
    n_folded = len(structure_data)
    avg_plddt = sum(d["plddt_mean"] for d in structure_data.values()) / n_folded if n_folded else 0

    threeDmol_script = '<script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>'

    stats_html = f"""
    <h2>ESMFold - Cross-Species Structure Comparison</h2>
    <div class="note">
      <b>ESMFold</b> predicts 3D structure directly from amino acid sequence.
      Comparing structures across species reveals which parts of the protein are
      structurally conserved (functional core) vs divergent (surface loops, species-specific adaptations).
    </div>
    <div class="stat-grid">
      <div class="stat"><div class="value">{n_folded}</div><div class="label">Structures</div></div>
      <div class="stat"><div class="value">{avg_plddt:.1f}</div><div class="label">Avg pLDDT</div></div>
      <div class="stat"><div class="value">{comparison['gene_name']}</div><div class="label">Gene</div></div>
    </div>
    """

    viewers_html = '<div class="structure-grid">'
    for species, sdata in structure_data.items():
        plddt_val = sdata["plddt_mean"]
        common = scores[species]["common_name"]

        if plddt_val >= 90:
            badge = '<span class="badge badge-success">Very High</span>'
        elif plddt_val >= 70:
            badge = '<span class="badge badge-info">Confident</span>'
        elif plddt_val >= 50:
            badge = '<span class="badge badge-warning">Low</span>'
        else:
            badge = '<span class="badge badge-danger">Disordered</span>'

        plddt_sparkline = _make_plddt_sparkline(sdata["plddt_per_residue"], width=300)
        pdb_escaped = sdata["pdb_str"].replace("\\", "\\\\").replace("`", "\\`").replace("$", "\\$")
        viewer_id = f"viewer_{hash(species) & 0xFFFFFF:06x}"

        viewers_html += f"""
        <div class="card" style="margin:0;">
          <h3 style="margin-top:0;">{species}
            <span style="font-size:0.7em;color:#6c757d;">({sdata['protein_length']} aa)</span>
            {badge}
          </h3>
          <p style="font-size:0.85em;color:#6c757d;margin:2px 0 8px;"><i>{common}</i></p>
          <div id="{viewer_id}" style="width:100%;max-width:320px;height:280px;border:1px solid #dbeafe;border-radius:8px;position:relative;"></div>
          <div style="margin-top:8px;">
            <b>Mean pLDDT:</b> {plddt_val:.1f} / 100
            <div style="margin-top:4px;">{plddt_sparkline}</div>
            <div style="font-size:0.75em;color:#9ca3af;margin-top:2px;">
              <span style="color:#0053d6;">&block; &gt;90</span>
              <span style="color:#65cbf3;">&block; 70-90</span>
              <span style="color:#ffdb13;">&block; 50-70</span>
              <span style="color:#ff7d45;">&block; &lt;50</span>
            </div>
          </div>
        </div>
        <script>
        (function() {{
          var pdb = `{pdb_escaped}`;
          function initViewer() {{
            if (typeof $3Dmol === 'undefined') {{ setTimeout(initViewer, 200); return; }}
            var viewer = $3Dmol.createViewer(document.getElementById("{viewer_id}"), {{backgroundColor: "white"}});
            viewer.addModel(pdb, "pdb");
            viewer.setStyle({{}}, {{cartoon: {{color: "spectrum"}}}});
            viewer.zoomTo();
            viewer.render();
            viewer.spin("y", 1);
          }}
          initViewer();
        }})();
        </script>
        """

    viewers_html += "</div>"

    # pLDDT comparison bar chart
    plddt_chart = _make_bar_chart(
        list(structure_data.keys()),
        {"Mean pLDDT": [d["plddt_mean"] for d in structure_data.values()]},
        title="Structure Confidence Comparison (pLDDT)",
        value_format=".1f",
        colors=["#0053d6"],
    )

    report_html = f"""
    {threeDmol_script}
    {stats_html}
    {viewers_html}
    <div class="chart-container">{plddt_chart}</div>
    """

    await flyte.report.replace.aio(_wrap_report(report_html), do_flush=True)

    return json.dumps(structure_data)

# ------------------------------------------------------------------
# Task 5: Generate summary
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def generate_summary(
    comparison_json: str,
    structures_json: str,
) -> str:
    """Generate comprehensive cross-species summary."""
    comparison = json.loads(comparison_json)
    structures = json.loads(structures_json)

    species = comparison["species"]
    scores = comparison["scores"]
    gene_name = comparison["gene_name"]
    dna_matrix = comparison["dna_identity_matrix"]
    protein_matrix = comparison["protein_identity_matrix"]

    html_parts = [
        f"<h2>{gene_name} - Cross-Species Evolution Summary</h2>",
        f'<div class="note">{comparison["description"]}</div>',
    ]

    # Key metrics
    # Average pairwise identity (exclude diagonal)
    n = len(species)
    dna_pairs = [dna_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    protein_pairs = [protein_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    avg_dna_id = sum(dna_pairs) / len(dna_pairs) if dna_pairs else 0
    avg_protein_id = sum(protein_pairs) / len(protein_pairs) if protein_pairs else 0
    avg_plddt = sum(d["plddt_mean"] for d in structures.values()) / len(structures) if structures else 0

    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{n}</div><div class="label">Species</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_dna_id:.0%}</div><div class="label">Avg DNA Identity</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_protein_id:.0%}</div><div class="label">Avg Protein Identity</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_plddt:.1f}</div><div class="label">Avg pLDDT</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{len(structures)}</div><div class="label">Structures Folded</div></div>')
    html_parts.append("</div>")

    # Full comparison table
    html_parts.append("<h3>Per-Species Detail</h3>")
    html_parts.append(
        "<table><tr><th>Species</th><th>Scientific Name</th><th>DNA (bp)</th>"
        "<th>Protein (aa)</th><th>GC%</th><th>Carbon LL</th><th>pLDDT</th></tr>"
    )
    for sp in species:
        s = scores[sp]
        plddt = structures.get(sp, {}).get("plddt_mean", "N/A")
        plddt_str = f"{plddt:.1f}" if isinstance(plddt, float) else plddt
        html_parts.append(
            f'<tr><td><b>{sp}</b></td><td><i>{s["common_name"]}</i></td>'
            f'<td>{s["length"]}</td><td>{s["protein_length"]}</td>'
            f'<td>{s["gc_content"]:.1%}</td><td>{s["log_likelihood"]:.2f}</td>'
            f'<td>{plddt_str}</td></tr>'
        )
    html_parts.append("</table>")

    # GC content comparison
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species,
        {"GC Content": [scores[s]["gc_content"] for s in species]},
        title="GC Content Across Species",
        value_format=".2f",
    ))
    html_parts.append("</div>")

    # DNA phylogenetic tree
    html_parts.append("<h3>Phylogenetic Relationships</h3>")
    html_parts.append(
        '<div class="note">'
        "Trees built from pairwise sequence identity using UPGMA clustering. "
        "Species that diverged more recently cluster together. DNA and protein trees "
        "may differ when synonymous mutations dominate."
        "</div>"
    )

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species, dna_matrix,
        title=f"{gene_name} - DNA Phylogenetic Tree",
    ))
    html_parts.append("</div>")

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species, protein_matrix,
        title=f"{gene_name} - Protein Phylogenetic Tree",
        color="#059669",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    summary = {
        "gene_name": gene_name,
        "n_species": n,
        "avg_dna_identity": round(avg_dna_id, 4),
        "avg_protein_identity": round(avg_protein_id, 4),
        "avg_plddt": round(avg_plddt, 1),
        "n_structures": len(structures),
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Pipeline orchestrator
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    gene_set: str = "insulin",
    model_name: str = "HuggingFaceBio/Carbon-3B",
    custom_json: str = "",
) -> tuple[str, str]:
    """
    End-to-end cross-species gene comparison pipeline.

    Returns (comparison JSON, structures JSON).

    1. Load homologous gene sequences across species
    2. Score with Carbon genomic language model
    3. Align sequences and compute pairwise similarity
    4. Fold translated proteins with ESMFold
    5. Generate comprehensive summary with phylogenetic trees
    """
    log.info(f"Starting cross-species gene comparison pipeline (gene_set={gene_set})...")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = [
            "Load Genes",
            "Carbon Scoring",
            "Sequence Alignment",
            "ESMFold Structures",
            "Generate Summary",
        ]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#2563eb;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#2563eb;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>Cross-Species Gene Comparison</h2>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    # Stage 1
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Loading homologous gene sequences...")),
        do_flush=True,
    )
    genes_dir = await load_genes(gene_set=gene_set, custom_json=custom_json)

    # Stage 2
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Scoring sequences with Carbon...")),
        do_flush=True,
    )
    scores_json = await score_sequences(genes_dir=genes_dir, model_name=model_name)

    # Stage 3
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Aligning sequences with Needleman-Wunsch...")),
        do_flush=True,
    )
    comparison_json = await align_and_compare(scores_json=scores_json, genes_dir=genes_dir)

    # Stage 4
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Folding proteins with ESMFold...")),
        do_flush=True,
    )
    structures_json = await fold_proteins(comparison_json=comparison_json)

    # Stage 5
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(5, "Generating summary report...")),
        do_flush=True,
    )
    summary_json = await generate_summary(
        comparison_json=comparison_json,
        structures_json=structures_json,
    )

    # Final report
    summary = json.loads(summary_json)
    comparison = json.loads(comparison_json)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{summary['gene_name']}</div><div class="label">Gene</div></div>
      <div class="stat"><div class="value">{summary['n_species']}</div><div class="label">Species</div></div>
      <div class="stat"><div class="value">{summary['avg_dna_identity']:.0%}</div><div class="label">Avg DNA Identity</div></div>
      <div class="stat"><div class="value">{summary['avg_protein_identity']:.0%}</div><div class="label">Avg Protein Identity</div></div>
      <div class="stat"><div class="value">{summary['avg_plddt']:.1f}</div><div class="label">Avg pLDDT</div></div>
      <div class="stat"><div class="value">{summary['n_structures']}</div><div class="label">3D Structures</div></div>
    </div>
    <div class="card">
      <b>Gene:</b> {summary['gene_name']} |
      <b>Species:</b> {', '.join(comparison['species'])} |
      <b>Model:</b> {model_name}
    </div>
    <div class="note">
      All 4 pipeline stages completed. View individual task reports for DNA/protein
      identity heatmaps, phylogenetic trees, interactive 3D protein structures with
      pLDDT confidence, Carbon log-likelihood scores, and evolutionary analysis.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)
    log.info("Pipeline complete.")
    return comparison_json, structures_json

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_gene_comparison/genomic_gene_comparison.py*

Dependencies are declared at the top of the file using the `uv` script style:

CODE0

## Orchestrate the pipeline

The top-level `pipeline` task chains four stages: load genes, Carbon scoring, sequence alignment, ESMFold folding, and a cross-species summary report.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="genomic-gene-comparison", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="genomic-gene-comparison-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="32Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="genomic-gene-comparison-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Homologous gene sets - same gene across species
# ------------------------------------------------------------------
# Full-length coding sequences from NCBI RefSeq (stop codon excluded).

GENE_SETS = {
    "insulin": {
        "gene_name": "Insulin",
        "description": "Insulin regulates blood sugar in all vertebrates. Highly conserved across 500M+ years of evolution - even fish insulin can lower blood sugar in humans. Comparing across species reveals which regions are functionally essential (conserved) vs free to drift.",
        "sequences": {
            "Human": {
                "dna": "ATGGCCCTGTGGATGCGCCTCCTGCCCCTGCTGGCGCTGCTGGCCCTCTGGGGACCTGACCCAGCCGCAGCCTTTGTGAACCAACACCTGTGCGGCTCACACCTGGTGGAAGCTCTCTACCTAGTGTGCGGGGAACGAGGCTTCTTCTACACACCCAAGACCCGCCGGGAGGCAGAGGACCTGCAGGTGGGGCAGGTGGAGCTGGGCGGGGGCCCTGGTGCAGGCAGCCTGCAGCCCTTGGCCCTGGAGGGGTCCCTGCAGAAGCGTGGCATTGTGGAACAATGCTGTACCAGCATCTGCTCCCTCTACCAGCTGGAGAACTACTGCAAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGGCCCTGTGGATGCGCTTCCTGCCCCTGCTGGCCCTGCTCTTCCTCTGGGAGTCCCACCCCACCCAGGCTTTTGTCAAGCAGCACCTTTGTGGTTCCCACCTGGTGGAGGCTCTCTACCTGGTGTGTGGGGAGCGTGGCTTCTTCTACACACCCATGTCCCGCCGTGAAGTGGAGGACCCACAAGTGGCACAACTGGAGCTGGGTGGAGGCCCGGGAGCAGGTGACCTTCAGACCTTGGCACTGGAGGTGGCCCAGCAGAAGCGTGGCATTGTAGATCAGTGCTGCACCAGCATCTGCTCCCTCTACCAGCTGGAGAACTACTGCAAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGCTCTCTGGATCCGATCACTGCCTCTTCTGGCTCTCCTTGTCTTTTCTGGCCCTGGAACCAGCTATGCAGCTGCCAACCAGCACCTCTGTGGCTCCCACTTGGTGGAGGCTCTCTACCTGGTGTGTGGAGAGCGTGGCTTCTTCTACTCCCCCAAAGCCCGACGGGATGTCGAGCAGCCCCTAGTGAGCAGTCCCTTGCGTGGCGAGGCAGGAGTGCTGCCTTTCCAGCAGGAGGAATACGAGAAAGTCAAGCGAGGGATTGTTGAGCAATGCTGCCATAACACGTGTTCCCTCTACCAACTGGAGAACTACTGCAAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGCAGTGTGGCTTCAGGCTGGTGCTCTGTTGGTCCTGTTGGTCGTGTCCAGTGTAAGCACTAACCCAGGCACACCGCAGCACCTGTGTGGATCTCATCTGGTCGATGCCCTTTATCTGGTCTGTGGCCCAACAGGCTTCTTCTACAACCCCAAGAGAGACGTTGAGCCCCTTCTGGGTTTCCTTCCTCCTAAATCTGCCCAGGAAACTGAGGTGGCTGACTTTGCATTTAAAGATCATGCCGAGCTGATAAGGAAGAGAGGCATTGTAGAGCAGTGCTGCCACAAACCCTGCAGCATCTTTGAGCTGCAGAACTACTGTAAC",
                "common_name": "Danio rerio",
            },
            "Frog": {
                "dna": "ATGGCTCTATGGATGCAGTGTCTGCCCCTGGTTCTTGTCCTCTTTTTCTCTACACCCAACACCGAAGCTCTAGTTAACCAGCACTTGTGTGGGTCTCACCTGGTAGAAGCCCTGTACTTAGTATGTGGGGATCGAGGCTTCTTCTACTACCCTAAGGTCAAACGGGACATGGAACAAGCACTTGTCAGTGGACCCCAGGATAATGAGTTGGATGGAATGCAGCTCCAGCCTCAGGAGTATCAGAAAATGAAGAGGGGGATTGTGGAGCAATGTTGCCACAGCACATGTTCTCTCTTCCAGCTGGAGAGTTACTGCAAC",
                "common_name": "Xenopus laevis",
            },
            "Cow": {
                "dna": "ATGGCCCTGTGGACACGCCTGGCGCCCCTGCTGGCCCTGCTGGCGCTCTGGGCCCCCGCCCCGGCCCGCGCCTTCGTCAACCAGCATCTGTGTGGCTCCCACCTGGTGGAGGCGCTGTACCTGGTGTGCGGAGAGCGCGGCTTCTTCTACACGCCCAAGGCCCGCCGGGAGGTGGAGGGCCCCCAGGTGGGGGCGCTGGAGCTGGCCGGAGGCCCGGGCGCGGGCGGCCTGGAGGGGCCCCCGCAGAAGCGTGGCATCGTGGAGCAGTGCTGTGCCAGCGTCTGCTCGCTCTACCAGCTGGAGAACTACTGTAAC",
                "common_name": "Bos taurus",
            },
        },
    },
    "hemoglobin": {
        "gene_name": "Hemoglobin Beta",
        "description": "Beta-globin carries oxygen from lungs to tissues. The most studied gene in molecular evolution - sequence differences power the 'molecular clock' hypothesis. Sickle cell mutation (E6V) in humans shows how a single base change creates devastating disease.",
        "sequences": {
            "Human": {
                "dna": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTGGTCTACCCTTGGACCCAGAGGTTCTTTGAGTCCTTTGGGGATCTGTCCACTCCTGATGCTGTTATGGGCAACCCTAAGGTGAAGGCTCATGGCAAGAAAGTGCTCGGTGCCTTTAGTGATGGCCTGGCTCACCTGGACAACCTCAAGGGCACCTTTGCCACACTGAGTGAGCTGCACTGTGACAAGCTGCACGTGGATCCTGAGAACTTCAGGCTCCTGGGCAACGTGCTGGTCTGTGTGCTGGCCCATCACTTTGGCAAAGAATTCACCCCACCAGTGCAGGCTGCCTATCAGAAAGTGGTGGCTGGTGTGGCTAATGCCCTGGCCCACAAGTATCAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGGTGCACCTGACTGATGCTGAGAAGGCTGCTGTCTCTGGCCTGTGGGGAAAGGTGAACGCCGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTTGTCTACCCTTGGACCCAGCGGTACTTTGATAGCTTTGGAGACCTATCCTCTGCCTCTGCTATCATGGGTAATGCCAAAGTGAAGGCCCATGGCAAGAAAGTGATAACTGCCTTTAACGATGGCCTGAATCACTTGGACAGCCTCAAGGGCACCTTTGCCAGCCTCAGTGAGCTCCACTGTGACAAGCTGCATGTGGATCCTGAGAACTTCAGGCTCCTGGGCAATATGATCGTGATTGTGCTGGGCCACCACCTGGGCAAGGATTTCACCCCCGCTGCACAGGCTGCCTTCCAGAAGGTGGTGGCTGGAGTGGCTGCTGCCCTGGCTCACAAGTACCAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGTGCACTGGACTGCTGAGGAGAAGCAGCTCATCACCGGCCTCTGGGGCAAGGTCAATGTGGCCGAATGTGGGGCTGAAGCCCTGGCCAGGCTGCTGATCGTCTACCCCTGGACCCAGAGGTTCTTTGCGTCCTTTGGGAACCTCTCCAGCCCCACTGCCATCCTTGGCAACCCCATGGTCCGCGCCCATGGCAAGAAAGTGCTCACCTCCTTTGGGGATGCTGTGAAGAACCTGGACAACATCAAGAACACCTTCTCCCAACTGTCCGAACTGCATTGTGACAAGCTGCATGTGGACCCCGAGAACTTCAGGCTCCTGGGTGACATCCTCATCATTGTCCTGGCCGCCCACTTCAGCAAGGACTTCACTCCTGAATGCCAGGCTGCCTGGCAGAAGCTGGTCCGCGTGGTGGCCCATGCCCTGGCTCGCAAGTACCAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGTTGAGTGGACAGATGCCGAGCGCACAGCCATCCTTGGCCTGTGGGGAAAGCTCAATATCGATGAAATCGGACCTCAGGCCCTATCCAGATGTCTGATCGTGTATCCCTGGACTCAGAGATATTTCGCCACATTCGGCAACCTGTCAAGCCCCGCTGCGATCATGGGTAACCCCAAAGTGGCAGCTCATGGGAGGACTGTGATGGGAGGTCTTGAGAGAGCCATCAAGAACATGGACAACGTCAAGAACACCTATGCCGCCCTCAGTGTGATGCACTCTGAGAAACTGCATGTGGATCCCGACAACTTCAGGCTTCTCGCTGATTGCATCACCGTTTGCGCTGCCATGAAGTTCGGCCAAGCTGGTTTCAATGCTGATGTCCAGGAGGCCTGGCAGAAGTTTCTGGCTGTGGTCGTTTCTGCTCTGTGCAGACAGTACCAC",
                "common_name": "Danio rerio",
            },
            "Frog": {
                "dna": "ATGGTTCATTGGACAGCTGAAGAGAAGGCCGCCATCACCTCTGTGTGGCAGGAGGTCAACCAGGAGCAAGATGGCCATGATGCACTCACAAGGCTGCTGGTTGTGTACCCCTGGACCCAGAGATACTTCAGCAGTTTTGGAAATCTCGGTAATGCCACAGCTATTGCTGGAAATGTCAAGGTGCGTGCCCATGGCAAGAAGGTTCTTTCAGCTGTTGGTGATGCCATCGCCCATCTTGACAACGTGAAGGGAACTCTCCATGACCTCAGTGTGGTCCACGCCTTCAAGCTCTATGTGGATCCTGAGAACTTCAAGCGTCTTGGTGAAGTTCTGGTCATTGTCTTGGCTTCCAAACTGGGATCAGCCTTTACTCCTCAAGTCCAGGGAGCCTGGGAGAAATTTGTTGCTGTTCTGGTTGATGCCCTCAGCCAAGGATACAAC",
                "common_name": "Xenopus laevis",
            },
            "Cow": {
                "dna": "ATGCTGACTGCTGAGGAGAAGGCTGCCGTCACCGCCTTTTGGGGCAAGGTGAAAGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTTGTCTACCCCTGGACTCAGAGGTTCTTTGAGTCCTTTGGGGACTTGTCCACTGCTGATGCTGTTATGAACAACCCTAAGGTGAAGGCCCATGGCAAGAAGGTGCTAGATTCCTTTAGTAATGGCATGAAGCATCTCGATGACCTCAAGGGCACCTTTGCTGCGCTGAGTGAGCTGCACTGTGATAAGCTGCATGTGGATCCTGAGAACTTCAAGCTCCTGGGCAACGTGCTAGTGGTTGTGCTGGCTCGCAATTTTGGCAAGGAATTCACCCCGGTGCTGCAGGCTGACTTTCAGAAGGTGGTGGCTGGTGTGGCCAATGCCCTGGCCCACAGATATCAT",
                "common_name": "Bos taurus",
            },
        },
    },
    "p53": {
        "gene_name": "p53 (TP53)",
        "description": "The 'guardian of the genome' - p53 detects DNA damage and triggers repair or cell death. Mutated in >50% of human cancers. Elephants have 20 copies of p53 (humans have 1), which may explain their extremely low cancer rates despite their size (Peto's paradox).",
        "sequences": {
            "Human": {
                "dna": "ATGGAGGAGCCGCAGTCAGATCCTAGCGTCGAGCCCCCTCTGAGTCAGGAAACATTTTCAGACCTATGGAAACTACTTCCTGAAAACAACGTTCTGTCCCCCTTGCCGTCCCAAGCAATGGATGATTTGATGCTGTCCCCGGACGATATTGAACAATGGTTCACTGAAGACCCAGGTCCAGATGAAGCTCCCAGAATGCCAGAGGCTGCTCCCCCCGTGGCCCCTGCACCAGCAGCTCCTACACCGGCGGCCCCTGCACCAGCCCCCTCCTGGCCCCTGTCATCTTCTGTCCCTTCCCAGAAAACCTACCAGGGCAGCTACGGTTTCCGTCTGGGCTTCTTGCATTCTGGGACAGCCAAGTCTGTGACTTGCACGTACTCCCCTGCCCTCAACAAGATGTTTTGCCAACTGGCCAAGACCTGCCCTGTGCAGCTGTGGGTTGATTCCACACCCCCGCCCGGCACCCGCGTCCGCGCCATGGCCATCTACAAGCAGTCACAGCACATGACGGAGGTTGTGAGGCGCTGCCCCCACCATGAGCGCTGCTCAGATAGCGATGGTCTGGCCCCTCCTCAGCATCTTATCCGAGTGGAAGGAAATTTGCGTGTGGAGTATTTGGATGACAGAAACACTTTTCGACATAGTGTGGTGGTGCCCTATGAGCCGCCTGAGGTTGGCTCTGACTGTACCACCATCCACTACAACTACATGTGTAACAGTTCCTGCATGGGCGGCATGAACCGGAGGCCCATCCTCACCATCATCACACTGGAAGACTCCAGTGGTAATCTACTGGGACGGAACAGCTTTGAGGTGCGTGTTTGTGCCTGTCCTGGGAGAGACCGGCGCACAGAGGAAGAGAATCTCCGCAAGAAAGGGGAGCCTCACCACGAGCTGCCCCCAGGGAGCACTAAGCGAGCACTGCCCAACAACACCAGCTCCTCTCCCCAGCCAAAGAAGAAACCACTGGATGGAGAATATTTCACCCTTCAGATCCGTGGGCGTGAGCGCTTCGAGATGTTCCGAGAGCTGAATGAGGCCTTGGAACTCAAGGATGCCCAGGCTGGGAAGGAGCCAGGGGGGAGCAGGGCTCACTCCAGCCACCTGAAGTCCAAAAAGGGTCAGTCTACCTCCCGCCATAAAAAACTCATGTTCAAGACAGAAGGGCCTGACTCAGAC",
                "common_name": "Homo sapiens",
            },
            "Mouse": {
                "dna": "ATGACTGCCATGGAGGAGTCACAGTCGGATATCAGCCTCGAGCTCCCTCTGAGCCAGGAGACATTTTCAGGCTTATGGAAACTACTTCCTCCAGAAGATATCCTGCCATCACCTCACTGCATGGACGATCTGTTGCTGCCCCAGGATGTTGAGGAGTTTTTTGAAGGCCCAAGTGAAGCCCTCCGAGTGTCAGGAGCTCCTGCAGCACAGGACCCTGTCACCGAGACCCCTGGGCCAGTGGCCCCTGCCCCAGCCACTCCATGGCCCCTGTCATCTTTTGTCCCTTCTCAAAAAACTTACCAGGGCAACTATGGCTTCCACCTGGGCTTCCTGCAGTCTGGGACAGCCAAGTCTGTTATGTGCACGTACTCTCCTCCCCTCAATAAGCTATTCTGCCAGCTGGCGAAGACGTGCCCTGTGCAGTTGTGGGTCAGCGCCACACCTCCAGCTGGGAGCCGTGTCCGCGCCATGGCCATCTACAAGAAGTCACAGCACATGACGGAGGTCGTGAGACGCTGCCCCCACCATGAGCGCTGCTCCGATGGTGATGGCCTGGCTCCTCCCCAGCATCTTATCCGGGTGGAAGGAAATTTGTATCCCGAGTATCTGGAAGACAGGCAGACTTTTCGCCACAGCGTGGTGGTACCTTATGAGCCACCCGAGGCCGGCTCTGAGTATACCACCATCCACTACAAGTACATGTGTAATAGCTCCTGCATGGGGGGCATGAACCGCCGACCTATCCTTACCATCATCACACTGGAAGACTCCAGTGGGAACCTTCTGGGACGGGACAGCTTTGAGGTTCGTGTTTGTGCCTGCCCTGGGAGAGACCGCCGTACAGAAGAAGAAAATTTCCGCAAAAAGGAAGTCCTTTGCCCTGAACTGCCCCCAGGGAGCGCAAAGAGAGCGCTGCCCACCTGCACAAGCGCCTCTCCCCCGCAAAAGAAAAAACCACTTGATGGAGAGTATTTCACCCTCAAGATCCGCGGGCGTAAACGCTTCGAGATGTTCCGGGAGCTGAATGAGGCCTTAGAGTTAAAGGATGCCCATGCTACAGAGGAGTCTGGAGACAGCAGGGCTCACTCCAGCTACCTGAAGACCAAGAAGGGCCAGTCTACTTCCCGCCATAAAAAAACAATGGTCAAGAAAGTGGGGCCTGACTCAGAC",
                "common_name": "Mus musculus",
            },
            "Chicken": {
                "dna": "ATGGCGGAGGAGATGGAACCATTGCTGGAACCCACTGAGGTCTTCATGGACCTCTGGAGCATGCTCCCCTATAGCATGCAACAGCTGCCCCTCCCTGAGGATCACAGCAACTGGCAGGAGCTGAGCCCCCTGGAACCCAGCGACCCCCCCCCACCACCGCCACCACCACCTCTGCCATTGGCCGCCGCCGCCCCCCCCCCATTAAACCCCCCCACCCCCCCCCGCGCTGCCCCCTCCCCGGTGGTCCCATCCACGGAGGATTATGGGGGGGACTTCGACTTCCGGGTGGGGTTCGTGGAGGCGGGCACAGCCAAATCGGTCACCTGCACTTACTCCCCGGTGCTGAATAAGGTCTATTGCCGCCTGGCCAAGCCGTGCCCGGTGCAGGTGAGGGTGGGGGTGGCGCCCCCCCCCGGTTCCTCCCTCCGCGCCGTGGCCGTCTATAAGAAATCAGAGCACGTGGCCGAAGTGGTGCGGCGCTGCCCCCACCACGAGCGCTGCGGGGGGGGCACCGACGGCCTGGCCCCCGCACAGCACCTCATCCGGGTGGAGGGGAACCCCCAGGCGCGTTACCACGACGACGAGACCACCAAACGGCACAGCGTCGTCGTCCCCTATGAGCCCCCCGAGGTGGGCTCTGACTGTACCACGGTGCTGTACAACTTCATGTGCAACAGTTCCTGCATGGGGGGGATGAACCGCCGCCCCATCCTCACCATCCTTACACTGGAGGGGCCGGGGGGGCAGCTGTTGGGGCGGCGCTGCTTCGAGGTGCGCGTGTGCGCATGTCCGGGGAGGGACCGCAAGATCGAGGAGGAGAACTTCCGCAAGAGGGGCGGGGCCGGGGGCGTGGCTAAGCGAGCCATGTCGCCCCCAACCGAAGCCCCCGAGCCCCCCAAGAAGCGCGTGCTGAACCCCGACAATGAGATATTCTACCTGCAGGTGCGCGGGCGCCGCCGCTATGAGATGCTGAAGGAGATCAATGAGGCGCTGCAGCTCGCCGAGGGGGGGTCCGCACCGCGGCCTTCCAAAGGCCGCCGTGTGAAGGTGGAGGGACCCCAACCCAGCTGCGGGAAGAAACTGCTGCAAAAAGGCTCGGAC",
                "common_name": "Gallus gallus",
            },
            "Zebrafish": {
                "dna": "ATGGCGCAAAACGACAGCCAAGAGTTCGCGGAGCTCTGGGAGAAGAATTTGATTATTCAGCCCCCAGGTGGTGGCTCTTGCTGGGACATCATTAATGATGAGGAGTACTTGCCGGGATCGTTTGACCCCAATTTTTTTGAAAATGTGCTTGAAGAACAGCCTCAGCCATCCACTCTCCCACCAACATCCACTGTTCCGGAGACAAGCGACTATCCCGGCGATCATGGATTTAGGCTCAGGTTCCCGCAGTCTGGCACAGCAAAATCTGTAACTTGCACTTATTCACCGGACCTGAATAAACTCTTCTGTCAGCTGGCAAAAACTTGCCCCGTTCAAATGGTGGTGGACGTTGCCCCTCCACAGGGCTCCGTGGTTCGAGCCACTGCCATCTATAAGAAGTCCGAGCATGTGGCTGAAGTGGTCCGCAGATGCCCCCATCATGAGCGAACCCCGGATGGAGATAACTTGGCGCCTGCTGGTCATTTGATAAGAGTGGAGGGCAATCAGCGAGCAAATTACAGGGAAGATAACATCACTTTAAGGCATAGTGTTTTTGTCCCATATGAAGCACCACAGCTTGGTGCTGAATGGACAACTGTGCTACTAAACTACATGTGCAATAGCAGCTGCATGGGGGGGATGAACCGCAGGCCCATCCTCACAATCATCACTCTGGAGACTCAGGAAGGTCAGTTGCTGGGCCGGAGGTCTTTTGAGGTGCGTGTGTGTGCATGTCCAGGCAGAGACAGGAAAACTGAGGAGAGCAACTTCAAGAAAGACCAAGAGACCAAAACCATGGCCAAAACCACCACTGGGACCAAACGTAGTTTGGTGAAAGAATCTTCTTCAGCTACATTACGACCTGAGGGGAGCAAAAAGGCCAAGGGCTCCAGCAGCGATGAGGAGATCTTTACCCTGCAGGTGAGGGGCAGGGAGCGTTATGAAATTTTAAAGAAATTGAACGACAGTCTGGAGTTAAGTGATGTGGTGCCTGCCTCAGATGCTGAAAAGTATCGTCAGAAATTCATGACAAAAAACAAAAAAGAGAATCGTGAATCATCTGAGCCCAAACAGGGAAAGAAGCTGATGGTGAAGGACGAAGGAAGAAGCGACTCTGAT",
                "common_name": "Danio rerio",
            },
            "Elephant": {
                "dna": "ATGGAGGAGCCCCAGTCAGATCTCAGCACCGAGCTCCCTCTGAGTCAAGAGACGTTTTCATACTTATGGGAACTCCTTCCTGAGAATCCGGTTCTGTCCCCCACACTACCCCCGGCAGTGGAGGTCATGGACGATCTGCTACTCTCAGAAGACACTGCAAACTGGCTAGAAAGCCAAGTTGAGGCTCAGGGAATGTCCACAACCCCTGCACCAGCCACCCCTACACCGGTGGCCCCCGCACCAGCCACCTCCTGGACCCTGTCATCTTCCGTCCCTTCCCAAAAGACCTACCCTGGCACCTATGGTTTCCGTCTGGGCTTCCTACATTCTGGGACAGCCAAGTCCGTCACCTGCACGTACTCCCCTGACCTTAACAAGCTGTTTTGCCAGCTGGCAAAAACCTGCCCAGTGCAGCTGTGGGTCGCCTCACCACCCCCGCCCGGCACCCGTGTTCGCACCATGGCCATCTACAAGAAGTCAGAGCATATGACGGAGGTCGTCAAGCGCTGCCCCCACCATGAGCGCTGCTCTGACTCTAGCGATGGCCTGGCCCCTCCTCAGCACCTCATCCGGGTGGAAGGAAACCTGCGTGCTGAGTATCTGGAGGACAGCATCACTCTCCGACACAGTGTGGTGGTGCCCTACGAGCCGCCCGAGGTTGGGTCTGACTGTACCACCATCCACTTCAACTTCATGTGTAACAGCTCCTGCATGGGGGGCATGAACCGGCGGCCCATCCTCACCATCATCACACTGGAAGACTCCAGTGGTAATCTGCTGGGACGTAACAGCTTTGAGGTGCGCATTTGTGCCTGTCCTGGAAGAGACAGACGTACAGAAGAAGAAAATTTCCACAAGAAGGGAGAGCCTTGCCCAGAGCCGCCACCCCCTGGGAGGAGCACTAAGCGAGCACTGCCCACCAACACCAGCTCCTCTACCCAGCCAAAGAAGAAGCCACTGGATGAAGAATATTTCACCCTTCAGATCCGTGGGCGTGAACGCTTCAAGATGTTCCTAGAGCTAAATGAGGCCTTGGAGCTGAAGGATGCCCAGGCTGGGAAGGAGCCAGAGGGGAGCCGGGCTCACTCCAGCCCTTCGAAGTCTAAGAAGGGACAGTCTACCTCCCGCCATAAAAAACCAATGTTCAAGAGAGAGGGACCTGACTCAGAC",
                "common_name": "Loxodonta africana",
            },
            "Dog": {
                "dna": "ATGGAGGAGTCGCAGTCAGAGCTCAATATCGACCCCCCTCTGAGCCAGGAGACATTTTCAGAATTGTGGAACCTGCTTCCTGAAAACAATGTTCTGTCTTCGGAGCTGTGCCCAGCAGTGGATGAGCTGCTGCTCCCAGAGAGCGTCGTGAACTGGCTAGACGAAGACTCAGATGATGCTCCCAGGATGCCAGCCACTTCTGCCCCCACAGCCCCTGGACCGGCCCCCTCGTGGCCCCTATCATCCTCTGTCCCTTCCCCGAAGACCTACCCTGGCACCTATGGGTTCCGTTTGGGGTTCCTGCATTCCGGGACAGCCAAGTCTGTTACTTGGACGTACTCCCCTCTCCTCAACAAGTTGTTTTGCCAGCTGGCGAAGACCTGCCCCGTGCAGCTGTGGGTCAGCTCCCCACCCCCACCCAATACCTGCGTCCGCGCTATGGCCATCTATAAGAAGTCGGAGTTCGTGACCGAGGTTGTGCGGCGCTGCCCCCACCATGAACGCTGCTCTGACAGTAGTGACGGTCTTGCCCCTCCTCAGCATCTCATCCGAGTGGAAGGAAATTTGCGGGCCAAGTACCTGGACGACAGAAACACTTTTCGACACAGTGTGGTGGTGCCTTATGAGCCACCCGAGGTTGGCTCTGACTATACCACCATCCACTACAACTACATGTGTAACAGTTCCTGCATGGGAGGCATGAACCGGCGGCCCATCCTCACTATCATCACCCTGGAAGACTCCAGTGGAAACGTGCTGGGACGCAACAGCTTTGAGGTACGCGTTTGTGCCTGTCCCGGGAGAGACCGCCGGACTGAGGAGGAGAATTTCCACAAGAAGGGGGAGCCTTGTCCTGAGCCACCCCCCGGGAGTACCAAGCGAGCACTGCCTCCCAGCACCAGCTCCTCTCCCCCGCAAAAGAAGAAGCCACTAGATGGAGAATATTTCACCCTTCAGATCCGTGGGCGTGAACGCTATGAGATGTTCAGGAATCTGAATGAAGCCTTGGAGCTGAAGGATGCCCAGAGTGGAAAGGAGCCAGGGGGAAGCAGGGCTCACTCCAGCCACCTGAAGGCAAAGAAGGGGCAATCTACCTCTCGCCATAAAAAACTGATGTTCAAGAGAGAAGGGCTTGACTCAGAC",
                "common_name": "Canis lupus familiaris",
            },
        },
    },
}

# Standard genetic code
CODON_TABLE = {
    "TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L", "CTT": "L", "CTC": "L",
    "CTA": "L", "CTG": "L", "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M",
    "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V", "TCT": "S", "TCC": "S",
    "TCA": "S", "TCG": "S", "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
    "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T", "GCT": "A", "GCC": "A",
    "GCA": "A", "GCG": "A", "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*",
    "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q", "AAT": "N", "AAC": "N",
    "AAA": "K", "AAG": "K", "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E",
    "TGT": "C", "TGC": "C", "TGA": "*", "TGG": "W", "CGT": "R", "CGC": "R",
    "CGA": "R", "CGG": "R", "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
    "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G",
}

BASE_COLORS = {"A": "#2ecc71", "T": "#e74c3c", "G": "#f39c12", "C": "#3498db"}

def _translate(dna: str) -> str:
    """Translate DNA to protein in reading frame 0."""
    dna = dna.upper()
    protein = []
    for i in range(0, len(dna) - 2, 3):
        codon = dna[i:i + 3]
        aa = CODON_TABLE.get(codon, "X")
        if aa == "*":
            break
        protein.append(aa)
    return "".join(protein)

def _gc_content(seq: str) -> float:
    if not seq:
        return 0.0
    return sum(1 for b in seq.upper() if b in "GC") / len(seq)

def _sequence_identity(seq1: str, seq2: str, match: int = 2, mismatch: int = -1, gap: int = -2) -> float:
    """Percent identity via Needleman-Wunsch global alignment."""
    if not seq1 or not seq2:
        return 0.0
    n, m = len(seq1), len(seq2)

    # Build score matrix
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        dp[i][0] = dp[i - 1][0] + gap
    for j in range(1, m + 1):
        dp[0][j] = dp[0][j - 1] + gap
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            s = match if seq1[i - 1] == seq2[j - 1] else mismatch
            dp[i][j] = max(dp[i - 1][j - 1] + s, dp[i - 1][j] + gap, dp[i][j - 1] + gap)

    # Traceback to count matches and alignment length
    i, j = n, m
    matches = 0
    aligned = 0
    while i > 0 or j > 0:
        if i > 0 and j > 0:
            s = match if seq1[i - 1] == seq2[j - 1] else mismatch
            if dp[i][j] == dp[i - 1][j - 1] + s:
                if seq1[i - 1] == seq2[j - 1]:
                    matches += 1
                aligned += 1
                i -= 1
                j -= 1
                continue
        if i > 0 and dp[i][j] == dp[i - 1][j] + gap:
            aligned += 1
            i -= 1
        else:
            aligned += 1
            j -= 1

    return matches / aligned if aligned else 0.0

# ------------------------------------------------------------------
# Report styling
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #1e3a5f; border-bottom: 2px solid #2563eb; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #1e40af; margin-top: 20px; }
  .report .card { background: #eff6ff; border: 1px solid #bfdbfe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #dbeafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #1e3a5f; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #1e3a5f; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dbeafe; }
  .report tr:nth-child(even) { background: #eff6ff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-warning { background: #fef3c7; color: #92400e; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #dbeafe; color: #1e40af; }
  .report .chart-container { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #eff6ff; border-left: 4px solid #2563eb; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .structure-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(340px, 1fr)); gap: 16px; margin: 12px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    width: int = 600,
    height: int = 500,
    value_format: str = ".1f",
    color_scale: str = "blue",
) -> str:
    """Generate an SVG heatmap."""
    n_rows = len(matrix)
    n_cols = len(matrix[0]) if matrix else 0
    if not n_rows or not n_cols:
        return ""

    show_values = n_rows <= 10 and n_cols <= 10
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    if color_scale == "blue":
        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 30))
            g = int(255 - t * (255 - 58))
            b = int(255 - t * (255 - 95))
            return f"rgb({r},{g},{b})"
    else:  # green
        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 6))
            g = int(255 - t * (255 - 95))
            b = int(255 - t * (255 - 70))
            return f"rgb({r},{g},{b})"

    ml = max(80, max(len(l) for l in row_labels) * 7 + 10) if row_labels else 80
    mr = 20
    mt = 80
    mb = 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>')

    for j, label in enumerate(col_labels):
        cx = ml + j * cell_w + cell_w / 2
        svg.append(f'<text x="{cx:.1f}" y="{mt - 8}" text-anchor="start" font-size="10" fill="#374151" transform="rotate(-45, {cx:.1f}, {mt - 8})">{label}</text>')

    for i, row_label in enumerate(row_labels):
        ry = mt + i * cell_h + cell_h / 2
        svg.append(f'<text x="{ml - 8}" y="{ry + 4:.1f}" text-anchor="end" font-size="10" fill="#374151">{row_label}</text>')
        for j in range(n_cols):
            val = matrix[i][j]
            color = get_color(val)
            cx = ml + j * cell_w
            cy = mt + i * cell_h
            svg.append(f'<rect x="{cx:.1f}" y="{cy:.1f}" width="{cell_w:.1f}" height="{cell_h:.1f}" fill="{color}" stroke="#fff" stroke-width="1"/>')
            if show_values:
                t = (val - v_min) / v_range
                text_color = "#fff" if t > 0.55 else "#1a1a2e"
                fs = min(10, int(cell_w / 4), int(cell_h / 2.5))
                fs = max(7, fs)
                svg.append(f'<text x="{cx + cell_w / 2:.1f}" y="{cy + cell_h / 2 + 3:.1f}" text-anchor="middle" font-size="{fs}" fill="{text_color}">{val:{value_format}}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_dendrogram(
    names: list[str],
    matrix: list[list[float]],
    title: str = "",
    width: int = 700,
    height: int = 350,
    color: str = "#2563eb",
) -> str:
    """Generate an SVG dendrogram from a similarity matrix using UPGMA."""
    n = len(names)
    if n < 2:
        return ""

    dist = [[1.0 - matrix[i][j] for j in range(n)] for i in range(n)]

    clusters = [{"members": [i], "height": 0.0, "left": None, "right": None} for i in range(n)]
    active = list(range(n))

    while len(active) > 1:
        best_d = float("inf")
        bi, bj = 0, 1
        for ii in range(len(active)):
            for jj in range(ii + 1, len(active)):
                ci, cj = active[ii], active[jj]
                d = 0
                count = 0
                for mi in clusters[ci]["members"]:
                    for mj in clusters[cj]["members"]:
                        d += dist[mi][mj]
                        count += 1
                avg_d = d / count if count else 0
                if avg_d < best_d:
                    best_d = avg_d
                    bi, bj = ii, jj

        ci, cj = active[bi], active[bj]
        new_cluster = {
            "members": clusters[ci]["members"] + clusters[cj]["members"],
            "height": best_d,
            "left": clusters[ci],
            "right": clusters[cj],
        }
        clusters.append(new_cluster)
        new_idx = len(clusters) - 1
        active.pop(bj)
        active.pop(bi)
        active.append(new_idx)

    root = clusters[active[0]]

    max_label_len = max((len(n) for n in names), default=0)
    ml, mr, mt, mb = max(50, max_label_len * 5 + 10), 30, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb
    max_h = root["height"] or 1

    leaf_positions = {}
    leaf_counter = [0]

    def assign_leaves(node):
        if node["left"] is None and node["right"] is None:
            leaf_positions[node["members"][0]] = leaf_counter[0]
            leaf_counter[0] += 1
        else:
            if node["left"]:
                assign_leaves(node["left"])
            if node["right"]:
                assign_leaves(node["right"])

    assign_leaves(root)
    n_leaves = len(leaf_positions)
    leaf_spacing = cw / max(n_leaves - 1, 1)

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="13" font-weight="600" fill="#1a1a2e">{title}</text>')

    def get_x(node):
        if node["left"] is None and node["right"] is None:
            return ml + leaf_positions[node["members"][0]] * leaf_spacing
        return (get_x(node["left"]) + get_x(node["right"])) / 2

    def get_y(h):
        return mt + ch - (h / max_h) * ch

    def draw_node(node):
        if node["left"] is None and node["right"] is None:
            return
        lx = get_x(node["left"])
        rx = get_x(node["right"])
        ly = get_y(node["left"]["height"])
        ry = get_y(node["right"]["height"])
        my = get_y(node["height"])

        svg.append(f'<line x1="{lx:.1f}" y1="{ly:.1f}" x2="{lx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')
        svg.append(f'<line x1="{rx:.1f}" y1="{ry:.1f}" x2="{rx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')
        svg.append(f'<line x1="{lx:.1f}" y1="{my:.1f}" x2="{rx:.1f}" y2="{my:.1f}" stroke="{color}" stroke-width="2"/>')

        if node["left"]:
            draw_node(node["left"])
        if node["right"]:
            draw_node(node["right"])

    draw_node(root)

    for idx, pos in leaf_positions.items():
        x = ml + pos * leaf_spacing
        svg.append(
            f'<text x="{x:.1f}" y="{mt + ch + 14}" text-anchor="start" font-size="10" fill="#374151" '
            f'transform="rotate(40, {x:.1f}, {mt + ch + 14})">{names[idx]}</text>'
        )

    for i in range(5):
        d = max_h * i / 4
        y = get_y(d)
        svg.append(f'<text x="{ml - 4}" y="{y + 3:.1f}" text-anchor="end" font-size="9" fill="#9ca3af">{d:.3f}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    value_format: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart."""
    if not labels:
        return ""

    default_colors = ["#2563eb", "#059669", "#f59e0b", "#dc2626", "#7c3aed"]
    colors = colors or default_colors

    ml, mr, mt, mb = 60, 20, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_min = min(all_vals) if all_vals else 0
    y_max = max(all_vals) if all_vals else 1
    if y_min >= 0:
        y_min_plot = 0
        y_max_plot = y_max * 1.15 or 1
    else:
        y_range = y_max - y_min or 1
        y_min_plot = y_min - y_range * 0.05
        y_max_plot = y_max + y_range * 0.15

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - ((v - y_min_plot) / (y_max_plot - y_min_plot)) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    for i in range(6):
        y_tick = y_min_plot + (y_max_plot - y_min_plot) * i / 5
        py = sy(y_tick)
        svg.append(f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" stroke="#e9ecef" stroke-width="1"/>')
        svg.append(f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" font-size="11" fill="#6c757d">{y_tick:{value_format}}</text>')

    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" height="{max(0, bh):.1f}" fill="{color}" rx="2"/>')
            svg.append(f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" text-anchor="middle" font-size="9" fill="#1a1a2e">{val:{value_format}}</text>')
        lx = gx + n_series * bar_width / 2
        svg.append(f'<text x="{lx:.1f}" y="{mt + ch + 14}" text-anchor="start" font-size="10" fill="#6c757d" transform="rotate(35, {lx:.1f}, {mt + ch + 14})">{label}</text>')

    if title:
        svg.append(f'<text x="{width / 2}" y="22" text-anchor="middle" font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>')

    if n_series > 1:
        lx = ml + cw - len(series) * 110
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(f'<rect x="{lx + si * 110}" y="{mt + ch + 55}" width="12" height="12" rx="2" fill="{color}"/>')
            svg.append(f'<text x="{lx + si * 110 + 16}" y="{mt + ch + 66}" font-size="11" fill="#1a1a2e">{name}</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_plddt_sparkline(values: list[float], width: int = 400, height: int = 50) -> str:
    """pLDDT sparkline with AlphaFold-style coloring."""
    if not values or len(values) < 2:
        return ""

    pad = 4
    cw = width - 2 * pad
    ch = height - 2 * pad
    seg_w = cw / len(values)

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
    ]

    for i, v in enumerate(values):
        x = pad + i * seg_w
        bar_h = (v / 100) * ch
        y = pad + ch - bar_h

        if v >= 90:
            color = "#0053d6"
        elif v >= 70:
            color = "#65cbf3"
        elif v >= 50:
            color = "#ffdb13"
        else:
            color = "#ff7d45"

        svg.append(f'<rect x="{x:.1f}" y="{y:.1f}" width="{max(seg_w, 1):.1f}" height="{bar_h:.1f}" fill="{color}"/>')

    ref_y = pad + ch - (70 / 100) * ch
    svg.append(f'<line x1="{pad}" y1="{ref_y:.1f}" x2="{pad + cw}" y2="{ref_y:.1f}" stroke="#adb5bd" stroke-width="0.5" stroke-dasharray="3,2"/>')

    svg.append("</svg>")
    return "\n".join(svg)

def _outputs_to_pdb(outputs, sequence: str) -> str:
    """Convert ESMFold outputs to PDB format string."""
    import numpy as np

    pos = outputs.positions[0]
    if pos.dim() == 4:
        pos = pos[-1]
    positions = pos.cpu().numpy()
    atom_names = ["N", "CA", "C", "O"]
    aa_3letter = {
        "A": "ALA", "R": "ARG", "N": "ASN", "D": "ASP", "C": "CYS",
        "Q": "GLN", "E": "GLU", "G": "GLY", "H": "HIS", "I": "ILE",
        "L": "LEU", "K": "LYS", "M": "MET", "F": "PHE", "P": "PRO",
        "S": "SER", "T": "THR", "W": "TRP", "Y": "TYR", "V": "VAL",
    }

    pdb_lines = []
    atom_idx = 1
    for res_idx, aa in enumerate(sequence):
        res_name = aa_3letter.get(aa, "UNK")
        for atom_i, atom_name in enumerate(atom_names):
            if atom_i >= positions.shape[1]:
                break
            x, y, z = positions[res_idx, atom_i]
            if any(math.isnan(c) for c in (x, y, z)):
                continue
            pdb_lines.append(
                f"ATOM  {atom_idx:5d}  {atom_name:<3s} {res_name} A{res_idx + 1:4d}    "
                f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {atom_name[0]:>2s}"
            )
            atom_idx += 1
    pdb_lines.append("END")
    return "\n".join(pdb_lines)

# ------------------------------------------------------------------
# Task 1: Load gene set
# ------------------------------------------------------------------

@cpu_env.task()
async def load_genes(
    gene_set: str = "insulin",
    custom_json: str = "",
) -> flyte.io.Dir:
    """Load a set of homologous genes from different species."""
    if custom_json:
        data = json.loads(custom_json)
    elif gene_set in GENE_SETS:
        data = GENE_SETS[gene_set]
    else:
        available = ", ".join(GENE_SETS.keys())
        raise ValueError(f"Unknown gene set '{gene_set}'. Available: {available}")

    log.info(f"Loaded gene set: {data['gene_name']} - {len(data['sequences'])} species")

    out_dir = tempfile.mkdtemp(prefix="gene_compare_")
    with open(os.path.join(out_dir, "genes.json"), "w") as f:
        json.dump(data, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Score sequences with Carbon
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def score_sequences(
    genes_dir: flyte.io.Dir,
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> str:
    """Score each species' gene with Carbon-3B genomic language model.

    Returns per-species log-likelihood scores and sequence metadata.
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info(f"Loading Carbon model: {model_name}")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, trust_remote_code=True,
        dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    ).to(device)
    model.eval()

    genes_path = await genes_dir.download()
    with open(os.path.join(genes_path, "genes.json")) as f:
        data = json.load(f)

    species_names = list(data["sequences"].keys())
    n = len(species_names)

    scores = {}
    for i, species in enumerate(species_names):
        await flyte.report.replace.aio(_wrap_report(
            f"<h2>Carbon Scoring</h2>"
            f"<p>Scoring {species} ({i + 1}/{n})...</p>"
        ), do_flush=True)

        dna = data["sequences"][species]["dna"]
        prompt = f"<dna>{dna}"
        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            output = model(**inputs, labels=inputs["input_ids"])
            loss = output.loss.item()
            ll = -loss * inputs["input_ids"].shape[1]

        protein = _translate(dna)
        scores[species] = {
            "log_likelihood": round(ll, 4),
            "loss": round(loss, 4),
            "gc_content": round(_gc_content(dna), 4),
            "length": len(dna),
            "protein": protein,
            "protein_length": len(protein),
            "common_name": data["sequences"][species]["common_name"],
        }
        log.info(f"  {species} ({data['sequences'][species]['common_name']}): LL={ll:.2f}, GC={_gc_content(dna):.1%}")

    # Report
    html_parts = [
        f"<h2>{data['gene_name']} - Carbon Scoring</h2>",
        f'<div class="note">{data["description"]}</div>',
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{n}</div><div class="label">Species</div></div>',
        f'<div class="stat"><div class="value">{data["gene_name"]}</div><div class="label">Gene</div></div>',
        f'<div class="stat"><div class="value">{model_name.split("/")[-1]}</div><div class="label">Model</div></div>',
        "</div>",
    ]

    html_parts.append(
        "<table><tr><th>Species</th><th>Scientific Name</th><th>DNA Length</th>"
        "<th>GC%</th><th>Protein Length</th><th>Carbon LL</th></tr>"
    )
    for species in species_names:
        s = scores[species]
        html_parts.append(
            f'<tr><td><b>{species}</b></td><td><i>{s["common_name"]}</i></td>'
            f'<td>{s["length"]}bp</td><td>{s["gc_content"]:.1%}</td>'
            f'<td>{s["protein_length"]}aa</td>'
            f'<td>{s["log_likelihood"]:.2f}</td></tr>'
        )
    html_parts.append("</table>")

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species_names,
        {"Log-Likelihood": [scores[s]["log_likelihood"] for s in species_names]},
        title="Carbon Log-Likelihood per Species",
        value_format=".1f",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    result = {
        "gene_name": data["gene_name"],
        "description": data["description"],
        "species": species_names,
        "scores": scores,
    }
    return json.dumps(result)

# ------------------------------------------------------------------
# Task 3: Align sequences and compute similarity
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def align_and_compare(
    scores_json: str,
    genes_dir: flyte.io.Dir,
) -> str:
    """Align sequences with Needleman-Wunsch and compute pairwise identity.

    Translates DNA to protein, builds DNA and protein identity matrices,
    and generates phylogenetic trees from sequence divergence.
    """
    scores_data = json.loads(scores_json)
    species_names = scores_data["species"]
    scores = scores_data["scores"]
    gene_name = scores_data["gene_name"]
    n = len(species_names)

    genes_path = await genes_dir.download()
    with open(os.path.join(genes_path, "genes.json")) as f:
        data = json.load(f)

    await flyte.report.replace.aio(_wrap_report(
        f"<h2>{gene_name} - Sequence Alignment</h2>"
        f"<p>Aligning {n} species with Needleman-Wunsch...</p>"
    ), do_flush=True)

    # Pairwise DNA identity matrix
    identity_matrix = []
    for sp1 in species_names:
        row = []
        for sp2 in species_names:
            dna1 = data["sequences"][sp1]["dna"]
            dna2 = data["sequences"][sp2]["dna"]
            identity = _sequence_identity(dna1, dna2)
            row.append(round(identity, 4))
        identity_matrix.append(row)

    # Pairwise protein identity matrix
    protein_matrix = []
    for sp1 in species_names:
        row = []
        for sp2 in species_names:
            identity = _sequence_identity(scores[sp1]["protein"], scores[sp2]["protein"])
            row.append(round(identity, 4))
        protein_matrix.append(row)

    # Average pairwise identities (exclude diagonal)
    dna_pairs = [identity_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    prot_pairs = [protein_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    avg_dna = sum(dna_pairs) / len(dna_pairs) if dna_pairs else 0
    avg_prot = sum(prot_pairs) / len(prot_pairs) if prot_pairs else 0

    # Most/least similar pair
    best_pair = max(range(len(dna_pairs)), key=lambda k: dna_pairs[k])
    worst_pair = min(range(len(dna_pairs)), key=lambda k: dna_pairs[k])
    pair_indices = [(i, j) for i in range(n) for j in range(i + 1, n)]
    best_sp = f"{species_names[pair_indices[best_pair][0]]}-{species_names[pair_indices[best_pair][1]]}"
    worst_sp = f"{species_names[pair_indices[worst_pair][0]]}-{species_names[pair_indices[worst_pair][1]]}"

    # Report
    html_parts = [
        f"<h2>{gene_name} - Sequence Alignment</h2>",
        f'<div class="note">Pairwise alignment using Needleman-Wunsch (match=2, mismatch=-1, gap=-2). '
        f"Identity is computed as matches / aligned length from the optimal global alignment.</div>",
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{n}</div><div class="label">Species Aligned</div></div>',
        f'<div class="stat"><div class="value">{n * (n - 1) // 2}</div><div class="label">Pairwise Alignments</div></div>',
        f'<div class="stat"><div class="value">{avg_dna:.0%}</div><div class="label">Avg DNA Identity</div></div>',
        f'<div class="stat"><div class="value">{avg_prot:.0%}</div><div class="label">Avg Protein Identity</div></div>',
        f'<div class="stat"><div class="value">{best_sp}</div><div class="label">Most Similar</div></div>',
        f'<div class="stat"><div class="value">{worst_sp}</div><div class="label">Most Divergent</div></div>',
        "</div>",
    ]

    # DNA identity heatmap
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        identity_matrix, species_names, species_names,
        title="Pairwise DNA Sequence Identity (%)",
        value_format=".0%",
    ))
    html_parts.append("</div>")

    # Protein identity heatmap
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        protein_matrix, species_names, species_names,
        title="Pairwise Protein Sequence Identity (%)",
        value_format=".0%",
        color_scale="green",
    ))
    html_parts.append("</div>")

    # DNA phylogenetic tree
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species_names, identity_matrix,
        title=f"{gene_name} - Phylogenetic Tree (DNA Identity)",
    ))
    html_parts.append("</div>")

    # Protein phylogenetic tree
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species_names, protein_matrix,
        title=f"{gene_name} - Phylogenetic Tree (Protein Identity)",
        color="#059669",
    ))
    html_parts.append("</div>")

    # DNA vs Protein conservation comparison
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species_names,
        {
            "DNA vs Human": [identity_matrix[0][j] for j in range(n)],
            "Protein vs Human": [protein_matrix[0][j] for j in range(n)],
        },
        title=f"Conservation vs {species_names[0]} (DNA and Protein)",
        value_format=".0%",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    result = {
        "gene_name": gene_name,
        "description": scores_data["description"],
        "species": species_names,
        "scores": scores,
        "dna_identity_matrix": identity_matrix,
        "protein_identity_matrix": protein_matrix,
    }
    return json.dumps(result)

# ------------------------------------------------------------------
# Task 4: Fold proteins with ESMFold
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def fold_proteins(
    comparison_json: str,
    max_length: int = 400,
) -> str:
    """Fold each species' translated protein with ESMFold for 3D comparison.

    Returns PDB strings and pLDDT confidence scores for each species.
    """
    import torch
    import numpy as np
    from transformers import AutoTokenizer, EsmForProteinFolding

    comparison = json.loads(comparison_json)
    species_names = comparison["species"]
    scores = comparison["scores"]

    log.info("Loading ESMFold model...")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
    model = model.to(device)
    model.eval()

    structure_data = {}
    n = len(species_names)

    for idx, species in enumerate(species_names):
        protein = scores[species]["protein"]

        if len(protein) > max_length:
            log.info(f"Skipping {species} ({len(protein)} aa > {max_length} max)")
            continue

        log.info(f"ESMFold [{idx + 1}/{n}]: {species} ({len(protein)} aa)")
        await flyte.report.replace.aio(_wrap_report(
            f"<h2>ESMFold - 3D Structure Prediction</h2>"
            f"<p>Folding {species} ({idx + 1}/{n}): {len(protein)} residues...</p>"
        ), do_flush=True)

        inputs = tokenizer(protein, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            outputs = model(**inputs)

        pdb_str = _outputs_to_pdb(outputs, protein)

        plddt_raw = outputs.plddt[0].cpu().numpy()
        if plddt_raw.ndim == 2:
            plddt_raw = plddt_raw[-1]
        plddt = plddt_raw.flatten()[:len(protein)]
        if plddt.max() <= 1.0:
            plddt = plddt * 100
        plddt_mean = float(np.mean(plddt))

        structure_data[species] = {
            "pdb_str": pdb_str,
            "plddt_mean": round(plddt_mean, 1),
            "plddt_per_residue": [round(float(v), 1) for v in plddt[:len(protein)]],
            "protein_length": len(protein),
        }
        log.info(f"  → mean pLDDT: {plddt_mean:.1f}")

    # Report with 3D viewers
    n_folded = len(structure_data)
    avg_plddt = sum(d["plddt_mean"] for d in structure_data.values()) / n_folded if n_folded else 0

    threeDmol_script = '<script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>'

    stats_html = f"""
    <h2>ESMFold - Cross-Species Structure Comparison</h2>
    <div class="note">
      <b>ESMFold</b> predicts 3D structure directly from amino acid sequence.
      Comparing structures across species reveals which parts of the protein are
      structurally conserved (functional core) vs divergent (surface loops, species-specific adaptations).
    </div>
    <div class="stat-grid">
      <div class="stat"><div class="value">{n_folded}</div><div class="label">Structures</div></div>
      <div class="stat"><div class="value">{avg_plddt:.1f}</div><div class="label">Avg pLDDT</div></div>
      <div class="stat"><div class="value">{comparison['gene_name']}</div><div class="label">Gene</div></div>
    </div>
    """

    viewers_html = '<div class="structure-grid">'
    for species, sdata in structure_data.items():
        plddt_val = sdata["plddt_mean"]
        common = scores[species]["common_name"]

        if plddt_val >= 90:
            badge = '<span class="badge badge-success">Very High</span>'
        elif plddt_val >= 70:
            badge = '<span class="badge badge-info">Confident</span>'
        elif plddt_val >= 50:
            badge = '<span class="badge badge-warning">Low</span>'
        else:
            badge = '<span class="badge badge-danger">Disordered</span>'

        plddt_sparkline = _make_plddt_sparkline(sdata["plddt_per_residue"], width=300)
        pdb_escaped = sdata["pdb_str"].replace("\\", "\\\\").replace("`", "\\`").replace("$", "\\$")
        viewer_id = f"viewer_{hash(species) & 0xFFFFFF:06x}"

        viewers_html += f"""
        <div class="card" style="margin:0;">
          <h3 style="margin-top:0;">{species}
            <span style="font-size:0.7em;color:#6c757d;">({sdata['protein_length']} aa)</span>
            {badge}
          </h3>
          <p style="font-size:0.85em;color:#6c757d;margin:2px 0 8px;"><i>{common}</i></p>
          <div id="{viewer_id}" style="width:100%;max-width:320px;height:280px;border:1px solid #dbeafe;border-radius:8px;position:relative;"></div>
          <div style="margin-top:8px;">
            <b>Mean pLDDT:</b> {plddt_val:.1f} / 100
            <div style="margin-top:4px;">{plddt_sparkline}</div>
            <div style="font-size:0.75em;color:#9ca3af;margin-top:2px;">
              <span style="color:#0053d6;">&block; &gt;90</span>
              <span style="color:#65cbf3;">&block; 70-90</span>
              <span style="color:#ffdb13;">&block; 50-70</span>
              <span style="color:#ff7d45;">&block; &lt;50</span>
            </div>
          </div>
        </div>
        <script>
        (function() {{
          var pdb = `{pdb_escaped}`;
          function initViewer() {{
            if (typeof $3Dmol === 'undefined') {{ setTimeout(initViewer, 200); return; }}
            var viewer = $3Dmol.createViewer(document.getElementById("{viewer_id}"), {{backgroundColor: "white"}});
            viewer.addModel(pdb, "pdb");
            viewer.setStyle({{}}, {{cartoon: {{color: "spectrum"}}}});
            viewer.zoomTo();
            viewer.render();
            viewer.spin("y", 1);
          }}
          initViewer();
        }})();
        </script>
        """

    viewers_html += "</div>"

    # pLDDT comparison bar chart
    plddt_chart = _make_bar_chart(
        list(structure_data.keys()),
        {"Mean pLDDT": [d["plddt_mean"] for d in structure_data.values()]},
        title="Structure Confidence Comparison (pLDDT)",
        value_format=".1f",
        colors=["#0053d6"],
    )

    report_html = f"""
    {threeDmol_script}
    {stats_html}
    {viewers_html}
    <div class="chart-container">{plddt_chart}</div>
    """

    await flyte.report.replace.aio(_wrap_report(report_html), do_flush=True)

    return json.dumps(structure_data)

# ------------------------------------------------------------------
# Task 5: Generate summary
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def generate_summary(
    comparison_json: str,
    structures_json: str,
) -> str:
    """Generate comprehensive cross-species summary."""
    comparison = json.loads(comparison_json)
    structures = json.loads(structures_json)

    species = comparison["species"]
    scores = comparison["scores"]
    gene_name = comparison["gene_name"]
    dna_matrix = comparison["dna_identity_matrix"]
    protein_matrix = comparison["protein_identity_matrix"]

    html_parts = [
        f"<h2>{gene_name} - Cross-Species Evolution Summary</h2>",
        f'<div class="note">{comparison["description"]}</div>',
    ]

    # Key metrics
    # Average pairwise identity (exclude diagonal)
    n = len(species)
    dna_pairs = [dna_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    protein_pairs = [protein_matrix[i][j] for i in range(n) for j in range(i + 1, n)]
    avg_dna_id = sum(dna_pairs) / len(dna_pairs) if dna_pairs else 0
    avg_protein_id = sum(protein_pairs) / len(protein_pairs) if protein_pairs else 0
    avg_plddt = sum(d["plddt_mean"] for d in structures.values()) / len(structures) if structures else 0

    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{n}</div><div class="label">Species</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_dna_id:.0%}</div><div class="label">Avg DNA Identity</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_protein_id:.0%}</div><div class="label">Avg Protein Identity</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{avg_plddt:.1f}</div><div class="label">Avg pLDDT</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{len(structures)}</div><div class="label">Structures Folded</div></div>')
    html_parts.append("</div>")

    # Full comparison table
    html_parts.append("<h3>Per-Species Detail</h3>")
    html_parts.append(
        "<table><tr><th>Species</th><th>Scientific Name</th><th>DNA (bp)</th>"
        "<th>Protein (aa)</th><th>GC%</th><th>Carbon LL</th><th>pLDDT</th></tr>"
    )
    for sp in species:
        s = scores[sp]
        plddt = structures.get(sp, {}).get("plddt_mean", "N/A")
        plddt_str = f"{plddt:.1f}" if isinstance(plddt, float) else plddt
        html_parts.append(
            f'<tr><td><b>{sp}</b></td><td><i>{s["common_name"]}</i></td>'
            f'<td>{s["length"]}</td><td>{s["protein_length"]}</td>'
            f'<td>{s["gc_content"]:.1%}</td><td>{s["log_likelihood"]:.2f}</td>'
            f'<td>{plddt_str}</td></tr>'
        )
    html_parts.append("</table>")

    # GC content comparison
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_bar_chart(
        species,
        {"GC Content": [scores[s]["gc_content"] for s in species]},
        title="GC Content Across Species",
        value_format=".2f",
    ))
    html_parts.append("</div>")

    # DNA phylogenetic tree
    html_parts.append("<h3>Phylogenetic Relationships</h3>")
    html_parts.append(
        '<div class="note">'
        "Trees built from pairwise sequence identity using UPGMA clustering. "
        "Species that diverged more recently cluster together. DNA and protein trees "
        "may differ when synonymous mutations dominate."
        "</div>"
    )

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species, dna_matrix,
        title=f"{gene_name} - DNA Phylogenetic Tree",
    ))
    html_parts.append("</div>")

    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_dendrogram(
        species, protein_matrix,
        title=f"{gene_name} - Protein Phylogenetic Tree",
        color="#059669",
    ))
    html_parts.append("</div>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    summary = {
        "gene_name": gene_name,
        "n_species": n,
        "avg_dna_identity": round(avg_dna_id, 4),
        "avg_protein_identity": round(avg_protein_id, 4),
        "avg_plddt": round(avg_plddt, 1),
        "n_structures": len(structures),
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Pipeline orchestrator
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    gene_set: str = "insulin",
    model_name: str = "HuggingFaceBio/Carbon-3B",
    custom_json: str = "",
) -> tuple[str, str]:
    """
    End-to-end cross-species gene comparison pipeline.

    Returns (comparison JSON, structures JSON).

    1. Load homologous gene sequences across species
    2. Score with Carbon genomic language model
    3. Align sequences and compute pairwise similarity
    4. Fold translated proteins with ESMFold
    5. Generate comprehensive summary with phylogenetic trees
    """
    log.info(f"Starting cross-species gene comparison pipeline (gene_set={gene_set})...")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = [
            "Load Genes",
            "Carbon Scoring",
            "Sequence Alignment",
            "ESMFold Structures",
            "Generate Summary",
        ]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#2563eb;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#2563eb;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>Cross-Species Gene Comparison</h2>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    # Stage 1
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Loading homologous gene sequences...")),
        do_flush=True,
    )
    genes_dir = await load_genes(gene_set=gene_set, custom_json=custom_json)

    # Stage 2
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Scoring sequences with Carbon...")),
        do_flush=True,
    )
    scores_json = await score_sequences(genes_dir=genes_dir, model_name=model_name)

    # Stage 3
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Aligning sequences with Needleman-Wunsch...")),
        do_flush=True,
    )
    comparison_json = await align_and_compare(scores_json=scores_json, genes_dir=genes_dir)

    # Stage 4
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Folding proteins with ESMFold...")),
        do_flush=True,
    )
    structures_json = await fold_proteins(comparison_json=comparison_json)

    # Stage 5
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(5, "Generating summary report...")),
        do_flush=True,
    )
    summary_json = await generate_summary(
        comparison_json=comparison_json,
        structures_json=structures_json,
    )

    # Final report
    summary = json.loads(summary_json)
    comparison = json.loads(comparison_json)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{summary['gene_name']}</div><div class="label">Gene</div></div>
      <div class="stat"><div class="value">{summary['n_species']}</div><div class="label">Species</div></div>
      <div class="stat"><div class="value">{summary['avg_dna_identity']:.0%}</div><div class="label">Avg DNA Identity</div></div>
      <div class="stat"><div class="value">{summary['avg_protein_identity']:.0%}</div><div class="label">Avg Protein Identity</div></div>
      <div class="stat"><div class="value">{summary['avg_plddt']:.1f}</div><div class="label">Avg pLDDT</div></div>
      <div class="stat"><div class="value">{summary['n_structures']}</div><div class="label">3D Structures</div></div>
    </div>
    <div class="card">
      <b>Gene:</b> {summary['gene_name']} |
      <b>Species:</b> {', '.join(comparison['species'])} |
      <b>Model:</b> {model_name}
    </div>
    <div class="note">
      All 4 pipeline stages completed. View individual task reports for DNA/protein
      identity heatmaps, phylogenetic trees, interactive 3D protein structures with
      pLDDT confidence, Carbon log-likelihood scores, and evolutionary analysis.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)
    log.info("Pipeline complete.")
    return comparison_json, structures_json

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
CODE1
cd v2/tutorials/genomic_gene_comparison
uv run --script genomic_gene_comparison.py
CODE2
flyte run genomic_gene_comparison.py pipeline --gene_set hemoglobin
```

This example needs a GPU for Carbon and ESMFold. Open the run URL and check each task's report tab for heatmaps, dendrograms, and interactive 3D viewers.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/genomic-variant-effect ===

# Genomic variant effect prediction

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic_variant_effect).

This tutorial demonstrates zero-shot variant effect prediction (VEP) with HuggingFace [Carbon](https://huggingface.co/HuggingFaceBio/Carbon-3B). The pipeline loads clinically relevant variants across genes such as BRCA2, TP53, CFTR, KRAS, and HBB, scores each mutation with a log-likelihood ratio, and produces rich HTML reports with DNA tracks, lollipop plots, confusion matrices, and ranked pathogenicity tables.

Flyte provides:

- **GPU-backed inference** for Carbon scoring with live progress reports.
- **CPU analysis tasks** for visualization and accuracy metrics without holding a GPU.
- **End-to-end orchestration** from variant loading through summary reporting.

## Define the task environments

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="genomic-variant-effect", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="genomic-variant-effect-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="genomic-variant-effect-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default gene variants — clinically relevant mutations
# ------------------------------------------------------------------
# Each entry: gene name -> { "sequence": reference DNA, "variants": [{ "pos": 0-indexed, "ref": base, "alt": base, "name": "...", "known_effect": "..." }] }
# Sequences are short windows (~120-200bp) around the variant site for tractable inference.

DEFAULT_GENE_VARIANTS = {
    "BRCA2 (Breast Cancer)": {
        "description": "Tumor suppressor critical for DNA repair via homologous recombination. Mutations dramatically increase breast and ovarian cancer risk.",
        "sequence": "ATGGCCTCGAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAG",
        "variants": [
            {"pos": 12, "ref": "A", "alt": "T", "name": "c.37A>T", "known_effect": "pathogenic", "clinical": "Nonsense mutation — truncates protein early"},
            {"pos": 18, "ref": "G", "alt": "A", "name": "c.55G>A", "known_effect": "benign", "clinical": "Synonymous — no amino acid change"},
            {"pos": 30, "ref": "C", "alt": "T", "name": "c.91C>T", "known_effect": "pathogenic", "clinical": "Missense in DNA-binding domain"},
            {"pos": 45, "ref": "G", "alt": "C", "name": "c.136G>C", "known_effect": "uncertain", "clinical": "Variant of uncertain significance (VUS)"},
        ],
    },
    "TP53 (Tumor Suppressor)": {
        "description": "Guardian of the genome. Activates DNA repair, cell cycle arrest, and apoptosis. Mutated in >50% of human cancers.",
        "sequence": "ATGGAGGAGCCGCAGTCAGATCCTAGCGTGAGTTTGCACCCTTCAGAGACAGAAACCACTGGATTGGAGACTACTTCCTGAAACAACGTTCTGTCCCCCTTGCCGTCCCAAGCAATGGATGAT",
        "variants": [
            {"pos": 15, "ref": "C", "alt": "T", "name": "R175H", "known_effect": "pathogenic", "clinical": "Hotspot — gain-of-function, dominant negative. Most common TP53 mutation in cancer"},
            {"pos": 36, "ref": "T", "alt": "C", "name": "P72R", "known_effect": "benign", "clinical": "Common polymorphism — subtle effect on apoptosis efficiency"},
            {"pos": 54, "ref": "C", "alt": "A", "name": "G245S", "known_effect": "pathogenic", "clinical": "Contact mutant — disrupts DNA binding"},
            {"pos": 72, "ref": "T", "alt": "G", "name": "R248W", "known_effect": "pathogenic", "clinical": "Structural mutant — destabilizes DNA-binding loop"},
            {"pos": 90, "ref": "C", "alt": "T", "name": "R273H", "known_effect": "pathogenic", "clinical": "Contact mutant — directly contacts DNA bases"},
        ],
    },
    "CFTR (Cystic Fibrosis)": {
        "description": "Chloride channel protein. Mutations cause cystic fibrosis — the most common lethal genetic disease in people of European descent.",
        "sequence": "ATGCAGAGGTCGCCTCTGGAAAAGGCCAGCGTTGTCTCCAAACTTTTTTTCAGCTGGACCAGACCAATTTTGAGGAAAGGATACAGACAGCGCCTGGAATTGTCAGACATATACCAAATCCCTTC",
        "variants": [
            {"pos": 9, "ref": "G", "alt": "A", "name": "G85E", "known_effect": "pathogenic", "clinical": "Disrupts chloride channel processing"},
            {"pos": 24, "ref": "C", "alt": "T", "name": "R117H", "known_effect": "pathogenic", "clinical": "Reduces channel conductance — milder CF phenotype"},
            {"pos": 48, "ref": "T", "alt": "C", "name": "I148T", "known_effect": "benign", "clinical": "Previously misclassified — now known benign polymorphism"},
            {"pos": 66, "ref": "A", "alt": "G", "name": "R334W", "known_effect": "pathogenic", "clinical": "Gating mutation — channel opens less frequently"},
        ],
    },
    "KRAS (Oncogene)": {
        "description": "GTPase signal switch. KRAS mutations are the most common oncogenic driver — found in ~25% of all human cancers, especially pancreatic, colorectal, and lung.",
        "sequence": "ATGACTGAATATAAACTTGTGGTAGTTGGAGCTGGTGGCGTAGGCAAGAGTGCCTTGACGATACAGCTAATTCAGAATCATTTTGTGGACGAATATGATCCAACAATAGAGGATTCCTACAGGAA",
        "variants": [
            {"pos": 34, "ref": "G", "alt": "T", "name": "G12V", "known_effect": "pathogenic", "clinical": "Locks KRAS in active state — constitutive proliferation signal"},
            {"pos": 35, "ref": "G", "alt": "A", "name": "G12D", "known_effect": "pathogenic", "clinical": "Most common KRAS mutation in pancreatic cancer"},
            {"pos": 37, "ref": "G", "alt": "T", "name": "G13D", "known_effect": "pathogenic", "clinical": "Constitutively active — common in colorectal cancer"},
            {"pos": 60, "ref": "C", "alt": "A", "name": "Q61K", "known_effect": "pathogenic", "clinical": "Impairs GTP hydrolysis — locked ON state"},
        ],
    },
    "HBB (Sickle Cell)": {
        "description": "Beta-globin subunit of hemoglobin. The sickle cell mutation (E6V) is the most well-known single-base disease variant in humans.",
        "sequence": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTGGTCTACCCTTGGACCCAGAGG",
        "variants": [
            {"pos": 17, "ref": "A", "alt": "T", "name": "E6V (HbS)", "known_effect": "pathogenic", "clinical": "THE sickle cell mutation — causes hemoglobin polymerization under low O2"},
            {"pos": 19, "ref": "G", "alt": "A", "name": "E6K (HbC)", "known_effect": "pathogenic", "clinical": "Hemoglobin C disease — milder than sickle cell but causes crystal formation"},
            {"pos": 36, "ref": "G", "alt": "A", "name": "E26K", "known_effect": "benign", "clinical": "Hemoglobin E — most common Hb variant worldwide, mild effect"},
            {"pos": 78, "ref": "C", "alt": "T", "name": "Q39X", "known_effect": "pathogenic", "clinical": "Nonsense — causes beta-thalassemia (no functional beta-globin)"},
        ],
    },
}

# DNA base colors (classic genomics color scheme)
BASE_COLORS = {"A": "#2ecc71", "T": "#e74c3c", "G": "#f39c12", "C": "#3498db"}
BASE_COMPLEMENT = {"A": "T", "T": "A", "G": "C", "C": "G"}

# Pathogenicity color scheme
EFFECT_COLORS = {
    "pathogenic": "#dc2626",
    "benign": "#059669",
    "uncertain": "#f59e0b",
}
EFFECT_BADGES = {
    "pathogenic": "badge-danger",
    "benign": "badge-success",
    "uncertain": "badge-warning",
}

# ------------------------------------------------------------------
# Report styling — genomics-themed deep blues and teals
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #1e3a5f; border-bottom: 2px solid #2563eb; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #1e40af; margin-top: 20px; }
  .report .card { background: #eff6ff; border: 1px solid #bfdbfe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #dbeafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #1e3a5f; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #1e3a5f; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dbeafe; }
  .report tr:nth-child(even) { background: #eff6ff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-warning { background: #fef3c7; color: #92400e; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #dbeafe; color: #1e40af; }
  .report .chart-container { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #eff6ff; border-left: 4px solid #2563eb; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .gene-card { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .dna-track { font-family: 'SF Mono', 'Fira Code', monospace; letter-spacing: 1px; }
</style>
"""

def _wrap_report(html: str) -> str:
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    value_format: str = ".2f",
) -> str:
    """Generate an SVG grouped bar chart."""
    if not labels:
        return ""

    default_colors = ["#2563eb", "#1e3a5f", "#3b82f6", "#60a5fa", "#93c5fd"]
    colors = colors or default_colors

    ml, mr, mt, mb = 70, 20, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(abs(v) for v in all_vals) if all_vals else 1
    y_min = min(all_vals) if all_vals else 0
    # For VEP scores (negative = more damaging), we need to handle negative values
    if y_min >= 0:
        y_min_plot = 0
        y_max_plot = y_max * 1.15 or 1
    else:
        y_max_plot = max(y_max * 1.15, 0.1)
        y_min_plot = y_min * 1.15

    y_range = y_max_plot - y_min_plot or 1

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v - y_min_plot) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_min_plot + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_format}}</text>'
        )

    # Zero line
    if y_min_plot < 0 < y_max_plot:
        zy = sy(0)
        svg.append(
            f'<line x1="{ml}" y1="{zy:.1f}" x2="{ml + cw}" y2="{zy:.1f}" '
            f'stroke="#374151" stroke-width="1.5"/>'
        )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            if val >= 0:
                by = sy(val)
                bh = sy(0) - by if y_min_plot < 0 else mt + ch - by
            else:
                by = sy(0) if y_min_plot < 0 else mt + ch
                bh = sy(val) - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{max(0, bh):.1f}" fill="{color}" rx="2"/>'
            )
            text_y = by - 4 if val >= 0 else by + bh + 12
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{text_y:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_format}}</text>'
            )
        # Rotated group label
        lx = gx + n_series * bar_width / 2
        svg.append(
            f'<text x="{lx:.1f}" y="{mt + ch + 14}" '
            f'text-anchor="end" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {lx:.1f}, {mt + ch + 14})">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 110
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 110}" y="{mt + ch + 55}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 110 + 16}" y="{mt + ch + 66}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    width: int = 700,
    height: int = 500,
    value_format: str = ".2f",
    diverging: bool = False,
) -> str:
    """Generate an SVG heatmap. If diverging=True, uses red-white-blue scale centered at 0."""
    n_rows = len(matrix)
    n_cols = len(matrix[0]) if matrix else 0
    if not n_rows or not n_cols:
        return ""

    show_values = n_rows <= 10 and n_cols <= 12

    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)

    if diverging:
        abs_max = max(abs(v_min), abs(v_max)) or 1

        def get_color(v):
            t = v / abs_max  # -1 to 1
            if t < 0:
                # White to red (negative = damaging)
                r = 255
                g = int(255 * (1 + t))
                b = int(255 * (1 + t))
            else:
                # White to blue (positive = benign)
                r = int(255 * (1 - t))
                g = int(255 * (1 - t))
                b = 255
            return f"rgb({r},{g},{b})"
    else:
        v_range = v_max - v_min or 1

        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 30))
            g = int(255 - t * (255 - 58))
            b = int(255 - t * (255 - 95))
            return f"rgb({r},{g},{b})"

    # Layout
    ml = max(140, max(len(l) for l in row_labels) * 7 + 20) if row_labels else 140
    mr = 20
    mt = 80 if col_labels else 40
    mb = 30
    cw = width - ml - mr
    ch = height - mt - mb

    cell_w = cw / n_cols
    cell_h = ch / n_rows

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Column labels (rotated)
    for j, label in enumerate(col_labels):
        cx = ml + j * cell_w + cell_w / 2
        svg.append(
            f'<text x="{cx:.1f}" y="{mt - 8}" text-anchor="end" '
            f'font-size="10" fill="#374151" '
            f'transform="rotate(-45, {cx:.1f}, {mt - 8})">{label}</text>'
        )

    # Row labels + cells
    for i, row_label in enumerate(row_labels):
        ry = mt + i * cell_h + cell_h / 2
        svg.append(
            f'<text x="{ml - 8}" y="{ry + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#374151">{row_label}</text>'
        )
        for j in range(n_cols):
            val = matrix[i][j]
            color = get_color(val)
            cx = ml + j * cell_w
            cy = mt + i * cell_h
            svg.append(
                f'<rect x="{cx:.1f}" y="{cy:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{color}" stroke="#fff" stroke-width="1"/>'
            )
            if show_values:
                if diverging:
                    t = abs(val) / (max(abs(v_min), abs(v_max)) or 1)
                else:
                    t = (val - v_min) / (v_max - v_min or 1)
                text_color = "#fff" if t > 0.55 else "#1a1a2e"
                font_size = min(10, int(cell_w / 4), int(cell_h / 2.5))
                font_size = max(7, font_size)
                svg.append(
                    f'<text x="{cx + cell_w / 2:.1f}" y="{cy + cell_h / 2 + 3:.1f}" '
                    f'text-anchor="middle" font-size="{font_size}" '
                    f'fill="{text_color}">{val:{value_format}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_dna_track(
    sequence: str,
    variants: list[dict],
    gene_name: str = "",
    width: int = 900,
) -> str:
    """Render a color-coded DNA sequence track with variant positions highlighted."""
    chars_per_line = 60
    char_w = 11
    line_h = 22
    label_w = 50
    n_lines = (len(sequence) + chars_per_line - 1) // chars_per_line

    # Extra space for variant annotations
    variant_positions = {v["pos"] for v in variants}
    svg_h = n_lines * line_h + 60

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {svg_h}" '
        f'style="width:100%;max-width:{width}px;height:auto;font-family:monospace;">',
        f'<rect width="{width}" height="{svg_h}" fill="#fff" rx="6"/>',
    ]

    if gene_name:
        svg.append(
            f'<text x="{width / 2}" y="16" text-anchor="middle" '
            f'font-size="12" font-weight="600" fill="#1e3a5f">{gene_name}</text>'
        )

    y_offset = 28

    for line_idx in range(n_lines):
        start = line_idx * chars_per_line
        end = min(start + chars_per_line, len(sequence))
        chunk = sequence[start:end]
        y = y_offset + line_idx * line_h

        # Position label
        svg.append(
            f'<text x="{label_w - 8}" y="{y}" text-anchor="end" '
            f'font-size="9" fill="#9ca3af">{start + 1}</text>'
        )

        for ci, base in enumerate(chunk):
            abs_pos = start + ci
            x = label_w + ci * char_w
            color = BASE_COLORS.get(base, "#6b7280")
            is_variant = abs_pos in variant_positions

            if is_variant:
                # Red highlight for variant positions
                svg.append(
                    f'<rect x="{x - 1}" y="{y - 13}" width="{char_w + 2}" height="{line_h}" '
                    f'fill="#fee2e2" stroke="#dc2626" stroke-width="1.5" rx="2"/>'
                )
                # Small triangle marker above
                svg.append(
                    f'<polygon points="{x + char_w / 2:.1f},{y - 16} {x + char_w / 2 - 3:.1f},{y - 20} {x + char_w / 2 + 3:.1f},{y - 20}" '
                    f'fill="#dc2626"/>'
                )
            else:
                # Subtle background
                svg.append(
                    f'<rect x="{x}" y="{y - 12}" width="{char_w}" height="{line_h - 2}" '
                    f'fill="{color}" opacity="0.08" rx="1"/>'
                )

            svg.append(
                f'<text x="{x + char_w / 2:.1f}" y="{y}" text-anchor="middle" '
                f'font-size="11" font-weight="{"700" if is_variant else "500"}" '
                f'fill="{color}">{base}</text>'
            )

            # Spacer every 10 bases
            if (abs_pos + 1) % 10 == 0 and ci < len(chunk) - 1:
                svg.append(
                    f'<line x1="{x + char_w + 1}" y1="{y - 11}" '
                    f'x2="{x + char_w + 1}" y2="{y + 4}" '
                    f'stroke="#e5e7eb" stroke-width="0.5"/>'
                )

    # Legend
    ly = svg_h - 12
    lx = label_w
    for base, color in BASE_COLORS.items():
        svg.append(f'<rect x="{lx}" y="{ly - 8}" width="8" height="8" rx="1" fill="{color}"/>')
        svg.append(f'<text x="{lx + 12}" y="{ly}" font-size="9" fill="#6b7280">{base}</text>')
        lx += 30
    lx += 10
    svg.append(f'<rect x="{lx}" y="{ly - 8}" width="8" height="8" rx="1" fill="#fee2e2" stroke="#dc2626" stroke-width="1"/>')
    svg.append(f'<text x="{lx + 12}" y="{ly}" font-size="9" fill="#6b7280">Variant site</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_lollipop_plot(
    variants: list[dict],
    seq_length: int,
    title: str = "",
    width: int = 700,
    height: int = 200,
) -> str:
    """Generate a lollipop plot showing variant positions along a gene with effect scores."""
    if not variants:
        return ""

    ml, mr, mt, mb = 50, 30, 40, 40
    cw = width - ml - mr
    ch = height - mt - mb

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Gene track (horizontal bar)
    track_y = mt + ch * 0.7
    svg.append(
        f'<rect x="{ml}" y="{track_y - 4}" width="{cw}" height="8" '
        f'fill="#dbeafe" rx="4" stroke="#93c5fd" stroke-width="1"/>'
    )

    # Position labels at ends
    svg.append(
        f'<text x="{ml}" y="{track_y + 20}" text-anchor="middle" '
        f'font-size="9" fill="#9ca3af">1</text>'
    )
    svg.append(
        f'<text x="{ml + cw}" y="{track_y + 20}" text-anchor="middle" '
        f'font-size="9" fill="#9ca3af">{seq_length}</text>'
    )

    # Lollipops
    scores = [v.get("score", 0) for v in variants]
    max_score = max(abs(s) for s in scores) if scores else 1

    for v in variants:
        pos = v["pos"]
        score = v.get("score", 0)
        effect = v.get("known_effect", "uncertain")
        color = EFFECT_COLORS.get(effect, "#6b7280")

        x = ml + (pos / max(seq_length - 1, 1)) * cw
        stem_h = max(20, abs(score) / max_score * (ch * 0.6))
        circle_y = track_y - stem_h - 6

        # Stem
        svg.append(
            f'<line x1="{x:.1f}" y1="{track_y - 4}" x2="{x:.1f}" y2="{circle_y + 5:.1f}" '
            f'stroke="{color}" stroke-width="2"/>'
        )
        # Circle
        svg.append(
            f'<circle cx="{x:.1f}" cy="{circle_y:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5"/>'
        )
        # Label
        svg.append(
            f'<text x="{x:.1f}" y="{circle_y - 9:.1f}" text-anchor="middle" '
            f'font-size="8" fill="#374151" font-weight="600">{v.get("name", "")}</text>'
        )

    # Legend
    lx = ml
    for effect, color in EFFECT_COLORS.items():
        svg.append(f'<circle cx="{lx + 4}" cy="{height - 10}" r="4" fill="{color}"/>')
        svg.append(
            f'<text x="{lx + 12}" y="{height - 6}" font-size="9" fill="#374151">'
            f'{effect.title()}</text>'
        )
        lx += len(effect) * 7 + 24

    svg.append("</svg>")
    return "\n".join(svg)

def _make_score_comparison_chart(
    gene_results: dict,
    width: int = 800,
    height: int = 350,
) -> str:
    """Generate a dot plot comparing VEP scores across all genes, colored by known effect."""
    all_variants = []
    for gene_name, gene_data in gene_results.items():
        for v in gene_data["variants"]:
            all_variants.append({
                "gene": gene_name.split("(")[0].strip(),
                "name": v["name"],
                "score": v.get("score", 0),
                "known_effect": v["known_effect"],
            })

    if not all_variants:
        return ""

    ml, mr, mt, mb = 120, 30, 40, 30
    cw = width - ml - mr
    ch = height - mt - mb

    scores = [v["score"] for v in all_variants]
    s_min, s_max = min(scores), max(scores)
    s_pad = (s_max - s_min) * 0.1 or 0.5
    s_min -= s_pad
    s_max += s_pad
    s_range = s_max - s_min or 1

    def sx(s):
        return ml + (s - s_min) / s_range * cw

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
        f'<text x="{width / 2}" y="22" text-anchor="middle" '
        f'font-size="14" font-weight="600" fill="#1a1a2e">'
        f'Variant Effect Scores — All Genes</text>',
    ]

    # Grid
    for i in range(6):
        gx = s_min + s_range * i / 5
        px = sx(gx)
        svg.append(
            f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
            f'stroke="#f3f4f6" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 16}" text-anchor="middle" '
            f'font-size="10" fill="#9ca3af">{gx:.2f}</text>'
        )

    # Zero line
    if s_min < 0 < s_max:
        zx = sx(0)
        svg.append(
            f'<line x1="{zx:.1f}" y1="{mt}" x2="{zx:.1f}" y2="{mt + ch}" '
            f'stroke="#374151" stroke-width="1" stroke-dasharray="4,3"/>'
        )

    # Rows per variant
    row_h = ch / max(len(all_variants), 1)
    for i, v in enumerate(all_variants):
        y = mt + i * row_h + row_h / 2
        color = EFFECT_COLORS.get(v["known_effect"], "#6b7280")
        px = sx(v["score"])

        # Label
        label = f'{v["gene"]} {v["name"]}'
        svg.append(
            f'<text x="{ml - 8}" y="{y + 3:.1f}" text-anchor="end" '
            f'font-size="9" fill="#374151">{label}</text>'
        )
        # Connector line
        svg.append(
            f'<line x1="{ml}" y1="{y:.1f}" x2="{px:.1f}" y2="{y:.1f}" '
            f'stroke="{color}" stroke-width="1" opacity="0.3"/>'
        )
        # Dot
        svg.append(
            f'<circle cx="{px:.1f}" cy="{y:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1"/>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate gene variants
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def load_variants(
    variants_json: str = "",
) -> flyte.io.Dir:
    """Load gene variant definitions, validate sequences, and save to a temp directory."""
    if variants_json:
        genes = json.loads(variants_json)
    else:
        genes = DEFAULT_GENE_VARIANTS

    # Validate
    valid_bases = set("ATGC")
    for gene_name, gene_data in genes.items():
        seq = gene_data["sequence"].upper()
        invalid = set(seq) - valid_bases
        if invalid:
            log.warning(f"{gene_name}: invalid bases {invalid} — removing them")
            seq = "".join(b for b in seq if b in valid_bases)
            gene_data["sequence"] = seq

        for v in gene_data["variants"]:
            pos = v["pos"]
            if pos < 0 or pos >= len(seq):
                log.warning(f"{gene_name} variant {v['name']}: position {pos} out of range [0, {len(seq)})")
            elif seq[pos] != v["ref"]:
                log.warning(f"{gene_name} variant {v['name']}: expected ref={v['ref']} at pos {pos}, found {seq[pos]}")

    total_variants = sum(len(g["variants"]) for g in genes.values())
    log.info(f"Loaded {len(genes)} genes with {total_variants} variants")

    out_dir = tempfile.mkdtemp(prefix="genomic_vep_")
    with open(os.path.join(out_dir, "genes.json"), "w") as f:
        json.dump(genes, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Run Carbon model for variant effect scoring
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def score_variants(
    variants_dir: flyte.io.Dir,
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> str:
    """Score each variant using Carbon's log-likelihood ratio.

    For each variant, we compute:
        score = log P(alt_sequence) - log P(ref_sequence)

    A negative score means the model considers the variant less likely than
    the reference — suggestive of a damaging/pathogenic effect. A score near
    zero means the model sees little difference (likely benign).
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info(f"Loading Carbon model: {model_name}")

    # Load model
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        log.warning("Running on CPU — inference will be slow. GPU recommended for production.")

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    ).to(device)
    model.eval()

    # Load variants
    variants_path = await variants_dir.download()
    with open(os.path.join(variants_path, "genes.json")) as f:
        genes = json.load(f)

    results = {}
    total_variants = sum(len(g["variants"]) for g in genes.values())
    scored = 0

    progress_html = """
    <h2>Carbon Variant Effect Scoring</h2>
    <div class="card">
        <b>Model:</b> {model}<br>
        <b>Device:</b> {device}<br>
        <b>Progress:</b> {scored}/{total} variants scored
    </div>
    """

    for gene_name, gene_data in genes.items():
        ref_seq = gene_data["sequence"]
        gene_results = {
            "description": gene_data.get("description", ""),
            "sequence": ref_seq,
            "variants": [],
        }

        # Score reference sequence
        ref_prompt = f"<dna>{ref_seq}"
        ref_inputs = tokenizer(ref_prompt, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            ref_output = model(**ref_inputs, labels=ref_inputs["input_ids"])
            ref_loss = ref_output.loss.item()
            ref_ll = -ref_loss * ref_inputs["input_ids"].shape[1]

        for variant in gene_data["variants"]:
            scored += 1
            await flyte.report.replace.aio(
                _wrap_report(progress_html.format(
                    model=model_name, device=device, scored=scored, total=total_variants
                )),
                do_flush=True,
            )

            # Create mutant sequence
            pos = variant["pos"]
            alt_seq = ref_seq[:pos] + variant["alt"] + ref_seq[pos + 1:]

            # Score mutant
            alt_prompt = f"<dna>{alt_seq}"
            alt_inputs = tokenizer(alt_prompt, return_tensors="pt", add_special_tokens=False).to(device)

            with torch.no_grad():
                alt_output = model(**alt_inputs, labels=alt_inputs["input_ids"])
                alt_loss = alt_output.loss.item()
                alt_ll = -alt_loss * alt_inputs["input_ids"].shape[1]

            # VEP score: positive = model prefers alt (likely benign), negative = model prefers ref (likely pathogenic)
            vep_score = alt_ll - ref_ll

            gene_results["variants"].append({
                **variant,
                "score": round(vep_score, 4),
                "ref_ll": round(ref_ll, 4),
                "alt_ll": round(alt_ll, 4),
            })

            log.info(
                f"  {gene_name} | {variant['name']}: score={vep_score:.4f} "
                f"(known: {variant['known_effect']})"
            )

        results[gene_name] = gene_results

    # Generate scoring report
    html_parts = [
        "<h2>Carbon Variant Effect Scoring</h2>",
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{len(genes)}</div><div class="label">Genes</div></div>',
        f'<div class="stat"><div class="value">{total_variants}</div><div class="label">Variants Scored</div></div>',
        f'<div class="stat"><div class="value">{model_name.split("/")[-1]}</div><div class="label">Model</div></div>',
        f'<div class="stat"><div class="value">{device.upper()}</div><div class="label">Device</div></div>',
        "</div>",
    ]

    # Per-gene tables
    for gene_name, gene_data in results.items():
        html_parts.append(f'<h3>{gene_name}</h3>')
        html_parts.append(f'<div class="note">{gene_data["description"]}</div>')
        html_parts.append("<table><tr><th>Variant</th><th>Ref</th><th>Alt</th><th>VEP Score</th><th>Known Effect</th><th>Clinical</th></tr>")

        for v in gene_data["variants"]:
            badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
            direction = "damaging" if v["score"] < -0.1 else "neutral" if abs(v["score"]) <= 0.1 else "tolerated"
            html_parts.append(
                f'<tr>'
                f'<td><b>{v["name"]}</b></td>'
                f'<td style="color:{BASE_COLORS.get(v["ref"], "#333")};font-weight:700">{v["ref"]}</td>'
                f'<td style="color:{BASE_COLORS.get(v["alt"], "#333")};font-weight:700">{v["alt"]}</td>'
                f'<td><b>{v["score"]:.4f}</b> <span style="font-size:0.8em;color:#6c757d">({direction})</span></td>'
                f'<td><span class="badge {badge}">{v["known_effect"]}</span></td>'
                f'<td style="font-size:0.85em">{v.get("clinical", "")}</td>'
                f'</tr>'
            )
        html_parts.append("</table>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    return json.dumps(results)

# ------------------------------------------------------------------
# Task 3: Analyze and visualize variant effects
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def analyze_effects(
    scores_json: str,
    variants_dir: flyte.io.Dir,
) -> str:
    """Analyze VEP scores: classification accuracy, gene-level summaries, and rich visualizations."""
    results = json.loads(scores_json)

    variants_path = await variants_dir.download()
    with open(os.path.join(variants_path, "genes.json")) as f:
        genes = json.load(f)

    html_parts = ["<h2>Variant Effect Analysis</h2>"]

    # ------------------------------------------------------------------
    # Overall accuracy: does the model's score direction match known labels?
    # ------------------------------------------------------------------
    all_variants = []
    correct = 0
    total_known = 0
    true_pos = 0
    false_pos = 0
    true_neg = 0
    false_neg = 0

    for gene_name, gene_data in results.items():
        for v in gene_data["variants"]:
            all_variants.append({**v, "gene": gene_name})
            if v["known_effect"] in ("pathogenic", "benign"):
                total_known += 1
                predicted_pathogenic = v["score"] < -0.05
                actual_pathogenic = v["known_effect"] == "pathogenic"
                if predicted_pathogenic == actual_pathogenic:
                    correct += 1
                if predicted_pathogenic and actual_pathogenic:
                    true_pos += 1
                elif predicted_pathogenic and not actual_pathogenic:
                    false_pos += 1
                elif not predicted_pathogenic and actual_pathogenic:
                    false_neg += 1
                else:
                    true_neg += 1

    accuracy = correct / total_known if total_known else 0
    precision = true_pos / (true_pos + false_pos) if (true_pos + false_pos) else 0
    recall = true_pos / (true_pos + false_neg) if (true_pos + false_neg) else 0

    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{accuracy:.0%}</div><div class="label">Direction Accuracy</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{precision:.0%}</div><div class="label">Precision (Pathogenic)</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{recall:.0%}</div><div class="label">Recall (Pathogenic)</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{len(all_variants)}</div><div class="label">Total Variants</div></div>')
    html_parts.append("</div>")

    html_parts.append(
        '<div class="note">'
        "<b>How to read VEP scores:</b> Negative scores mean Carbon considers the variant "
        "less likely than the reference sequence — suggestive of a damaging effect. Scores near "
        "zero indicate the model sees little difference (likely benign). The magnitude indicates "
        "confidence."
        "</div>"
    )

    # ------------------------------------------------------------------
    # Cross-gene score comparison dot plot
    # ------------------------------------------------------------------
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_score_comparison_chart(results))
    html_parts.append("</div>")

    # ------------------------------------------------------------------
    # Per-gene visualizations
    # ------------------------------------------------------------------
    for gene_name, gene_data in results.items():
        short_name = gene_name.split("(")[0].strip()
        html_parts.append(f'<h3>{gene_name}</h3>')
        html_parts.append(f'<div class="note">{gene_data["description"]}</div>')

        # DNA track with variant positions highlighted
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_dna_track(
            gene_data["sequence"],
            gene_data["variants"],
            gene_name=f"{short_name} Reference Sequence",
        ))
        html_parts.append("</div>")

        # Lollipop plot
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_lollipop_plot(
            gene_data["variants"],
            len(gene_data["sequence"]),
            title=f"{short_name} — Variant Positions & Effect Scores",
        ))
        html_parts.append("</div>")

        # Score bar chart for this gene
        variant_names = [v["name"] for v in gene_data["variants"]]
        variant_scores = [v["score"] for v in gene_data["variants"]]
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_bar_chart(
            variant_names,
            {"VEP Score": variant_scores},
            title=f"{short_name} — Log-Likelihood Ratio Scores",
            colors=[EFFECT_COLORS.get(v["known_effect"], "#6b7280") for v in gene_data["variants"]],
            value_format=".3f",
        ))
        html_parts.append("</div>")

        # Variant detail cards
        for v in gene_data["variants"]:
            badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
            score_color = "#dc2626" if v["score"] < -0.1 else "#059669" if v["score"] > 0.05 else "#f59e0b"
            html_parts.append(
                f'<div class="gene-card">'
                f'<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:8px;">'
                f'<b style="font-size:1.1em">{v["name"]}</b>'
                f'<span class="badge {badge}">{v["known_effect"]}</span>'
                f'</div>'
                f'<div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:8px;">'
                f'<div><span style="color:#6c757d;font-size:0.85em">Ref base:</span> '
                f'<b style="color:{BASE_COLORS.get(v["ref"], "#333")}">{v["ref"]}</b></div>'
                f'<div><span style="color:#6c757d;font-size:0.85em">Alt base:</span> '
                f'<b style="color:{BASE_COLORS.get(v["alt"], "#333")}">{v["alt"]}</b></div>'
                f'<div><span style="color:#6c757d;font-size:0.85em">VEP Score:</span> '
                f'<b style="color:{score_color}">{v["score"]:.4f}</b></div>'
                f'</div>'
                f'<div style="margin-top:8px;font-size:0.9em;color:#374151">{v.get("clinical", "")}</div>'
                f'</div>'
            )

    # ------------------------------------------------------------------
    # Confusion matrix as heatmap
    # ------------------------------------------------------------------
    html_parts.append("<h3>Classification Performance</h3>")
    html_parts.append(
        '<div class="note">'
        "Using a simple threshold (score &lt; -0.05 = predicted pathogenic). "
        "This is zero-shot — no training on these specific variants."
        "</div>"
    )

    conf_matrix = [[true_pos, false_neg], [false_pos, true_neg]]
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        conf_matrix,
        ["Actual Pathogenic", "Actual Benign"],
        ["Predicted Pathogenic", "Predicted Benign"],
        title="Confusion Matrix (Known Variants Only)",
        value_format=".0f",
        width=400,
        height=300,
    ))
    html_parts.append("</div>")

    # ------------------------------------------------------------------
    # Score distribution by known effect
    # ------------------------------------------------------------------
    html_parts.append("<h3>Score Distribution by Known Effect</h3>")

    pathogenic_scores = [v["score"] for v in all_variants if v["known_effect"] == "pathogenic"]
    benign_scores = [v["score"] for v in all_variants if v["known_effect"] == "benign"]
    uncertain_scores = [v["score"] for v in all_variants if v["known_effect"] == "uncertain"]

    stats_html = '<div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:12px;">'
    for label, scores, color in [
        ("Pathogenic", pathogenic_scores, "#dc2626"),
        ("Benign", benign_scores, "#059669"),
        ("Uncertain", uncertain_scores, "#f59e0b"),
    ]:
        if scores:
            mean_s = sum(scores) / len(scores)
            min_s = min(scores)
            max_s = max(scores)
            stats_html += (
                f'<div class="gene-card" style="border-left:4px solid {color}">'
                f'<b style="color:{color}">{label}</b> (n={len(scores)})<br>'
                f'Mean: {mean_s:.4f}<br>'
                f'Range: [{min_s:.4f}, {max_s:.4f}]'
                f'</div>'
            )
    stats_html += "</div>"
    html_parts.append(stats_html)

    # Summary
    analysis = {
        "total_variants": len(all_variants),
        "total_known": total_known,
        "accuracy": round(accuracy, 4),
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "true_pos": true_pos,
        "false_pos": false_pos,
        "true_neg": true_neg,
        "false_neg": false_neg,
        "pathogenic_mean_score": round(sum(pathogenic_scores) / len(pathogenic_scores), 4) if pathogenic_scores else None,
        "benign_mean_score": round(sum(benign_scores) / len(benign_scores), 4) if benign_scores else None,
    }

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)
    return json.dumps(analysis)

# ------------------------------------------------------------------
# Task 4: Generate comprehensive summary report
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def generate_summary(
    scores_json: str,
    analysis_json: str,
) -> str:
    """Generate the final summary report combining all results."""
    results = json.loads(scores_json)
    analysis = json.loads(analysis_json)

    html_parts = [
        "<h2>Genomic Variant Effect Prediction — Summary</h2>",
        '<div class="note">'
        "This pipeline uses <b>HuggingFace Carbon</b>, an autoregressive genomic foundation model "
        "trained on 1 trillion tokens of DNA sequence, to perform <b>zero-shot variant effect "
        "prediction</b>. No fine-tuning or labeled training data was used — the model scores "
        "variants purely based on its learned understanding of DNA sequence grammar."
        "</div>",
    ]

    # Key metrics
    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{len(results)}</div><div class="label">Genes Analyzed</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["total_variants"]}</div><div class="label">Variants Scored</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["accuracy"]:.0%}</div><div class="label">Direction Accuracy</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["precision"]:.0%}</div><div class="label">Precision</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["recall"]:.0%}</div><div class="label">Recall</div></div>')
    html_parts.append("</div>")

    # Gene summary table
    html_parts.append("<h3>Per-Gene Summary</h3>")
    html_parts.append(
        "<table><tr><th>Gene</th><th>Variants</th><th>Mean Score</th>"
        "<th>Pathogenic</th><th>Benign</th><th>Uncertain</th></tr>"
    )

    for gene_name, gene_data in results.items():
        variants = gene_data["variants"]
        scores = [v["score"] for v in variants]
        mean_score = sum(scores) / len(scores) if scores else 0
        n_path = sum(1 for v in variants if v["known_effect"] == "pathogenic")
        n_benign = sum(1 for v in variants if v["known_effect"] == "benign")
        n_unc = sum(1 for v in variants if v["known_effect"] == "uncertain")
        short = gene_name.split("(")[0].strip()

        html_parts.append(
            f"<tr><td><b>{short}</b></td><td>{len(variants)}</td>"
            f"<td>{mean_score:.4f}</td>"
            f'<td><span class="badge badge-danger">{n_path}</span></td>'
            f'<td><span class="badge badge-success">{n_benign}</span></td>'
            f'<td><span class="badge badge-warning">{n_unc}</span></td></tr>'
        )
    html_parts.append("</table>")

    # Cross-gene heatmap: gene x metric
    gene_names = [g.split("(")[0].strip() for g in results.keys()]
    metrics = ["Mean Score", "Min Score", "Max Score", "# Variants"]
    matrix = []
    for gene_data in results.values():
        scores = [v["score"] for v in gene_data["variants"]]
        matrix.append([
            sum(scores) / len(scores) if scores else 0,
            min(scores) if scores else 0,
            max(scores) if scores else 0,
            len(scores),
        ])

    html_parts.append("<h3>Gene-Level Metrics</h3>")
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        matrix,
        gene_names,
        metrics,
        title="Gene-Level VEP Score Summary",
        value_format=".2f",
        width=600,
        height=350,
    ))
    html_parts.append("</div>")

    # All variants ranked by score (most damaging first)
    html_parts.append("<h3>All Variants Ranked by Impact</h3>")
    all_vars_sorted = []
    for gene_name, gene_data in results.items():
        for v in gene_data["variants"]:
            all_vars_sorted.append({**v, "gene": gene_name.split("(")[0].strip()})
    all_vars_sorted.sort(key=lambda x: x["score"])

    html_parts.append(
        "<table><tr><th>#</th><th>Gene</th><th>Variant</th><th>Score</th>"
        "<th>Known</th><th>Clinical Significance</th></tr>"
    )
    for i, v in enumerate(all_vars_sorted):
        badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
        score_color = "#dc2626" if v["score"] < -0.1 else "#059669" if v["score"] > 0.05 else "#f59e0b"
        html_parts.append(
            f'<tr><td>{i + 1}</td><td><b>{v["gene"]}</b></td><td>{v["name"]}</td>'
            f'<td style="color:{score_color};font-weight:700">{v["score"]:.4f}</td>'
            f'<td><span class="badge {badge}">{v["known_effect"]}</span></td>'
            f'<td style="font-size:0.85em">{v.get("clinical", "")}</td></tr>'
        )
    html_parts.append("</table>")

    # Method note
    html_parts.append(
        '<div class="note">'
        "<b>Method:</b> Zero-shot variant effect prediction using log-likelihood ratio scoring. "
        "For each variant, we compute score = log P(mutant sequence | Carbon) - log P(reference sequence | Carbon). "
        "Negative scores indicate the model considers the mutant less probable than the reference, "
        "which correlates with pathogenicity. This approach requires no fine-tuning and generalizes "
        "across genes and variant types.<br><br>"
        "<b>Limitations:</b> These are short sequence windows — real clinical VEP would use longer "
        "genomic context (Carbon supports up to 786kbp). The threshold for pathogenicity classification "
        "(-0.05) is a simple heuristic; clinical use requires calibrated thresholds per gene."
        "</div>"
    )

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)
    return json.dumps({"status": "complete", "analysis": analysis})

# ------------------------------------------------------------------
# Pipeline orchestrator
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    variants_json: str = "",
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> tuple[str, str]:
    """
    End-to-end genomic variant effect prediction pipeline.

    Returns (scores JSON, analysis JSON).

    1. Load and validate gene variants
    2. Score variants with Carbon (log-likelihood ratio)
    3. Analyze effects — accuracy, visualizations, classification
    4. Generate comprehensive summary report
    """
    log.info("Starting genomic variant effect prediction pipeline...")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = [
            "Load Variants",
            "Carbon Scoring",
            "Analyze Effects",
            "Generate Summary",
        ]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#2563eb;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#2563eb;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>Genomic Variant Effect Prediction</h2>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    # Stage 1: Load variants
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Loading and validating gene variants...")),
        do_flush=True,
    )
    var_dir = await load_variants(variants_json=variants_json)

    # Stage 2: Score with Carbon
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Running Carbon model for variant effect scoring...")),
        do_flush=True,
    )
    scores_json = await score_variants(variants_dir=var_dir, model_name=model_name)

    # Stage 3: Analyze effects
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Analyzing variant effects and generating visualizations...")),
        do_flush=True,
    )
    analysis_json = await analyze_effects(scores_json=scores_json, variants_dir=var_dir)

    # Stage 4: Summary
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Generating comprehensive summary report...")),
        do_flush=True,
    )
    summary_json = await generate_summary(scores_json=scores_json, analysis_json=analysis_json)

    # Final pipeline report
    analysis = json.loads(analysis_json)
    results = json.loads(scores_json)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(results)}</div><div class="label">Genes Analyzed</div></div>
      <div class="stat"><div class="value">{analysis['total_variants']}</div><div class="label">Variants Scored</div></div>
      <div class="stat"><div class="value">{analysis['accuracy']:.0%}</div><div class="label">Direction Accuracy</div></div>
      <div class="stat"><div class="value">{analysis['precision']:.0%}</div><div class="label">Precision</div></div>
      <div class="stat"><div class="value">{analysis['recall']:.0%}</div><div class="label">Recall</div></div>
    </div>
    <div class="card">
      <b>Model:</b> HuggingFace Carbon |
      <b>Method:</b> Zero-shot log-likelihood ratio scoring |
      <b>Genes:</b> {', '.join(g.split('(')[0].strip() for g in results.keys())}
    </div>
    <div class="note">
      All 4 pipeline stages completed successfully. View individual task reports for detailed
      visualizations including DNA sequence tracks, variant lollipop plots, VEP score charts,
      confusion matrices, and ranked variant tables.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)

    log.info("Pipeline complete.")
    return scores_json, analysis_json

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_variant_effect/genomic_variant_effect.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "numpy",
# ]
# ///
```

## Orchestrate the pipeline

The `pipeline` task loads variants, scores them with Carbon, analyzes classification accuracy against known labels, and generates a summary report.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="genomic-variant-effect", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="genomic-variant-effect-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="genomic-variant-effect-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default gene variants — clinically relevant mutations
# ------------------------------------------------------------------
# Each entry: gene name -> { "sequence": reference DNA, "variants": [{ "pos": 0-indexed, "ref": base, "alt": base, "name": "...", "known_effect": "..." }] }
# Sequences are short windows (~120-200bp) around the variant site for tractable inference.

DEFAULT_GENE_VARIANTS = {
    "BRCA2 (Breast Cancer)": {
        "description": "Tumor suppressor critical for DNA repair via homologous recombination. Mutations dramatically increase breast and ovarian cancer risk.",
        "sequence": "ATGGCCTCGAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAGCAG",
        "variants": [
            {"pos": 12, "ref": "A", "alt": "T", "name": "c.37A>T", "known_effect": "pathogenic", "clinical": "Nonsense mutation — truncates protein early"},
            {"pos": 18, "ref": "G", "alt": "A", "name": "c.55G>A", "known_effect": "benign", "clinical": "Synonymous — no amino acid change"},
            {"pos": 30, "ref": "C", "alt": "T", "name": "c.91C>T", "known_effect": "pathogenic", "clinical": "Missense in DNA-binding domain"},
            {"pos": 45, "ref": "G", "alt": "C", "name": "c.136G>C", "known_effect": "uncertain", "clinical": "Variant of uncertain significance (VUS)"},
        ],
    },
    "TP53 (Tumor Suppressor)": {
        "description": "Guardian of the genome. Activates DNA repair, cell cycle arrest, and apoptosis. Mutated in >50% of human cancers.",
        "sequence": "ATGGAGGAGCCGCAGTCAGATCCTAGCGTGAGTTTGCACCCTTCAGAGACAGAAACCACTGGATTGGAGACTACTTCCTGAAACAACGTTCTGTCCCCCTTGCCGTCCCAAGCAATGGATGAT",
        "variants": [
            {"pos": 15, "ref": "C", "alt": "T", "name": "R175H", "known_effect": "pathogenic", "clinical": "Hotspot — gain-of-function, dominant negative. Most common TP53 mutation in cancer"},
            {"pos": 36, "ref": "T", "alt": "C", "name": "P72R", "known_effect": "benign", "clinical": "Common polymorphism — subtle effect on apoptosis efficiency"},
            {"pos": 54, "ref": "C", "alt": "A", "name": "G245S", "known_effect": "pathogenic", "clinical": "Contact mutant — disrupts DNA binding"},
            {"pos": 72, "ref": "T", "alt": "G", "name": "R248W", "known_effect": "pathogenic", "clinical": "Structural mutant — destabilizes DNA-binding loop"},
            {"pos": 90, "ref": "C", "alt": "T", "name": "R273H", "known_effect": "pathogenic", "clinical": "Contact mutant — directly contacts DNA bases"},
        ],
    },
    "CFTR (Cystic Fibrosis)": {
        "description": "Chloride channel protein. Mutations cause cystic fibrosis — the most common lethal genetic disease in people of European descent.",
        "sequence": "ATGCAGAGGTCGCCTCTGGAAAAGGCCAGCGTTGTCTCCAAACTTTTTTTCAGCTGGACCAGACCAATTTTGAGGAAAGGATACAGACAGCGCCTGGAATTGTCAGACATATACCAAATCCCTTC",
        "variants": [
            {"pos": 9, "ref": "G", "alt": "A", "name": "G85E", "known_effect": "pathogenic", "clinical": "Disrupts chloride channel processing"},
            {"pos": 24, "ref": "C", "alt": "T", "name": "R117H", "known_effect": "pathogenic", "clinical": "Reduces channel conductance — milder CF phenotype"},
            {"pos": 48, "ref": "T", "alt": "C", "name": "I148T", "known_effect": "benign", "clinical": "Previously misclassified — now known benign polymorphism"},
            {"pos": 66, "ref": "A", "alt": "G", "name": "R334W", "known_effect": "pathogenic", "clinical": "Gating mutation — channel opens less frequently"},
        ],
    },
    "KRAS (Oncogene)": {
        "description": "GTPase signal switch. KRAS mutations are the most common oncogenic driver — found in ~25% of all human cancers, especially pancreatic, colorectal, and lung.",
        "sequence": "ATGACTGAATATAAACTTGTGGTAGTTGGAGCTGGTGGCGTAGGCAAGAGTGCCTTGACGATACAGCTAATTCAGAATCATTTTGTGGACGAATATGATCCAACAATAGAGGATTCCTACAGGAA",
        "variants": [
            {"pos": 34, "ref": "G", "alt": "T", "name": "G12V", "known_effect": "pathogenic", "clinical": "Locks KRAS in active state — constitutive proliferation signal"},
            {"pos": 35, "ref": "G", "alt": "A", "name": "G12D", "known_effect": "pathogenic", "clinical": "Most common KRAS mutation in pancreatic cancer"},
            {"pos": 37, "ref": "G", "alt": "T", "name": "G13D", "known_effect": "pathogenic", "clinical": "Constitutively active — common in colorectal cancer"},
            {"pos": 60, "ref": "C", "alt": "A", "name": "Q61K", "known_effect": "pathogenic", "clinical": "Impairs GTP hydrolysis — locked ON state"},
        ],
    },
    "HBB (Sickle Cell)": {
        "description": "Beta-globin subunit of hemoglobin. The sickle cell mutation (E6V) is the most well-known single-base disease variant in humans.",
        "sequence": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTGGTCTACCCTTGGACCCAGAGG",
        "variants": [
            {"pos": 17, "ref": "A", "alt": "T", "name": "E6V (HbS)", "known_effect": "pathogenic", "clinical": "THE sickle cell mutation — causes hemoglobin polymerization under low O2"},
            {"pos": 19, "ref": "G", "alt": "A", "name": "E6K (HbC)", "known_effect": "pathogenic", "clinical": "Hemoglobin C disease — milder than sickle cell but causes crystal formation"},
            {"pos": 36, "ref": "G", "alt": "A", "name": "E26K", "known_effect": "benign", "clinical": "Hemoglobin E — most common Hb variant worldwide, mild effect"},
            {"pos": 78, "ref": "C", "alt": "T", "name": "Q39X", "known_effect": "pathogenic", "clinical": "Nonsense — causes beta-thalassemia (no functional beta-globin)"},
        ],
    },
}

# DNA base colors (classic genomics color scheme)
BASE_COLORS = {"A": "#2ecc71", "T": "#e74c3c", "G": "#f39c12", "C": "#3498db"}
BASE_COMPLEMENT = {"A": "T", "T": "A", "G": "C", "C": "G"}

# Pathogenicity color scheme
EFFECT_COLORS = {
    "pathogenic": "#dc2626",
    "benign": "#059669",
    "uncertain": "#f59e0b",
}
EFFECT_BADGES = {
    "pathogenic": "badge-danger",
    "benign": "badge-success",
    "uncertain": "badge-warning",
}

# ------------------------------------------------------------------
# Report styling — genomics-themed deep blues and teals
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #1e3a5f; border-bottom: 2px solid #2563eb; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #1e40af; margin-top: 20px; }
  .report .card { background: #eff6ff; border: 1px solid #bfdbfe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #dbeafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #1e3a5f; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #1e3a5f; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dbeafe; }
  .report tr:nth-child(even) { background: #eff6ff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-warning { background: #fef3c7; color: #92400e; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #dbeafe; color: #1e40af; }
  .report .chart-container { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #eff6ff; border-left: 4px solid #2563eb; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .gene-card { background: #fff; border: 1px solid #dbeafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .dna-track { font-family: 'SF Mono', 'Fira Code', monospace; letter-spacing: 1px; }
</style>
"""

def _wrap_report(html: str) -> str:
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    value_format: str = ".2f",
) -> str:
    """Generate an SVG grouped bar chart."""
    if not labels:
        return ""

    default_colors = ["#2563eb", "#1e3a5f", "#3b82f6", "#60a5fa", "#93c5fd"]
    colors = colors or default_colors

    ml, mr, mt, mb = 70, 20, 40, 80
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(abs(v) for v in all_vals) if all_vals else 1
    y_min = min(all_vals) if all_vals else 0
    # For VEP scores (negative = more damaging), we need to handle negative values
    if y_min >= 0:
        y_min_plot = 0
        y_max_plot = y_max * 1.15 or 1
    else:
        y_max_plot = max(y_max * 1.15, 0.1)
        y_min_plot = y_min * 1.15

    y_range = y_max_plot - y_min_plot or 1

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v - y_min_plot) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_min_plot + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_format}}</text>'
        )

    # Zero line
    if y_min_plot < 0 < y_max_plot:
        zy = sy(0)
        svg.append(
            f'<line x1="{ml}" y1="{zy:.1f}" x2="{ml + cw}" y2="{zy:.1f}" '
            f'stroke="#374151" stroke-width="1.5"/>'
        )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            if val >= 0:
                by = sy(val)
                bh = sy(0) - by if y_min_plot < 0 else mt + ch - by
            else:
                by = sy(0) if y_min_plot < 0 else mt + ch
                bh = sy(val) - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{max(0, bh):.1f}" fill="{color}" rx="2"/>'
            )
            text_y = by - 4 if val >= 0 else by + bh + 12
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{text_y:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_format}}</text>'
            )
        # Rotated group label
        lx = gx + n_series * bar_width / 2
        svg.append(
            f'<text x="{lx:.1f}" y="{mt + ch + 14}" '
            f'text-anchor="end" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {lx:.1f}, {mt + ch + 14})">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 110
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 110}" y="{mt + ch + 55}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 110 + 16}" y="{mt + ch + 66}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    width: int = 700,
    height: int = 500,
    value_format: str = ".2f",
    diverging: bool = False,
) -> str:
    """Generate an SVG heatmap. If diverging=True, uses red-white-blue scale centered at 0."""
    n_rows = len(matrix)
    n_cols = len(matrix[0]) if matrix else 0
    if not n_rows or not n_cols:
        return ""

    show_values = n_rows <= 10 and n_cols <= 12

    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)

    if diverging:
        abs_max = max(abs(v_min), abs(v_max)) or 1

        def get_color(v):
            t = v / abs_max  # -1 to 1
            if t < 0:
                # White to red (negative = damaging)
                r = 255
                g = int(255 * (1 + t))
                b = int(255 * (1 + t))
            else:
                # White to blue (positive = benign)
                r = int(255 * (1 - t))
                g = int(255 * (1 - t))
                b = 255
            return f"rgb({r},{g},{b})"
    else:
        v_range = v_max - v_min or 1

        def get_color(v):
            t = (v - v_min) / v_range
            r = int(255 - t * (255 - 30))
            g = int(255 - t * (255 - 58))
            b = int(255 - t * (255 - 95))
            return f"rgb({r},{g},{b})"

    # Layout
    ml = max(140, max(len(l) for l in row_labels) * 7 + 20) if row_labels else 140
    mr = 20
    mt = 80 if col_labels else 40
    mb = 30
    cw = width - ml - mr
    ch = height - mt - mb

    cell_w = cw / n_cols
    cell_h = ch / n_rows

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Column labels (rotated)
    for j, label in enumerate(col_labels):
        cx = ml + j * cell_w + cell_w / 2
        svg.append(
            f'<text x="{cx:.1f}" y="{mt - 8}" text-anchor="end" '
            f'font-size="10" fill="#374151" '
            f'transform="rotate(-45, {cx:.1f}, {mt - 8})">{label}</text>'
        )

    # Row labels + cells
    for i, row_label in enumerate(row_labels):
        ry = mt + i * cell_h + cell_h / 2
        svg.append(
            f'<text x="{ml - 8}" y="{ry + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#374151">{row_label}</text>'
        )
        for j in range(n_cols):
            val = matrix[i][j]
            color = get_color(val)
            cx = ml + j * cell_w
            cy = mt + i * cell_h
            svg.append(
                f'<rect x="{cx:.1f}" y="{cy:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{color}" stroke="#fff" stroke-width="1"/>'
            )
            if show_values:
                if diverging:
                    t = abs(val) / (max(abs(v_min), abs(v_max)) or 1)
                else:
                    t = (val - v_min) / (v_max - v_min or 1)
                text_color = "#fff" if t > 0.55 else "#1a1a2e"
                font_size = min(10, int(cell_w / 4), int(cell_h / 2.5))
                font_size = max(7, font_size)
                svg.append(
                    f'<text x="{cx + cell_w / 2:.1f}" y="{cy + cell_h / 2 + 3:.1f}" '
                    f'text-anchor="middle" font-size="{font_size}" '
                    f'fill="{text_color}">{val:{value_format}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_dna_track(
    sequence: str,
    variants: list[dict],
    gene_name: str = "",
    width: int = 900,
) -> str:
    """Render a color-coded DNA sequence track with variant positions highlighted."""
    chars_per_line = 60
    char_w = 11
    line_h = 22
    label_w = 50
    n_lines = (len(sequence) + chars_per_line - 1) // chars_per_line

    # Extra space for variant annotations
    variant_positions = {v["pos"] for v in variants}
    svg_h = n_lines * line_h + 60

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {svg_h}" '
        f'style="width:100%;max-width:{width}px;height:auto;font-family:monospace;">',
        f'<rect width="{width}" height="{svg_h}" fill="#fff" rx="6"/>',
    ]

    if gene_name:
        svg.append(
            f'<text x="{width / 2}" y="16" text-anchor="middle" '
            f'font-size="12" font-weight="600" fill="#1e3a5f">{gene_name}</text>'
        )

    y_offset = 28

    for line_idx in range(n_lines):
        start = line_idx * chars_per_line
        end = min(start + chars_per_line, len(sequence))
        chunk = sequence[start:end]
        y = y_offset + line_idx * line_h

        # Position label
        svg.append(
            f'<text x="{label_w - 8}" y="{y}" text-anchor="end" '
            f'font-size="9" fill="#9ca3af">{start + 1}</text>'
        )

        for ci, base in enumerate(chunk):
            abs_pos = start + ci
            x = label_w + ci * char_w
            color = BASE_COLORS.get(base, "#6b7280")
            is_variant = abs_pos in variant_positions

            if is_variant:
                # Red highlight for variant positions
                svg.append(
                    f'<rect x="{x - 1}" y="{y - 13}" width="{char_w + 2}" height="{line_h}" '
                    f'fill="#fee2e2" stroke="#dc2626" stroke-width="1.5" rx="2"/>'
                )
                # Small triangle marker above
                svg.append(
                    f'<polygon points="{x + char_w / 2:.1f},{y - 16} {x + char_w / 2 - 3:.1f},{y - 20} {x + char_w / 2 + 3:.1f},{y - 20}" '
                    f'fill="#dc2626"/>'
                )
            else:
                # Subtle background
                svg.append(
                    f'<rect x="{x}" y="{y - 12}" width="{char_w}" height="{line_h - 2}" '
                    f'fill="{color}" opacity="0.08" rx="1"/>'
                )

            svg.append(
                f'<text x="{x + char_w / 2:.1f}" y="{y}" text-anchor="middle" '
                f'font-size="11" font-weight="{"700" if is_variant else "500"}" '
                f'fill="{color}">{base}</text>'
            )

            # Spacer every 10 bases
            if (abs_pos + 1) % 10 == 0 and ci < len(chunk) - 1:
                svg.append(
                    f'<line x1="{x + char_w + 1}" y1="{y - 11}" '
                    f'x2="{x + char_w + 1}" y2="{y + 4}" '
                    f'stroke="#e5e7eb" stroke-width="0.5"/>'
                )

    # Legend
    ly = svg_h - 12
    lx = label_w
    for base, color in BASE_COLORS.items():
        svg.append(f'<rect x="{lx}" y="{ly - 8}" width="8" height="8" rx="1" fill="{color}"/>')
        svg.append(f'<text x="{lx + 12}" y="{ly}" font-size="9" fill="#6b7280">{base}</text>')
        lx += 30
    lx += 10
    svg.append(f'<rect x="{lx}" y="{ly - 8}" width="8" height="8" rx="1" fill="#fee2e2" stroke="#dc2626" stroke-width="1"/>')
    svg.append(f'<text x="{lx + 12}" y="{ly}" font-size="9" fill="#6b7280">Variant site</text>')

    svg.append("</svg>")
    return "\n".join(svg)

def _make_lollipop_plot(
    variants: list[dict],
    seq_length: int,
    title: str = "",
    width: int = 700,
    height: int = 200,
) -> str:
    """Generate a lollipop plot showing variant positions along a gene with effect scores."""
    if not variants:
        return ""

    ml, mr, mt, mb = 50, 30, 40, 40
    cw = width - ml - mr
    ch = height - mt - mb

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Gene track (horizontal bar)
    track_y = mt + ch * 0.7
    svg.append(
        f'<rect x="{ml}" y="{track_y - 4}" width="{cw}" height="8" '
        f'fill="#dbeafe" rx="4" stroke="#93c5fd" stroke-width="1"/>'
    )

    # Position labels at ends
    svg.append(
        f'<text x="{ml}" y="{track_y + 20}" text-anchor="middle" '
        f'font-size="9" fill="#9ca3af">1</text>'
    )
    svg.append(
        f'<text x="{ml + cw}" y="{track_y + 20}" text-anchor="middle" '
        f'font-size="9" fill="#9ca3af">{seq_length}</text>'
    )

    # Lollipops
    scores = [v.get("score", 0) for v in variants]
    max_score = max(abs(s) for s in scores) if scores else 1

    for v in variants:
        pos = v["pos"]
        score = v.get("score", 0)
        effect = v.get("known_effect", "uncertain")
        color = EFFECT_COLORS.get(effect, "#6b7280")

        x = ml + (pos / max(seq_length - 1, 1)) * cw
        stem_h = max(20, abs(score) / max_score * (ch * 0.6))
        circle_y = track_y - stem_h - 6

        # Stem
        svg.append(
            f'<line x1="{x:.1f}" y1="{track_y - 4}" x2="{x:.1f}" y2="{circle_y + 5:.1f}" '
            f'stroke="{color}" stroke-width="2"/>'
        )
        # Circle
        svg.append(
            f'<circle cx="{x:.1f}" cy="{circle_y:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5"/>'
        )
        # Label
        svg.append(
            f'<text x="{x:.1f}" y="{circle_y - 9:.1f}" text-anchor="middle" '
            f'font-size="8" fill="#374151" font-weight="600">{v.get("name", "")}</text>'
        )

    # Legend
    lx = ml
    for effect, color in EFFECT_COLORS.items():
        svg.append(f'<circle cx="{lx + 4}" cy="{height - 10}" r="4" fill="{color}"/>')
        svg.append(
            f'<text x="{lx + 12}" y="{height - 6}" font-size="9" fill="#374151">'
            f'{effect.title()}</text>'
        )
        lx += len(effect) * 7 + 24

    svg.append("</svg>")
    return "\n".join(svg)

def _make_score_comparison_chart(
    gene_results: dict,
    width: int = 800,
    height: int = 350,
) -> str:
    """Generate a dot plot comparing VEP scores across all genes, colored by known effect."""
    all_variants = []
    for gene_name, gene_data in gene_results.items():
        for v in gene_data["variants"]:
            all_variants.append({
                "gene": gene_name.split("(")[0].strip(),
                "name": v["name"],
                "score": v.get("score", 0),
                "known_effect": v["known_effect"],
            })

    if not all_variants:
        return ""

    ml, mr, mt, mb = 120, 30, 40, 30
    cw = width - ml - mr
    ch = height - mt - mb

    scores = [v["score"] for v in all_variants]
    s_min, s_max = min(scores), max(scores)
    s_pad = (s_max - s_min) * 0.1 or 0.5
    s_min -= s_pad
    s_max += s_pad
    s_range = s_max - s_min or 1

    def sx(s):
        return ml + (s - s_min) / s_range * cw

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
        f'<text x="{width / 2}" y="22" text-anchor="middle" '
        f'font-size="14" font-weight="600" fill="#1a1a2e">'
        f'Variant Effect Scores — All Genes</text>',
    ]

    # Grid
    for i in range(6):
        gx = s_min + s_range * i / 5
        px = sx(gx)
        svg.append(
            f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
            f'stroke="#f3f4f6" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 16}" text-anchor="middle" '
            f'font-size="10" fill="#9ca3af">{gx:.2f}</text>'
        )

    # Zero line
    if s_min < 0 < s_max:
        zx = sx(0)
        svg.append(
            f'<line x1="{zx:.1f}" y1="{mt}" x2="{zx:.1f}" y2="{mt + ch}" '
            f'stroke="#374151" stroke-width="1" stroke-dasharray="4,3"/>'
        )

    # Rows per variant
    row_h = ch / max(len(all_variants), 1)
    for i, v in enumerate(all_variants):
        y = mt + i * row_h + row_h / 2
        color = EFFECT_COLORS.get(v["known_effect"], "#6b7280")
        px = sx(v["score"])

        # Label
        label = f'{v["gene"]} {v["name"]}'
        svg.append(
            f'<text x="{ml - 8}" y="{y + 3:.1f}" text-anchor="end" '
            f'font-size="9" fill="#374151">{label}</text>'
        )
        # Connector line
        svg.append(
            f'<line x1="{ml}" y1="{y:.1f}" x2="{px:.1f}" y2="{y:.1f}" '
            f'stroke="{color}" stroke-width="1" opacity="0.3"/>'
        )
        # Dot
        svg.append(
            f'<circle cx="{px:.1f}" cy="{y:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1"/>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate gene variants
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def load_variants(
    variants_json: str = "",
) -> flyte.io.Dir:
    """Load gene variant definitions, validate sequences, and save to a temp directory."""
    if variants_json:
        genes = json.loads(variants_json)
    else:
        genes = DEFAULT_GENE_VARIANTS

    # Validate
    valid_bases = set("ATGC")
    for gene_name, gene_data in genes.items():
        seq = gene_data["sequence"].upper()
        invalid = set(seq) - valid_bases
        if invalid:
            log.warning(f"{gene_name}: invalid bases {invalid} — removing them")
            seq = "".join(b for b in seq if b in valid_bases)
            gene_data["sequence"] = seq

        for v in gene_data["variants"]:
            pos = v["pos"]
            if pos < 0 or pos >= len(seq):
                log.warning(f"{gene_name} variant {v['name']}: position {pos} out of range [0, {len(seq)})")
            elif seq[pos] != v["ref"]:
                log.warning(f"{gene_name} variant {v['name']}: expected ref={v['ref']} at pos {pos}, found {seq[pos]}")

    total_variants = sum(len(g["variants"]) for g in genes.values())
    log.info(f"Loaded {len(genes)} genes with {total_variants} variants")

    out_dir = tempfile.mkdtemp(prefix="genomic_vep_")
    with open(os.path.join(out_dir, "genes.json"), "w") as f:
        json.dump(genes, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Run Carbon model for variant effect scoring
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def score_variants(
    variants_dir: flyte.io.Dir,
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> str:
    """Score each variant using Carbon's log-likelihood ratio.

    For each variant, we compute:
        score = log P(alt_sequence) - log P(ref_sequence)

    A negative score means the model considers the variant less likely than
    the reference — suggestive of a damaging/pathogenic effect. A score near
    zero means the model sees little difference (likely benign).
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info(f"Loading Carbon model: {model_name}")

    # Load model
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        log.warning("Running on CPU — inference will be slow. GPU recommended for production.")

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    ).to(device)
    model.eval()

    # Load variants
    variants_path = await variants_dir.download()
    with open(os.path.join(variants_path, "genes.json")) as f:
        genes = json.load(f)

    results = {}
    total_variants = sum(len(g["variants"]) for g in genes.values())
    scored = 0

    progress_html = """
    <h2>Carbon Variant Effect Scoring</h2>
    <div class="card">
        <b>Model:</b> {model}<br>
        <b>Device:</b> {device}<br>
        <b>Progress:</b> {scored}/{total} variants scored
    </div>
    """

    for gene_name, gene_data in genes.items():
        ref_seq = gene_data["sequence"]
        gene_results = {
            "description": gene_data.get("description", ""),
            "sequence": ref_seq,
            "variants": [],
        }

        # Score reference sequence
        ref_prompt = f"<dna>{ref_seq}"
        ref_inputs = tokenizer(ref_prompt, return_tensors="pt", add_special_tokens=False).to(device)

        with torch.no_grad():
            ref_output = model(**ref_inputs, labels=ref_inputs["input_ids"])
            ref_loss = ref_output.loss.item()
            ref_ll = -ref_loss * ref_inputs["input_ids"].shape[1]

        for variant in gene_data["variants"]:
            scored += 1
            await flyte.report.replace.aio(
                _wrap_report(progress_html.format(
                    model=model_name, device=device, scored=scored, total=total_variants
                )),
                do_flush=True,
            )

            # Create mutant sequence
            pos = variant["pos"]
            alt_seq = ref_seq[:pos] + variant["alt"] + ref_seq[pos + 1:]

            # Score mutant
            alt_prompt = f"<dna>{alt_seq}"
            alt_inputs = tokenizer(alt_prompt, return_tensors="pt", add_special_tokens=False).to(device)

            with torch.no_grad():
                alt_output = model(**alt_inputs, labels=alt_inputs["input_ids"])
                alt_loss = alt_output.loss.item()
                alt_ll = -alt_loss * alt_inputs["input_ids"].shape[1]

            # VEP score: positive = model prefers alt (likely benign), negative = model prefers ref (likely pathogenic)
            vep_score = alt_ll - ref_ll

            gene_results["variants"].append({
                **variant,
                "score": round(vep_score, 4),
                "ref_ll": round(ref_ll, 4),
                "alt_ll": round(alt_ll, 4),
            })

            log.info(
                f"  {gene_name} | {variant['name']}: score={vep_score:.4f} "
                f"(known: {variant['known_effect']})"
            )

        results[gene_name] = gene_results

    # Generate scoring report
    html_parts = [
        "<h2>Carbon Variant Effect Scoring</h2>",
        '<div class="stat-grid">',
        f'<div class="stat"><div class="value">{len(genes)}</div><div class="label">Genes</div></div>',
        f'<div class="stat"><div class="value">{total_variants}</div><div class="label">Variants Scored</div></div>',
        f'<div class="stat"><div class="value">{model_name.split("/")[-1]}</div><div class="label">Model</div></div>',
        f'<div class="stat"><div class="value">{device.upper()}</div><div class="label">Device</div></div>',
        "</div>",
    ]

    # Per-gene tables
    for gene_name, gene_data in results.items():
        html_parts.append(f'<h3>{gene_name}</h3>')
        html_parts.append(f'<div class="note">{gene_data["description"]}</div>')
        html_parts.append("<table><tr><th>Variant</th><th>Ref</th><th>Alt</th><th>VEP Score</th><th>Known Effect</th><th>Clinical</th></tr>")

        for v in gene_data["variants"]:
            badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
            direction = "damaging" if v["score"] < -0.1 else "neutral" if abs(v["score"]) <= 0.1 else "tolerated"
            html_parts.append(
                f'<tr>'
                f'<td><b>{v["name"]}</b></td>'
                f'<td style="color:{BASE_COLORS.get(v["ref"], "#333")};font-weight:700">{v["ref"]}</td>'
                f'<td style="color:{BASE_COLORS.get(v["alt"], "#333")};font-weight:700">{v["alt"]}</td>'
                f'<td><b>{v["score"]:.4f}</b> <span style="font-size:0.8em;color:#6c757d">({direction})</span></td>'
                f'<td><span class="badge {badge}">{v["known_effect"]}</span></td>'
                f'<td style="font-size:0.85em">{v.get("clinical", "")}</td>'
                f'</tr>'
            )
        html_parts.append("</table>")

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)

    return json.dumps(results)

# ------------------------------------------------------------------
# Task 3: Analyze and visualize variant effects
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def analyze_effects(
    scores_json: str,
    variants_dir: flyte.io.Dir,
) -> str:
    """Analyze VEP scores: classification accuracy, gene-level summaries, and rich visualizations."""
    results = json.loads(scores_json)

    variants_path = await variants_dir.download()
    with open(os.path.join(variants_path, "genes.json")) as f:
        genes = json.load(f)

    html_parts = ["<h2>Variant Effect Analysis</h2>"]

    # ------------------------------------------------------------------
    # Overall accuracy: does the model's score direction match known labels?
    # ------------------------------------------------------------------
    all_variants = []
    correct = 0
    total_known = 0
    true_pos = 0
    false_pos = 0
    true_neg = 0
    false_neg = 0

    for gene_name, gene_data in results.items():
        for v in gene_data["variants"]:
            all_variants.append({**v, "gene": gene_name})
            if v["known_effect"] in ("pathogenic", "benign"):
                total_known += 1
                predicted_pathogenic = v["score"] < -0.05
                actual_pathogenic = v["known_effect"] == "pathogenic"
                if predicted_pathogenic == actual_pathogenic:
                    correct += 1
                if predicted_pathogenic and actual_pathogenic:
                    true_pos += 1
                elif predicted_pathogenic and not actual_pathogenic:
                    false_pos += 1
                elif not predicted_pathogenic and actual_pathogenic:
                    false_neg += 1
                else:
                    true_neg += 1

    accuracy = correct / total_known if total_known else 0
    precision = true_pos / (true_pos + false_pos) if (true_pos + false_pos) else 0
    recall = true_pos / (true_pos + false_neg) if (true_pos + false_neg) else 0

    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{accuracy:.0%}</div><div class="label">Direction Accuracy</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{precision:.0%}</div><div class="label">Precision (Pathogenic)</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{recall:.0%}</div><div class="label">Recall (Pathogenic)</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{len(all_variants)}</div><div class="label">Total Variants</div></div>')
    html_parts.append("</div>")

    html_parts.append(
        '<div class="note">'
        "<b>How to read VEP scores:</b> Negative scores mean Carbon considers the variant "
        "less likely than the reference sequence — suggestive of a damaging effect. Scores near "
        "zero indicate the model sees little difference (likely benign). The magnitude indicates "
        "confidence."
        "</div>"
    )

    # ------------------------------------------------------------------
    # Cross-gene score comparison dot plot
    # ------------------------------------------------------------------
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_score_comparison_chart(results))
    html_parts.append("</div>")

    # ------------------------------------------------------------------
    # Per-gene visualizations
    # ------------------------------------------------------------------
    for gene_name, gene_data in results.items():
        short_name = gene_name.split("(")[0].strip()
        html_parts.append(f'<h3>{gene_name}</h3>')
        html_parts.append(f'<div class="note">{gene_data["description"]}</div>')

        # DNA track with variant positions highlighted
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_dna_track(
            gene_data["sequence"],
            gene_data["variants"],
            gene_name=f"{short_name} Reference Sequence",
        ))
        html_parts.append("</div>")

        # Lollipop plot
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_lollipop_plot(
            gene_data["variants"],
            len(gene_data["sequence"]),
            title=f"{short_name} — Variant Positions & Effect Scores",
        ))
        html_parts.append("</div>")

        # Score bar chart for this gene
        variant_names = [v["name"] for v in gene_data["variants"]]
        variant_scores = [v["score"] for v in gene_data["variants"]]
        html_parts.append('<div class="chart-container">')
        html_parts.append(_make_bar_chart(
            variant_names,
            {"VEP Score": variant_scores},
            title=f"{short_name} — Log-Likelihood Ratio Scores",
            colors=[EFFECT_COLORS.get(v["known_effect"], "#6b7280") for v in gene_data["variants"]],
            value_format=".3f",
        ))
        html_parts.append("</div>")

        # Variant detail cards
        for v in gene_data["variants"]:
            badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
            score_color = "#dc2626" if v["score"] < -0.1 else "#059669" if v["score"] > 0.05 else "#f59e0b"
            html_parts.append(
                f'<div class="gene-card">'
                f'<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:8px;">'
                f'<b style="font-size:1.1em">{v["name"]}</b>'
                f'<span class="badge {badge}">{v["known_effect"]}</span>'
                f'</div>'
                f'<div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:8px;">'
                f'<div><span style="color:#6c757d;font-size:0.85em">Ref base:</span> '
                f'<b style="color:{BASE_COLORS.get(v["ref"], "#333")}">{v["ref"]}</b></div>'
                f'<div><span style="color:#6c757d;font-size:0.85em">Alt base:</span> '
                f'<b style="color:{BASE_COLORS.get(v["alt"], "#333")}">{v["alt"]}</b></div>'
                f'<div><span style="color:#6c757d;font-size:0.85em">VEP Score:</span> '
                f'<b style="color:{score_color}">{v["score"]:.4f}</b></div>'
                f'</div>'
                f'<div style="margin-top:8px;font-size:0.9em;color:#374151">{v.get("clinical", "")}</div>'
                f'</div>'
            )

    # ------------------------------------------------------------------
    # Confusion matrix as heatmap
    # ------------------------------------------------------------------
    html_parts.append("<h3>Classification Performance</h3>")
    html_parts.append(
        '<div class="note">'
        "Using a simple threshold (score &lt; -0.05 = predicted pathogenic). "
        "This is zero-shot — no training on these specific variants."
        "</div>"
    )

    conf_matrix = [[true_pos, false_neg], [false_pos, true_neg]]
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        conf_matrix,
        ["Actual Pathogenic", "Actual Benign"],
        ["Predicted Pathogenic", "Predicted Benign"],
        title="Confusion Matrix (Known Variants Only)",
        value_format=".0f",
        width=400,
        height=300,
    ))
    html_parts.append("</div>")

    # ------------------------------------------------------------------
    # Score distribution by known effect
    # ------------------------------------------------------------------
    html_parts.append("<h3>Score Distribution by Known Effect</h3>")

    pathogenic_scores = [v["score"] for v in all_variants if v["known_effect"] == "pathogenic"]
    benign_scores = [v["score"] for v in all_variants if v["known_effect"] == "benign"]
    uncertain_scores = [v["score"] for v in all_variants if v["known_effect"] == "uncertain"]

    stats_html = '<div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:12px;">'
    for label, scores, color in [
        ("Pathogenic", pathogenic_scores, "#dc2626"),
        ("Benign", benign_scores, "#059669"),
        ("Uncertain", uncertain_scores, "#f59e0b"),
    ]:
        if scores:
            mean_s = sum(scores) / len(scores)
            min_s = min(scores)
            max_s = max(scores)
            stats_html += (
                f'<div class="gene-card" style="border-left:4px solid {color}">'
                f'<b style="color:{color}">{label}</b> (n={len(scores)})<br>'
                f'Mean: {mean_s:.4f}<br>'
                f'Range: [{min_s:.4f}, {max_s:.4f}]'
                f'</div>'
            )
    stats_html += "</div>"
    html_parts.append(stats_html)

    # Summary
    analysis = {
        "total_variants": len(all_variants),
        "total_known": total_known,
        "accuracy": round(accuracy, 4),
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "true_pos": true_pos,
        "false_pos": false_pos,
        "true_neg": true_neg,
        "false_neg": false_neg,
        "pathogenic_mean_score": round(sum(pathogenic_scores) / len(pathogenic_scores), 4) if pathogenic_scores else None,
        "benign_mean_score": round(sum(benign_scores) / len(benign_scores), 4) if benign_scores else None,
    }

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)
    return json.dumps(analysis)

# ------------------------------------------------------------------
# Task 4: Generate comprehensive summary report
# ------------------------------------------------------------------

@cpu_env.task(report=True)
async def generate_summary(
    scores_json: str,
    analysis_json: str,
) -> str:
    """Generate the final summary report combining all results."""
    results = json.loads(scores_json)
    analysis = json.loads(analysis_json)

    html_parts = [
        "<h2>Genomic Variant Effect Prediction — Summary</h2>",
        '<div class="note">'
        "This pipeline uses <b>HuggingFace Carbon</b>, an autoregressive genomic foundation model "
        "trained on 1 trillion tokens of DNA sequence, to perform <b>zero-shot variant effect "
        "prediction</b>. No fine-tuning or labeled training data was used — the model scores "
        "variants purely based on its learned understanding of DNA sequence grammar."
        "</div>",
    ]

    # Key metrics
    html_parts.append('<div class="stat-grid">')
    html_parts.append(f'<div class="stat"><div class="value">{len(results)}</div><div class="label">Genes Analyzed</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["total_variants"]}</div><div class="label">Variants Scored</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["accuracy"]:.0%}</div><div class="label">Direction Accuracy</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["precision"]:.0%}</div><div class="label">Precision</div></div>')
    html_parts.append(f'<div class="stat"><div class="value">{analysis["recall"]:.0%}</div><div class="label">Recall</div></div>')
    html_parts.append("</div>")

    # Gene summary table
    html_parts.append("<h3>Per-Gene Summary</h3>")
    html_parts.append(
        "<table><tr><th>Gene</th><th>Variants</th><th>Mean Score</th>"
        "<th>Pathogenic</th><th>Benign</th><th>Uncertain</th></tr>"
    )

    for gene_name, gene_data in results.items():
        variants = gene_data["variants"]
        scores = [v["score"] for v in variants]
        mean_score = sum(scores) / len(scores) if scores else 0
        n_path = sum(1 for v in variants if v["known_effect"] == "pathogenic")
        n_benign = sum(1 for v in variants if v["known_effect"] == "benign")
        n_unc = sum(1 for v in variants if v["known_effect"] == "uncertain")
        short = gene_name.split("(")[0].strip()

        html_parts.append(
            f"<tr><td><b>{short}</b></td><td>{len(variants)}</td>"
            f"<td>{mean_score:.4f}</td>"
            f'<td><span class="badge badge-danger">{n_path}</span></td>'
            f'<td><span class="badge badge-success">{n_benign}</span></td>'
            f'<td><span class="badge badge-warning">{n_unc}</span></td></tr>'
        )
    html_parts.append("</table>")

    # Cross-gene heatmap: gene x metric
    gene_names = [g.split("(")[0].strip() for g in results.keys()]
    metrics = ["Mean Score", "Min Score", "Max Score", "# Variants"]
    matrix = []
    for gene_data in results.values():
        scores = [v["score"] for v in gene_data["variants"]]
        matrix.append([
            sum(scores) / len(scores) if scores else 0,
            min(scores) if scores else 0,
            max(scores) if scores else 0,
            len(scores),
        ])

    html_parts.append("<h3>Gene-Level Metrics</h3>")
    html_parts.append('<div class="chart-container">')
    html_parts.append(_make_heatmap(
        matrix,
        gene_names,
        metrics,
        title="Gene-Level VEP Score Summary",
        value_format=".2f",
        width=600,
        height=350,
    ))
    html_parts.append("</div>")

    # All variants ranked by score (most damaging first)
    html_parts.append("<h3>All Variants Ranked by Impact</h3>")
    all_vars_sorted = []
    for gene_name, gene_data in results.items():
        for v in gene_data["variants"]:
            all_vars_sorted.append({**v, "gene": gene_name.split("(")[0].strip()})
    all_vars_sorted.sort(key=lambda x: x["score"])

    html_parts.append(
        "<table><tr><th>#</th><th>Gene</th><th>Variant</th><th>Score</th>"
        "<th>Known</th><th>Clinical Significance</th></tr>"
    )
    for i, v in enumerate(all_vars_sorted):
        badge = EFFECT_BADGES.get(v["known_effect"], "badge-info")
        score_color = "#dc2626" if v["score"] < -0.1 else "#059669" if v["score"] > 0.05 else "#f59e0b"
        html_parts.append(
            f'<tr><td>{i + 1}</td><td><b>{v["gene"]}</b></td><td>{v["name"]}</td>'
            f'<td style="color:{score_color};font-weight:700">{v["score"]:.4f}</td>'
            f'<td><span class="badge {badge}">{v["known_effect"]}</span></td>'
            f'<td style="font-size:0.85em">{v.get("clinical", "")}</td></tr>'
        )
    html_parts.append("</table>")

    # Method note
    html_parts.append(
        '<div class="note">'
        "<b>Method:</b> Zero-shot variant effect prediction using log-likelihood ratio scoring. "
        "For each variant, we compute score = log P(mutant sequence | Carbon) - log P(reference sequence | Carbon). "
        "Negative scores indicate the model considers the mutant less probable than the reference, "
        "which correlates with pathogenicity. This approach requires no fine-tuning and generalizes "
        "across genes and variant types.<br><br>"
        "<b>Limitations:</b> These are short sequence windows — real clinical VEP would use longer "
        "genomic context (Carbon supports up to 786kbp). The threshold for pathogenicity classification "
        "(-0.05) is a simple heuristic; clinical use requires calibrated thresholds per gene."
        "</div>"
    )

    await flyte.report.replace.aio(_wrap_report("\n".join(html_parts)), do_flush=True)
    return json.dumps({"status": "complete", "analysis": analysis})

# ------------------------------------------------------------------
# Pipeline orchestrator
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    variants_json: str = "",
    model_name: str = "HuggingFaceBio/Carbon-3B",
) -> tuple[str, str]:
    """
    End-to-end genomic variant effect prediction pipeline.

    Returns (scores JSON, analysis JSON).

    1. Load and validate gene variants
    2. Score variants with Carbon (log-likelihood ratio)
    3. Analyze effects — accuracy, visualizations, classification
    4. Generate comprehensive summary report
    """
    log.info("Starting genomic variant effect prediction pipeline...")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = [
            "Load Variants",
            "Carbon Scoring",
            "Analyze Effects",
            "Generate Summary",
        ]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#2563eb;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#2563eb;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>Genomic Variant Effect Prediction</h2>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    # Stage 1: Load variants
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Loading and validating gene variants...")),
        do_flush=True,
    )
    var_dir = await load_variants(variants_json=variants_json)

    # Stage 2: Score with Carbon
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Running Carbon model for variant effect scoring...")),
        do_flush=True,
    )
    scores_json = await score_variants(variants_dir=var_dir, model_name=model_name)

    # Stage 3: Analyze effects
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Analyzing variant effects and generating visualizations...")),
        do_flush=True,
    )
    analysis_json = await analyze_effects(scores_json=scores_json, variants_dir=var_dir)

    # Stage 4: Summary
    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Generating comprehensive summary report...")),
        do_flush=True,
    )
    summary_json = await generate_summary(scores_json=scores_json, analysis_json=analysis_json)

    # Final pipeline report
    analysis = json.loads(analysis_json)
    results = json.loads(scores_json)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(results)}</div><div class="label">Genes Analyzed</div></div>
      <div class="stat"><div class="value">{analysis['total_variants']}</div><div class="label">Variants Scored</div></div>
      <div class="stat"><div class="value">{analysis['accuracy']:.0%}</div><div class="label">Direction Accuracy</div></div>
      <div class="stat"><div class="value">{analysis['precision']:.0%}</div><div class="label">Precision</div></div>
      <div class="stat"><div class="value">{analysis['recall']:.0%}</div><div class="label">Recall</div></div>
    </div>
    <div class="card">
      <b>Model:</b> HuggingFace Carbon |
      <b>Method:</b> Zero-shot log-likelihood ratio scoring |
      <b>Genes:</b> {', '.join(g.split('(')[0].strip() for g in results.keys())}
    </div>
    <div class="note">
      All 4 pipeline stages completed successfully. View individual task reports for detailed
      visualizations including DNA sequence tracks, variant lollipop plots, VEP score charts,
      confusion matrices, and ranked variant tables.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)

    log.info("Pipeline complete.")
    return scores_json, analysis_json

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic_variant_effect/genomic_variant_effect.py*

## Run the workflow

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic_variant_effect):

```
cd v2/tutorials/genomic_variant_effect
uv run --script genomic_variant_effect.py
```

Use a smaller Carbon model for faster iteration:

```
flyte run genomic_variant_effect.py pipeline --model_name HuggingFaceBio/Carbon-500M
```

Negative VEP scores indicate the model prefers the reference allele over the alternate, a signal correlated with pathogenicity in this zero-shot setup.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/drug-molecule-screening ===

# Drug molecule screening agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/drug_molecule_screening).

This tutorial builds an **agentic** virtual drug-screening workflow on Flyte. A medicinal-chemistry agent interprets your therapeutic goal in plain language, derives screening criteria, and composes durable RDKit stage tasks, while the scientific core (property computation, Lipinski filters, Tanimoto similarity, ranking, and HTML reports) stays in trusted, deterministic tools.

The pattern follows how cheminformatics agents like ChemCrow and PharmAgents are built: **the LLM plans and reflects; RDKit computes.**

Flyte provides:

- **Flyte-native agent orchestration** via `flyte.ai.agents.Agent` (see **Agents > Build an agent > Flyte-native agents**)
- **Typed agent tool I/O**: Flyte 2.5.4+ passes `flyte.io.Dir`, `File`, and `DataFrame` between agent tool calls so the LLM can compose multi-step pipelines directly
- **Cached molecule loading** so repeated runs skip re-parsing SMILES
- **Report-enabled stage tasks** that stream property charts, similarity matrices, and candidate spotlights as each step completes
- **Hybrid iteration**: the agent re-runs `screen_candidates` and `generate_report` with adjusted criteria when the funnel is too narrow, reusing cached `molecule_dir` and `properties_json`

> [!NOTE] Prerequisites
> Create an Anthropic API key secret (the key name must match the `TaskEnvironment`):
>
> ```
> flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
> ```
>
> See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

## Define the task environment

The pipeline runs on CPU with RDKit, LiteLLM, and system libraries for 2D structure rendering.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.4",
#    "litellm",
#    "rdkit",
#    "numpy",
#    "scikit-learn",
#    "pillow",
# ]
# main = "pipeline"
# params = ""
# ///
"""Virtual drug molecule screening — compute properties, apply Lipinski filters, rank candidates."""

import base64
import io
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report
from flyte.ai.agents import Agent, tool

MODEL = os.getenv("DRUG_SCREENING_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="drug-molecule-screening", pre=True).with_apt_packages(
    "libxrender1", "libxext6", "libexpat1",
)

env = flyte.TaskEnvironment(
    name="drug-molecule-screening",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default molecule library — real SMILES for well-known drugs
# ------------------------------------------------------------------

DEFAULT_MOLECULES = {
    "Aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
    "Ibuprofen": "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
    "Caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
    "Penicillin G": "CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C",
    "Metformin": "CN(C)C(=N)NC(=N)N",
    "Paracetamol": "CC(=O)NC1=CC=C(C=C1)O",
    "Diazepam": "ClC1=CC2=C(C=C1)N(C(=O)CN=C2C3=CC=CC=C3)C",
    "Omeprazole": "CC1=CN=C(C(=C1OC)C)CS(=O)C2=NC3=CC=CC=C3N2",
    "Atorvastatin": "CC(C)C1=C(C(=C(N1CCC(CC(CC(=O)O)O)O)C2=CC=C(C=C2)F)C3=CC=CC=C3)C(=O)NC4=CC=CC=C4",
    "Methotrexate": "CN(CC1=CN=C2N=C(N=C(N)C2=N1)N)C3=CC=C(C=C3)C(=O)NC(CCC(=O)O)C(=O)O",
    "Doxorubicin": "CC1C(C(CC(O1)OC2CC(CC3=C2C(=C4C(=C3O)C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)CO)O)N)O",
    "Tamoxifen": "CCC(=C(C1=CC=CC=C1)C2=CC=C(C=C2)OCCN(C)C)C3=CC=CC=C3",
    "Lopinavir": "CC1=C(C(=CC=C1)C)OCC(=O)NC(CC2=CC=CC=C2)C(CC(CC3=CC=CC=C3)NC(=O)C(C(C)C)N4CCCNC4=O)O",
    "Remdesivir": "CCC(CC)COC(=O)C(C)NP(=O)(OCC1C(C(C(O1)C2=CC=C3N2N=CN=C3N)O)O)OC4=CC=CC=C4",
    "Erlotinib": "COCCOC1=CC2=C(C=C1OCCOC)C(=NC=N2)NC3=CC=CC(=C3)C#C",
}

# ------------------------------------------------------------------
# Report styling — pharma blue/cyan theme
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #0e4f6e; border-bottom: 2px solid #0891b2; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #155e75; margin-top: 20px; }
  .report .card { background: #ecfeff; border: 1px solid #a5f3fc; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #cffafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0e4f6e; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0e4f6e; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #cffafe; }
  .report tr:nth-child(even) { background: #ecfeff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #cffafe; color: #155e75; }
  .report .chart-container { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #ecfeff; border-left: 4px solid #0891b2; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .molecule-card { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .molecule-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 12px; margin: 16px 0; }
  .report .funnel { text-align: center; margin: 24px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _mol_to_data_uri(mol, size: tuple[int, int] = (300, 300)) -> str:
    """Convert an RDKit molecule to a PNG base64 data URI."""
    from rdkit.Chem import Draw

    img = Draw.MolToImage(mol, size=size)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 340,
    y_max_cap: float | None = None,
    horizontal: bool = False,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels.
        series: Dict mapping series name to list of values.
        title: Chart title.
        colors: Colors for each series.
        width/height: SVG dimensions.
        y_max_cap: Cap the y-axis at this value.
        horizontal: If True, draw horizontal bars.
        value_fmt: Format string for value labels.

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#a5f3fc", "#155e75"]
    colors = colors or default_colors

    if horizontal:
        return _make_horizontal_bar_chart(labels, series, title, colors, width, height, value_fmt)

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_fmt}}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_fmt}}</text>'
            )
        # Truncate long labels
        disp_label = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 16}" '
            f'text-anchor="middle" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {gx + n_series * bar_width / 2:.1f}, {mt + ch + 16})">'
            f'{disp_label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 100
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 100}" y="{mt + ch + 40}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 51}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_horizontal_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 400,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG horizontal bar chart (sorted)."""
    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0"]
    colors = colors or default_colors

    n = len(labels)
    row_height = max(22, min(35, (height - 80) // max(n, 1)))
    actual_height = max(height, 80 + n * row_height)
    ml, mr, mt, mb = 120, 60, 40, 20
    cw = width - ml - mr
    ch = actual_height - mt - mb

    # Use first series
    first_key = list(series.keys())[0]
    vals = series[first_key]
    x_max = max(vals) * 1.15 if vals else 1

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {actual_height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{actual_height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    bar_h = row_height * 0.65
    for i, (label, val) in enumerate(zip(labels, vals)):
        y = mt + i * row_height
        bw = (val / x_max) * cw if x_max else 0
        color = colors[i % len(colors)]
        # Label
        disp = label if len(label) <= 14 else label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + bar_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#1a1a2e">{disp}</text>'
        )
        # Bar
        svg.append(
            f'<rect x="{ml}" y="{y:.1f}" width="{bw:.1f}" height="{bar_h:.1f}" '
            f'fill="{color}" rx="3"/>'
        )
        # Value
        svg.append(
            f'<text x="{ml + bw + 6:.1f}" y="{y + bar_h / 2 + 4:.1f}" '
            f'font-size="11" fill="#0e4f6e" font-weight="600">{val:{value_fmt}}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    color_scale: str = "cyan",
    width: int = 700,
    height: int = 500,
    value_fmt: str = ".2f",
) -> str:
    """Generate an SVG heatmap.

    Args:
        matrix: 2D list of values (rows x cols).
        row_labels: Labels for rows.
        col_labels: Labels for columns.
        title: Chart title.
        color_scale: Color scheme ("cyan", "red", "green").
        width/height: SVG dimensions.
        value_fmt: Format string for cell values.

    Returns:
        SVG string.
    """
    if not matrix or not matrix[0]:
        return ""

    n_rows = len(matrix)
    n_cols = len(matrix[0])

    ml, mr, mt, mb = 110, 20, 70, 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    # Flatten to find range
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    def color_for(v):
        t = (v - v_min) / v_range
        if color_scale == "cyan":
            # White to deep teal
            r = int(255 - t * (255 - 14))
            g = int(255 - t * (255 - 79))
            b = int(255 - t * (255 - 110))
        elif color_scale == "red":
            r = int(255 - t * 50)
            g = int(255 - t * 200)
            b = int(255 - t * 200)
        else:  # green
            r = int(255 - t * 200)
            g = int(255 - t * 50)
            b = int(255 - t * 200)
        return f"rgb({r},{g},{b})"

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Column labels (rotated)
    for ci, label in enumerate(col_labels):
        x = ml + ci * cell_w + cell_w / 2
        disp = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{x:.1f}" y="{mt - 8}" text-anchor="end" font-size="10" '
            f'fill="#1a1a2e" transform="rotate(-45, {x:.1f}, {mt - 8})">{disp}</text>'
        )

    # Row labels + cells
    for ri, (row_label, row_vals) in enumerate(zip(row_labels, matrix)):
        y = mt + ri * cell_h
        disp = row_label if len(row_label) <= 14 else row_label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + cell_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#1a1a2e">{disp}</text>'
        )
        for ci, val in enumerate(row_vals):
            x = ml + ci * cell_w
            fill = color_for(val)
            svg.append(
                f'<rect x="{x:.1f}" y="{y:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{fill}" stroke="#fff" stroke-width="1"/>'
            )
            # Text color: dark on light, light on dark
            t = (val - v_min) / v_range
            txt_color = "#fff" if t > 0.55 else "#1a1a2e"
            # Only show text if cells are large enough
            if cell_w > 30 and cell_h > 18:
                svg.append(
                    f'<text x="{x + cell_w / 2:.1f}" y="{y + cell_h / 2 + 4:.1f}" '
                    f'text-anchor="middle" font-size="9" fill="{txt_color}">'
                    f'{val:{value_fmt}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_scatter_plot(
    points: list[dict],
    x_label: str = "MW",
    y_label: str = "LogP",
    title: str = "",
    reference_lines: list[dict] | None = None,
    width: int = 700,
    height: int = 400,
) -> str:
    """Generate an SVG scatter plot.

    Args:
        points: List of dicts with "x", "y", "label" keys.
        x_label/y_label: Axis labels.
        title: Chart title.
        reference_lines: List of dicts with "axis" ("x"/"y"), "value", "label".
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not points:
        return ""

    ml, mr, mt, mb = 60, 30, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [p["x"] for p in points]
    y_vals = [p["y"] for p in points]
    x_min, x_max = min(x_vals) * 0.9, max(x_vals) * 1.1
    y_min, y_max = min(y_vals) - 1, max(y_vals) + 1

    # Extend ranges to include reference lines
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                x_max = max(x_max, rl["value"] * 1.1)
            else:
                y_max = max(y_max, rl["value"] * 1.1)

    x_range = x_max - x_min or 1
    y_range = y_max - y_min or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid
    for i in range(6):
        y_tick = y_min + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.1f}</text>'
        )

    for i in range(6):
        x_tick = x_min + x_range * i / 5
        px = sx(x_tick)
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
            f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Reference lines (Lipinski boundaries)
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                px = sx(rl["value"])
                svg.append(
                    f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{px + 4:.1f}" y="{mt + 14}" font-size="10" '
                    f'fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )
            else:
                py = sy(rl["value"])
                svg.append(
                    f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{ml + cw - 4:.1f}" y="{py - 6:.1f}" text-anchor="end" '
                    f'font-size="10" fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )

    # Drug-like zone shading (MW<=500 and LogP<=5 quadrant)
    if reference_lines:
        mw_line = next((rl for rl in reference_lines if rl["axis"] == "x"), None)
        logp_line = next((rl for rl in reference_lines if rl["axis"] == "y"), None)
        if mw_line and logp_line:
            zx1 = sx(x_min)
            zx2 = sx(min(mw_line["value"], x_max))
            zy1 = sy(min(logp_line["value"], y_max))
            zy2 = sy(y_min)
            svg.append(
                f'<rect x="{zx1:.1f}" y="{zy1:.1f}" '
                f'width="{zx2 - zx1:.1f}" height="{zy2 - zy1:.1f}" '
                f'fill="#0891b2" opacity="0.06" rx="4"/>'
            )
            svg.append(
                f'<text x="{zx1 + 8:.1f}" y="{zy2 - 8:.1f}" font-size="11" '
                f'fill="#0891b2" font-weight="600" opacity="0.6">Drug-like Zone</text>'
            )

    # Points
    point_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#155e75", "#0284c7",
                    "#059669", "#0d9488", "#0369a1", "#047857", "#115e59",
                    "#0c4a6e", "#064e3b", "#1e3a5f", "#134e4a", "#075985"]
    for i, pt in enumerate(points):
        px, py = sx(pt["x"]), sy(pt["y"])
        color = point_colors[i % len(point_colors)]
        svg.append(
            f'<circle cx="{px:.1f}" cy="{py:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5" opacity="0.85"/>'
        )
        # Label offset to avoid overlap
        offset_x = 8
        offset_y = -8 if i % 2 == 0 else 14
        label = pt["label"] if len(pt["label"]) <= 12 else pt["label"][:10] + ".."
        svg.append(
            f'<text x="{px + offset_x:.1f}" y="{py + offset_y:.1f}" '
            f'font-size="9" fill="#1a1a2e">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Axis labels
    if x_label:
        svg.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        svg.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_funnel(
    stages: list[dict],
    title: str = "",
    width: int = 600,
    height: int = 400,
) -> str:
    """Generate an SVG funnel visualization.

    Args:
        stages: List of dicts with "label", "count", "total" keys.
        title: Chart title.
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not stages:
        return ""

    n = len(stages)
    mt = 50
    mb = 20
    available_h = height - mt - mb
    stage_h = available_h / n
    cx = width / 2

    # Color gradient from light cyan to deep teal
    colors = []
    for i in range(n):
        t = i / max(n - 1, 1)
        r = int(207 - t * (207 - 14))
        g = int(250 - t * (250 - 79))
        b = int(254 - t * (254 - 110))
        colors.append(f"rgb({r},{g},{b})")

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{cx}" y="28" text-anchor="middle" '
            f'font-size="16" font-weight="700" fill="#0e4f6e">{title}</text>'
        )

    max_count = stages[0]["count"] if stages else 1
    max_width = width * 0.75

    for i, stage in enumerate(stages):
        y_top = mt + i * stage_h
        y_bot = y_top + stage_h

        # Width proportional to count
        w_top = max_width * (stage["count"] / max_count) if i == 0 else prev_w_bot
        if i < n - 1:
            w_bot = max_width * (stages[i + 1]["count"] / max_count)
        else:
            w_bot = max_width * (stage["count"] / max_count) * 0.7

        prev_w_bot = w_bot

        # Trapezoid
        x1_top = cx - w_top / 2
        x2_top = cx + w_top / 2
        x1_bot = cx - w_bot / 2
        x2_bot = cx + w_bot / 2

        svg.append(
            f'<polygon points="{x1_top:.1f},{y_top:.1f} {x2_top:.1f},{y_top:.1f} '
            f'{x2_bot:.1f},{y_bot:.1f} {x1_bot:.1f},{y_bot:.1f}" '
            f'fill="{colors[i]}" stroke="#fff" stroke-width="2"/>'
        )

        # Text: dark on light, white on dark
        t = i / max(n - 1, 1)
        txt_color = "#0e4f6e" if t < 0.5 else "#fff"
        y_mid = (y_top + y_bot) / 2

        svg.append(
            f'<text x="{cx}" y="{y_mid - 4:.1f}" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="{txt_color}">{stage["label"]}</text>'
        )
        svg.append(
            f'<text x="{cx}" y="{y_mid + 14:.1f}" text-anchor="middle" '
            f'font-size="12" fill="{txt_color}" opacity="0.85">'
            f'{stage["count"]} / {stage["total"]}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate molecules
# ------------------------------------------------------------------

@tool
@env.task(cache="auto")
async def load_molecules(
    molecules_json: str = "",
) -> flyte.io.Dir:
    """Parse SMILES strings, validate with RDKit, generate 2D depictions.

    Args:
        molecules_json: JSON string mapping molecule names to SMILES.
            Defaults to a curated library of ~15 well-known drugs.

    Returns:
        flyte.io.Dir containing molecule data (JSON + PNG depictions).
        Pass this directory to compute_properties and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Draw

    if molecules_json.strip():
        molecules = json.loads(molecules_json)
    else:
        molecules = DEFAULT_MOLECULES

    out_dir = tempfile.mkdtemp(prefix="mol_library_")
    results = []
    valid_count = 0
    invalid_count = 0

    log.info(f"Parsing {len(molecules)} molecules...")

    for name, smiles in molecules.items():
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            log.warning(f"  [INVALID] {name}: {smiles}")
            invalid_count += 1
            continue

        valid_count += 1

        # Generate 2D depiction as PNG
        img = Draw.MolToImage(mol, size=(300, 300))
        img_path = os.path.join(out_dir, f"{name.replace(' ', '_')}.png")
        img.save(img_path)

        results.append({
            "name": name,
            "smiles": smiles,
            "valid": True,
            "image_file": os.path.basename(img_path),
        })

    # Save molecule manifest
    manifest = {
        "total": len(molecules),
        "valid": valid_count,
        "invalid": invalid_count,
        "molecules": results,
    }
    manifest_path = os.path.join(out_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    log.info(f"Loaded {valid_count} valid molecules ({invalid_count} invalid)")

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Compute physicochemical properties
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def compute_properties(
    molecule_dir: flyte.io.Dir,
) -> str:
    """Compute drug-likeness properties for all molecules.

    Computes MW, LogP, HBD, HBA, TPSA, rotatable bonds, formal charge,
    ring count, QED, and Lipinski Rule of Five compliance.

    Args:
        molecule_dir: Directory from load_molecules.

    Returns:
        JSON string with all computed properties. Pass to screen_candidates
        and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Descriptors, Lipinski
    from rdkit.Chem.QED import qed

    # --- Loading report ---
    await flyte.report.replace.aio(
        _wrap_report("<h2>Computing Molecular Properties...</h2>"
                      "<p>Analyzing physicochemical descriptors for all molecules.</p>"),
        do_flush=True,
    )

    mol_dir = await molecule_dir.download()
    with open(os.path.join(mol_dir, "manifest.json")) as f:
        manifest = json.load(f)

    molecules_data = []
    lipinski_pass = 0

    for mol_info in manifest["molecules"]:
        mol = Chem.MolFromSmiles(mol_info["smiles"])
        if mol is None:
            continue

        mw = Descriptors.MolWt(mol)
        logp = Descriptors.MolLogP(mol)
        hbd = Lipinski.NumHDonors(mol)
        hba = Lipinski.NumHAcceptors(mol)
        tpsa = Descriptors.TPSA(mol)
        rotatable = Lipinski.NumRotatableBonds(mol)
        formal_charge = Chem.GetFormalCharge(mol)
        num_rings = Lipinski.RingCount(mol)
        qed_score = qed(mol)

        # Lipinski Rule of Five
        lipinski = {
            "mw_ok": mw <= 500,
            "logp_ok": logp <= 5,
            "hbd_ok": hbd <= 5,
            "hba_ok": hba <= 10,
        }
        lipinski_all = all(lipinski.values())
        if lipinski_all:
            lipinski_pass += 1

        # Read image for data URI
        img_path = os.path.join(mol_dir, mol_info["image_file"])
        data_uri = ""
        if os.path.exists(img_path):
            with open(img_path, "rb") as img_f:
                b64 = base64.b64encode(img_f.read()).decode()
                data_uri = f"data:image/png;base64,{b64}"

        molecules_data.append({
            "name": mol_info["name"],
            "smiles": mol_info["smiles"],
            "mw": round(mw, 2),
            "logp": round(logp, 2),
            "hbd": hbd,
            "hba": hba,
            "tpsa": round(tpsa, 2),
            "rotatable_bonds": rotatable,
            "formal_charge": formal_charge,
            "num_rings": num_rings,
            "qed": round(qed_score, 4),
            "lipinski": lipinski,
            "lipinski_pass": lipinski_all,
            "image_data_uri": data_uri,
        })

    total = len(molecules_data)
    avg_mw = sum(m["mw"] for m in molecules_data) / total if total else 0
    avg_logp = sum(m["logp"] for m in molecules_data) / total if total else 0
    lipinski_rate = lipinski_pass / total * 100 if total else 0

    # ---- Build report ----
    html_parts = []

    # Header
    html_parts.append("<h2>Molecular Properties Analysis</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Total Molecules"),
        (f"{lipinski_rate:.0f}%", "Lipinski Pass Rate"),
        (f"{avg_mw:.1f}", "Avg. MW (Da)"),
        (f"{avg_logp:.2f}", "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Molecule gallery
    html_parts.append("<h3>Molecule Library</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in molecules_data:
        if m["image_data_uri"]:
            badge_class = "badge-success" if m["lipinski_pass"] else "badge-danger"
            badge_text = "Lipinski Pass" if m["lipinski_pass"] else "Lipinski Fail"
            html_parts.append(
                f'<div class="molecule-card" style="text-align:center;">'
                f'<img src="{m["image_data_uri"]}" style="width:160px;height:160px;object-fit:contain;"/>'
                f'<div style="font-weight:600;margin-top:6px;color:#0e4f6e;">{m["name"]}</div>'
                f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f}</div>'
                f'<div><span class="badge {badge_class}">{badge_text}</span></div>'
                f'</div>'
            )
    html_parts.append("</div>")

    # MW bar chart (horizontal, sorted)
    sorted_by_mw = sorted(molecules_data, key=lambda m: m["mw"], reverse=True)
    mw_labels = [m["name"] for m in sorted_by_mw]
    mw_vals = [m["mw"] for m in sorted_by_mw]
    mw_chart = _make_bar_chart(
        mw_labels, {"MW (Da)": mw_vals},
        title="Molecular Weight Distribution",
        horizontal=True,
        width=700, height=max(300, len(mw_labels) * 30 + 80),
        value_fmt=".1f",
    )
    html_parts.append("<h3>Molecular Weight</h3>")
    html_parts.append(f'<div class="chart-container">{mw_chart}</div>')

    # LogP vs MW scatter plot
    scatter_points = [
        {"x": m["mw"], "y": m["logp"], "label": m["name"]}
        for m in molecules_data
    ]
    scatter_chart = _make_scatter_plot(
        scatter_points,
        x_label="Molecular Weight (Da)",
        y_label="LogP",
        title="LogP vs. Molecular Weight (Lipinski Boundaries)",
        reference_lines=[
            {"axis": "x", "value": 500, "label": "MW = 500"},
            {"axis": "y", "value": 5, "label": "LogP = 5"},
        ],
        width=700,
        height=420,
    )
    html_parts.append("<h3>Lipinski Space</h3>")
    html_parts.append(f'<div class="chart-container">{scatter_chart}</div>')

    # Property heatmap (molecules x properties)
    prop_names = ["MW", "LogP", "HBD", "HBA", "TPSA", "Rot. Bonds"]
    # Normalize each property to 0-1 for heatmap
    raw_matrix = []
    for m in molecules_data:
        raw_matrix.append([m["mw"], m["logp"], m["hbd"], m["hba"], m["tpsa"], m["rotatable_bonds"]])

    # Normalize per column
    n_props = len(prop_names)
    col_min = [min(row[c] for row in raw_matrix) for c in range(n_props)]
    col_max = [max(row[c] for row in raw_matrix) for c in range(n_props)]
    norm_matrix = []
    for row in raw_matrix:
        norm_row = []
        for c in range(n_props):
            rng = col_max[c] - col_min[c]
            norm_row.append((row[c] - col_min[c]) / rng if rng else 0.5)
        norm_matrix.append(norm_row)

    heatmap_labels = [m["name"] for m in molecules_data]
    heatmap = _make_heatmap(
        norm_matrix, heatmap_labels, prop_names,
        title="Normalized Property Heatmap",
        color_scale="cyan",
        width=700,
        height=max(400, len(heatmap_labels) * 28 + 100),
    )
    html_parts.append("<h3>Property Heatmap</h3>")
    html_parts.append(f'<div class="chart-container">{heatmap}</div>')

    # Lipinski compliance table
    html_parts.append("<h3>Lipinski Rule of Five Compliance</h3>")
    html_parts.append("<table><tr><th>Molecule</th><th>MW &le; 500</th>"
                      "<th>LogP &le; 5</th><th>HBD &le; 5</th>"
                      "<th>HBA &le; 10</th><th>Overall</th></tr>")
    for m in molecules_data:
        lip = m["lipinski"]

        def _badge(ok):
            if ok:
                return '<span class="badge badge-success">Pass</span>'
            return '<span class="badge badge-danger">Fail</span>'

        overall_badge = _badge(m["lipinski_pass"])
        html_parts.append(
            f'<tr><td><strong>{m["name"]}</strong></td>'
            f'<td>{_badge(lip["mw_ok"])}</td>'
            f'<td>{_badge(lip["logp_ok"])}</td>'
            f'<td>{_badge(lip["hbd_ok"])}</td>'
            f'<td>{_badge(lip["hba_ok"])}</td>'
            f'<td>{overall_badge}</td></tr>'
        )
    html_parts.append("</table>")

    # QED bar chart
    sorted_by_qed = sorted(molecules_data, key=lambda m: m["qed"], reverse=True)
    qed_labels = [m["name"] for m in sorted_by_qed]
    qed_vals = [m["qed"] for m in sorted_by_qed]
    qed_chart = _make_bar_chart(
        qed_labels, {"QED Score": qed_vals},
        title="Drug-likeness (QED Score)",
        horizontal=True,
        width=700, height=max(300, len(qed_labels) * 30 + 80),
        value_fmt=".3f",
        colors=["#06d6a0"],
    )
    html_parts.append("<h3>Drug-likeness (QED)</h3>")
    html_parts.append(f'<div class="chart-container">{qed_chart}</div>')

    # Flush full report
    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # Return properties as JSON (strip image data URIs to reduce size)
    output = {
        "total": total,
        "lipinski_pass_count": lipinski_pass,
        "lipinski_pass_rate": round(lipinski_rate, 2),
        "avg_mw": round(avg_mw, 2),
        "avg_logp": round(avg_logp, 2),
        "molecules": [
            {k: v for k, v in m.items() if k != "image_data_uri"}
            for m in molecules_data
        ],
    }
    return json.dumps(output)

# ------------------------------------------------------------------
# Task 3: Screen candidates against target profile
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def screen_candidates(
    properties_json: str,
    target_profile: str = "",
) -> str:
    """Screen molecules against a target drug profile and rank candidates.

    Scores each molecule on how well it matches the target profile, computes
    pairwise Tanimoto similarity, and produces a ranked list.

    Args:
        properties_json: JSON from compute_properties.
        target_profile: JSON string with desired property ranges
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5.0]}).

    Returns:
        JSON string with ranked_molecules, similarity_matrix, similarity_labels,
        funnel, and target_profile. Pass the full return value verbatim to
        generate_report along with molecule_dir and properties_json.
    """
    from rdkit import Chem, DataStructs
    from rdkit.Chem import AllChem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Screening Candidates...</h2>"
                      "<p>Evaluating molecules against the target drug profile.</p>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    molecules = props["molecules"]

    # Default target profile
    if target_profile.strip():
        profile = json.loads(target_profile)
    else:
        profile = {
            "mw": [150, 500],
            "logp": [-0.5, 5.0],
            "hbd": [0, 5],
            "hba": [0, 10],
            "tpsa": [20, 140],
        }

    # --- Screening ---
    funnel_total = len(molecules)
    pass_mw = 0
    pass_logp = 0
    pass_lipinski = 0
    final_candidates = 0

    scored = []
    for m in molecules:
        score = 0
        max_score = 0
        criteria = {}

        # Check each profile criterion
        checks = [
            ("mw", m["mw"]),
            ("logp", m["logp"]),
            ("hbd", m["hbd"]),
            ("hba", m["hba"]),
            ("tpsa", m["tpsa"]),
        ]

        for key, val in checks:
            if key in profile:
                lo, hi = profile[key]
                max_score += 1
                in_range = lo <= val <= hi
                criteria[key] = in_range
                if in_range:
                    score += 1
                    # Bonus: closer to midpoint = higher score
                    mid = (lo + hi) / 2
                    rng = (hi - lo) / 2
                    dist = abs(val - mid) / rng if rng else 0
                    score += max(0, 0.5 * (1 - dist))

        # QED bonus
        score += m["qed"] * 2
        max_score += 2

        # Lipinski bonus
        if m["lipinski_pass"]:
            score += 1
        max_score += 1

        normalized_score = score / max_score if max_score else 0

        # Funnel tracking — cascading filter (each stage requires passing the previous)
        mw_ok = criteria.get("mw", True)
        logp_ok = criteria.get("logp", True)
        if mw_ok:
            pass_mw += 1
            if logp_ok:
                pass_logp += 1
                if m["lipinski_pass"]:
                    pass_lipinski += 1
                    if all(criteria.values()):
                        final_candidates += 1

        scored.append({
            **m,
            "screening_score": round(normalized_score, 4),
            "criteria_met": criteria,
            "all_criteria_met": all(criteria.values()),
        })

    # Sort by score descending
    scored.sort(key=lambda m: m["screening_score"], reverse=True)

    # --- Tanimoto similarity matrix ---
    fps = []
    valid_names = []
    for m in scored:
        mol = Chem.MolFromSmiles(m["smiles"])
        if mol:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
            fps.append(fp)
            valid_names.append(m["name"])

    similarity_matrix = []
    for i in range(len(fps)):
        row = []
        for j in range(len(fps)):
            sim = DataStructs.TanimotoSimilarity(fps[i], fps[j])
            row.append(round(sim, 3))
        similarity_matrix.append(row)

    # ---- Build report ----
    html_parts = []
    html_parts.append("<h2>Candidate Screening Results</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(funnel_total), "Total Screened"),
        (str(pass_lipinski), "Lipinski Passes"),
        (str(final_candidates), "All Criteria Met"),
        (f"{scored[0]['screening_score']:.3f}" if scored else "N/A", "Top Score"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Screening funnel
    funnel_stages = [
        {"label": "Total Molecules", "count": funnel_total, "total": funnel_total},
        {"label": "Pass MW Filter", "count": pass_mw, "total": funnel_total},
        {"label": "Pass LogP Filter", "count": pass_logp, "total": funnel_total},
        {"label": "Lipinski Compliant", "count": pass_lipinski, "total": funnel_total},
        {"label": "All Criteria Met", "count": final_candidates, "total": funnel_total},
    ]
    funnel_svg = _make_funnel(
        funnel_stages,
        title="Screening Funnel",
        width=600,
        height=380,
    )
    html_parts.append("<h3>Screening Funnel</h3>")
    html_parts.append(f'<div class="chart-container" style="text-align:center;">{funnel_svg}</div>')

    # Ranked candidates table
    html_parts.append("<h3>Ranked Candidates</h3>")
    html_parts.append(
        "<table><tr><th>Rank</th><th>Molecule</th><th>Score</th>"
        "<th>MW</th><th>LogP</th><th>QED</th><th>Lipinski</th><th>All Criteria</th></tr>"
    )
    for rank, m in enumerate(scored, 1):
        lip_badge = ('<span class="badge badge-success">Pass</span>'
                     if m["lipinski_pass"]
                     else '<span class="badge badge-danger">Fail</span>')
        crit_badge = ('<span class="badge badge-success">Pass</span>'
                      if m["all_criteria_met"]
                      else '<span class="badge badge-danger">Fail</span>')
        # Highlight top 3
        row_style = ' style="background:#ecfeff;font-weight:600;"' if rank <= 3 else ""
        html_parts.append(
            f"<tr{row_style}><td>{rank}</td><td>{m['name']}</td>"
            f"<td>{m['screening_score']:.3f}</td>"
            f"<td>{m['mw']:.1f}</td><td>{m['logp']:.2f}</td>"
            f"<td>{m['qed']:.3f}</td><td>{lip_badge}</td><td>{crit_badge}</td></tr>"
        )
    html_parts.append("</table>")

    # Top 5 candidate cards with structures
    html_parts.append("<h3>Top 5 Candidates</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in scored[:5]:
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(250, 250)) if mol else ""
        badge_class = "badge-success" if m["all_criteria_met"] else "badge-info"
        badge_text = "All Criteria Met" if m["all_criteria_met"] else "Partial Match"
        html_parts.append(
            f'<div class="molecule-card" style="text-align:center;">'
            f'<img src="{img_uri}" style="width:140px;height:140px;object-fit:contain;"/>'
            f'<div style="font-weight:700;margin-top:6px;color:#0e4f6e;font-size:1.05em;">{m["name"]}</div>'
            f'<div style="font-size:0.85em;color:#155e75;margin:4px 0;">Score: {m["screening_score"]:.3f}</div>'
            f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f} | QED: {m["qed"]:.3f}</div>'
            f'<div style="margin-top:4px;"><span class="badge {badge_class}">{badge_text}</span></div>'
            f'</div>'
        )
    html_parts.append("</div>")

    # Tanimoto similarity heatmap
    if similarity_matrix:
        sim_heatmap = _make_heatmap(
            similarity_matrix, valid_names, valid_names,
            title="Pairwise Tanimoto Similarity (Morgan Fingerprints)",
            color_scale="cyan",
            width=700,
            height=max(500, len(valid_names) * 32 + 100),
        )
        html_parts.append("<h3>Chemical Similarity</h3>")
        html_parts.append(f'<div class="chart-container">{sim_heatmap}</div>')

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    output = {
        "ranked_molecules": scored,
        "similarity_matrix": similarity_matrix,
        "similarity_labels": valid_names,
        "funnel": funnel_stages,
        "target_profile": profile,
    }
    return json.dumps(output)

def _parse_screening_json(screening_json: str) -> dict:
    """Parse screening JSON from screen_candidates, with safe defaults.

    The agent must pass the exact tool return value. Partial or hand-built JSON
    is tolerated for optional similarity fields only.
    """
    screening = json.loads(screening_json)
    if "ranked_molecules" not in screening:
        raise ValueError(
            "screening_json must be the exact JSON string returned by "
            "screen_candidates (missing 'ranked_molecules'). Do not construct, "
            "truncate, or summarize tool output."
        )
    screening.setdefault("similarity_matrix", [])
    screening.setdefault("similarity_labels", [])
    return screening

# ------------------------------------------------------------------
# Task 4: Generate final comprehensive report
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def generate_report(
    molecule_dir: flyte.io.Dir,
    properties_json: str,
    screening_json: str,
) -> str:
    """Generate a comprehensive drug screening report.

    Produces an executive summary, top candidate spotlight cards, property
    distributions, chemical diversity analysis, and final recommendation.

    Args:
        molecule_dir: Directory from load_molecules.
        properties_json: JSON from compute_properties.
        screening_json: Exact verbatim JSON string returned by screen_candidates
            (must include ranked_molecules, similarity_matrix, similarity_labels).
            Do not construct or summarize this payload yourself.

    Returns:
        JSON summary with total_screened, lipinski_passes, all_criteria_met,
        top_candidate, top_score, and top_3 ranked molecules.
    """
    from rdkit import Chem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Generating Final Report...</h2>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    screening = _parse_screening_json(screening_json)
    ranked = screening["ranked_molecules"]
    sim_matrix = screening["similarity_matrix"]
    sim_labels = screening["similarity_labels"]

    total = props["total"]
    lipinski_pass = props["lipinski_pass_count"]
    all_criteria = sum(1 for m in ranked if m["all_criteria_met"])
    top = ranked[0] if ranked else None

    html_parts = []

    # --- Executive Summary ---
    html_parts.append("<h2>Drug Molecule Screening Report</h2>")
    top_name = top["name"] if top else "N/A"
    top_score = f'{top["screening_score"]:.3f}' if top else "N/A"
    html_parts.append(
        f'<div class="card">'
        f'<h3 style="margin-top:0;color:#0e4f6e;">Executive Summary</h3>'
        f'<p style="font-size:1.05em;">'
        f'<strong>{total}</strong> molecules were screened against the target drug profile. '
        f'<strong>{lipinski_pass}</strong> passed Lipinski\'s Rule of Five, and '
        f'<strong>{all_criteria}</strong> met all screening criteria. '
        f'The top candidate is <strong style="color:#0891b2;">{top_name}</strong> '
        f'with a screening score of <strong>{top_score}</strong>.</p>'
        f'</div>'
    )

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Molecules Screened"),
        (str(lipinski_pass), "Lipinski Passes"),
        (str(all_criteria), "All Criteria Met"),
        (top_score, "Top Score"),
        (f'{props["avg_mw"]:.0f} Da', "Avg. Molecular Weight"),
        (f'{props["avg_logp"]:.2f}', "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # --- Top 3 Candidate Spotlights ---
    html_parts.append("<h2>Top Candidate Spotlights</h2>")

    for rank, m in enumerate(ranked[:3], 1):
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(300, 300)) if mol else ""

        medal = ["gold", "silver", "#cd7f32"][rank - 1]
        medal_emoji = ["1st", "2nd", "3rd"][rank - 1]

        lip_badges = ""
        for rule, key in [("MW", "mw_ok"), ("LogP", "logp_ok"),
                          ("HBD", "hbd_ok"), ("HBA", "hba_ok")]:
            ok = m["lipinski"].get(key, False)
            cls = "badge-success" if ok else "badge-danger"
            lip_badges += f'<span class="badge {cls}" style="margin:2px;">{rule}</span> '

        html_parts.append(
            f'<div class="molecule-card" style="display:flex;gap:20px;align-items:flex-start;flex-wrap:wrap;">'
            f'<div style="text-align:center;min-width:180px;">'
            f'<div style="font-size:1.6em;font-weight:800;color:{medal};">{medal_emoji}</div>'
            f'<img src="{img_uri}" style="width:200px;height:200px;object-fit:contain;border-radius:8px;'
            f'border:2px solid #a5f3fc;"/>'
            f'<div style="font-weight:700;font-size:1.1em;color:#0e4f6e;margin-top:8px;">{m["name"]}</div>'
            f'</div>'
            f'<div style="flex:1;min-width:280px;">'
            f'<table style="margin:0;">'
            f'<tr><td><strong>SMILES</strong></td><td style="font-family:monospace;font-size:0.8em;word-break:break-all;">{m["smiles"]}</td></tr>'
            f'<tr><td><strong>Screening Score</strong></td><td style="font-weight:700;color:#0891b2;font-size:1.1em;">{m["screening_score"]:.3f}</td></tr>'
            f'<tr><td><strong>Molecular Weight</strong></td><td>{m["mw"]:.1f} Da</td></tr>'
            f'<tr><td><strong>LogP</strong></td><td>{m["logp"]:.2f}</td></tr>'
            f'<tr><td><strong>H-Bond Donors</strong></td><td>{m["hbd"]}</td></tr>'
            f'<tr><td><strong>H-Bond Acceptors</strong></td><td>{m["hba"]}</td></tr>'
            f'<tr><td><strong>TPSA</strong></td><td>{m["tpsa"]:.1f} A&sup2;</td></tr>'
            f'<tr><td><strong>Rotatable Bonds</strong></td><td>{m["rotatable_bonds"]}</td></tr>'
            f'<tr><td><strong>QED</strong></td><td>{m["qed"]:.4f}</td></tr>'
            f'<tr><td><strong>Lipinski Compliance</strong></td><td>{lip_badges}</td></tr>'
            f'</table>'
            f'</div>'
            f'</div>'
        )

    # --- Property Distribution (box-plot style as bars with min/max/median) ---
    html_parts.append("<h2>Property Distributions</h2>")

    prop_keys = [("mw", "Molecular Weight (Da)"), ("logp", "LogP"),
                 ("tpsa", "TPSA"), ("qed", "QED Score")]
    for key, label in prop_keys:
        vals = sorted([m[key] for m in ranked])
        n = len(vals)
        if n == 0:
            continue
        v_min = vals[0]
        v_max = vals[-1]
        median = vals[n // 2] if n % 2 == 1 else (vals[n // 2 - 1] + vals[n // 2]) / 2
        q1 = vals[n // 4] if n >= 4 else v_min
        q3 = vals[3 * n // 4] if n >= 4 else v_max

        # Simple horizontal box-plot as SVG
        box_w = 500
        box_h = 50
        margin_l = 10
        v_range = v_max - v_min or 1

        def sx(v):
            return margin_l + ((v - v_min) / v_range) * (box_w - 2 * margin_l)

        box_svg = (
            f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {box_w} {box_h}" '
            f'style="width:100%;max-width:{box_w}px;height:auto;">'
            f'<rect width="{box_w}" height="{box_h}" fill="#fff" rx="4"/>'
            # Whisker line
            f'<line x1="{sx(v_min):.1f}" y1="25" x2="{sx(v_max):.1f}" y2="25" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Min whisker
            f'<line x1="{sx(v_min):.1f}" y1="18" x2="{sx(v_min):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Max whisker
            f'<line x1="{sx(v_max):.1f}" y1="18" x2="{sx(v_max):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # IQR box
            f'<rect x="{sx(q1):.1f}" y="14" width="{sx(q3) - sx(q1):.1f}" height="22" '
            f'fill="#a5f3fc" stroke="#0891b2" stroke-width="1.5" rx="3"/>'
            # Median line
            f'<line x1="{sx(median):.1f}" y1="12" x2="{sx(median):.1f}" y2="38" '
            f'stroke="#0e4f6e" stroke-width="2"/>'
            # Labels
            f'<text x="{sx(v_min):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_min:.1f}</text>'
            f'<text x="{sx(median):.1f}" y="10" text-anchor="middle" font-size="9" fill="#0e4f6e" font-weight="600">{median:.1f}</text>'
            f'<text x="{sx(v_max):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_max:.1f}</text>'
            f'</svg>'
        )
        html_parts.append(
            f'<div style="margin:8px 0;"><strong style="color:#155e75;">{label}</strong>'
            f'<div class="chart-container" style="padding:8px;">{box_svg}</div></div>'
        )

    # --- Chemical Diversity ---
    html_parts.append("<h2>Chemical Diversity Analysis</h2>")

    if sim_matrix and len(sim_matrix) > 1:
        # Compute average pairwise similarity (off-diagonal)
        n_mols = len(sim_matrix)
        off_diag = []
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                off_diag.append(sim_matrix[i][j])

        avg_sim = sum(off_diag) / len(off_diag) if off_diag else 0
        max_sim = max(off_diag) if off_diag else 0
        min_sim = min(off_diag) if off_diag else 0

        # Find most similar pair
        best_i, best_j = 0, 1
        best_val = 0
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                if sim_matrix[i][j] > best_val:
                    best_val = sim_matrix[i][j]
                    best_i, best_j = i, j

        html_parts.append('<div class="stat-grid">')
        html_parts.append(
            f'<div class="stat"><div class="value">{avg_sim:.3f}</div>'
            f'<div class="label">Avg. Pairwise Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{min_sim:.3f}</div>'
            f'<div class="label">Min Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{max_sim:.3f}</div>'
            f'<div class="label">Max Similarity</div></div>'
        )
        html_parts.append("</div>")

        diversity_text = "highly diverse" if avg_sim < 0.3 else "moderately diverse" if avg_sim < 0.5 else "relatively similar"
        html_parts.append(
            f'<div class="note">'
            f'The library is <strong>{diversity_text}</strong> (avg. Tanimoto = {avg_sim:.3f}). '
            f'The most similar pair is <strong>{sim_labels[best_i]}</strong> and '
            f'<strong>{sim_labels[best_j]}</strong> (similarity = {best_val:.3f}).</div>'
        )

    # --- Recommendation ---
    html_parts.append("<h2>Recommendation</h2>")
    if top:
        html_parts.append(
            f'<div class="card">'
            f'<h3 style="margin-top:0;color:#0891b2;">Top Candidate: {top["name"]}</h3>'
            f'<p>Based on the virtual screening analysis, <strong>{top["name"]}</strong> '
            f'achieved the highest composite screening score of <strong>{top["screening_score"]:.3f}</strong>. '
        )

        reasons = []
        if top["lipinski_pass"]:
            reasons.append("full Lipinski Rule of Five compliance")
        if top["qed"] > 0.5:
            reasons.append(f"high drug-likeness (QED = {top['qed']:.3f})")
        if top.get("all_criteria_met"):
            reasons.append("all target profile criteria met")
        if top["mw"] <= 500:
            reasons.append(f"favorable molecular weight ({top['mw']:.1f} Da)")

        if reasons:
            html_parts.append(
                f'This candidate stands out due to: {", ".join(reasons)}.</p>'
            )
        else:
            html_parts.append("</p>")

        # Runner-up mentions
        if len(ranked) >= 2:
            html_parts.append(
                f'<p style="font-size:0.9em;color:#6c757d;">Runner-up candidates: '
            )
            runners = []
            for m in ranked[1:4]:
                runners.append(f'{m["name"]} (score: {m["screening_score"]:.3f})')
            html_parts.append(", ".join(runners) + ".</p>")

        html_parts.append("</div>")

    # Final note
    html_parts.append(
        '<div class="note">'
        "This is a virtual screening analysis. All candidates should undergo "
        "further computational validation (molecular dynamics, docking) and "
        "experimental testing before advancing to clinical trials.</div>"
    )

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # JSON summary
    summary = {
        "total_screened": total,
        "lipinski_passes": lipinski_pass,
        "all_criteria_met": all_criteria,
        "top_candidate": top["name"] if top else None,
        "top_score": top["screening_score"] if top else None,
        "top_3": [
            {"name": m["name"], "score": m["screening_score"]}
            for m in ranked[:3]
        ],
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Agent
# ------------------------------------------------------------------

# {{docs-fragment agent}}
SCREENING_AGENT_INSTRUCTIONS = """\
You are a medicinal chemistry screening strategist. You orchestrate a virtual \
screening pipeline using durable Flyte tools. You NEVER invent molecular \
properties — only RDKit tools compute them.

Workflow:
1. If target_profile is not provided in the user message, derive a JSON \
target_profile from the therapeutic brief. Valid keys: mw, logp, hbd, hba, tpsa \
(each [min, max]). Ground choices in oral bioavailability / kinase / CNS rules \
as appropriate to the brief.
2. First pass (always): load_molecules → compute_properties → \
screen_candidates → generate_report. Pass tool outputs between steps exactly \
(molecule_dir from load_molecules into compute_properties and generate_report; \
properties_json from compute_properties into screen_candidates and \
generate_report; screening_json must be the complete, unmodified string \
returned by screen_candidates — never rebuild or summarize JSON yourself).
3. Read the JSON summary returned by generate_report. Reflect:
   - If all_criteria_met == 0: relax exactly ONE profile bound by ~10–20% \
and re-run screen_candidates then generate_report only, reusing the same \
molecule_dir and properties_json from the first pass.
   - If all molecules pass but diversity is a stated goal: note high similarity \
in your summary; do not re-run unless brief asks for stricter filters.
   - Maximum ONE rescreen iteration.
4. Finish with plain text: top candidate, rationale tied to computed metrics \
from the tool JSON, funnel interpretation, and suggested next steps (docking, \
ADMET lab tests).

If the user supplies an explicit target_profile JSON, use it as-is.

Do NOT ask the user for SMILES or molecule lists when molecules_json is empty — \
the default library is loaded automatically.
"""

screening_agent = Agent(
    name="drug-screening-agent",
    instructions=SCREENING_AGENT_INSTRUCTIONS,
    model=MODEL,
    tools=[
        load_molecules,
        compute_properties,
        screen_candidates,
        generate_report,
    ],
    max_turns=12,
)
# {{/docs-fragment agent}}

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def pipeline(
    brief: str = "Screen the default drug library for orally bioavailable small molecules.",
    molecules_json: str = "",
    target_profile: str = "",
) -> str:
    """Agentic virtual drug molecule screening pipeline.

    A medicinal-chemistry agent interprets the screening brief, derives or
    applies a target profile, orchestrates the RDKit screening stages, and
    optionally re-screens when funnel results are too narrow.

    Args:
        brief: Natural-language therapeutic goal (e.g. oral kinase inhibitors,
            CNS-penetrant small molecules).
        molecules_json: JSON mapping molecule names to SMILES strings.
            Defaults to a curated library of ~15 well-known drugs.
        target_profile: Optional JSON with desired property ranges that
            overrides agent-derived criteria
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5]}).

    Returns:
        Agent summary with screening rationale and key results.
    """
    prompt_parts = [
        f"Screening brief: {brief}",
        'Use molecules_json="" for the built-in default library unless provided below.',
        "Compose the four stage tools in order: load_molecules → compute_properties "
        "→ screen_candidates → generate_report. Pass each tool's full return value "
        "verbatim to the next step (especially screening_json). Re-run "
        "screen_candidates and generate_report at most once if the funnel is too narrow.",
    ]
    if molecules_json.strip():
        prompt_parts.append(f"molecules_json: {molecules_json}")
    if target_profile.strip():
        prompt_parts.append(f"Use this target_profile exactly: {target_profile}")

    result = await screening_agent.run.aio("\n".join(prompt_parts))
    return result.summary or result.error or ""

# {{/docs-fragment pipeline}}

# ------------------------------------------------------------------
# Rescreen demo — tight profile + explicit rescreen instructions
# ------------------------------------------------------------------

# Initial profile is deliberately strict (narrow MW + low LogP cap) so
# all_criteria_met is typically 0 on the default library; the brief then
# forces a single rescreen with a widened LogP window.
RESCREEN_DEMO_TARGET_PROFILE = (
    '{"mw": [150, 200], "logp": [-0.5, 1.0], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_TARGET_PROFILE_RESCREEN = (
    '{"mw": [150, 200], "logp": [-0.5, 3.5], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_BRIEF = f"""\
Two-round agentic screening demo on the default library.

**Round 1 (strict profile):** load_molecules → compute_properties → \
screen_candidates → generate_report using the initial target_profile exactly.

**Round 2 (required — do not skip):** call screen_candidates then generate_report \
again, reusing the same molecule_dir and properties_json from round 1, with this \
relaxed target_profile (wider LogP window only): \
{RESCREEN_DEMO_TARGET_PROFILE_RESCREEN}

Pass every tool return value verbatim to the next step. After both rounds, \
summarize how the funnel and top candidates changed between round 1 and round 2."""

# {{docs-fragment rescreen_demo}}
@env.task(report=True)
async def rescreen_demo() -> str:
    """Example run with a two-round execution graph (rescreen).

    Round 1 uses a strict CNS-like profile; round 2 always re-runs
    screen_candidates and generate_report with a widened LogP window,
    reusing cached molecule_dir and properties_json.
    """
    return await pipeline(
        brief=RESCREEN_DEMO_BRIEF,
        target_profile=RESCREEN_DEMO_TARGET_PROFILE,
    )

# {{/docs-fragment rescreen_demo}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/drug_molecule_screening/drug_molecule_screening.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.4",
#    "litellm",
#    "rdkit",
#    "numpy",
#    "scikit-learn",
#    "pillow",
# ]
# ///
```

## Define the screening agent

The agent receives a natural-language brief and composes four stage tools in order. Each tool is a durable Flyte task with its own `report=True` surface in the Flyte UI.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.4",
#    "litellm",
#    "rdkit",
#    "numpy",
#    "scikit-learn",
#    "pillow",
# ]
# main = "pipeline"
# params = ""
# ///
"""Virtual drug molecule screening — compute properties, apply Lipinski filters, rank candidates."""

import base64
import io
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report
from flyte.ai.agents import Agent, tool

MODEL = os.getenv("DRUG_SCREENING_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="drug-molecule-screening", pre=True).with_apt_packages(
    "libxrender1", "libxext6", "libexpat1",
)

env = flyte.TaskEnvironment(
    name="drug-molecule-screening",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default molecule library — real SMILES for well-known drugs
# ------------------------------------------------------------------

DEFAULT_MOLECULES = {
    "Aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
    "Ibuprofen": "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
    "Caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
    "Penicillin G": "CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C",
    "Metformin": "CN(C)C(=N)NC(=N)N",
    "Paracetamol": "CC(=O)NC1=CC=C(C=C1)O",
    "Diazepam": "ClC1=CC2=C(C=C1)N(C(=O)CN=C2C3=CC=CC=C3)C",
    "Omeprazole": "CC1=CN=C(C(=C1OC)C)CS(=O)C2=NC3=CC=CC=C3N2",
    "Atorvastatin": "CC(C)C1=C(C(=C(N1CCC(CC(CC(=O)O)O)O)C2=CC=C(C=C2)F)C3=CC=CC=C3)C(=O)NC4=CC=CC=C4",
    "Methotrexate": "CN(CC1=CN=C2N=C(N=C(N)C2=N1)N)C3=CC=C(C=C3)C(=O)NC(CCC(=O)O)C(=O)O",
    "Doxorubicin": "CC1C(C(CC(O1)OC2CC(CC3=C2C(=C4C(=C3O)C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)CO)O)N)O",
    "Tamoxifen": "CCC(=C(C1=CC=CC=C1)C2=CC=C(C=C2)OCCN(C)C)C3=CC=CC=C3",
    "Lopinavir": "CC1=C(C(=CC=C1)C)OCC(=O)NC(CC2=CC=CC=C2)C(CC(CC3=CC=CC=C3)NC(=O)C(C(C)C)N4CCCNC4=O)O",
    "Remdesivir": "CCC(CC)COC(=O)C(C)NP(=O)(OCC1C(C(C(O1)C2=CC=C3N2N=CN=C3N)O)O)OC4=CC=CC=C4",
    "Erlotinib": "COCCOC1=CC2=C(C=C1OCCOC)C(=NC=N2)NC3=CC=CC(=C3)C#C",
}

# ------------------------------------------------------------------
# Report styling — pharma blue/cyan theme
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #0e4f6e; border-bottom: 2px solid #0891b2; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #155e75; margin-top: 20px; }
  .report .card { background: #ecfeff; border: 1px solid #a5f3fc; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #cffafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0e4f6e; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0e4f6e; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #cffafe; }
  .report tr:nth-child(even) { background: #ecfeff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #cffafe; color: #155e75; }
  .report .chart-container { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #ecfeff; border-left: 4px solid #0891b2; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .molecule-card { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .molecule-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 12px; margin: 16px 0; }
  .report .funnel { text-align: center; margin: 24px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _mol_to_data_uri(mol, size: tuple[int, int] = (300, 300)) -> str:
    """Convert an RDKit molecule to a PNG base64 data URI."""
    from rdkit.Chem import Draw

    img = Draw.MolToImage(mol, size=size)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 340,
    y_max_cap: float | None = None,
    horizontal: bool = False,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels.
        series: Dict mapping series name to list of values.
        title: Chart title.
        colors: Colors for each series.
        width/height: SVG dimensions.
        y_max_cap: Cap the y-axis at this value.
        horizontal: If True, draw horizontal bars.
        value_fmt: Format string for value labels.

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#a5f3fc", "#155e75"]
    colors = colors or default_colors

    if horizontal:
        return _make_horizontal_bar_chart(labels, series, title, colors, width, height, value_fmt)

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_fmt}}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_fmt}}</text>'
            )
        # Truncate long labels
        disp_label = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 16}" '
            f'text-anchor="middle" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {gx + n_series * bar_width / 2:.1f}, {mt + ch + 16})">'
            f'{disp_label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 100
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 100}" y="{mt + ch + 40}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 51}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_horizontal_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 400,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG horizontal bar chart (sorted)."""
    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0"]
    colors = colors or default_colors

    n = len(labels)
    row_height = max(22, min(35, (height - 80) // max(n, 1)))
    actual_height = max(height, 80 + n * row_height)
    ml, mr, mt, mb = 120, 60, 40, 20
    cw = width - ml - mr
    ch = actual_height - mt - mb

    # Use first series
    first_key = list(series.keys())[0]
    vals = series[first_key]
    x_max = max(vals) * 1.15 if vals else 1

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {actual_height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{actual_height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    bar_h = row_height * 0.65
    for i, (label, val) in enumerate(zip(labels, vals)):
        y = mt + i * row_height
        bw = (val / x_max) * cw if x_max else 0
        color = colors[i % len(colors)]
        # Label
        disp = label if len(label) <= 14 else label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + bar_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#1a1a2e">{disp}</text>'
        )
        # Bar
        svg.append(
            f'<rect x="{ml}" y="{y:.1f}" width="{bw:.1f}" height="{bar_h:.1f}" '
            f'fill="{color}" rx="3"/>'
        )
        # Value
        svg.append(
            f'<text x="{ml + bw + 6:.1f}" y="{y + bar_h / 2 + 4:.1f}" '
            f'font-size="11" fill="#0e4f6e" font-weight="600">{val:{value_fmt}}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    color_scale: str = "cyan",
    width: int = 700,
    height: int = 500,
    value_fmt: str = ".2f",
) -> str:
    """Generate an SVG heatmap.

    Args:
        matrix: 2D list of values (rows x cols).
        row_labels: Labels for rows.
        col_labels: Labels for columns.
        title: Chart title.
        color_scale: Color scheme ("cyan", "red", "green").
        width/height: SVG dimensions.
        value_fmt: Format string for cell values.

    Returns:
        SVG string.
    """
    if not matrix or not matrix[0]:
        return ""

    n_rows = len(matrix)
    n_cols = len(matrix[0])

    ml, mr, mt, mb = 110, 20, 70, 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    # Flatten to find range
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    def color_for(v):
        t = (v - v_min) / v_range
        if color_scale == "cyan":
            # White to deep teal
            r = int(255 - t * (255 - 14))
            g = int(255 - t * (255 - 79))
            b = int(255 - t * (255 - 110))
        elif color_scale == "red":
            r = int(255 - t * 50)
            g = int(255 - t * 200)
            b = int(255 - t * 200)
        else:  # green
            r = int(255 - t * 200)
            g = int(255 - t * 50)
            b = int(255 - t * 200)
        return f"rgb({r},{g},{b})"

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Column labels (rotated)
    for ci, label in enumerate(col_labels):
        x = ml + ci * cell_w + cell_w / 2
        disp = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{x:.1f}" y="{mt - 8}" text-anchor="end" font-size="10" '
            f'fill="#1a1a2e" transform="rotate(-45, {x:.1f}, {mt - 8})">{disp}</text>'
        )

    # Row labels + cells
    for ri, (row_label, row_vals) in enumerate(zip(row_labels, matrix)):
        y = mt + ri * cell_h
        disp = row_label if len(row_label) <= 14 else row_label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + cell_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#1a1a2e">{disp}</text>'
        )
        for ci, val in enumerate(row_vals):
            x = ml + ci * cell_w
            fill = color_for(val)
            svg.append(
                f'<rect x="{x:.1f}" y="{y:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{fill}" stroke="#fff" stroke-width="1"/>'
            )
            # Text color: dark on light, light on dark
            t = (val - v_min) / v_range
            txt_color = "#fff" if t > 0.55 else "#1a1a2e"
            # Only show text if cells are large enough
            if cell_w > 30 and cell_h > 18:
                svg.append(
                    f'<text x="{x + cell_w / 2:.1f}" y="{y + cell_h / 2 + 4:.1f}" '
                    f'text-anchor="middle" font-size="9" fill="{txt_color}">'
                    f'{val:{value_fmt}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_scatter_plot(
    points: list[dict],
    x_label: str = "MW",
    y_label: str = "LogP",
    title: str = "",
    reference_lines: list[dict] | None = None,
    width: int = 700,
    height: int = 400,
) -> str:
    """Generate an SVG scatter plot.

    Args:
        points: List of dicts with "x", "y", "label" keys.
        x_label/y_label: Axis labels.
        title: Chart title.
        reference_lines: List of dicts with "axis" ("x"/"y"), "value", "label".
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not points:
        return ""

    ml, mr, mt, mb = 60, 30, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [p["x"] for p in points]
    y_vals = [p["y"] for p in points]
    x_min, x_max = min(x_vals) * 0.9, max(x_vals) * 1.1
    y_min, y_max = min(y_vals) - 1, max(y_vals) + 1

    # Extend ranges to include reference lines
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                x_max = max(x_max, rl["value"] * 1.1)
            else:
                y_max = max(y_max, rl["value"] * 1.1)

    x_range = x_max - x_min or 1
    y_range = y_max - y_min or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid
    for i in range(6):
        y_tick = y_min + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.1f}</text>'
        )

    for i in range(6):
        x_tick = x_min + x_range * i / 5
        px = sx(x_tick)
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
            f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Reference lines (Lipinski boundaries)
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                px = sx(rl["value"])
                svg.append(
                    f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{px + 4:.1f}" y="{mt + 14}" font-size="10" '
                    f'fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )
            else:
                py = sy(rl["value"])
                svg.append(
                    f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{ml + cw - 4:.1f}" y="{py - 6:.1f}" text-anchor="end" '
                    f'font-size="10" fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )

    # Drug-like zone shading (MW<=500 and LogP<=5 quadrant)
    if reference_lines:
        mw_line = next((rl for rl in reference_lines if rl["axis"] == "x"), None)
        logp_line = next((rl for rl in reference_lines if rl["axis"] == "y"), None)
        if mw_line and logp_line:
            zx1 = sx(x_min)
            zx2 = sx(min(mw_line["value"], x_max))
            zy1 = sy(min(logp_line["value"], y_max))
            zy2 = sy(y_min)
            svg.append(
                f'<rect x="{zx1:.1f}" y="{zy1:.1f}" '
                f'width="{zx2 - zx1:.1f}" height="{zy2 - zy1:.1f}" '
                f'fill="#0891b2" opacity="0.06" rx="4"/>'
            )
            svg.append(
                f'<text x="{zx1 + 8:.1f}" y="{zy2 - 8:.1f}" font-size="11" '
                f'fill="#0891b2" font-weight="600" opacity="0.6">Drug-like Zone</text>'
            )

    # Points
    point_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#155e75", "#0284c7",
                    "#059669", "#0d9488", "#0369a1", "#047857", "#115e59",
                    "#0c4a6e", "#064e3b", "#1e3a5f", "#134e4a", "#075985"]
    for i, pt in enumerate(points):
        px, py = sx(pt["x"]), sy(pt["y"])
        color = point_colors[i % len(point_colors)]
        svg.append(
            f'<circle cx="{px:.1f}" cy="{py:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5" opacity="0.85"/>'
        )
        # Label offset to avoid overlap
        offset_x = 8
        offset_y = -8 if i % 2 == 0 else 14
        label = pt["label"] if len(pt["label"]) <= 12 else pt["label"][:10] + ".."
        svg.append(
            f'<text x="{px + offset_x:.1f}" y="{py + offset_y:.1f}" '
            f'font-size="9" fill="#1a1a2e">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Axis labels
    if x_label:
        svg.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        svg.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_funnel(
    stages: list[dict],
    title: str = "",
    width: int = 600,
    height: int = 400,
) -> str:
    """Generate an SVG funnel visualization.

    Args:
        stages: List of dicts with "label", "count", "total" keys.
        title: Chart title.
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not stages:
        return ""

    n = len(stages)
    mt = 50
    mb = 20
    available_h = height - mt - mb
    stage_h = available_h / n
    cx = width / 2

    # Color gradient from light cyan to deep teal
    colors = []
    for i in range(n):
        t = i / max(n - 1, 1)
        r = int(207 - t * (207 - 14))
        g = int(250 - t * (250 - 79))
        b = int(254 - t * (254 - 110))
        colors.append(f"rgb({r},{g},{b})")

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{cx}" y="28" text-anchor="middle" '
            f'font-size="16" font-weight="700" fill="#0e4f6e">{title}</text>'
        )

    max_count = stages[0]["count"] if stages else 1
    max_width = width * 0.75

    for i, stage in enumerate(stages):
        y_top = mt + i * stage_h
        y_bot = y_top + stage_h

        # Width proportional to count
        w_top = max_width * (stage["count"] / max_count) if i == 0 else prev_w_bot
        if i < n - 1:
            w_bot = max_width * (stages[i + 1]["count"] / max_count)
        else:
            w_bot = max_width * (stage["count"] / max_count) * 0.7

        prev_w_bot = w_bot

        # Trapezoid
        x1_top = cx - w_top / 2
        x2_top = cx + w_top / 2
        x1_bot = cx - w_bot / 2
        x2_bot = cx + w_bot / 2

        svg.append(
            f'<polygon points="{x1_top:.1f},{y_top:.1f} {x2_top:.1f},{y_top:.1f} '
            f'{x2_bot:.1f},{y_bot:.1f} {x1_bot:.1f},{y_bot:.1f}" '
            f'fill="{colors[i]}" stroke="#fff" stroke-width="2"/>'
        )

        # Text: dark on light, white on dark
        t = i / max(n - 1, 1)
        txt_color = "#0e4f6e" if t < 0.5 else "#fff"
        y_mid = (y_top + y_bot) / 2

        svg.append(
            f'<text x="{cx}" y="{y_mid - 4:.1f}" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="{txt_color}">{stage["label"]}</text>'
        )
        svg.append(
            f'<text x="{cx}" y="{y_mid + 14:.1f}" text-anchor="middle" '
            f'font-size="12" fill="{txt_color}" opacity="0.85">'
            f'{stage["count"]} / {stage["total"]}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate molecules
# ------------------------------------------------------------------

@tool
@env.task(cache="auto")
async def load_molecules(
    molecules_json: str = "",
) -> flyte.io.Dir:
    """Parse SMILES strings, validate with RDKit, generate 2D depictions.

    Args:
        molecules_json: JSON string mapping molecule names to SMILES.
            Defaults to a curated library of ~15 well-known drugs.

    Returns:
        flyte.io.Dir containing molecule data (JSON + PNG depictions).
        Pass this directory to compute_properties and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Draw

    if molecules_json.strip():
        molecules = json.loads(molecules_json)
    else:
        molecules = DEFAULT_MOLECULES

    out_dir = tempfile.mkdtemp(prefix="mol_library_")
    results = []
    valid_count = 0
    invalid_count = 0

    log.info(f"Parsing {len(molecules)} molecules...")

    for name, smiles in molecules.items():
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            log.warning(f"  [INVALID] {name}: {smiles}")
            invalid_count += 1
            continue

        valid_count += 1

        # Generate 2D depiction as PNG
        img = Draw.MolToImage(mol, size=(300, 300))
        img_path = os.path.join(out_dir, f"{name.replace(' ', '_')}.png")
        img.save(img_path)

        results.append({
            "name": name,
            "smiles": smiles,
            "valid": True,
            "image_file": os.path.basename(img_path),
        })

    # Save molecule manifest
    manifest = {
        "total": len(molecules),
        "valid": valid_count,
        "invalid": invalid_count,
        "molecules": results,
    }
    manifest_path = os.path.join(out_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    log.info(f"Loaded {valid_count} valid molecules ({invalid_count} invalid)")

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Compute physicochemical properties
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def compute_properties(
    molecule_dir: flyte.io.Dir,
) -> str:
    """Compute drug-likeness properties for all molecules.

    Computes MW, LogP, HBD, HBA, TPSA, rotatable bonds, formal charge,
    ring count, QED, and Lipinski Rule of Five compliance.

    Args:
        molecule_dir: Directory from load_molecules.

    Returns:
        JSON string with all computed properties. Pass to screen_candidates
        and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Descriptors, Lipinski
    from rdkit.Chem.QED import qed

    # --- Loading report ---
    await flyte.report.replace.aio(
        _wrap_report("<h2>Computing Molecular Properties...</h2>"
                      "<p>Analyzing physicochemical descriptors for all molecules.</p>"),
        do_flush=True,
    )

    mol_dir = await molecule_dir.download()
    with open(os.path.join(mol_dir, "manifest.json")) as f:
        manifest = json.load(f)

    molecules_data = []
    lipinski_pass = 0

    for mol_info in manifest["molecules"]:
        mol = Chem.MolFromSmiles(mol_info["smiles"])
        if mol is None:
            continue

        mw = Descriptors.MolWt(mol)
        logp = Descriptors.MolLogP(mol)
        hbd = Lipinski.NumHDonors(mol)
        hba = Lipinski.NumHAcceptors(mol)
        tpsa = Descriptors.TPSA(mol)
        rotatable = Lipinski.NumRotatableBonds(mol)
        formal_charge = Chem.GetFormalCharge(mol)
        num_rings = Lipinski.RingCount(mol)
        qed_score = qed(mol)

        # Lipinski Rule of Five
        lipinski = {
            "mw_ok": mw <= 500,
            "logp_ok": logp <= 5,
            "hbd_ok": hbd <= 5,
            "hba_ok": hba <= 10,
        }
        lipinski_all = all(lipinski.values())
        if lipinski_all:
            lipinski_pass += 1

        # Read image for data URI
        img_path = os.path.join(mol_dir, mol_info["image_file"])
        data_uri = ""
        if os.path.exists(img_path):
            with open(img_path, "rb") as img_f:
                b64 = base64.b64encode(img_f.read()).decode()
                data_uri = f"data:image/png;base64,{b64}"

        molecules_data.append({
            "name": mol_info["name"],
            "smiles": mol_info["smiles"],
            "mw": round(mw, 2),
            "logp": round(logp, 2),
            "hbd": hbd,
            "hba": hba,
            "tpsa": round(tpsa, 2),
            "rotatable_bonds": rotatable,
            "formal_charge": formal_charge,
            "num_rings": num_rings,
            "qed": round(qed_score, 4),
            "lipinski": lipinski,
            "lipinski_pass": lipinski_all,
            "image_data_uri": data_uri,
        })

    total = len(molecules_data)
    avg_mw = sum(m["mw"] for m in molecules_data) / total if total else 0
    avg_logp = sum(m["logp"] for m in molecules_data) / total if total else 0
    lipinski_rate = lipinski_pass / total * 100 if total else 0

    # ---- Build report ----
    html_parts = []

    # Header
    html_parts.append("<h2>Molecular Properties Analysis</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Total Molecules"),
        (f"{lipinski_rate:.0f}%", "Lipinski Pass Rate"),
        (f"{avg_mw:.1f}", "Avg. MW (Da)"),
        (f"{avg_logp:.2f}", "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Molecule gallery
    html_parts.append("<h3>Molecule Library</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in molecules_data:
        if m["image_data_uri"]:
            badge_class = "badge-success" if m["lipinski_pass"] else "badge-danger"
            badge_text = "Lipinski Pass" if m["lipinski_pass"] else "Lipinski Fail"
            html_parts.append(
                f'<div class="molecule-card" style="text-align:center;">'
                f'<img src="{m["image_data_uri"]}" style="width:160px;height:160px;object-fit:contain;"/>'
                f'<div style="font-weight:600;margin-top:6px;color:#0e4f6e;">{m["name"]}</div>'
                f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f}</div>'
                f'<div><span class="badge {badge_class}">{badge_text}</span></div>'
                f'</div>'
            )
    html_parts.append("</div>")

    # MW bar chart (horizontal, sorted)
    sorted_by_mw = sorted(molecules_data, key=lambda m: m["mw"], reverse=True)
    mw_labels = [m["name"] for m in sorted_by_mw]
    mw_vals = [m["mw"] for m in sorted_by_mw]
    mw_chart = _make_bar_chart(
        mw_labels, {"MW (Da)": mw_vals},
        title="Molecular Weight Distribution",
        horizontal=True,
        width=700, height=max(300, len(mw_labels) * 30 + 80),
        value_fmt=".1f",
    )
    html_parts.append("<h3>Molecular Weight</h3>")
    html_parts.append(f'<div class="chart-container">{mw_chart}</div>')

    # LogP vs MW scatter plot
    scatter_points = [
        {"x": m["mw"], "y": m["logp"], "label": m["name"]}
        for m in molecules_data
    ]
    scatter_chart = _make_scatter_plot(
        scatter_points,
        x_label="Molecular Weight (Da)",
        y_label="LogP",
        title="LogP vs. Molecular Weight (Lipinski Boundaries)",
        reference_lines=[
            {"axis": "x", "value": 500, "label": "MW = 500"},
            {"axis": "y", "value": 5, "label": "LogP = 5"},
        ],
        width=700,
        height=420,
    )
    html_parts.append("<h3>Lipinski Space</h3>")
    html_parts.append(f'<div class="chart-container">{scatter_chart}</div>')

    # Property heatmap (molecules x properties)
    prop_names = ["MW", "LogP", "HBD", "HBA", "TPSA", "Rot. Bonds"]
    # Normalize each property to 0-1 for heatmap
    raw_matrix = []
    for m in molecules_data:
        raw_matrix.append([m["mw"], m["logp"], m["hbd"], m["hba"], m["tpsa"], m["rotatable_bonds"]])

    # Normalize per column
    n_props = len(prop_names)
    col_min = [min(row[c] for row in raw_matrix) for c in range(n_props)]
    col_max = [max(row[c] for row in raw_matrix) for c in range(n_props)]
    norm_matrix = []
    for row in raw_matrix:
        norm_row = []
        for c in range(n_props):
            rng = col_max[c] - col_min[c]
            norm_row.append((row[c] - col_min[c]) / rng if rng else 0.5)
        norm_matrix.append(norm_row)

    heatmap_labels = [m["name"] for m in molecules_data]
    heatmap = _make_heatmap(
        norm_matrix, heatmap_labels, prop_names,
        title="Normalized Property Heatmap",
        color_scale="cyan",
        width=700,
        height=max(400, len(heatmap_labels) * 28 + 100),
    )
    html_parts.append("<h3>Property Heatmap</h3>")
    html_parts.append(f'<div class="chart-container">{heatmap}</div>')

    # Lipinski compliance table
    html_parts.append("<h3>Lipinski Rule of Five Compliance</h3>")
    html_parts.append("<table><tr><th>Molecule</th><th>MW &le; 500</th>"
                      "<th>LogP &le; 5</th><th>HBD &le; 5</th>"
                      "<th>HBA &le; 10</th><th>Overall</th></tr>")
    for m in molecules_data:
        lip = m["lipinski"]

        def _badge(ok):
            if ok:
                return '<span class="badge badge-success">Pass</span>'
            return '<span class="badge badge-danger">Fail</span>'

        overall_badge = _badge(m["lipinski_pass"])
        html_parts.append(
            f'<tr><td><strong>{m["name"]}</strong></td>'
            f'<td>{_badge(lip["mw_ok"])}</td>'
            f'<td>{_badge(lip["logp_ok"])}</td>'
            f'<td>{_badge(lip["hbd_ok"])}</td>'
            f'<td>{_badge(lip["hba_ok"])}</td>'
            f'<td>{overall_badge}</td></tr>'
        )
    html_parts.append("</table>")

    # QED bar chart
    sorted_by_qed = sorted(molecules_data, key=lambda m: m["qed"], reverse=True)
    qed_labels = [m["name"] for m in sorted_by_qed]
    qed_vals = [m["qed"] for m in sorted_by_qed]
    qed_chart = _make_bar_chart(
        qed_labels, {"QED Score": qed_vals},
        title="Drug-likeness (QED Score)",
        horizontal=True,
        width=700, height=max(300, len(qed_labels) * 30 + 80),
        value_fmt=".3f",
        colors=["#06d6a0"],
    )
    html_parts.append("<h3>Drug-likeness (QED)</h3>")
    html_parts.append(f'<div class="chart-container">{qed_chart}</div>')

    # Flush full report
    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # Return properties as JSON (strip image data URIs to reduce size)
    output = {
        "total": total,
        "lipinski_pass_count": lipinski_pass,
        "lipinski_pass_rate": round(lipinski_rate, 2),
        "avg_mw": round(avg_mw, 2),
        "avg_logp": round(avg_logp, 2),
        "molecules": [
            {k: v for k, v in m.items() if k != "image_data_uri"}
            for m in molecules_data
        ],
    }
    return json.dumps(output)

# ------------------------------------------------------------------
# Task 3: Screen candidates against target profile
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def screen_candidates(
    properties_json: str,
    target_profile: str = "",
) -> str:
    """Screen molecules against a target drug profile and rank candidates.

    Scores each molecule on how well it matches the target profile, computes
    pairwise Tanimoto similarity, and produces a ranked list.

    Args:
        properties_json: JSON from compute_properties.
        target_profile: JSON string with desired property ranges
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5.0]}).

    Returns:
        JSON string with ranked_molecules, similarity_matrix, similarity_labels,
        funnel, and target_profile. Pass the full return value verbatim to
        generate_report along with molecule_dir and properties_json.
    """
    from rdkit import Chem, DataStructs
    from rdkit.Chem import AllChem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Screening Candidates...</h2>"
                      "<p>Evaluating molecules against the target drug profile.</p>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    molecules = props["molecules"]

    # Default target profile
    if target_profile.strip():
        profile = json.loads(target_profile)
    else:
        profile = {
            "mw": [150, 500],
            "logp": [-0.5, 5.0],
            "hbd": [0, 5],
            "hba": [0, 10],
            "tpsa": [20, 140],
        }

    # --- Screening ---
    funnel_total = len(molecules)
    pass_mw = 0
    pass_logp = 0
    pass_lipinski = 0
    final_candidates = 0

    scored = []
    for m in molecules:
        score = 0
        max_score = 0
        criteria = {}

        # Check each profile criterion
        checks = [
            ("mw", m["mw"]),
            ("logp", m["logp"]),
            ("hbd", m["hbd"]),
            ("hba", m["hba"]),
            ("tpsa", m["tpsa"]),
        ]

        for key, val in checks:
            if key in profile:
                lo, hi = profile[key]
                max_score += 1
                in_range = lo <= val <= hi
                criteria[key] = in_range
                if in_range:
                    score += 1
                    # Bonus: closer to midpoint = higher score
                    mid = (lo + hi) / 2
                    rng = (hi - lo) / 2
                    dist = abs(val - mid) / rng if rng else 0
                    score += max(0, 0.5 * (1 - dist))

        # QED bonus
        score += m["qed"] * 2
        max_score += 2

        # Lipinski bonus
        if m["lipinski_pass"]:
            score += 1
        max_score += 1

        normalized_score = score / max_score if max_score else 0

        # Funnel tracking — cascading filter (each stage requires passing the previous)
        mw_ok = criteria.get("mw", True)
        logp_ok = criteria.get("logp", True)
        if mw_ok:
            pass_mw += 1
            if logp_ok:
                pass_logp += 1
                if m["lipinski_pass"]:
                    pass_lipinski += 1
                    if all(criteria.values()):
                        final_candidates += 1

        scored.append({
            **m,
            "screening_score": round(normalized_score, 4),
            "criteria_met": criteria,
            "all_criteria_met": all(criteria.values()),
        })

    # Sort by score descending
    scored.sort(key=lambda m: m["screening_score"], reverse=True)

    # --- Tanimoto similarity matrix ---
    fps = []
    valid_names = []
    for m in scored:
        mol = Chem.MolFromSmiles(m["smiles"])
        if mol:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
            fps.append(fp)
            valid_names.append(m["name"])

    similarity_matrix = []
    for i in range(len(fps)):
        row = []
        for j in range(len(fps)):
            sim = DataStructs.TanimotoSimilarity(fps[i], fps[j])
            row.append(round(sim, 3))
        similarity_matrix.append(row)

    # ---- Build report ----
    html_parts = []
    html_parts.append("<h2>Candidate Screening Results</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(funnel_total), "Total Screened"),
        (str(pass_lipinski), "Lipinski Passes"),
        (str(final_candidates), "All Criteria Met"),
        (f"{scored[0]['screening_score']:.3f}" if scored else "N/A", "Top Score"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Screening funnel
    funnel_stages = [
        {"label": "Total Molecules", "count": funnel_total, "total": funnel_total},
        {"label": "Pass MW Filter", "count": pass_mw, "total": funnel_total},
        {"label": "Pass LogP Filter", "count": pass_logp, "total": funnel_total},
        {"label": "Lipinski Compliant", "count": pass_lipinski, "total": funnel_total},
        {"label": "All Criteria Met", "count": final_candidates, "total": funnel_total},
    ]
    funnel_svg = _make_funnel(
        funnel_stages,
        title="Screening Funnel",
        width=600,
        height=380,
    )
    html_parts.append("<h3>Screening Funnel</h3>")
    html_parts.append(f'<div class="chart-container" style="text-align:center;">{funnel_svg}</div>')

    # Ranked candidates table
    html_parts.append("<h3>Ranked Candidates</h3>")
    html_parts.append(
        "<table><tr><th>Rank</th><th>Molecule</th><th>Score</th>"
        "<th>MW</th><th>LogP</th><th>QED</th><th>Lipinski</th><th>All Criteria</th></tr>"
    )
    for rank, m in enumerate(scored, 1):
        lip_badge = ('<span class="badge badge-success">Pass</span>'
                     if m["lipinski_pass"]
                     else '<span class="badge badge-danger">Fail</span>')
        crit_badge = ('<span class="badge badge-success">Pass</span>'
                      if m["all_criteria_met"]
                      else '<span class="badge badge-danger">Fail</span>')
        # Highlight top 3
        row_style = ' style="background:#ecfeff;font-weight:600;"' if rank <= 3 else ""
        html_parts.append(
            f"<tr{row_style}><td>{rank}</td><td>{m['name']}</td>"
            f"<td>{m['screening_score']:.3f}</td>"
            f"<td>{m['mw']:.1f}</td><td>{m['logp']:.2f}</td>"
            f"<td>{m['qed']:.3f}</td><td>{lip_badge}</td><td>{crit_badge}</td></tr>"
        )
    html_parts.append("</table>")

    # Top 5 candidate cards with structures
    html_parts.append("<h3>Top 5 Candidates</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in scored[:5]:
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(250, 250)) if mol else ""
        badge_class = "badge-success" if m["all_criteria_met"] else "badge-info"
        badge_text = "All Criteria Met" if m["all_criteria_met"] else "Partial Match"
        html_parts.append(
            f'<div class="molecule-card" style="text-align:center;">'
            f'<img src="{img_uri}" style="width:140px;height:140px;object-fit:contain;"/>'
            f'<div style="font-weight:700;margin-top:6px;color:#0e4f6e;font-size:1.05em;">{m["name"]}</div>'
            f'<div style="font-size:0.85em;color:#155e75;margin:4px 0;">Score: {m["screening_score"]:.3f}</div>'
            f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f} | QED: {m["qed"]:.3f}</div>'
            f'<div style="margin-top:4px;"><span class="badge {badge_class}">{badge_text}</span></div>'
            f'</div>'
        )
    html_parts.append("</div>")

    # Tanimoto similarity heatmap
    if similarity_matrix:
        sim_heatmap = _make_heatmap(
            similarity_matrix, valid_names, valid_names,
            title="Pairwise Tanimoto Similarity (Morgan Fingerprints)",
            color_scale="cyan",
            width=700,
            height=max(500, len(valid_names) * 32 + 100),
        )
        html_parts.append("<h3>Chemical Similarity</h3>")
        html_parts.append(f'<div class="chart-container">{sim_heatmap}</div>')

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    output = {
        "ranked_molecules": scored,
        "similarity_matrix": similarity_matrix,
        "similarity_labels": valid_names,
        "funnel": funnel_stages,
        "target_profile": profile,
    }
    return json.dumps(output)

def _parse_screening_json(screening_json: str) -> dict:
    """Parse screening JSON from screen_candidates, with safe defaults.

    The agent must pass the exact tool return value. Partial or hand-built JSON
    is tolerated for optional similarity fields only.
    """
    screening = json.loads(screening_json)
    if "ranked_molecules" not in screening:
        raise ValueError(
            "screening_json must be the exact JSON string returned by "
            "screen_candidates (missing 'ranked_molecules'). Do not construct, "
            "truncate, or summarize tool output."
        )
    screening.setdefault("similarity_matrix", [])
    screening.setdefault("similarity_labels", [])
    return screening

# ------------------------------------------------------------------
# Task 4: Generate final comprehensive report
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def generate_report(
    molecule_dir: flyte.io.Dir,
    properties_json: str,
    screening_json: str,
) -> str:
    """Generate a comprehensive drug screening report.

    Produces an executive summary, top candidate spotlight cards, property
    distributions, chemical diversity analysis, and final recommendation.

    Args:
        molecule_dir: Directory from load_molecules.
        properties_json: JSON from compute_properties.
        screening_json: Exact verbatim JSON string returned by screen_candidates
            (must include ranked_molecules, similarity_matrix, similarity_labels).
            Do not construct or summarize this payload yourself.

    Returns:
        JSON summary with total_screened, lipinski_passes, all_criteria_met,
        top_candidate, top_score, and top_3 ranked molecules.
    """
    from rdkit import Chem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Generating Final Report...</h2>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    screening = _parse_screening_json(screening_json)
    ranked = screening["ranked_molecules"]
    sim_matrix = screening["similarity_matrix"]
    sim_labels = screening["similarity_labels"]

    total = props["total"]
    lipinski_pass = props["lipinski_pass_count"]
    all_criteria = sum(1 for m in ranked if m["all_criteria_met"])
    top = ranked[0] if ranked else None

    html_parts = []

    # --- Executive Summary ---
    html_parts.append("<h2>Drug Molecule Screening Report</h2>")
    top_name = top["name"] if top else "N/A"
    top_score = f'{top["screening_score"]:.3f}' if top else "N/A"
    html_parts.append(
        f'<div class="card">'
        f'<h3 style="margin-top:0;color:#0e4f6e;">Executive Summary</h3>'
        f'<p style="font-size:1.05em;">'
        f'<strong>{total}</strong> molecules were screened against the target drug profile. '
        f'<strong>{lipinski_pass}</strong> passed Lipinski\'s Rule of Five, and '
        f'<strong>{all_criteria}</strong> met all screening criteria. '
        f'The top candidate is <strong style="color:#0891b2;">{top_name}</strong> '
        f'with a screening score of <strong>{top_score}</strong>.</p>'
        f'</div>'
    )

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Molecules Screened"),
        (str(lipinski_pass), "Lipinski Passes"),
        (str(all_criteria), "All Criteria Met"),
        (top_score, "Top Score"),
        (f'{props["avg_mw"]:.0f} Da', "Avg. Molecular Weight"),
        (f'{props["avg_logp"]:.2f}', "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # --- Top 3 Candidate Spotlights ---
    html_parts.append("<h2>Top Candidate Spotlights</h2>")

    for rank, m in enumerate(ranked[:3], 1):
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(300, 300)) if mol else ""

        medal = ["gold", "silver", "#cd7f32"][rank - 1]
        medal_emoji = ["1st", "2nd", "3rd"][rank - 1]

        lip_badges = ""
        for rule, key in [("MW", "mw_ok"), ("LogP", "logp_ok"),
                          ("HBD", "hbd_ok"), ("HBA", "hba_ok")]:
            ok = m["lipinski"].get(key, False)
            cls = "badge-success" if ok else "badge-danger"
            lip_badges += f'<span class="badge {cls}" style="margin:2px;">{rule}</span> '

        html_parts.append(
            f'<div class="molecule-card" style="display:flex;gap:20px;align-items:flex-start;flex-wrap:wrap;">'
            f'<div style="text-align:center;min-width:180px;">'
            f'<div style="font-size:1.6em;font-weight:800;color:{medal};">{medal_emoji}</div>'
            f'<img src="{img_uri}" style="width:200px;height:200px;object-fit:contain;border-radius:8px;'
            f'border:2px solid #a5f3fc;"/>'
            f'<div style="font-weight:700;font-size:1.1em;color:#0e4f6e;margin-top:8px;">{m["name"]}</div>'
            f'</div>'
            f'<div style="flex:1;min-width:280px;">'
            f'<table style="margin:0;">'
            f'<tr><td><strong>SMILES</strong></td><td style="font-family:monospace;font-size:0.8em;word-break:break-all;">{m["smiles"]}</td></tr>'
            f'<tr><td><strong>Screening Score</strong></td><td style="font-weight:700;color:#0891b2;font-size:1.1em;">{m["screening_score"]:.3f}</td></tr>'
            f'<tr><td><strong>Molecular Weight</strong></td><td>{m["mw"]:.1f} Da</td></tr>'
            f'<tr><td><strong>LogP</strong></td><td>{m["logp"]:.2f}</td></tr>'
            f'<tr><td><strong>H-Bond Donors</strong></td><td>{m["hbd"]}</td></tr>'
            f'<tr><td><strong>H-Bond Acceptors</strong></td><td>{m["hba"]}</td></tr>'
            f'<tr><td><strong>TPSA</strong></td><td>{m["tpsa"]:.1f} A&sup2;</td></tr>'
            f'<tr><td><strong>Rotatable Bonds</strong></td><td>{m["rotatable_bonds"]}</td></tr>'
            f'<tr><td><strong>QED</strong></td><td>{m["qed"]:.4f}</td></tr>'
            f'<tr><td><strong>Lipinski Compliance</strong></td><td>{lip_badges}</td></tr>'
            f'</table>'
            f'</div>'
            f'</div>'
        )

    # --- Property Distribution (box-plot style as bars with min/max/median) ---
    html_parts.append("<h2>Property Distributions</h2>")

    prop_keys = [("mw", "Molecular Weight (Da)"), ("logp", "LogP"),
                 ("tpsa", "TPSA"), ("qed", "QED Score")]
    for key, label in prop_keys:
        vals = sorted([m[key] for m in ranked])
        n = len(vals)
        if n == 0:
            continue
        v_min = vals[0]
        v_max = vals[-1]
        median = vals[n // 2] if n % 2 == 1 else (vals[n // 2 - 1] + vals[n // 2]) / 2
        q1 = vals[n // 4] if n >= 4 else v_min
        q3 = vals[3 * n // 4] if n >= 4 else v_max

        # Simple horizontal box-plot as SVG
        box_w = 500
        box_h = 50
        margin_l = 10
        v_range = v_max - v_min or 1

        def sx(v):
            return margin_l + ((v - v_min) / v_range) * (box_w - 2 * margin_l)

        box_svg = (
            f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {box_w} {box_h}" '
            f'style="width:100%;max-width:{box_w}px;height:auto;">'
            f'<rect width="{box_w}" height="{box_h}" fill="#fff" rx="4"/>'
            # Whisker line
            f'<line x1="{sx(v_min):.1f}" y1="25" x2="{sx(v_max):.1f}" y2="25" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Min whisker
            f'<line x1="{sx(v_min):.1f}" y1="18" x2="{sx(v_min):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Max whisker
            f'<line x1="{sx(v_max):.1f}" y1="18" x2="{sx(v_max):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # IQR box
            f'<rect x="{sx(q1):.1f}" y="14" width="{sx(q3) - sx(q1):.1f}" height="22" '
            f'fill="#a5f3fc" stroke="#0891b2" stroke-width="1.5" rx="3"/>'
            # Median line
            f'<line x1="{sx(median):.1f}" y1="12" x2="{sx(median):.1f}" y2="38" '
            f'stroke="#0e4f6e" stroke-width="2"/>'
            # Labels
            f'<text x="{sx(v_min):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_min:.1f}</text>'
            f'<text x="{sx(median):.1f}" y="10" text-anchor="middle" font-size="9" fill="#0e4f6e" font-weight="600">{median:.1f}</text>'
            f'<text x="{sx(v_max):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_max:.1f}</text>'
            f'</svg>'
        )
        html_parts.append(
            f'<div style="margin:8px 0;"><strong style="color:#155e75;">{label}</strong>'
            f'<div class="chart-container" style="padding:8px;">{box_svg}</div></div>'
        )

    # --- Chemical Diversity ---
    html_parts.append("<h2>Chemical Diversity Analysis</h2>")

    if sim_matrix and len(sim_matrix) > 1:
        # Compute average pairwise similarity (off-diagonal)
        n_mols = len(sim_matrix)
        off_diag = []
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                off_diag.append(sim_matrix[i][j])

        avg_sim = sum(off_diag) / len(off_diag) if off_diag else 0
        max_sim = max(off_diag) if off_diag else 0
        min_sim = min(off_diag) if off_diag else 0

        # Find most similar pair
        best_i, best_j = 0, 1
        best_val = 0
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                if sim_matrix[i][j] > best_val:
                    best_val = sim_matrix[i][j]
                    best_i, best_j = i, j

        html_parts.append('<div class="stat-grid">')
        html_parts.append(
            f'<div class="stat"><div class="value">{avg_sim:.3f}</div>'
            f'<div class="label">Avg. Pairwise Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{min_sim:.3f}</div>'
            f'<div class="label">Min Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{max_sim:.3f}</div>'
            f'<div class="label">Max Similarity</div></div>'
        )
        html_parts.append("</div>")

        diversity_text = "highly diverse" if avg_sim < 0.3 else "moderately diverse" if avg_sim < 0.5 else "relatively similar"
        html_parts.append(
            f'<div class="note">'
            f'The library is <strong>{diversity_text}</strong> (avg. Tanimoto = {avg_sim:.3f}). '
            f'The most similar pair is <strong>{sim_labels[best_i]}</strong> and '
            f'<strong>{sim_labels[best_j]}</strong> (similarity = {best_val:.3f}).</div>'
        )

    # --- Recommendation ---
    html_parts.append("<h2>Recommendation</h2>")
    if top:
        html_parts.append(
            f'<div class="card">'
            f'<h3 style="margin-top:0;color:#0891b2;">Top Candidate: {top["name"]}</h3>'
            f'<p>Based on the virtual screening analysis, <strong>{top["name"]}</strong> '
            f'achieved the highest composite screening score of <strong>{top["screening_score"]:.3f}</strong>. '
        )

        reasons = []
        if top["lipinski_pass"]:
            reasons.append("full Lipinski Rule of Five compliance")
        if top["qed"] > 0.5:
            reasons.append(f"high drug-likeness (QED = {top['qed']:.3f})")
        if top.get("all_criteria_met"):
            reasons.append("all target profile criteria met")
        if top["mw"] <= 500:
            reasons.append(f"favorable molecular weight ({top['mw']:.1f} Da)")

        if reasons:
            html_parts.append(
                f'This candidate stands out due to: {", ".join(reasons)}.</p>'
            )
        else:
            html_parts.append("</p>")

        # Runner-up mentions
        if len(ranked) >= 2:
            html_parts.append(
                f'<p style="font-size:0.9em;color:#6c757d;">Runner-up candidates: '
            )
            runners = []
            for m in ranked[1:4]:
                runners.append(f'{m["name"]} (score: {m["screening_score"]:.3f})')
            html_parts.append(", ".join(runners) + ".</p>")

        html_parts.append("</div>")

    # Final note
    html_parts.append(
        '<div class="note">'
        "This is a virtual screening analysis. All candidates should undergo "
        "further computational validation (molecular dynamics, docking) and "
        "experimental testing before advancing to clinical trials.</div>"
    )

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # JSON summary
    summary = {
        "total_screened": total,
        "lipinski_passes": lipinski_pass,
        "all_criteria_met": all_criteria,
        "top_candidate": top["name"] if top else None,
        "top_score": top["screening_score"] if top else None,
        "top_3": [
            {"name": m["name"], "score": m["screening_score"]}
            for m in ranked[:3]
        ],
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Agent
# ------------------------------------------------------------------

# {{docs-fragment agent}}
SCREENING_AGENT_INSTRUCTIONS = """\
You are a medicinal chemistry screening strategist. You orchestrate a virtual \
screening pipeline using durable Flyte tools. You NEVER invent molecular \
properties — only RDKit tools compute them.

Workflow:
1. If target_profile is not provided in the user message, derive a JSON \
target_profile from the therapeutic brief. Valid keys: mw, logp, hbd, hba, tpsa \
(each [min, max]). Ground choices in oral bioavailability / kinase / CNS rules \
as appropriate to the brief.
2. First pass (always): load_molecules → compute_properties → \
screen_candidates → generate_report. Pass tool outputs between steps exactly \
(molecule_dir from load_molecules into compute_properties and generate_report; \
properties_json from compute_properties into screen_candidates and \
generate_report; screening_json must be the complete, unmodified string \
returned by screen_candidates — never rebuild or summarize JSON yourself).
3. Read the JSON summary returned by generate_report. Reflect:
   - If all_criteria_met == 0: relax exactly ONE profile bound by ~10–20% \
and re-run screen_candidates then generate_report only, reusing the same \
molecule_dir and properties_json from the first pass.
   - If all molecules pass but diversity is a stated goal: note high similarity \
in your summary; do not re-run unless brief asks for stricter filters.
   - Maximum ONE rescreen iteration.
4. Finish with plain text: top candidate, rationale tied to computed metrics \
from the tool JSON, funnel interpretation, and suggested next steps (docking, \
ADMET lab tests).

If the user supplies an explicit target_profile JSON, use it as-is.

Do NOT ask the user for SMILES or molecule lists when molecules_json is empty — \
the default library is loaded automatically.
"""

screening_agent = Agent(
    name="drug-screening-agent",
    instructions=SCREENING_AGENT_INSTRUCTIONS,
    model=MODEL,
    tools=[
        load_molecules,
        compute_properties,
        screen_candidates,
        generate_report,
    ],
    max_turns=12,
)
# {{/docs-fragment agent}}

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def pipeline(
    brief: str = "Screen the default drug library for orally bioavailable small molecules.",
    molecules_json: str = "",
    target_profile: str = "",
) -> str:
    """Agentic virtual drug molecule screening pipeline.

    A medicinal-chemistry agent interprets the screening brief, derives or
    applies a target profile, orchestrates the RDKit screening stages, and
    optionally re-screens when funnel results are too narrow.

    Args:
        brief: Natural-language therapeutic goal (e.g. oral kinase inhibitors,
            CNS-penetrant small molecules).
        molecules_json: JSON mapping molecule names to SMILES strings.
            Defaults to a curated library of ~15 well-known drugs.
        target_profile: Optional JSON with desired property ranges that
            overrides agent-derived criteria
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5]}).

    Returns:
        Agent summary with screening rationale and key results.
    """
    prompt_parts = [
        f"Screening brief: {brief}",
        'Use molecules_json="" for the built-in default library unless provided below.',
        "Compose the four stage tools in order: load_molecules → compute_properties "
        "→ screen_candidates → generate_report. Pass each tool's full return value "
        "verbatim to the next step (especially screening_json). Re-run "
        "screen_candidates and generate_report at most once if the funnel is too narrow.",
    ]
    if molecules_json.strip():
        prompt_parts.append(f"molecules_json: {molecules_json}")
    if target_profile.strip():
        prompt_parts.append(f"Use this target_profile exactly: {target_profile}")

    result = await screening_agent.run.aio("\n".join(prompt_parts))
    return result.summary or result.error or ""

# {{/docs-fragment pipeline}}

# ------------------------------------------------------------------
# Rescreen demo — tight profile + explicit rescreen instructions
# ------------------------------------------------------------------

# Initial profile is deliberately strict (narrow MW + low LogP cap) so
# all_criteria_met is typically 0 on the default library; the brief then
# forces a single rescreen with a widened LogP window.
RESCREEN_DEMO_TARGET_PROFILE = (
    '{"mw": [150, 200], "logp": [-0.5, 1.0], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_TARGET_PROFILE_RESCREEN = (
    '{"mw": [150, 200], "logp": [-0.5, 3.5], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_BRIEF = f"""\
Two-round agentic screening demo on the default library.

**Round 1 (strict profile):** load_molecules → compute_properties → \
screen_candidates → generate_report using the initial target_profile exactly.

**Round 2 (required — do not skip):** call screen_candidates then generate_report \
again, reusing the same molecule_dir and properties_json from round 1, with this \
relaxed target_profile (wider LogP window only): \
{RESCREEN_DEMO_TARGET_PROFILE_RESCREEN}

Pass every tool return value verbatim to the next step. After both rounds, \
summarize how the funnel and top candidates changed between round 1 and round 2."""

# {{docs-fragment rescreen_demo}}
@env.task(report=True)
async def rescreen_demo() -> str:
    """Example run with a two-round execution graph (rescreen).

    Round 1 uses a strict CNS-like profile; round 2 always re-runs
    screen_candidates and generate_report with a widened LogP window,
    reusing cached molecule_dir and properties_json.
    """
    return await pipeline(
        brief=RESCREEN_DEMO_BRIEF,
        target_profile=RESCREEN_DEMO_TARGET_PROFILE,
    )

# {{/docs-fragment rescreen_demo}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/drug_molecule_screening/drug_molecule_screening.py*

## Run the agentic pipeline

The `pipeline` task delegates to the screening agent:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.4",
#    "litellm",
#    "rdkit",
#    "numpy",
#    "scikit-learn",
#    "pillow",
# ]
# main = "pipeline"
# params = ""
# ///
"""Virtual drug molecule screening — compute properties, apply Lipinski filters, rank candidates."""

import base64
import io
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report
from flyte.ai.agents import Agent, tool

MODEL = os.getenv("DRUG_SCREENING_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="drug-molecule-screening", pre=True).with_apt_packages(
    "libxrender1", "libxext6", "libexpat1",
)

env = flyte.TaskEnvironment(
    name="drug-molecule-screening",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default molecule library — real SMILES for well-known drugs
# ------------------------------------------------------------------

DEFAULT_MOLECULES = {
    "Aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
    "Ibuprofen": "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
    "Caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
    "Penicillin G": "CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C",
    "Metformin": "CN(C)C(=N)NC(=N)N",
    "Paracetamol": "CC(=O)NC1=CC=C(C=C1)O",
    "Diazepam": "ClC1=CC2=C(C=C1)N(C(=O)CN=C2C3=CC=CC=C3)C",
    "Omeprazole": "CC1=CN=C(C(=C1OC)C)CS(=O)C2=NC3=CC=CC=C3N2",
    "Atorvastatin": "CC(C)C1=C(C(=C(N1CCC(CC(CC(=O)O)O)O)C2=CC=C(C=C2)F)C3=CC=CC=C3)C(=O)NC4=CC=CC=C4",
    "Methotrexate": "CN(CC1=CN=C2N=C(N=C(N)C2=N1)N)C3=CC=C(C=C3)C(=O)NC(CCC(=O)O)C(=O)O",
    "Doxorubicin": "CC1C(C(CC(O1)OC2CC(CC3=C2C(=C4C(=C3O)C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)CO)O)N)O",
    "Tamoxifen": "CCC(=C(C1=CC=CC=C1)C2=CC=C(C=C2)OCCN(C)C)C3=CC=CC=C3",
    "Lopinavir": "CC1=C(C(=CC=C1)C)OCC(=O)NC(CC2=CC=CC=C2)C(CC(CC3=CC=CC=C3)NC(=O)C(C(C)C)N4CCCNC4=O)O",
    "Remdesivir": "CCC(CC)COC(=O)C(C)NP(=O)(OCC1C(C(C(O1)C2=CC=C3N2N=CN=C3N)O)O)OC4=CC=CC=C4",
    "Erlotinib": "COCCOC1=CC2=C(C=C1OCCOC)C(=NC=N2)NC3=CC=CC(=C3)C#C",
}

# ------------------------------------------------------------------
# Report styling — pharma blue/cyan theme
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #0e4f6e; border-bottom: 2px solid #0891b2; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #155e75; margin-top: 20px; }
  .report .card { background: #ecfeff; border: 1px solid #a5f3fc; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #cffafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0e4f6e; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0e4f6e; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #cffafe; }
  .report tr:nth-child(even) { background: #ecfeff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #cffafe; color: #155e75; }
  .report .chart-container { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #ecfeff; border-left: 4px solid #0891b2; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .molecule-card { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .molecule-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 12px; margin: 16px 0; }
  .report .funnel { text-align: center; margin: 24px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _mol_to_data_uri(mol, size: tuple[int, int] = (300, 300)) -> str:
    """Convert an RDKit molecule to a PNG base64 data URI."""
    from rdkit.Chem import Draw

    img = Draw.MolToImage(mol, size=size)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 340,
    y_max_cap: float | None = None,
    horizontal: bool = False,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels.
        series: Dict mapping series name to list of values.
        title: Chart title.
        colors: Colors for each series.
        width/height: SVG dimensions.
        y_max_cap: Cap the y-axis at this value.
        horizontal: If True, draw horizontal bars.
        value_fmt: Format string for value labels.

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#a5f3fc", "#155e75"]
    colors = colors or default_colors

    if horizontal:
        return _make_horizontal_bar_chart(labels, series, title, colors, width, height, value_fmt)

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_fmt}}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_fmt}}</text>'
            )
        # Truncate long labels
        disp_label = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 16}" '
            f'text-anchor="middle" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {gx + n_series * bar_width / 2:.1f}, {mt + ch + 16})">'
            f'{disp_label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 100
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 100}" y="{mt + ch + 40}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 51}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_horizontal_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 400,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG horizontal bar chart (sorted)."""
    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0"]
    colors = colors or default_colors

    n = len(labels)
    row_height = max(22, min(35, (height - 80) // max(n, 1)))
    actual_height = max(height, 80 + n * row_height)
    ml, mr, mt, mb = 120, 60, 40, 20
    cw = width - ml - mr
    ch = actual_height - mt - mb

    # Use first series
    first_key = list(series.keys())[0]
    vals = series[first_key]
    x_max = max(vals) * 1.15 if vals else 1

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {actual_height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{actual_height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    bar_h = row_height * 0.65
    for i, (label, val) in enumerate(zip(labels, vals)):
        y = mt + i * row_height
        bw = (val / x_max) * cw if x_max else 0
        color = colors[i % len(colors)]
        # Label
        disp = label if len(label) <= 14 else label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + bar_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#1a1a2e">{disp}</text>'
        )
        # Bar
        svg.append(
            f'<rect x="{ml}" y="{y:.1f}" width="{bw:.1f}" height="{bar_h:.1f}" '
            f'fill="{color}" rx="3"/>'
        )
        # Value
        svg.append(
            f'<text x="{ml + bw + 6:.1f}" y="{y + bar_h / 2 + 4:.1f}" '
            f'font-size="11" fill="#0e4f6e" font-weight="600">{val:{value_fmt}}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    color_scale: str = "cyan",
    width: int = 700,
    height: int = 500,
    value_fmt: str = ".2f",
) -> str:
    """Generate an SVG heatmap.

    Args:
        matrix: 2D list of values (rows x cols).
        row_labels: Labels for rows.
        col_labels: Labels for columns.
        title: Chart title.
        color_scale: Color scheme ("cyan", "red", "green").
        width/height: SVG dimensions.
        value_fmt: Format string for cell values.

    Returns:
        SVG string.
    """
    if not matrix or not matrix[0]:
        return ""

    n_rows = len(matrix)
    n_cols = len(matrix[0])

    ml, mr, mt, mb = 110, 20, 70, 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    # Flatten to find range
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    def color_for(v):
        t = (v - v_min) / v_range
        if color_scale == "cyan":
            # White to deep teal
            r = int(255 - t * (255 - 14))
            g = int(255 - t * (255 - 79))
            b = int(255 - t * (255 - 110))
        elif color_scale == "red":
            r = int(255 - t * 50)
            g = int(255 - t * 200)
            b = int(255 - t * 200)
        else:  # green
            r = int(255 - t * 200)
            g = int(255 - t * 50)
            b = int(255 - t * 200)
        return f"rgb({r},{g},{b})"

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Column labels (rotated)
    for ci, label in enumerate(col_labels):
        x = ml + ci * cell_w + cell_w / 2
        disp = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{x:.1f}" y="{mt - 8}" text-anchor="end" font-size="10" '
            f'fill="#1a1a2e" transform="rotate(-45, {x:.1f}, {mt - 8})">{disp}</text>'
        )

    # Row labels + cells
    for ri, (row_label, row_vals) in enumerate(zip(row_labels, matrix)):
        y = mt + ri * cell_h
        disp = row_label if len(row_label) <= 14 else row_label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + cell_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#1a1a2e">{disp}</text>'
        )
        for ci, val in enumerate(row_vals):
            x = ml + ci * cell_w
            fill = color_for(val)
            svg.append(
                f'<rect x="{x:.1f}" y="{y:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{fill}" stroke="#fff" stroke-width="1"/>'
            )
            # Text color: dark on light, light on dark
            t = (val - v_min) / v_range
            txt_color = "#fff" if t > 0.55 else "#1a1a2e"
            # Only show text if cells are large enough
            if cell_w > 30 and cell_h > 18:
                svg.append(
                    f'<text x="{x + cell_w / 2:.1f}" y="{y + cell_h / 2 + 4:.1f}" '
                    f'text-anchor="middle" font-size="9" fill="{txt_color}">'
                    f'{val:{value_fmt}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_scatter_plot(
    points: list[dict],
    x_label: str = "MW",
    y_label: str = "LogP",
    title: str = "",
    reference_lines: list[dict] | None = None,
    width: int = 700,
    height: int = 400,
) -> str:
    """Generate an SVG scatter plot.

    Args:
        points: List of dicts with "x", "y", "label" keys.
        x_label/y_label: Axis labels.
        title: Chart title.
        reference_lines: List of dicts with "axis" ("x"/"y"), "value", "label".
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not points:
        return ""

    ml, mr, mt, mb = 60, 30, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [p["x"] for p in points]
    y_vals = [p["y"] for p in points]
    x_min, x_max = min(x_vals) * 0.9, max(x_vals) * 1.1
    y_min, y_max = min(y_vals) - 1, max(y_vals) + 1

    # Extend ranges to include reference lines
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                x_max = max(x_max, rl["value"] * 1.1)
            else:
                y_max = max(y_max, rl["value"] * 1.1)

    x_range = x_max - x_min or 1
    y_range = y_max - y_min or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid
    for i in range(6):
        y_tick = y_min + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.1f}</text>'
        )

    for i in range(6):
        x_tick = x_min + x_range * i / 5
        px = sx(x_tick)
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
            f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Reference lines (Lipinski boundaries)
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                px = sx(rl["value"])
                svg.append(
                    f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{px + 4:.1f}" y="{mt + 14}" font-size="10" '
                    f'fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )
            else:
                py = sy(rl["value"])
                svg.append(
                    f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{ml + cw - 4:.1f}" y="{py - 6:.1f}" text-anchor="end" '
                    f'font-size="10" fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )

    # Drug-like zone shading (MW<=500 and LogP<=5 quadrant)
    if reference_lines:
        mw_line = next((rl for rl in reference_lines if rl["axis"] == "x"), None)
        logp_line = next((rl for rl in reference_lines if rl["axis"] == "y"), None)
        if mw_line and logp_line:
            zx1 = sx(x_min)
            zx2 = sx(min(mw_line["value"], x_max))
            zy1 = sy(min(logp_line["value"], y_max))
            zy2 = sy(y_min)
            svg.append(
                f'<rect x="{zx1:.1f}" y="{zy1:.1f}" '
                f'width="{zx2 - zx1:.1f}" height="{zy2 - zy1:.1f}" '
                f'fill="#0891b2" opacity="0.06" rx="4"/>'
            )
            svg.append(
                f'<text x="{zx1 + 8:.1f}" y="{zy2 - 8:.1f}" font-size="11" '
                f'fill="#0891b2" font-weight="600" opacity="0.6">Drug-like Zone</text>'
            )

    # Points
    point_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#155e75", "#0284c7",
                    "#059669", "#0d9488", "#0369a1", "#047857", "#115e59",
                    "#0c4a6e", "#064e3b", "#1e3a5f", "#134e4a", "#075985"]
    for i, pt in enumerate(points):
        px, py = sx(pt["x"]), sy(pt["y"])
        color = point_colors[i % len(point_colors)]
        svg.append(
            f'<circle cx="{px:.1f}" cy="{py:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5" opacity="0.85"/>'
        )
        # Label offset to avoid overlap
        offset_x = 8
        offset_y = -8 if i % 2 == 0 else 14
        label = pt["label"] if len(pt["label"]) <= 12 else pt["label"][:10] + ".."
        svg.append(
            f'<text x="{px + offset_x:.1f}" y="{py + offset_y:.1f}" '
            f'font-size="9" fill="#1a1a2e">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Axis labels
    if x_label:
        svg.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        svg.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_funnel(
    stages: list[dict],
    title: str = "",
    width: int = 600,
    height: int = 400,
) -> str:
    """Generate an SVG funnel visualization.

    Args:
        stages: List of dicts with "label", "count", "total" keys.
        title: Chart title.
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not stages:
        return ""

    n = len(stages)
    mt = 50
    mb = 20
    available_h = height - mt - mb
    stage_h = available_h / n
    cx = width / 2

    # Color gradient from light cyan to deep teal
    colors = []
    for i in range(n):
        t = i / max(n - 1, 1)
        r = int(207 - t * (207 - 14))
        g = int(250 - t * (250 - 79))
        b = int(254 - t * (254 - 110))
        colors.append(f"rgb({r},{g},{b})")

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{cx}" y="28" text-anchor="middle" '
            f'font-size="16" font-weight="700" fill="#0e4f6e">{title}</text>'
        )

    max_count = stages[0]["count"] if stages else 1
    max_width = width * 0.75

    for i, stage in enumerate(stages):
        y_top = mt + i * stage_h
        y_bot = y_top + stage_h

        # Width proportional to count
        w_top = max_width * (stage["count"] / max_count) if i == 0 else prev_w_bot
        if i < n - 1:
            w_bot = max_width * (stages[i + 1]["count"] / max_count)
        else:
            w_bot = max_width * (stage["count"] / max_count) * 0.7

        prev_w_bot = w_bot

        # Trapezoid
        x1_top = cx - w_top / 2
        x2_top = cx + w_top / 2
        x1_bot = cx - w_bot / 2
        x2_bot = cx + w_bot / 2

        svg.append(
            f'<polygon points="{x1_top:.1f},{y_top:.1f} {x2_top:.1f},{y_top:.1f} '
            f'{x2_bot:.1f},{y_bot:.1f} {x1_bot:.1f},{y_bot:.1f}" '
            f'fill="{colors[i]}" stroke="#fff" stroke-width="2"/>'
        )

        # Text: dark on light, white on dark
        t = i / max(n - 1, 1)
        txt_color = "#0e4f6e" if t < 0.5 else "#fff"
        y_mid = (y_top + y_bot) / 2

        svg.append(
            f'<text x="{cx}" y="{y_mid - 4:.1f}" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="{txt_color}">{stage["label"]}</text>'
        )
        svg.append(
            f'<text x="{cx}" y="{y_mid + 14:.1f}" text-anchor="middle" '
            f'font-size="12" fill="{txt_color}" opacity="0.85">'
            f'{stage["count"]} / {stage["total"]}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate molecules
# ------------------------------------------------------------------

@tool
@env.task(cache="auto")
async def load_molecules(
    molecules_json: str = "",
) -> flyte.io.Dir:
    """Parse SMILES strings, validate with RDKit, generate 2D depictions.

    Args:
        molecules_json: JSON string mapping molecule names to SMILES.
            Defaults to a curated library of ~15 well-known drugs.

    Returns:
        flyte.io.Dir containing molecule data (JSON + PNG depictions).
        Pass this directory to compute_properties and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Draw

    if molecules_json.strip():
        molecules = json.loads(molecules_json)
    else:
        molecules = DEFAULT_MOLECULES

    out_dir = tempfile.mkdtemp(prefix="mol_library_")
    results = []
    valid_count = 0
    invalid_count = 0

    log.info(f"Parsing {len(molecules)} molecules...")

    for name, smiles in molecules.items():
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            log.warning(f"  [INVALID] {name}: {smiles}")
            invalid_count += 1
            continue

        valid_count += 1

        # Generate 2D depiction as PNG
        img = Draw.MolToImage(mol, size=(300, 300))
        img_path = os.path.join(out_dir, f"{name.replace(' ', '_')}.png")
        img.save(img_path)

        results.append({
            "name": name,
            "smiles": smiles,
            "valid": True,
            "image_file": os.path.basename(img_path),
        })

    # Save molecule manifest
    manifest = {
        "total": len(molecules),
        "valid": valid_count,
        "invalid": invalid_count,
        "molecules": results,
    }
    manifest_path = os.path.join(out_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    log.info(f"Loaded {valid_count} valid molecules ({invalid_count} invalid)")

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Compute physicochemical properties
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def compute_properties(
    molecule_dir: flyte.io.Dir,
) -> str:
    """Compute drug-likeness properties for all molecules.

    Computes MW, LogP, HBD, HBA, TPSA, rotatable bonds, formal charge,
    ring count, QED, and Lipinski Rule of Five compliance.

    Args:
        molecule_dir: Directory from load_molecules.

    Returns:
        JSON string with all computed properties. Pass to screen_candidates
        and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Descriptors, Lipinski
    from rdkit.Chem.QED import qed

    # --- Loading report ---
    await flyte.report.replace.aio(
        _wrap_report("<h2>Computing Molecular Properties...</h2>"
                      "<p>Analyzing physicochemical descriptors for all molecules.</p>"),
        do_flush=True,
    )

    mol_dir = await molecule_dir.download()
    with open(os.path.join(mol_dir, "manifest.json")) as f:
        manifest = json.load(f)

    molecules_data = []
    lipinski_pass = 0

    for mol_info in manifest["molecules"]:
        mol = Chem.MolFromSmiles(mol_info["smiles"])
        if mol is None:
            continue

        mw = Descriptors.MolWt(mol)
        logp = Descriptors.MolLogP(mol)
        hbd = Lipinski.NumHDonors(mol)
        hba = Lipinski.NumHAcceptors(mol)
        tpsa = Descriptors.TPSA(mol)
        rotatable = Lipinski.NumRotatableBonds(mol)
        formal_charge = Chem.GetFormalCharge(mol)
        num_rings = Lipinski.RingCount(mol)
        qed_score = qed(mol)

        # Lipinski Rule of Five
        lipinski = {
            "mw_ok": mw <= 500,
            "logp_ok": logp <= 5,
            "hbd_ok": hbd <= 5,
            "hba_ok": hba <= 10,
        }
        lipinski_all = all(lipinski.values())
        if lipinski_all:
            lipinski_pass += 1

        # Read image for data URI
        img_path = os.path.join(mol_dir, mol_info["image_file"])
        data_uri = ""
        if os.path.exists(img_path):
            with open(img_path, "rb") as img_f:
                b64 = base64.b64encode(img_f.read()).decode()
                data_uri = f"data:image/png;base64,{b64}"

        molecules_data.append({
            "name": mol_info["name"],
            "smiles": mol_info["smiles"],
            "mw": round(mw, 2),
            "logp": round(logp, 2),
            "hbd": hbd,
            "hba": hba,
            "tpsa": round(tpsa, 2),
            "rotatable_bonds": rotatable,
            "formal_charge": formal_charge,
            "num_rings": num_rings,
            "qed": round(qed_score, 4),
            "lipinski": lipinski,
            "lipinski_pass": lipinski_all,
            "image_data_uri": data_uri,
        })

    total = len(molecules_data)
    avg_mw = sum(m["mw"] for m in molecules_data) / total if total else 0
    avg_logp = sum(m["logp"] for m in molecules_data) / total if total else 0
    lipinski_rate = lipinski_pass / total * 100 if total else 0

    # ---- Build report ----
    html_parts = []

    # Header
    html_parts.append("<h2>Molecular Properties Analysis</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Total Molecules"),
        (f"{lipinski_rate:.0f}%", "Lipinski Pass Rate"),
        (f"{avg_mw:.1f}", "Avg. MW (Da)"),
        (f"{avg_logp:.2f}", "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Molecule gallery
    html_parts.append("<h3>Molecule Library</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in molecules_data:
        if m["image_data_uri"]:
            badge_class = "badge-success" if m["lipinski_pass"] else "badge-danger"
            badge_text = "Lipinski Pass" if m["lipinski_pass"] else "Lipinski Fail"
            html_parts.append(
                f'<div class="molecule-card" style="text-align:center;">'
                f'<img src="{m["image_data_uri"]}" style="width:160px;height:160px;object-fit:contain;"/>'
                f'<div style="font-weight:600;margin-top:6px;color:#0e4f6e;">{m["name"]}</div>'
                f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f}</div>'
                f'<div><span class="badge {badge_class}">{badge_text}</span></div>'
                f'</div>'
            )
    html_parts.append("</div>")

    # MW bar chart (horizontal, sorted)
    sorted_by_mw = sorted(molecules_data, key=lambda m: m["mw"], reverse=True)
    mw_labels = [m["name"] for m in sorted_by_mw]
    mw_vals = [m["mw"] for m in sorted_by_mw]
    mw_chart = _make_bar_chart(
        mw_labels, {"MW (Da)": mw_vals},
        title="Molecular Weight Distribution",
        horizontal=True,
        width=700, height=max(300, len(mw_labels) * 30 + 80),
        value_fmt=".1f",
    )
    html_parts.append("<h3>Molecular Weight</h3>")
    html_parts.append(f'<div class="chart-container">{mw_chart}</div>')

    # LogP vs MW scatter plot
    scatter_points = [
        {"x": m["mw"], "y": m["logp"], "label": m["name"]}
        for m in molecules_data
    ]
    scatter_chart = _make_scatter_plot(
        scatter_points,
        x_label="Molecular Weight (Da)",
        y_label="LogP",
        title="LogP vs. Molecular Weight (Lipinski Boundaries)",
        reference_lines=[
            {"axis": "x", "value": 500, "label": "MW = 500"},
            {"axis": "y", "value": 5, "label": "LogP = 5"},
        ],
        width=700,
        height=420,
    )
    html_parts.append("<h3>Lipinski Space</h3>")
    html_parts.append(f'<div class="chart-container">{scatter_chart}</div>')

    # Property heatmap (molecules x properties)
    prop_names = ["MW", "LogP", "HBD", "HBA", "TPSA", "Rot. Bonds"]
    # Normalize each property to 0-1 for heatmap
    raw_matrix = []
    for m in molecules_data:
        raw_matrix.append([m["mw"], m["logp"], m["hbd"], m["hba"], m["tpsa"], m["rotatable_bonds"]])

    # Normalize per column
    n_props = len(prop_names)
    col_min = [min(row[c] for row in raw_matrix) for c in range(n_props)]
    col_max = [max(row[c] for row in raw_matrix) for c in range(n_props)]
    norm_matrix = []
    for row in raw_matrix:
        norm_row = []
        for c in range(n_props):
            rng = col_max[c] - col_min[c]
            norm_row.append((row[c] - col_min[c]) / rng if rng else 0.5)
        norm_matrix.append(norm_row)

    heatmap_labels = [m["name"] for m in molecules_data]
    heatmap = _make_heatmap(
        norm_matrix, heatmap_labels, prop_names,
        title="Normalized Property Heatmap",
        color_scale="cyan",
        width=700,
        height=max(400, len(heatmap_labels) * 28 + 100),
    )
    html_parts.append("<h3>Property Heatmap</h3>")
    html_parts.append(f'<div class="chart-container">{heatmap}</div>')

    # Lipinski compliance table
    html_parts.append("<h3>Lipinski Rule of Five Compliance</h3>")
    html_parts.append("<table><tr><th>Molecule</th><th>MW &le; 500</th>"
                      "<th>LogP &le; 5</th><th>HBD &le; 5</th>"
                      "<th>HBA &le; 10</th><th>Overall</th></tr>")
    for m in molecules_data:
        lip = m["lipinski"]

        def _badge(ok):
            if ok:
                return '<span class="badge badge-success">Pass</span>'
            return '<span class="badge badge-danger">Fail</span>'

        overall_badge = _badge(m["lipinski_pass"])
        html_parts.append(
            f'<tr><td><strong>{m["name"]}</strong></td>'
            f'<td>{_badge(lip["mw_ok"])}</td>'
            f'<td>{_badge(lip["logp_ok"])}</td>'
            f'<td>{_badge(lip["hbd_ok"])}</td>'
            f'<td>{_badge(lip["hba_ok"])}</td>'
            f'<td>{overall_badge}</td></tr>'
        )
    html_parts.append("</table>")

    # QED bar chart
    sorted_by_qed = sorted(molecules_data, key=lambda m: m["qed"], reverse=True)
    qed_labels = [m["name"] for m in sorted_by_qed]
    qed_vals = [m["qed"] for m in sorted_by_qed]
    qed_chart = _make_bar_chart(
        qed_labels, {"QED Score": qed_vals},
        title="Drug-likeness (QED Score)",
        horizontal=True,
        width=700, height=max(300, len(qed_labels) * 30 + 80),
        value_fmt=".3f",
        colors=["#06d6a0"],
    )
    html_parts.append("<h3>Drug-likeness (QED)</h3>")
    html_parts.append(f'<div class="chart-container">{qed_chart}</div>')

    # Flush full report
    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # Return properties as JSON (strip image data URIs to reduce size)
    output = {
        "total": total,
        "lipinski_pass_count": lipinski_pass,
        "lipinski_pass_rate": round(lipinski_rate, 2),
        "avg_mw": round(avg_mw, 2),
        "avg_logp": round(avg_logp, 2),
        "molecules": [
            {k: v for k, v in m.items() if k != "image_data_uri"}
            for m in molecules_data
        ],
    }
    return json.dumps(output)

# ------------------------------------------------------------------
# Task 3: Screen candidates against target profile
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def screen_candidates(
    properties_json: str,
    target_profile: str = "",
) -> str:
    """Screen molecules against a target drug profile and rank candidates.

    Scores each molecule on how well it matches the target profile, computes
    pairwise Tanimoto similarity, and produces a ranked list.

    Args:
        properties_json: JSON from compute_properties.
        target_profile: JSON string with desired property ranges
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5.0]}).

    Returns:
        JSON string with ranked_molecules, similarity_matrix, similarity_labels,
        funnel, and target_profile. Pass the full return value verbatim to
        generate_report along with molecule_dir and properties_json.
    """
    from rdkit import Chem, DataStructs
    from rdkit.Chem import AllChem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Screening Candidates...</h2>"
                      "<p>Evaluating molecules against the target drug profile.</p>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    molecules = props["molecules"]

    # Default target profile
    if target_profile.strip():
        profile = json.loads(target_profile)
    else:
        profile = {
            "mw": [150, 500],
            "logp": [-0.5, 5.0],
            "hbd": [0, 5],
            "hba": [0, 10],
            "tpsa": [20, 140],
        }

    # --- Screening ---
    funnel_total = len(molecules)
    pass_mw = 0
    pass_logp = 0
    pass_lipinski = 0
    final_candidates = 0

    scored = []
    for m in molecules:
        score = 0
        max_score = 0
        criteria = {}

        # Check each profile criterion
        checks = [
            ("mw", m["mw"]),
            ("logp", m["logp"]),
            ("hbd", m["hbd"]),
            ("hba", m["hba"]),
            ("tpsa", m["tpsa"]),
        ]

        for key, val in checks:
            if key in profile:
                lo, hi = profile[key]
                max_score += 1
                in_range = lo <= val <= hi
                criteria[key] = in_range
                if in_range:
                    score += 1
                    # Bonus: closer to midpoint = higher score
                    mid = (lo + hi) / 2
                    rng = (hi - lo) / 2
                    dist = abs(val - mid) / rng if rng else 0
                    score += max(0, 0.5 * (1 - dist))

        # QED bonus
        score += m["qed"] * 2
        max_score += 2

        # Lipinski bonus
        if m["lipinski_pass"]:
            score += 1
        max_score += 1

        normalized_score = score / max_score if max_score else 0

        # Funnel tracking — cascading filter (each stage requires passing the previous)
        mw_ok = criteria.get("mw", True)
        logp_ok = criteria.get("logp", True)
        if mw_ok:
            pass_mw += 1
            if logp_ok:
                pass_logp += 1
                if m["lipinski_pass"]:
                    pass_lipinski += 1
                    if all(criteria.values()):
                        final_candidates += 1

        scored.append({
            **m,
            "screening_score": round(normalized_score, 4),
            "criteria_met": criteria,
            "all_criteria_met": all(criteria.values()),
        })

    # Sort by score descending
    scored.sort(key=lambda m: m["screening_score"], reverse=True)

    # --- Tanimoto similarity matrix ---
    fps = []
    valid_names = []
    for m in scored:
        mol = Chem.MolFromSmiles(m["smiles"])
        if mol:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
            fps.append(fp)
            valid_names.append(m["name"])

    similarity_matrix = []
    for i in range(len(fps)):
        row = []
        for j in range(len(fps)):
            sim = DataStructs.TanimotoSimilarity(fps[i], fps[j])
            row.append(round(sim, 3))
        similarity_matrix.append(row)

    # ---- Build report ----
    html_parts = []
    html_parts.append("<h2>Candidate Screening Results</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(funnel_total), "Total Screened"),
        (str(pass_lipinski), "Lipinski Passes"),
        (str(final_candidates), "All Criteria Met"),
        (f"{scored[0]['screening_score']:.3f}" if scored else "N/A", "Top Score"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Screening funnel
    funnel_stages = [
        {"label": "Total Molecules", "count": funnel_total, "total": funnel_total},
        {"label": "Pass MW Filter", "count": pass_mw, "total": funnel_total},
        {"label": "Pass LogP Filter", "count": pass_logp, "total": funnel_total},
        {"label": "Lipinski Compliant", "count": pass_lipinski, "total": funnel_total},
        {"label": "All Criteria Met", "count": final_candidates, "total": funnel_total},
    ]
    funnel_svg = _make_funnel(
        funnel_stages,
        title="Screening Funnel",
        width=600,
        height=380,
    )
    html_parts.append("<h3>Screening Funnel</h3>")
    html_parts.append(f'<div class="chart-container" style="text-align:center;">{funnel_svg}</div>')

    # Ranked candidates table
    html_parts.append("<h3>Ranked Candidates</h3>")
    html_parts.append(
        "<table><tr><th>Rank</th><th>Molecule</th><th>Score</th>"
        "<th>MW</th><th>LogP</th><th>QED</th><th>Lipinski</th><th>All Criteria</th></tr>"
    )
    for rank, m in enumerate(scored, 1):
        lip_badge = ('<span class="badge badge-success">Pass</span>'
                     if m["lipinski_pass"]
                     else '<span class="badge badge-danger">Fail</span>')
        crit_badge = ('<span class="badge badge-success">Pass</span>'
                      if m["all_criteria_met"]
                      else '<span class="badge badge-danger">Fail</span>')
        # Highlight top 3
        row_style = ' style="background:#ecfeff;font-weight:600;"' if rank <= 3 else ""
        html_parts.append(
            f"<tr{row_style}><td>{rank}</td><td>{m['name']}</td>"
            f"<td>{m['screening_score']:.3f}</td>"
            f"<td>{m['mw']:.1f}</td><td>{m['logp']:.2f}</td>"
            f"<td>{m['qed']:.3f}</td><td>{lip_badge}</td><td>{crit_badge}</td></tr>"
        )
    html_parts.append("</table>")

    # Top 5 candidate cards with structures
    html_parts.append("<h3>Top 5 Candidates</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in scored[:5]:
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(250, 250)) if mol else ""
        badge_class = "badge-success" if m["all_criteria_met"] else "badge-info"
        badge_text = "All Criteria Met" if m["all_criteria_met"] else "Partial Match"
        html_parts.append(
            f'<div class="molecule-card" style="text-align:center;">'
            f'<img src="{img_uri}" style="width:140px;height:140px;object-fit:contain;"/>'
            f'<div style="font-weight:700;margin-top:6px;color:#0e4f6e;font-size:1.05em;">{m["name"]}</div>'
            f'<div style="font-size:0.85em;color:#155e75;margin:4px 0;">Score: {m["screening_score"]:.3f}</div>'
            f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f} | QED: {m["qed"]:.3f}</div>'
            f'<div style="margin-top:4px;"><span class="badge {badge_class}">{badge_text}</span></div>'
            f'</div>'
        )
    html_parts.append("</div>")

    # Tanimoto similarity heatmap
    if similarity_matrix:
        sim_heatmap = _make_heatmap(
            similarity_matrix, valid_names, valid_names,
            title="Pairwise Tanimoto Similarity (Morgan Fingerprints)",
            color_scale="cyan",
            width=700,
            height=max(500, len(valid_names) * 32 + 100),
        )
        html_parts.append("<h3>Chemical Similarity</h3>")
        html_parts.append(f'<div class="chart-container">{sim_heatmap}</div>')

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    output = {
        "ranked_molecules": scored,
        "similarity_matrix": similarity_matrix,
        "similarity_labels": valid_names,
        "funnel": funnel_stages,
        "target_profile": profile,
    }
    return json.dumps(output)

def _parse_screening_json(screening_json: str) -> dict:
    """Parse screening JSON from screen_candidates, with safe defaults.

    The agent must pass the exact tool return value. Partial or hand-built JSON
    is tolerated for optional similarity fields only.
    """
    screening = json.loads(screening_json)
    if "ranked_molecules" not in screening:
        raise ValueError(
            "screening_json must be the exact JSON string returned by "
            "screen_candidates (missing 'ranked_molecules'). Do not construct, "
            "truncate, or summarize tool output."
        )
    screening.setdefault("similarity_matrix", [])
    screening.setdefault("similarity_labels", [])
    return screening

# ------------------------------------------------------------------
# Task 4: Generate final comprehensive report
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def generate_report(
    molecule_dir: flyte.io.Dir,
    properties_json: str,
    screening_json: str,
) -> str:
    """Generate a comprehensive drug screening report.

    Produces an executive summary, top candidate spotlight cards, property
    distributions, chemical diversity analysis, and final recommendation.

    Args:
        molecule_dir: Directory from load_molecules.
        properties_json: JSON from compute_properties.
        screening_json: Exact verbatim JSON string returned by screen_candidates
            (must include ranked_molecules, similarity_matrix, similarity_labels).
            Do not construct or summarize this payload yourself.

    Returns:
        JSON summary with total_screened, lipinski_passes, all_criteria_met,
        top_candidate, top_score, and top_3 ranked molecules.
    """
    from rdkit import Chem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Generating Final Report...</h2>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    screening = _parse_screening_json(screening_json)
    ranked = screening["ranked_molecules"]
    sim_matrix = screening["similarity_matrix"]
    sim_labels = screening["similarity_labels"]

    total = props["total"]
    lipinski_pass = props["lipinski_pass_count"]
    all_criteria = sum(1 for m in ranked if m["all_criteria_met"])
    top = ranked[0] if ranked else None

    html_parts = []

    # --- Executive Summary ---
    html_parts.append("<h2>Drug Molecule Screening Report</h2>")
    top_name = top["name"] if top else "N/A"
    top_score = f'{top["screening_score"]:.3f}' if top else "N/A"
    html_parts.append(
        f'<div class="card">'
        f'<h3 style="margin-top:0;color:#0e4f6e;">Executive Summary</h3>'
        f'<p style="font-size:1.05em;">'
        f'<strong>{total}</strong> molecules were screened against the target drug profile. '
        f'<strong>{lipinski_pass}</strong> passed Lipinski\'s Rule of Five, and '
        f'<strong>{all_criteria}</strong> met all screening criteria. '
        f'The top candidate is <strong style="color:#0891b2;">{top_name}</strong> '
        f'with a screening score of <strong>{top_score}</strong>.</p>'
        f'</div>'
    )

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Molecules Screened"),
        (str(lipinski_pass), "Lipinski Passes"),
        (str(all_criteria), "All Criteria Met"),
        (top_score, "Top Score"),
        (f'{props["avg_mw"]:.0f} Da', "Avg. Molecular Weight"),
        (f'{props["avg_logp"]:.2f}', "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # --- Top 3 Candidate Spotlights ---
    html_parts.append("<h2>Top Candidate Spotlights</h2>")

    for rank, m in enumerate(ranked[:3], 1):
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(300, 300)) if mol else ""

        medal = ["gold", "silver", "#cd7f32"][rank - 1]
        medal_emoji = ["1st", "2nd", "3rd"][rank - 1]

        lip_badges = ""
        for rule, key in [("MW", "mw_ok"), ("LogP", "logp_ok"),
                          ("HBD", "hbd_ok"), ("HBA", "hba_ok")]:
            ok = m["lipinski"].get(key, False)
            cls = "badge-success" if ok else "badge-danger"
            lip_badges += f'<span class="badge {cls}" style="margin:2px;">{rule}</span> '

        html_parts.append(
            f'<div class="molecule-card" style="display:flex;gap:20px;align-items:flex-start;flex-wrap:wrap;">'
            f'<div style="text-align:center;min-width:180px;">'
            f'<div style="font-size:1.6em;font-weight:800;color:{medal};">{medal_emoji}</div>'
            f'<img src="{img_uri}" style="width:200px;height:200px;object-fit:contain;border-radius:8px;'
            f'border:2px solid #a5f3fc;"/>'
            f'<div style="font-weight:700;font-size:1.1em;color:#0e4f6e;margin-top:8px;">{m["name"]}</div>'
            f'</div>'
            f'<div style="flex:1;min-width:280px;">'
            f'<table style="margin:0;">'
            f'<tr><td><strong>SMILES</strong></td><td style="font-family:monospace;font-size:0.8em;word-break:break-all;">{m["smiles"]}</td></tr>'
            f'<tr><td><strong>Screening Score</strong></td><td style="font-weight:700;color:#0891b2;font-size:1.1em;">{m["screening_score"]:.3f}</td></tr>'
            f'<tr><td><strong>Molecular Weight</strong></td><td>{m["mw"]:.1f} Da</td></tr>'
            f'<tr><td><strong>LogP</strong></td><td>{m["logp"]:.2f}</td></tr>'
            f'<tr><td><strong>H-Bond Donors</strong></td><td>{m["hbd"]}</td></tr>'
            f'<tr><td><strong>H-Bond Acceptors</strong></td><td>{m["hba"]}</td></tr>'
            f'<tr><td><strong>TPSA</strong></td><td>{m["tpsa"]:.1f} A&sup2;</td></tr>'
            f'<tr><td><strong>Rotatable Bonds</strong></td><td>{m["rotatable_bonds"]}</td></tr>'
            f'<tr><td><strong>QED</strong></td><td>{m["qed"]:.4f}</td></tr>'
            f'<tr><td><strong>Lipinski Compliance</strong></td><td>{lip_badges}</td></tr>'
            f'</table>'
            f'</div>'
            f'</div>'
        )

    # --- Property Distribution (box-plot style as bars with min/max/median) ---
    html_parts.append("<h2>Property Distributions</h2>")

    prop_keys = [("mw", "Molecular Weight (Da)"), ("logp", "LogP"),
                 ("tpsa", "TPSA"), ("qed", "QED Score")]
    for key, label in prop_keys:
        vals = sorted([m[key] for m in ranked])
        n = len(vals)
        if n == 0:
            continue
        v_min = vals[0]
        v_max = vals[-1]
        median = vals[n // 2] if n % 2 == 1 else (vals[n // 2 - 1] + vals[n // 2]) / 2
        q1 = vals[n // 4] if n >= 4 else v_min
        q3 = vals[3 * n // 4] if n >= 4 else v_max

        # Simple horizontal box-plot as SVG
        box_w = 500
        box_h = 50
        margin_l = 10
        v_range = v_max - v_min or 1

        def sx(v):
            return margin_l + ((v - v_min) / v_range) * (box_w - 2 * margin_l)

        box_svg = (
            f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {box_w} {box_h}" '
            f'style="width:100%;max-width:{box_w}px;height:auto;">'
            f'<rect width="{box_w}" height="{box_h}" fill="#fff" rx="4"/>'
            # Whisker line
            f'<line x1="{sx(v_min):.1f}" y1="25" x2="{sx(v_max):.1f}" y2="25" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Min whisker
            f'<line x1="{sx(v_min):.1f}" y1="18" x2="{sx(v_min):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Max whisker
            f'<line x1="{sx(v_max):.1f}" y1="18" x2="{sx(v_max):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # IQR box
            f'<rect x="{sx(q1):.1f}" y="14" width="{sx(q3) - sx(q1):.1f}" height="22" '
            f'fill="#a5f3fc" stroke="#0891b2" stroke-width="1.5" rx="3"/>'
            # Median line
            f'<line x1="{sx(median):.1f}" y1="12" x2="{sx(median):.1f}" y2="38" '
            f'stroke="#0e4f6e" stroke-width="2"/>'
            # Labels
            f'<text x="{sx(v_min):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_min:.1f}</text>'
            f'<text x="{sx(median):.1f}" y="10" text-anchor="middle" font-size="9" fill="#0e4f6e" font-weight="600">{median:.1f}</text>'
            f'<text x="{sx(v_max):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_max:.1f}</text>'
            f'</svg>'
        )
        html_parts.append(
            f'<div style="margin:8px 0;"><strong style="color:#155e75;">{label}</strong>'
            f'<div class="chart-container" style="padding:8px;">{box_svg}</div></div>'
        )

    # --- Chemical Diversity ---
    html_parts.append("<h2>Chemical Diversity Analysis</h2>")

    if sim_matrix and len(sim_matrix) > 1:
        # Compute average pairwise similarity (off-diagonal)
        n_mols = len(sim_matrix)
        off_diag = []
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                off_diag.append(sim_matrix[i][j])

        avg_sim = sum(off_diag) / len(off_diag) if off_diag else 0
        max_sim = max(off_diag) if off_diag else 0
        min_sim = min(off_diag) if off_diag else 0

        # Find most similar pair
        best_i, best_j = 0, 1
        best_val = 0
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                if sim_matrix[i][j] > best_val:
                    best_val = sim_matrix[i][j]
                    best_i, best_j = i, j

        html_parts.append('<div class="stat-grid">')
        html_parts.append(
            f'<div class="stat"><div class="value">{avg_sim:.3f}</div>'
            f'<div class="label">Avg. Pairwise Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{min_sim:.3f}</div>'
            f'<div class="label">Min Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{max_sim:.3f}</div>'
            f'<div class="label">Max Similarity</div></div>'
        )
        html_parts.append("</div>")

        diversity_text = "highly diverse" if avg_sim < 0.3 else "moderately diverse" if avg_sim < 0.5 else "relatively similar"
        html_parts.append(
            f'<div class="note">'
            f'The library is <strong>{diversity_text}</strong> (avg. Tanimoto = {avg_sim:.3f}). '
            f'The most similar pair is <strong>{sim_labels[best_i]}</strong> and '
            f'<strong>{sim_labels[best_j]}</strong> (similarity = {best_val:.3f}).</div>'
        )

    # --- Recommendation ---
    html_parts.append("<h2>Recommendation</h2>")
    if top:
        html_parts.append(
            f'<div class="card">'
            f'<h3 style="margin-top:0;color:#0891b2;">Top Candidate: {top["name"]}</h3>'
            f'<p>Based on the virtual screening analysis, <strong>{top["name"]}</strong> '
            f'achieved the highest composite screening score of <strong>{top["screening_score"]:.3f}</strong>. '
        )

        reasons = []
        if top["lipinski_pass"]:
            reasons.append("full Lipinski Rule of Five compliance")
        if top["qed"] > 0.5:
            reasons.append(f"high drug-likeness (QED = {top['qed']:.3f})")
        if top.get("all_criteria_met"):
            reasons.append("all target profile criteria met")
        if top["mw"] <= 500:
            reasons.append(f"favorable molecular weight ({top['mw']:.1f} Da)")

        if reasons:
            html_parts.append(
                f'This candidate stands out due to: {", ".join(reasons)}.</p>'
            )
        else:
            html_parts.append("</p>")

        # Runner-up mentions
        if len(ranked) >= 2:
            html_parts.append(
                f'<p style="font-size:0.9em;color:#6c757d;">Runner-up candidates: '
            )
            runners = []
            for m in ranked[1:4]:
                runners.append(f'{m["name"]} (score: {m["screening_score"]:.3f})')
            html_parts.append(", ".join(runners) + ".</p>")

        html_parts.append("</div>")

    # Final note
    html_parts.append(
        '<div class="note">'
        "This is a virtual screening analysis. All candidates should undergo "
        "further computational validation (molecular dynamics, docking) and "
        "experimental testing before advancing to clinical trials.</div>"
    )

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # JSON summary
    summary = {
        "total_screened": total,
        "lipinski_passes": lipinski_pass,
        "all_criteria_met": all_criteria,
        "top_candidate": top["name"] if top else None,
        "top_score": top["screening_score"] if top else None,
        "top_3": [
            {"name": m["name"], "score": m["screening_score"]}
            for m in ranked[:3]
        ],
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Agent
# ------------------------------------------------------------------

# {{docs-fragment agent}}
SCREENING_AGENT_INSTRUCTIONS = """\
You are a medicinal chemistry screening strategist. You orchestrate a virtual \
screening pipeline using durable Flyte tools. You NEVER invent molecular \
properties — only RDKit tools compute them.

Workflow:
1. If target_profile is not provided in the user message, derive a JSON \
target_profile from the therapeutic brief. Valid keys: mw, logp, hbd, hba, tpsa \
(each [min, max]). Ground choices in oral bioavailability / kinase / CNS rules \
as appropriate to the brief.
2. First pass (always): load_molecules → compute_properties → \
screen_candidates → generate_report. Pass tool outputs between steps exactly \
(molecule_dir from load_molecules into compute_properties and generate_report; \
properties_json from compute_properties into screen_candidates and \
generate_report; screening_json must be the complete, unmodified string \
returned by screen_candidates — never rebuild or summarize JSON yourself).
3. Read the JSON summary returned by generate_report. Reflect:
   - If all_criteria_met == 0: relax exactly ONE profile bound by ~10–20% \
and re-run screen_candidates then generate_report only, reusing the same \
molecule_dir and properties_json from the first pass.
   - If all molecules pass but diversity is a stated goal: note high similarity \
in your summary; do not re-run unless brief asks for stricter filters.
   - Maximum ONE rescreen iteration.
4. Finish with plain text: top candidate, rationale tied to computed metrics \
from the tool JSON, funnel interpretation, and suggested next steps (docking, \
ADMET lab tests).

If the user supplies an explicit target_profile JSON, use it as-is.

Do NOT ask the user for SMILES or molecule lists when molecules_json is empty — \
the default library is loaded automatically.
"""

screening_agent = Agent(
    name="drug-screening-agent",
    instructions=SCREENING_AGENT_INSTRUCTIONS,
    model=MODEL,
    tools=[
        load_molecules,
        compute_properties,
        screen_candidates,
        generate_report,
    ],
    max_turns=12,
)
# {{/docs-fragment agent}}

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def pipeline(
    brief: str = "Screen the default drug library for orally bioavailable small molecules.",
    molecules_json: str = "",
    target_profile: str = "",
) -> str:
    """Agentic virtual drug molecule screening pipeline.

    A medicinal-chemistry agent interprets the screening brief, derives or
    applies a target profile, orchestrates the RDKit screening stages, and
    optionally re-screens when funnel results are too narrow.

    Args:
        brief: Natural-language therapeutic goal (e.g. oral kinase inhibitors,
            CNS-penetrant small molecules).
        molecules_json: JSON mapping molecule names to SMILES strings.
            Defaults to a curated library of ~15 well-known drugs.
        target_profile: Optional JSON with desired property ranges that
            overrides agent-derived criteria
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5]}).

    Returns:
        Agent summary with screening rationale and key results.
    """
    prompt_parts = [
        f"Screening brief: {brief}",
        'Use molecules_json="" for the built-in default library unless provided below.',
        "Compose the four stage tools in order: load_molecules → compute_properties "
        "→ screen_candidates → generate_report. Pass each tool's full return value "
        "verbatim to the next step (especially screening_json). Re-run "
        "screen_candidates and generate_report at most once if the funnel is too narrow.",
    ]
    if molecules_json.strip():
        prompt_parts.append(f"molecules_json: {molecules_json}")
    if target_profile.strip():
        prompt_parts.append(f"Use this target_profile exactly: {target_profile}")

    result = await screening_agent.run.aio("\n".join(prompt_parts))
    return result.summary or result.error or ""

# {{/docs-fragment pipeline}}

# ------------------------------------------------------------------
# Rescreen demo — tight profile + explicit rescreen instructions
# ------------------------------------------------------------------

# Initial profile is deliberately strict (narrow MW + low LogP cap) so
# all_criteria_met is typically 0 on the default library; the brief then
# forces a single rescreen with a widened LogP window.
RESCREEN_DEMO_TARGET_PROFILE = (
    '{"mw": [150, 200], "logp": [-0.5, 1.0], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_TARGET_PROFILE_RESCREEN = (
    '{"mw": [150, 200], "logp": [-0.5, 3.5], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_BRIEF = f"""\
Two-round agentic screening demo on the default library.

**Round 1 (strict profile):** load_molecules → compute_properties → \
screen_candidates → generate_report using the initial target_profile exactly.

**Round 2 (required — do not skip):** call screen_candidates then generate_report \
again, reusing the same molecule_dir and properties_json from round 1, with this \
relaxed target_profile (wider LogP window only): \
{RESCREEN_DEMO_TARGET_PROFILE_RESCREEN}

Pass every tool return value verbatim to the next step. After both rounds, \
summarize how the funnel and top candidates changed between round 1 and round 2."""

# {{docs-fragment rescreen_demo}}
@env.task(report=True)
async def rescreen_demo() -> str:
    """Example run with a two-round execution graph (rescreen).

    Round 1 uses a strict CNS-like profile; round 2 always re-runs
    screen_candidates and generate_report with a widened LogP window,
    reusing cached molecule_dir and properties_json.
    """
    return await pipeline(
        brief=RESCREEN_DEMO_BRIEF,
        target_profile=RESCREEN_DEMO_TARGET_PROFILE,
    )

# {{/docs-fragment rescreen_demo}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/drug_molecule_screening/drug_molecule_screening.py*

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/drug_molecule_screening):

```
cd v2/tutorials/drug_molecule_screening
uv run --script drug_molecule_screening.py
```

Pass a natural-language brief (the agent derives the target profile):

```
flyte run drug_molecule_screening.py pipeline \
  --brief "Find oral kinase inhibitor candidates under 400 Da with moderate LogP"
```

Or pass an explicit target profile to skip agent-derived criteria:

```
flyte run drug_molecule_screening.py pipeline \
  --target_profile '{"mw": [100, 400], "logp": [-0.5, 4.0]}'
```

### Two-round rescreen demo (complex execution graph)

The `rescreen_demo` task always runs two screening rounds: a strict first pass (`load_molecules` → `compute_properties` → `screen_candidates` → `generate_report`), then a second `screen_candidates` → `generate_report` with a widened LogP window reusing the same `molecule_dir` and `properties_json`. The Flyte UI shows six stage tasks instead of four.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.4",
#    "litellm",
#    "rdkit",
#    "numpy",
#    "scikit-learn",
#    "pillow",
# ]
# main = "pipeline"
# params = ""
# ///
"""Virtual drug molecule screening — compute properties, apply Lipinski filters, rank candidates."""

import base64
import io
import json
import logging
import math
import os
import tempfile

import flyte
import flyte.io
import flyte.report
from flyte.ai.agents import Agent, tool

MODEL = os.getenv("DRUG_SCREENING_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="drug-molecule-screening", pre=True).with_apt_packages(
    "libxrender1", "libxext6", "libexpat1",
)

env = flyte.TaskEnvironment(
    name="drug-molecule-screening",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Default molecule library — real SMILES for well-known drugs
# ------------------------------------------------------------------

DEFAULT_MOLECULES = {
    "Aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
    "Ibuprofen": "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
    "Caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
    "Penicillin G": "CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C",
    "Metformin": "CN(C)C(=N)NC(=N)N",
    "Paracetamol": "CC(=O)NC1=CC=C(C=C1)O",
    "Diazepam": "ClC1=CC2=C(C=C1)N(C(=O)CN=C2C3=CC=CC=C3)C",
    "Omeprazole": "CC1=CN=C(C(=C1OC)C)CS(=O)C2=NC3=CC=CC=C3N2",
    "Atorvastatin": "CC(C)C1=C(C(=C(N1CCC(CC(CC(=O)O)O)O)C2=CC=C(C=C2)F)C3=CC=CC=C3)C(=O)NC4=CC=CC=C4",
    "Methotrexate": "CN(CC1=CN=C2N=C(N=C(N)C2=N1)N)C3=CC=C(C=C3)C(=O)NC(CCC(=O)O)C(=O)O",
    "Doxorubicin": "CC1C(C(CC(O1)OC2CC(CC3=C2C(=C4C(=C3O)C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)CO)O)N)O",
    "Tamoxifen": "CCC(=C(C1=CC=CC=C1)C2=CC=C(C=C2)OCCN(C)C)C3=CC=CC=C3",
    "Lopinavir": "CC1=C(C(=CC=C1)C)OCC(=O)NC(CC2=CC=CC=C2)C(CC(CC3=CC=CC=C3)NC(=O)C(C(C)C)N4CCCNC4=O)O",
    "Remdesivir": "CCC(CC)COC(=O)C(C)NP(=O)(OCC1C(C(C(O1)C2=CC=C3N2N=CN=C3N)O)O)OC4=CC=CC=C4",
    "Erlotinib": "COCCOC1=CC2=C(C=C1OCCOC)C(=NC=N2)NC3=CC=CC(=C3)C#C",
}

# ------------------------------------------------------------------
# Report styling — pharma blue/cyan theme
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #0e4f6e; border-bottom: 2px solid #0891b2; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #155e75; margin-top: 20px; }
  .report .card { background: #ecfeff; border: 1px solid #a5f3fc; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #cffafe; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0e4f6e; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0e4f6e; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #cffafe; }
  .report tr:nth-child(even) { background: #ecfeff; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d1fae5; color: #065f46; }
  .report .badge-danger { background: #fee2e2; color: #991b1b; }
  .report .badge-info { background: #cffafe; color: #155e75; }
  .report .chart-container { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 16px 0; }
  .report .note { background: #ecfeff; border-left: 4px solid #0891b2; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .molecule-card { background: #fff; border: 1px solid #cffafe; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .molecule-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 12px; margin: 16px 0; }
  .report .funnel { text-align: center; margin: 24px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers
# ------------------------------------------------------------------

def _mol_to_data_uri(mol, size: tuple[int, int] = (300, 300)) -> str:
    """Convert an RDKit molecule to a PNG base64 data URI."""
    from rdkit.Chem import Draw

    img = Draw.MolToImage(mol, size=size)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 340,
    y_max_cap: float | None = None,
    horizontal: bool = False,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels.
        series: Dict mapping series name to list of values.
        title: Chart title.
        colors: Colors for each series.
        width/height: SVG dimensions.
        y_max_cap: Cap the y-axis at this value.
        horizontal: If True, draw horizontal bars.
        value_fmt: Format string for value labels.

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#a5f3fc", "#155e75"]
    colors = colors or default_colors

    if horizontal:
        return _make_horizontal_bar_chart(labels, series, title, colors, width, height, value_fmt)

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:{value_fmt}}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="9" fill="#1a1a2e">'
                f'{val:{value_fmt}}</text>'
            )
        # Truncate long labels
        disp_label = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 16}" '
            f'text-anchor="middle" font-size="10" fill="#6c757d" '
            f'transform="rotate(-35, {gx + n_series * bar_width / 2:.1f}, {mt + ch + 16})">'
            f'{disp_label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Legend
    if n_series > 1:
        lx = ml + cw - len(series) * 100
        for si, name in enumerate(series):
            color = colors[si % len(colors)]
            svg.append(
                f'<rect x="{lx + si * 100}" y="{mt + ch + 40}" width="12" '
                f'height="12" rx="2" fill="{color}"/>'
            )
            svg.append(
                f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 51}" font-size="11" '
                f'fill="#1a1a2e">{name}</text>'
            )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_horizontal_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 400,
    value_fmt: str = ".1f",
) -> str:
    """Generate an SVG horizontal bar chart (sorted)."""
    default_colors = ["#0891b2", "#0e4f6e", "#06d6a0"]
    colors = colors or default_colors

    n = len(labels)
    row_height = max(22, min(35, (height - 80) // max(n, 1)))
    actual_height = max(height, 80 + n * row_height)
    ml, mr, mt, mb = 120, 60, 40, 20
    cw = width - ml - mr
    ch = actual_height - mt - mb

    # Use first series
    first_key = list(series.keys())[0]
    vals = series[first_key]
    x_max = max(vals) * 1.15 if vals else 1

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {actual_height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{actual_height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    bar_h = row_height * 0.65
    for i, (label, val) in enumerate(zip(labels, vals)):
        y = mt + i * row_height
        bw = (val / x_max) * cw if x_max else 0
        color = colors[i % len(colors)]
        # Label
        disp = label if len(label) <= 14 else label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + bar_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#1a1a2e">{disp}</text>'
        )
        # Bar
        svg.append(
            f'<rect x="{ml}" y="{y:.1f}" width="{bw:.1f}" height="{bar_h:.1f}" '
            f'fill="{color}" rx="3"/>'
        )
        # Value
        svg.append(
            f'<text x="{ml + bw + 6:.1f}" y="{y + bar_h / 2 + 4:.1f}" '
            f'font-size="11" fill="#0e4f6e" font-weight="600">{val:{value_fmt}}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_heatmap(
    matrix: list[list[float]],
    row_labels: list[str],
    col_labels: list[str],
    title: str = "",
    color_scale: str = "cyan",
    width: int = 700,
    height: int = 500,
    value_fmt: str = ".2f",
) -> str:
    """Generate an SVG heatmap.

    Args:
        matrix: 2D list of values (rows x cols).
        row_labels: Labels for rows.
        col_labels: Labels for columns.
        title: Chart title.
        color_scale: Color scheme ("cyan", "red", "green").
        width/height: SVG dimensions.
        value_fmt: Format string for cell values.

    Returns:
        SVG string.
    """
    if not matrix or not matrix[0]:
        return ""

    n_rows = len(matrix)
    n_cols = len(matrix[0])

    ml, mr, mt, mb = 110, 20, 70, 20
    cw = width - ml - mr
    ch = height - mt - mb
    cell_w = cw / n_cols
    cell_h = ch / n_rows

    # Flatten to find range
    flat = [v for row in matrix for v in row]
    v_min = min(flat)
    v_max = max(flat)
    v_range = v_max - v_min or 1

    def color_for(v):
        t = (v - v_min) / v_range
        if color_scale == "cyan":
            # White to deep teal
            r = int(255 - t * (255 - 14))
            g = int(255 - t * (255 - 79))
            b = int(255 - t * (255 - 110))
        elif color_scale == "red":
            r = int(255 - t * 50)
            g = int(255 - t * 200)
            b = int(255 - t * 200)
        else:  # green
            r = int(255 - t * 200)
            g = int(255 - t * 50)
            b = int(255 - t * 200)
        return f"rgb({r},{g},{b})"

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Column labels (rotated)
    for ci, label in enumerate(col_labels):
        x = ml + ci * cell_w + cell_w / 2
        disp = label if len(label) <= 12 else label[:10] + ".."
        svg.append(
            f'<text x="{x:.1f}" y="{mt - 8}" text-anchor="end" font-size="10" '
            f'fill="#1a1a2e" transform="rotate(-45, {x:.1f}, {mt - 8})">{disp}</text>'
        )

    # Row labels + cells
    for ri, (row_label, row_vals) in enumerate(zip(row_labels, matrix)):
        y = mt + ri * cell_h
        disp = row_label if len(row_label) <= 14 else row_label[:12] + ".."
        svg.append(
            f'<text x="{ml - 8}" y="{y + cell_h / 2 + 4:.1f}" text-anchor="end" '
            f'font-size="10" fill="#1a1a2e">{disp}</text>'
        )
        for ci, val in enumerate(row_vals):
            x = ml + ci * cell_w
            fill = color_for(val)
            svg.append(
                f'<rect x="{x:.1f}" y="{y:.1f}" width="{cell_w:.1f}" '
                f'height="{cell_h:.1f}" fill="{fill}" stroke="#fff" stroke-width="1"/>'
            )
            # Text color: dark on light, light on dark
            t = (val - v_min) / v_range
            txt_color = "#fff" if t > 0.55 else "#1a1a2e"
            # Only show text if cells are large enough
            if cell_w > 30 and cell_h > 18:
                svg.append(
                    f'<text x="{x + cell_w / 2:.1f}" y="{y + cell_h / 2 + 4:.1f}" '
                    f'text-anchor="middle" font-size="9" fill="{txt_color}">'
                    f'{val:{value_fmt}}</text>'
                )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_scatter_plot(
    points: list[dict],
    x_label: str = "MW",
    y_label: str = "LogP",
    title: str = "",
    reference_lines: list[dict] | None = None,
    width: int = 700,
    height: int = 400,
) -> str:
    """Generate an SVG scatter plot.

    Args:
        points: List of dicts with "x", "y", "label" keys.
        x_label/y_label: Axis labels.
        title: Chart title.
        reference_lines: List of dicts with "axis" ("x"/"y"), "value", "label".
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not points:
        return ""

    ml, mr, mt, mb = 60, 30, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [p["x"] for p in points]
    y_vals = [p["y"] for p in points]
    x_min, x_max = min(x_vals) * 0.9, max(x_vals) * 1.1
    y_min, y_max = min(y_vals) - 1, max(y_vals) + 1

    # Extend ranges to include reference lines
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                x_max = max(x_max, rl["value"] * 1.1)
            else:
                y_max = max(y_max, rl["value"] * 1.1)

    x_range = x_max - x_min or 1
    y_range = y_max - y_min or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min) / y_range * ch

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid
    for i in range(6):
        y_tick = y_min + y_range * i / 5
        py = sy(y_tick)
        svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e0f2fe" stroke-width="1"/>'
        )
        svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.1f}</text>'
        )

    for i in range(6):
        x_tick = x_min + x_range * i / 5
        px = sx(x_tick)
        svg.append(
            f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
            f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
        )

    # Axes
    svg.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )
    svg.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#94a3b8" stroke-width="1.5"/>'
    )

    # Reference lines (Lipinski boundaries)
    if reference_lines:
        for rl in reference_lines:
            if rl["axis"] == "x":
                px = sx(rl["value"])
                svg.append(
                    f'<line x1="{px:.1f}" y1="{mt}" x2="{px:.1f}" y2="{mt + ch}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{px + 4:.1f}" y="{mt + 14}" font-size="10" '
                    f'fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )
            else:
                py = sy(rl["value"])
                svg.append(
                    f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
                    f'stroke="#ef4444" stroke-width="1.5" stroke-dasharray="6,4"/>'
                )
                svg.append(
                    f'<text x="{ml + cw - 4:.1f}" y="{py - 6:.1f}" text-anchor="end" '
                    f'font-size="10" fill="#ef4444" font-weight="600">{rl["label"]}</text>'
                )

    # Drug-like zone shading (MW<=500 and LogP<=5 quadrant)
    if reference_lines:
        mw_line = next((rl for rl in reference_lines if rl["axis"] == "x"), None)
        logp_line = next((rl for rl in reference_lines if rl["axis"] == "y"), None)
        if mw_line and logp_line:
            zx1 = sx(x_min)
            zx2 = sx(min(mw_line["value"], x_max))
            zy1 = sy(min(logp_line["value"], y_max))
            zy2 = sy(y_min)
            svg.append(
                f'<rect x="{zx1:.1f}" y="{zy1:.1f}" '
                f'width="{zx2 - zx1:.1f}" height="{zy2 - zy1:.1f}" '
                f'fill="#0891b2" opacity="0.06" rx="4"/>'
            )
            svg.append(
                f'<text x="{zx1 + 8:.1f}" y="{zy2 - 8:.1f}" font-size="11" '
                f'fill="#0891b2" font-weight="600" opacity="0.6">Drug-like Zone</text>'
            )

    # Points
    point_colors = ["#0891b2", "#0e4f6e", "#06d6a0", "#155e75", "#0284c7",
                    "#059669", "#0d9488", "#0369a1", "#047857", "#115e59",
                    "#0c4a6e", "#064e3b", "#1e3a5f", "#134e4a", "#075985"]
    for i, pt in enumerate(points):
        px, py = sx(pt["x"]), sy(pt["y"])
        color = point_colors[i % len(point_colors)]
        svg.append(
            f'<circle cx="{px:.1f}" cy="{py:.1f}" r="5" fill="{color}" '
            f'stroke="#fff" stroke-width="1.5" opacity="0.85"/>'
        )
        # Label offset to avoid overlap
        offset_x = 8
        offset_y = -8 if i % 2 == 0 else 14
        label = pt["label"] if len(pt["label"]) <= 12 else pt["label"][:10] + ".."
        svg.append(
            f'<text x="{px + offset_x:.1f}" y="{py + offset_y:.1f}" '
            f'font-size="9" fill="#1a1a2e">{label}</text>'
        )

    # Title
    if title:
        svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#0e4f6e">{title}</text>'
        )

    # Axis labels
    if x_label:
        svg.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        svg.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

def _make_funnel(
    stages: list[dict],
    title: str = "",
    width: int = 600,
    height: int = 400,
) -> str:
    """Generate an SVG funnel visualization.

    Args:
        stages: List of dicts with "label", "count", "total" keys.
        title: Chart title.
        width/height: SVG dimensions.

    Returns:
        SVG string.
    """
    if not stages:
        return ""

    n = len(stages)
    mt = 50
    mb = 20
    available_h = height - mt - mb
    stage_h = available_h / n
    cx = width / 2

    # Color gradient from light cyan to deep teal
    colors = []
    for i in range(n):
        t = i / max(n - 1, 1)
        r = int(207 - t * (207 - 14))
        g = int(250 - t * (250 - 79))
        b = int(254 - t * (254 - 110))
        colors.append(f"rgb({r},{g},{b})")

    svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    if title:
        svg.append(
            f'<text x="{cx}" y="28" text-anchor="middle" '
            f'font-size="16" font-weight="700" fill="#0e4f6e">{title}</text>'
        )

    max_count = stages[0]["count"] if stages else 1
    max_width = width * 0.75

    for i, stage in enumerate(stages):
        y_top = mt + i * stage_h
        y_bot = y_top + stage_h

        # Width proportional to count
        w_top = max_width * (stage["count"] / max_count) if i == 0 else prev_w_bot
        if i < n - 1:
            w_bot = max_width * (stages[i + 1]["count"] / max_count)
        else:
            w_bot = max_width * (stage["count"] / max_count) * 0.7

        prev_w_bot = w_bot

        # Trapezoid
        x1_top = cx - w_top / 2
        x2_top = cx + w_top / 2
        x1_bot = cx - w_bot / 2
        x2_bot = cx + w_bot / 2

        svg.append(
            f'<polygon points="{x1_top:.1f},{y_top:.1f} {x2_top:.1f},{y_top:.1f} '
            f'{x2_bot:.1f},{y_bot:.1f} {x1_bot:.1f},{y_bot:.1f}" '
            f'fill="{colors[i]}" stroke="#fff" stroke-width="2"/>'
        )

        # Text: dark on light, white on dark
        t = i / max(n - 1, 1)
        txt_color = "#0e4f6e" if t < 0.5 else "#fff"
        y_mid = (y_top + y_bot) / 2

        svg.append(
            f'<text x="{cx}" y="{y_mid - 4:.1f}" text-anchor="middle" '
            f'font-size="13" font-weight="600" fill="{txt_color}">{stage["label"]}</text>'
        )
        svg.append(
            f'<text x="{cx}" y="{y_mid + 14:.1f}" text-anchor="middle" '
            f'font-size="12" fill="{txt_color}" opacity="0.85">'
            f'{stage["count"]} / {stage["total"]}</text>'
        )

    svg.append("</svg>")
    return "\n".join(svg)

# ------------------------------------------------------------------
# Task 1: Load and validate molecules
# ------------------------------------------------------------------

@tool
@env.task(cache="auto")
async def load_molecules(
    molecules_json: str = "",
) -> flyte.io.Dir:
    """Parse SMILES strings, validate with RDKit, generate 2D depictions.

    Args:
        molecules_json: JSON string mapping molecule names to SMILES.
            Defaults to a curated library of ~15 well-known drugs.

    Returns:
        flyte.io.Dir containing molecule data (JSON + PNG depictions).
        Pass this directory to compute_properties and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Draw

    if molecules_json.strip():
        molecules = json.loads(molecules_json)
    else:
        molecules = DEFAULT_MOLECULES

    out_dir = tempfile.mkdtemp(prefix="mol_library_")
    results = []
    valid_count = 0
    invalid_count = 0

    log.info(f"Parsing {len(molecules)} molecules...")

    for name, smiles in molecules.items():
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            log.warning(f"  [INVALID] {name}: {smiles}")
            invalid_count += 1
            continue

        valid_count += 1

        # Generate 2D depiction as PNG
        img = Draw.MolToImage(mol, size=(300, 300))
        img_path = os.path.join(out_dir, f"{name.replace(' ', '_')}.png")
        img.save(img_path)

        results.append({
            "name": name,
            "smiles": smiles,
            "valid": True,
            "image_file": os.path.basename(img_path),
        })

    # Save molecule manifest
    manifest = {
        "total": len(molecules),
        "valid": valid_count,
        "invalid": invalid_count,
        "molecules": results,
    }
    manifest_path = os.path.join(out_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    log.info(f"Loaded {valid_count} valid molecules ({invalid_count} invalid)")

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Task 2: Compute physicochemical properties
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def compute_properties(
    molecule_dir: flyte.io.Dir,
) -> str:
    """Compute drug-likeness properties for all molecules.

    Computes MW, LogP, HBD, HBA, TPSA, rotatable bonds, formal charge,
    ring count, QED, and Lipinski Rule of Five compliance.

    Args:
        molecule_dir: Directory from load_molecules.

    Returns:
        JSON string with all computed properties. Pass to screen_candidates
        and generate_report.
    """
    from rdkit import Chem
    from rdkit.Chem import Descriptors, Lipinski
    from rdkit.Chem.QED import qed

    # --- Loading report ---
    await flyte.report.replace.aio(
        _wrap_report("<h2>Computing Molecular Properties...</h2>"
                      "<p>Analyzing physicochemical descriptors for all molecules.</p>"),
        do_flush=True,
    )

    mol_dir = await molecule_dir.download()
    with open(os.path.join(mol_dir, "manifest.json")) as f:
        manifest = json.load(f)

    molecules_data = []
    lipinski_pass = 0

    for mol_info in manifest["molecules"]:
        mol = Chem.MolFromSmiles(mol_info["smiles"])
        if mol is None:
            continue

        mw = Descriptors.MolWt(mol)
        logp = Descriptors.MolLogP(mol)
        hbd = Lipinski.NumHDonors(mol)
        hba = Lipinski.NumHAcceptors(mol)
        tpsa = Descriptors.TPSA(mol)
        rotatable = Lipinski.NumRotatableBonds(mol)
        formal_charge = Chem.GetFormalCharge(mol)
        num_rings = Lipinski.RingCount(mol)
        qed_score = qed(mol)

        # Lipinski Rule of Five
        lipinski = {
            "mw_ok": mw <= 500,
            "logp_ok": logp <= 5,
            "hbd_ok": hbd <= 5,
            "hba_ok": hba <= 10,
        }
        lipinski_all = all(lipinski.values())
        if lipinski_all:
            lipinski_pass += 1

        # Read image for data URI
        img_path = os.path.join(mol_dir, mol_info["image_file"])
        data_uri = ""
        if os.path.exists(img_path):
            with open(img_path, "rb") as img_f:
                b64 = base64.b64encode(img_f.read()).decode()
                data_uri = f"data:image/png;base64,{b64}"

        molecules_data.append({
            "name": mol_info["name"],
            "smiles": mol_info["smiles"],
            "mw": round(mw, 2),
            "logp": round(logp, 2),
            "hbd": hbd,
            "hba": hba,
            "tpsa": round(tpsa, 2),
            "rotatable_bonds": rotatable,
            "formal_charge": formal_charge,
            "num_rings": num_rings,
            "qed": round(qed_score, 4),
            "lipinski": lipinski,
            "lipinski_pass": lipinski_all,
            "image_data_uri": data_uri,
        })

    total = len(molecules_data)
    avg_mw = sum(m["mw"] for m in molecules_data) / total if total else 0
    avg_logp = sum(m["logp"] for m in molecules_data) / total if total else 0
    lipinski_rate = lipinski_pass / total * 100 if total else 0

    # ---- Build report ----
    html_parts = []

    # Header
    html_parts.append("<h2>Molecular Properties Analysis</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Total Molecules"),
        (f"{lipinski_rate:.0f}%", "Lipinski Pass Rate"),
        (f"{avg_mw:.1f}", "Avg. MW (Da)"),
        (f"{avg_logp:.2f}", "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Molecule gallery
    html_parts.append("<h3>Molecule Library</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in molecules_data:
        if m["image_data_uri"]:
            badge_class = "badge-success" if m["lipinski_pass"] else "badge-danger"
            badge_text = "Lipinski Pass" if m["lipinski_pass"] else "Lipinski Fail"
            html_parts.append(
                f'<div class="molecule-card" style="text-align:center;">'
                f'<img src="{m["image_data_uri"]}" style="width:160px;height:160px;object-fit:contain;"/>'
                f'<div style="font-weight:600;margin-top:6px;color:#0e4f6e;">{m["name"]}</div>'
                f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f}</div>'
                f'<div><span class="badge {badge_class}">{badge_text}</span></div>'
                f'</div>'
            )
    html_parts.append("</div>")

    # MW bar chart (horizontal, sorted)
    sorted_by_mw = sorted(molecules_data, key=lambda m: m["mw"], reverse=True)
    mw_labels = [m["name"] for m in sorted_by_mw]
    mw_vals = [m["mw"] for m in sorted_by_mw]
    mw_chart = _make_bar_chart(
        mw_labels, {"MW (Da)": mw_vals},
        title="Molecular Weight Distribution",
        horizontal=True,
        width=700, height=max(300, len(mw_labels) * 30 + 80),
        value_fmt=".1f",
    )
    html_parts.append("<h3>Molecular Weight</h3>")
    html_parts.append(f'<div class="chart-container">{mw_chart}</div>')

    # LogP vs MW scatter plot
    scatter_points = [
        {"x": m["mw"], "y": m["logp"], "label": m["name"]}
        for m in molecules_data
    ]
    scatter_chart = _make_scatter_plot(
        scatter_points,
        x_label="Molecular Weight (Da)",
        y_label="LogP",
        title="LogP vs. Molecular Weight (Lipinski Boundaries)",
        reference_lines=[
            {"axis": "x", "value": 500, "label": "MW = 500"},
            {"axis": "y", "value": 5, "label": "LogP = 5"},
        ],
        width=700,
        height=420,
    )
    html_parts.append("<h3>Lipinski Space</h3>")
    html_parts.append(f'<div class="chart-container">{scatter_chart}</div>')

    # Property heatmap (molecules x properties)
    prop_names = ["MW", "LogP", "HBD", "HBA", "TPSA", "Rot. Bonds"]
    # Normalize each property to 0-1 for heatmap
    raw_matrix = []
    for m in molecules_data:
        raw_matrix.append([m["mw"], m["logp"], m["hbd"], m["hba"], m["tpsa"], m["rotatable_bonds"]])

    # Normalize per column
    n_props = len(prop_names)
    col_min = [min(row[c] for row in raw_matrix) for c in range(n_props)]
    col_max = [max(row[c] for row in raw_matrix) for c in range(n_props)]
    norm_matrix = []
    for row in raw_matrix:
        norm_row = []
        for c in range(n_props):
            rng = col_max[c] - col_min[c]
            norm_row.append((row[c] - col_min[c]) / rng if rng else 0.5)
        norm_matrix.append(norm_row)

    heatmap_labels = [m["name"] for m in molecules_data]
    heatmap = _make_heatmap(
        norm_matrix, heatmap_labels, prop_names,
        title="Normalized Property Heatmap",
        color_scale="cyan",
        width=700,
        height=max(400, len(heatmap_labels) * 28 + 100),
    )
    html_parts.append("<h3>Property Heatmap</h3>")
    html_parts.append(f'<div class="chart-container">{heatmap}</div>')

    # Lipinski compliance table
    html_parts.append("<h3>Lipinski Rule of Five Compliance</h3>")
    html_parts.append("<table><tr><th>Molecule</th><th>MW &le; 500</th>"
                      "<th>LogP &le; 5</th><th>HBD &le; 5</th>"
                      "<th>HBA &le; 10</th><th>Overall</th></tr>")
    for m in molecules_data:
        lip = m["lipinski"]

        def _badge(ok):
            if ok:
                return '<span class="badge badge-success">Pass</span>'
            return '<span class="badge badge-danger">Fail</span>'

        overall_badge = _badge(m["lipinski_pass"])
        html_parts.append(
            f'<tr><td><strong>{m["name"]}</strong></td>'
            f'<td>{_badge(lip["mw_ok"])}</td>'
            f'<td>{_badge(lip["logp_ok"])}</td>'
            f'<td>{_badge(lip["hbd_ok"])}</td>'
            f'<td>{_badge(lip["hba_ok"])}</td>'
            f'<td>{overall_badge}</td></tr>'
        )
    html_parts.append("</table>")

    # QED bar chart
    sorted_by_qed = sorted(molecules_data, key=lambda m: m["qed"], reverse=True)
    qed_labels = [m["name"] for m in sorted_by_qed]
    qed_vals = [m["qed"] for m in sorted_by_qed]
    qed_chart = _make_bar_chart(
        qed_labels, {"QED Score": qed_vals},
        title="Drug-likeness (QED Score)",
        horizontal=True,
        width=700, height=max(300, len(qed_labels) * 30 + 80),
        value_fmt=".3f",
        colors=["#06d6a0"],
    )
    html_parts.append("<h3>Drug-likeness (QED)</h3>")
    html_parts.append(f'<div class="chart-container">{qed_chart}</div>')

    # Flush full report
    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # Return properties as JSON (strip image data URIs to reduce size)
    output = {
        "total": total,
        "lipinski_pass_count": lipinski_pass,
        "lipinski_pass_rate": round(lipinski_rate, 2),
        "avg_mw": round(avg_mw, 2),
        "avg_logp": round(avg_logp, 2),
        "molecules": [
            {k: v for k, v in m.items() if k != "image_data_uri"}
            for m in molecules_data
        ],
    }
    return json.dumps(output)

# ------------------------------------------------------------------
# Task 3: Screen candidates against target profile
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def screen_candidates(
    properties_json: str,
    target_profile: str = "",
) -> str:
    """Screen molecules against a target drug profile and rank candidates.

    Scores each molecule on how well it matches the target profile, computes
    pairwise Tanimoto similarity, and produces a ranked list.

    Args:
        properties_json: JSON from compute_properties.
        target_profile: JSON string with desired property ranges
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5.0]}).

    Returns:
        JSON string with ranked_molecules, similarity_matrix, similarity_labels,
        funnel, and target_profile. Pass the full return value verbatim to
        generate_report along with molecule_dir and properties_json.
    """
    from rdkit import Chem, DataStructs
    from rdkit.Chem import AllChem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Screening Candidates...</h2>"
                      "<p>Evaluating molecules against the target drug profile.</p>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    molecules = props["molecules"]

    # Default target profile
    if target_profile.strip():
        profile = json.loads(target_profile)
    else:
        profile = {
            "mw": [150, 500],
            "logp": [-0.5, 5.0],
            "hbd": [0, 5],
            "hba": [0, 10],
            "tpsa": [20, 140],
        }

    # --- Screening ---
    funnel_total = len(molecules)
    pass_mw = 0
    pass_logp = 0
    pass_lipinski = 0
    final_candidates = 0

    scored = []
    for m in molecules:
        score = 0
        max_score = 0
        criteria = {}

        # Check each profile criterion
        checks = [
            ("mw", m["mw"]),
            ("logp", m["logp"]),
            ("hbd", m["hbd"]),
            ("hba", m["hba"]),
            ("tpsa", m["tpsa"]),
        ]

        for key, val in checks:
            if key in profile:
                lo, hi = profile[key]
                max_score += 1
                in_range = lo <= val <= hi
                criteria[key] = in_range
                if in_range:
                    score += 1
                    # Bonus: closer to midpoint = higher score
                    mid = (lo + hi) / 2
                    rng = (hi - lo) / 2
                    dist = abs(val - mid) / rng if rng else 0
                    score += max(0, 0.5 * (1 - dist))

        # QED bonus
        score += m["qed"] * 2
        max_score += 2

        # Lipinski bonus
        if m["lipinski_pass"]:
            score += 1
        max_score += 1

        normalized_score = score / max_score if max_score else 0

        # Funnel tracking — cascading filter (each stage requires passing the previous)
        mw_ok = criteria.get("mw", True)
        logp_ok = criteria.get("logp", True)
        if mw_ok:
            pass_mw += 1
            if logp_ok:
                pass_logp += 1
                if m["lipinski_pass"]:
                    pass_lipinski += 1
                    if all(criteria.values()):
                        final_candidates += 1

        scored.append({
            **m,
            "screening_score": round(normalized_score, 4),
            "criteria_met": criteria,
            "all_criteria_met": all(criteria.values()),
        })

    # Sort by score descending
    scored.sort(key=lambda m: m["screening_score"], reverse=True)

    # --- Tanimoto similarity matrix ---
    fps = []
    valid_names = []
    for m in scored:
        mol = Chem.MolFromSmiles(m["smiles"])
        if mol:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
            fps.append(fp)
            valid_names.append(m["name"])

    similarity_matrix = []
    for i in range(len(fps)):
        row = []
        for j in range(len(fps)):
            sim = DataStructs.TanimotoSimilarity(fps[i], fps[j])
            row.append(round(sim, 3))
        similarity_matrix.append(row)

    # ---- Build report ----
    html_parts = []
    html_parts.append("<h2>Candidate Screening Results</h2>")

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(funnel_total), "Total Screened"),
        (str(pass_lipinski), "Lipinski Passes"),
        (str(final_candidates), "All Criteria Met"),
        (f"{scored[0]['screening_score']:.3f}" if scored else "N/A", "Top Score"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # Screening funnel
    funnel_stages = [
        {"label": "Total Molecules", "count": funnel_total, "total": funnel_total},
        {"label": "Pass MW Filter", "count": pass_mw, "total": funnel_total},
        {"label": "Pass LogP Filter", "count": pass_logp, "total": funnel_total},
        {"label": "Lipinski Compliant", "count": pass_lipinski, "total": funnel_total},
        {"label": "All Criteria Met", "count": final_candidates, "total": funnel_total},
    ]
    funnel_svg = _make_funnel(
        funnel_stages,
        title="Screening Funnel",
        width=600,
        height=380,
    )
    html_parts.append("<h3>Screening Funnel</h3>")
    html_parts.append(f'<div class="chart-container" style="text-align:center;">{funnel_svg}</div>')

    # Ranked candidates table
    html_parts.append("<h3>Ranked Candidates</h3>")
    html_parts.append(
        "<table><tr><th>Rank</th><th>Molecule</th><th>Score</th>"
        "<th>MW</th><th>LogP</th><th>QED</th><th>Lipinski</th><th>All Criteria</th></tr>"
    )
    for rank, m in enumerate(scored, 1):
        lip_badge = ('<span class="badge badge-success">Pass</span>'
                     if m["lipinski_pass"]
                     else '<span class="badge badge-danger">Fail</span>')
        crit_badge = ('<span class="badge badge-success">Pass</span>'
                      if m["all_criteria_met"]
                      else '<span class="badge badge-danger">Fail</span>')
        # Highlight top 3
        row_style = ' style="background:#ecfeff;font-weight:600;"' if rank <= 3 else ""
        html_parts.append(
            f"<tr{row_style}><td>{rank}</td><td>{m['name']}</td>"
            f"<td>{m['screening_score']:.3f}</td>"
            f"<td>{m['mw']:.1f}</td><td>{m['logp']:.2f}</td>"
            f"<td>{m['qed']:.3f}</td><td>{lip_badge}</td><td>{crit_badge}</td></tr>"
        )
    html_parts.append("</table>")

    # Top 5 candidate cards with structures
    html_parts.append("<h3>Top 5 Candidates</h3>")
    html_parts.append('<div class="molecule-grid">')
    for m in scored[:5]:
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(250, 250)) if mol else ""
        badge_class = "badge-success" if m["all_criteria_met"] else "badge-info"
        badge_text = "All Criteria Met" if m["all_criteria_met"] else "Partial Match"
        html_parts.append(
            f'<div class="molecule-card" style="text-align:center;">'
            f'<img src="{img_uri}" style="width:140px;height:140px;object-fit:contain;"/>'
            f'<div style="font-weight:700;margin-top:6px;color:#0e4f6e;font-size:1.05em;">{m["name"]}</div>'
            f'<div style="font-size:0.85em;color:#155e75;margin:4px 0;">Score: {m["screening_score"]:.3f}</div>'
            f'<div style="font-size:0.8em;color:#6c757d;">MW: {m["mw"]:.1f} | LogP: {m["logp"]:.2f} | QED: {m["qed"]:.3f}</div>'
            f'<div style="margin-top:4px;"><span class="badge {badge_class}">{badge_text}</span></div>'
            f'</div>'
        )
    html_parts.append("</div>")

    # Tanimoto similarity heatmap
    if similarity_matrix:
        sim_heatmap = _make_heatmap(
            similarity_matrix, valid_names, valid_names,
            title="Pairwise Tanimoto Similarity (Morgan Fingerprints)",
            color_scale="cyan",
            width=700,
            height=max(500, len(valid_names) * 32 + 100),
        )
        html_parts.append("<h3>Chemical Similarity</h3>")
        html_parts.append(f'<div class="chart-container">{sim_heatmap}</div>')

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    output = {
        "ranked_molecules": scored,
        "similarity_matrix": similarity_matrix,
        "similarity_labels": valid_names,
        "funnel": funnel_stages,
        "target_profile": profile,
    }
    return json.dumps(output)

def _parse_screening_json(screening_json: str) -> dict:
    """Parse screening JSON from screen_candidates, with safe defaults.

    The agent must pass the exact tool return value. Partial or hand-built JSON
    is tolerated for optional similarity fields only.
    """
    screening = json.loads(screening_json)
    if "ranked_molecules" not in screening:
        raise ValueError(
            "screening_json must be the exact JSON string returned by "
            "screen_candidates (missing 'ranked_molecules'). Do not construct, "
            "truncate, or summarize tool output."
        )
    screening.setdefault("similarity_matrix", [])
    screening.setdefault("similarity_labels", [])
    return screening

# ------------------------------------------------------------------
# Task 4: Generate final comprehensive report
# ------------------------------------------------------------------

@tool
@env.task(report=True)
async def generate_report(
    molecule_dir: flyte.io.Dir,
    properties_json: str,
    screening_json: str,
) -> str:
    """Generate a comprehensive drug screening report.

    Produces an executive summary, top candidate spotlight cards, property
    distributions, chemical diversity analysis, and final recommendation.

    Args:
        molecule_dir: Directory from load_molecules.
        properties_json: JSON from compute_properties.
        screening_json: Exact verbatim JSON string returned by screen_candidates
            (must include ranked_molecules, similarity_matrix, similarity_labels).
            Do not construct or summarize this payload yourself.

    Returns:
        JSON summary with total_screened, lipinski_passes, all_criteria_met,
        top_candidate, top_score, and top_3 ranked molecules.
    """
    from rdkit import Chem

    await flyte.report.replace.aio(
        _wrap_report("<h2>Generating Final Report...</h2>"),
        do_flush=True,
    )

    props = json.loads(properties_json)
    screening = _parse_screening_json(screening_json)
    ranked = screening["ranked_molecules"]
    sim_matrix = screening["similarity_matrix"]
    sim_labels = screening["similarity_labels"]

    total = props["total"]
    lipinski_pass = props["lipinski_pass_count"]
    all_criteria = sum(1 for m in ranked if m["all_criteria_met"])
    top = ranked[0] if ranked else None

    html_parts = []

    # --- Executive Summary ---
    html_parts.append("<h2>Drug Molecule Screening Report</h2>")
    top_name = top["name"] if top else "N/A"
    top_score = f'{top["screening_score"]:.3f}' if top else "N/A"
    html_parts.append(
        f'<div class="card">'
        f'<h3 style="margin-top:0;color:#0e4f6e;">Executive Summary</h3>'
        f'<p style="font-size:1.05em;">'
        f'<strong>{total}</strong> molecules were screened against the target drug profile. '
        f'<strong>{lipinski_pass}</strong> passed Lipinski\'s Rule of Five, and '
        f'<strong>{all_criteria}</strong> met all screening criteria. '
        f'The top candidate is <strong style="color:#0891b2;">{top_name}</strong> '
        f'with a screening score of <strong>{top_score}</strong>.</p>'
        f'</div>'
    )

    # Stat grid
    html_parts.append('<div class="stat-grid">')
    for val, label in [
        (str(total), "Molecules Screened"),
        (str(lipinski_pass), "Lipinski Passes"),
        (str(all_criteria), "All Criteria Met"),
        (top_score, "Top Score"),
        (f'{props["avg_mw"]:.0f} Da', "Avg. Molecular Weight"),
        (f'{props["avg_logp"]:.2f}', "Avg. LogP"),
    ]:
        html_parts.append(
            f'<div class="stat"><div class="value">{val}</div>'
            f'<div class="label">{label}</div></div>'
        )
    html_parts.append("</div>")

    # --- Top 3 Candidate Spotlights ---
    html_parts.append("<h2>Top Candidate Spotlights</h2>")

    for rank, m in enumerate(ranked[:3], 1):
        mol = Chem.MolFromSmiles(m["smiles"])
        img_uri = _mol_to_data_uri(mol, size=(300, 300)) if mol else ""

        medal = ["gold", "silver", "#cd7f32"][rank - 1]
        medal_emoji = ["1st", "2nd", "3rd"][rank - 1]

        lip_badges = ""
        for rule, key in [("MW", "mw_ok"), ("LogP", "logp_ok"),
                          ("HBD", "hbd_ok"), ("HBA", "hba_ok")]:
            ok = m["lipinski"].get(key, False)
            cls = "badge-success" if ok else "badge-danger"
            lip_badges += f'<span class="badge {cls}" style="margin:2px;">{rule}</span> '

        html_parts.append(
            f'<div class="molecule-card" style="display:flex;gap:20px;align-items:flex-start;flex-wrap:wrap;">'
            f'<div style="text-align:center;min-width:180px;">'
            f'<div style="font-size:1.6em;font-weight:800;color:{medal};">{medal_emoji}</div>'
            f'<img src="{img_uri}" style="width:200px;height:200px;object-fit:contain;border-radius:8px;'
            f'border:2px solid #a5f3fc;"/>'
            f'<div style="font-weight:700;font-size:1.1em;color:#0e4f6e;margin-top:8px;">{m["name"]}</div>'
            f'</div>'
            f'<div style="flex:1;min-width:280px;">'
            f'<table style="margin:0;">'
            f'<tr><td><strong>SMILES</strong></td><td style="font-family:monospace;font-size:0.8em;word-break:break-all;">{m["smiles"]}</td></tr>'
            f'<tr><td><strong>Screening Score</strong></td><td style="font-weight:700;color:#0891b2;font-size:1.1em;">{m["screening_score"]:.3f}</td></tr>'
            f'<tr><td><strong>Molecular Weight</strong></td><td>{m["mw"]:.1f} Da</td></tr>'
            f'<tr><td><strong>LogP</strong></td><td>{m["logp"]:.2f}</td></tr>'
            f'<tr><td><strong>H-Bond Donors</strong></td><td>{m["hbd"]}</td></tr>'
            f'<tr><td><strong>H-Bond Acceptors</strong></td><td>{m["hba"]}</td></tr>'
            f'<tr><td><strong>TPSA</strong></td><td>{m["tpsa"]:.1f} A&sup2;</td></tr>'
            f'<tr><td><strong>Rotatable Bonds</strong></td><td>{m["rotatable_bonds"]}</td></tr>'
            f'<tr><td><strong>QED</strong></td><td>{m["qed"]:.4f}</td></tr>'
            f'<tr><td><strong>Lipinski Compliance</strong></td><td>{lip_badges}</td></tr>'
            f'</table>'
            f'</div>'
            f'</div>'
        )

    # --- Property Distribution (box-plot style as bars with min/max/median) ---
    html_parts.append("<h2>Property Distributions</h2>")

    prop_keys = [("mw", "Molecular Weight (Da)"), ("logp", "LogP"),
                 ("tpsa", "TPSA"), ("qed", "QED Score")]
    for key, label in prop_keys:
        vals = sorted([m[key] for m in ranked])
        n = len(vals)
        if n == 0:
            continue
        v_min = vals[0]
        v_max = vals[-1]
        median = vals[n // 2] if n % 2 == 1 else (vals[n // 2 - 1] + vals[n // 2]) / 2
        q1 = vals[n // 4] if n >= 4 else v_min
        q3 = vals[3 * n // 4] if n >= 4 else v_max

        # Simple horizontal box-plot as SVG
        box_w = 500
        box_h = 50
        margin_l = 10
        v_range = v_max - v_min or 1

        def sx(v):
            return margin_l + ((v - v_min) / v_range) * (box_w - 2 * margin_l)

        box_svg = (
            f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {box_w} {box_h}" '
            f'style="width:100%;max-width:{box_w}px;height:auto;">'
            f'<rect width="{box_w}" height="{box_h}" fill="#fff" rx="4"/>'
            # Whisker line
            f'<line x1="{sx(v_min):.1f}" y1="25" x2="{sx(v_max):.1f}" y2="25" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Min whisker
            f'<line x1="{sx(v_min):.1f}" y1="18" x2="{sx(v_min):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # Max whisker
            f'<line x1="{sx(v_max):.1f}" y1="18" x2="{sx(v_max):.1f}" y2="32" '
            f'stroke="#94a3b8" stroke-width="1.5"/>'
            # IQR box
            f'<rect x="{sx(q1):.1f}" y="14" width="{sx(q3) - sx(q1):.1f}" height="22" '
            f'fill="#a5f3fc" stroke="#0891b2" stroke-width="1.5" rx="3"/>'
            # Median line
            f'<line x1="{sx(median):.1f}" y1="12" x2="{sx(median):.1f}" y2="38" '
            f'stroke="#0e4f6e" stroke-width="2"/>'
            # Labels
            f'<text x="{sx(v_min):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_min:.1f}</text>'
            f'<text x="{sx(median):.1f}" y="10" text-anchor="middle" font-size="9" fill="#0e4f6e" font-weight="600">{median:.1f}</text>'
            f'<text x="{sx(v_max):.1f}" y="46" text-anchor="middle" font-size="9" fill="#6c757d">{v_max:.1f}</text>'
            f'</svg>'
        )
        html_parts.append(
            f'<div style="margin:8px 0;"><strong style="color:#155e75;">{label}</strong>'
            f'<div class="chart-container" style="padding:8px;">{box_svg}</div></div>'
        )

    # --- Chemical Diversity ---
    html_parts.append("<h2>Chemical Diversity Analysis</h2>")

    if sim_matrix and len(sim_matrix) > 1:
        # Compute average pairwise similarity (off-diagonal)
        n_mols = len(sim_matrix)
        off_diag = []
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                off_diag.append(sim_matrix[i][j])

        avg_sim = sum(off_diag) / len(off_diag) if off_diag else 0
        max_sim = max(off_diag) if off_diag else 0
        min_sim = min(off_diag) if off_diag else 0

        # Find most similar pair
        best_i, best_j = 0, 1
        best_val = 0
        for i in range(n_mols):
            for j in range(i + 1, n_mols):
                if sim_matrix[i][j] > best_val:
                    best_val = sim_matrix[i][j]
                    best_i, best_j = i, j

        html_parts.append('<div class="stat-grid">')
        html_parts.append(
            f'<div class="stat"><div class="value">{avg_sim:.3f}</div>'
            f'<div class="label">Avg. Pairwise Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{min_sim:.3f}</div>'
            f'<div class="label">Min Similarity</div></div>'
        )
        html_parts.append(
            f'<div class="stat"><div class="value">{max_sim:.3f}</div>'
            f'<div class="label">Max Similarity</div></div>'
        )
        html_parts.append("</div>")

        diversity_text = "highly diverse" if avg_sim < 0.3 else "moderately diverse" if avg_sim < 0.5 else "relatively similar"
        html_parts.append(
            f'<div class="note">'
            f'The library is <strong>{diversity_text}</strong> (avg. Tanimoto = {avg_sim:.3f}). '
            f'The most similar pair is <strong>{sim_labels[best_i]}</strong> and '
            f'<strong>{sim_labels[best_j]}</strong> (similarity = {best_val:.3f}).</div>'
        )

    # --- Recommendation ---
    html_parts.append("<h2>Recommendation</h2>")
    if top:
        html_parts.append(
            f'<div class="card">'
            f'<h3 style="margin-top:0;color:#0891b2;">Top Candidate: {top["name"]}</h3>'
            f'<p>Based on the virtual screening analysis, <strong>{top["name"]}</strong> '
            f'achieved the highest composite screening score of <strong>{top["screening_score"]:.3f}</strong>. '
        )

        reasons = []
        if top["lipinski_pass"]:
            reasons.append("full Lipinski Rule of Five compliance")
        if top["qed"] > 0.5:
            reasons.append(f"high drug-likeness (QED = {top['qed']:.3f})")
        if top.get("all_criteria_met"):
            reasons.append("all target profile criteria met")
        if top["mw"] <= 500:
            reasons.append(f"favorable molecular weight ({top['mw']:.1f} Da)")

        if reasons:
            html_parts.append(
                f'This candidate stands out due to: {", ".join(reasons)}.</p>'
            )
        else:
            html_parts.append("</p>")

        # Runner-up mentions
        if len(ranked) >= 2:
            html_parts.append(
                f'<p style="font-size:0.9em;color:#6c757d;">Runner-up candidates: '
            )
            runners = []
            for m in ranked[1:4]:
                runners.append(f'{m["name"]} (score: {m["screening_score"]:.3f})')
            html_parts.append(", ".join(runners) + ".</p>")

        html_parts.append("</div>")

    # Final note
    html_parts.append(
        '<div class="note">'
        "This is a virtual screening analysis. All candidates should undergo "
        "further computational validation (molecular dynamics, docking) and "
        "experimental testing before advancing to clinical trials.</div>"
    )

    await flyte.report.replace.aio(
        _wrap_report("\n".join(html_parts)),
        do_flush=True,
    )

    # JSON summary
    summary = {
        "total_screened": total,
        "lipinski_passes": lipinski_pass,
        "all_criteria_met": all_criteria,
        "top_candidate": top["name"] if top else None,
        "top_score": top["screening_score"] if top else None,
        "top_3": [
            {"name": m["name"], "score": m["screening_score"]}
            for m in ranked[:3]
        ],
    }
    return json.dumps(summary)

# ------------------------------------------------------------------
# Agent
# ------------------------------------------------------------------

# {{docs-fragment agent}}
SCREENING_AGENT_INSTRUCTIONS = """\
You are a medicinal chemistry screening strategist. You orchestrate a virtual \
screening pipeline using durable Flyte tools. You NEVER invent molecular \
properties — only RDKit tools compute them.

Workflow:
1. If target_profile is not provided in the user message, derive a JSON \
target_profile from the therapeutic brief. Valid keys: mw, logp, hbd, hba, tpsa \
(each [min, max]). Ground choices in oral bioavailability / kinase / CNS rules \
as appropriate to the brief.
2. First pass (always): load_molecules → compute_properties → \
screen_candidates → generate_report. Pass tool outputs between steps exactly \
(molecule_dir from load_molecules into compute_properties and generate_report; \
properties_json from compute_properties into screen_candidates and \
generate_report; screening_json must be the complete, unmodified string \
returned by screen_candidates — never rebuild or summarize JSON yourself).
3. Read the JSON summary returned by generate_report. Reflect:
   - If all_criteria_met == 0: relax exactly ONE profile bound by ~10–20% \
and re-run screen_candidates then generate_report only, reusing the same \
molecule_dir and properties_json from the first pass.
   - If all molecules pass but diversity is a stated goal: note high similarity \
in your summary; do not re-run unless brief asks for stricter filters.
   - Maximum ONE rescreen iteration.
4. Finish with plain text: top candidate, rationale tied to computed metrics \
from the tool JSON, funnel interpretation, and suggested next steps (docking, \
ADMET lab tests).

If the user supplies an explicit target_profile JSON, use it as-is.

Do NOT ask the user for SMILES or molecule lists when molecules_json is empty — \
the default library is loaded automatically.
"""

screening_agent = Agent(
    name="drug-screening-agent",
    instructions=SCREENING_AGENT_INSTRUCTIONS,
    model=MODEL,
    tools=[
        load_molecules,
        compute_properties,
        screen_candidates,
        generate_report,
    ],
    max_turns=12,
)
# {{/docs-fragment agent}}

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def pipeline(
    brief: str = "Screen the default drug library for orally bioavailable small molecules.",
    molecules_json: str = "",
    target_profile: str = "",
) -> str:
    """Agentic virtual drug molecule screening pipeline.

    A medicinal-chemistry agent interprets the screening brief, derives or
    applies a target profile, orchestrates the RDKit screening stages, and
    optionally re-screens when funnel results are too narrow.

    Args:
        brief: Natural-language therapeutic goal (e.g. oral kinase inhibitors,
            CNS-penetrant small molecules).
        molecules_json: JSON mapping molecule names to SMILES strings.
            Defaults to a curated library of ~15 well-known drugs.
        target_profile: Optional JSON with desired property ranges that
            overrides agent-derived criteria
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5]}).

    Returns:
        Agent summary with screening rationale and key results.
    """
    prompt_parts = [
        f"Screening brief: {brief}",
        'Use molecules_json="" for the built-in default library unless provided below.',
        "Compose the four stage tools in order: load_molecules → compute_properties "
        "→ screen_candidates → generate_report. Pass each tool's full return value "
        "verbatim to the next step (especially screening_json). Re-run "
        "screen_candidates and generate_report at most once if the funnel is too narrow.",
    ]
    if molecules_json.strip():
        prompt_parts.append(f"molecules_json: {molecules_json}")
    if target_profile.strip():
        prompt_parts.append(f"Use this target_profile exactly: {target_profile}")

    result = await screening_agent.run.aio("\n".join(prompt_parts))
    return result.summary or result.error or ""

# {{/docs-fragment pipeline}}

# ------------------------------------------------------------------
# Rescreen demo — tight profile + explicit rescreen instructions
# ------------------------------------------------------------------

# Initial profile is deliberately strict (narrow MW + low LogP cap) so
# all_criteria_met is typically 0 on the default library; the brief then
# forces a single rescreen with a widened LogP window.
RESCREEN_DEMO_TARGET_PROFILE = (
    '{"mw": [150, 200], "logp": [-0.5, 1.0], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_TARGET_PROFILE_RESCREEN = (
    '{"mw": [150, 200], "logp": [-0.5, 3.5], "hbd": [0, 1], '
    '"hba": [0, 3], "tpsa": [20, 45]}'
)
RESCREEN_DEMO_BRIEF = f"""\
Two-round agentic screening demo on the default library.

**Round 1 (strict profile):** load_molecules → compute_properties → \
screen_candidates → generate_report using the initial target_profile exactly.

**Round 2 (required — do not skip):** call screen_candidates then generate_report \
again, reusing the same molecule_dir and properties_json from round 1, with this \
relaxed target_profile (wider LogP window only): \
{RESCREEN_DEMO_TARGET_PROFILE_RESCREEN}

Pass every tool return value verbatim to the next step. After both rounds, \
summarize how the funnel and top candidates changed between round 1 and round 2."""

# {{docs-fragment rescreen_demo}}
@env.task(report=True)
async def rescreen_demo() -> str:
    """Example run with a two-round execution graph (rescreen).

    Round 1 uses a strict CNS-like profile; round 2 always re-runs
    screen_candidates and generate_report with a widened LogP window,
    reusing cached molecule_dir and properties_json.
    """
    return await pipeline(
        brief=RESCREEN_DEMO_BRIEF,
        target_profile=RESCREEN_DEMO_TARGET_PROFILE,
    )

# {{/docs-fragment rescreen_demo}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/drug_molecule_screening/drug_molecule_screening.py*

```
flyte run drug_molecule_screening.py rescreen_demo
```

Or pass the same inputs to `pipeline` directly:

```
flyte run drug_molecule_screening.py pipeline \
  --brief "Screen the default library. If all_criteria_met is 0 after generate_report, re-run screen_candidates and generate_report with target_profile {\"mw\": [150, 200], \"logp\": [-0.5, 3.5], \"hbd\": [0, 1], \"hba\": [0, 3], \"tpsa\": [20, 45]}." \
  --target_profile '{"mw": [150, 200], "logp": [-0.5, 1.0], "hbd": [0, 1], "hba": [0, 3], "tpsa": [20, 45]}'
```

Open the run URL and follow the report panel for funnel charts, property distributions, top-candidate spotlights, and the agent's final screening summary. A successful rescreen demo shows two rounds of `screen_candidates` and `generate_report` in the action tree.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/geospatial ===

# Geospatial

Tutorials for satellite imagery, remote sensing, and earth and atmospheric modeling workloads.

### **Geospatial > GPU-accelerated climate modeling**

Run ensemble atmospheric simulations on H200 GPUs with multi-source data ingestion and real-time extreme event detection.

## Subpages
- **Geospatial > GPU-accelerated climate modeling**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/geospatial/climate-modeling ===

# GPU-accelerated climate modeling

Climate modeling is hard for two reasons: data and compute. Satellite imagery arrives continuously from multiple sources. Reanalysis datasets have to be pulled from remote APIs. Weather station data shows up in different formats and schemas. And once all of that is finally in one place, running atmospheric physics simulations demands serious GPU compute.

In practice, many climate workflows are held together with scripts, cron jobs, and a lot of manual babysitting. Data ingestion breaks without warning. GPU jobs run overnight with little visibility into what's happening. When something interesting shows up in a simulation, like a developing hurricane, no one notices until the job finishes hours later.

In this tutorial, we build a production-grade climate modeling pipeline using Flyte. We ingest data from three different sources in parallel, combine it with Dask, run ensemble atmospheric simulations on H200 GPUs, detect extreme weather events as they emerge, and visualize everything in a live dashboard. The entire pipeline is orchestrated, cached, and fault-tolerant, so it can run reliably at scale.

![Report](../../../_static/images/tutorials/climate-modeling/report.png)

> [!NOTE]
> Full code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/climate_modeling/simulation.py).

## Overview

We're building an ensemble weather forecasting system. Ensemble forecasting runs the same simulation multiple times with slightly different initial conditions. This quantifies forecast uncertainty. Instead of saying "the temperature will be 25°C", we can say "the temperature will be 24-26°C with 90% confidence".

The pipeline has five stages:

1. **Data ingestion**: Pull satellite imagery from NOAA GOES, reanalysis data from ERA5, and surface observations from weather stations in parallel.
2. **Preprocessing**: Fuse the datasets, interpolate to a common grid, and run quality control using Dask for distributed computation.
3. **GPU simulation**: Run ensemble atmospheric physics on H200 GPUs. Each ensemble member evolves independently. PyTorch handles the tensor operations; `torch.compile` optimizes the kernels.
4. **Event detection**: Monitor for hurricanes (high wind + low pressure) and heatwaves during simulation. When extreme events are detected, the pipeline can adaptively refine the grid resolution.
5. **Real-time reporting**: Stream metrics to a live Flyte Reports dashboard showing convergence and detected events.

This workflow is a good example of where Flyte shines!

- **Parallel data ingestion**: Three different data sources, three different APIs, all running concurrently. Flyte's async task execution handles this naturally.
- **Resource heterogeneity**: Data ingestion needs CPU and network. Preprocessing needs a Dask cluster. Simulation needs GPUs. Flyte provisions exactly what each stage needs.
- **Caching**: ERA5 data fetches can take minutes. Run the pipeline twice with the same date range, and Flyte skips the fetch entirely.
- **Adaptive workflows**: When a hurricane is detected, we can dynamically refine the simulation. Flyte makes this kind of conditional logic straightforward.

## Implementation

### Dependencies and container image

```
import asyncio
import gc
import io
import json
import os
import tempfile
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Literal

import flyte
import numpy as np
import pandas as pd
import xarray as xr
from flyte.io import File
from flyteplugins.dask import Dask, Scheduler, WorkerGroup
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

The key imports include `xarray` for multi-dimensional climate data, `flyteplugins.dask` for distributed preprocessing, and `flyte` for orchestration.

```
climate_image = (
    flyte.Image.from_debian_base(name="climate_modeling_h200")
    .with_apt_packages(
        "libnetcdf-dev",  # NetCDF for climate data
        "libhdf5-dev",  # HDF5 for large datasets
        "libeccodes-dev",  # GRIB format support (ECMWF's native format)
        "libudunits2-dev",  # Unit conversions
    )
    .with_pip_packages(
        "numpy==2.3.5",
        "pandas==2.3.3",
        "xarray==2025.11.0",
        "torch==2.9.1",
        "netCDF4==1.7.3",
        "s3fs==2025.10.0",
        "aiohttp==3.13.2",
        "ecmwf-datastores-client==0.4.1",
        "h5netcdf==1.7.3",
        "cfgrib==0.9.15.1",
        "pyarrow==22.0.0",
        "scipy==1.15.1",
        "flyteplugins-dask>=2.0.0b33",
        "nvidia-ml-py3==7.352.0",
    )
    .with_env_vars({"PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512"})
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

Climate data comes in specialized formats such as NetCDF, HDF5, and GRIB. The container image includes libraries to work with all of them, along with PyTorch for GPU computation and the ECMWF client for accessing ERA5 data.

### Simulation parameters and data structures

```
@dataclass
class SimulationParams:
    grid_resolution_km: float = 10.0
    time_step_minutes: int = 10
    simulation_hours: int = 240
    physics_model: Literal["WRF", "MPAS", "CAM"] = "WRF"
    boundary_layer_scheme: str = "YSU"
    microphysics_scheme: str = "Thompson"
    radiation_scheme: str = "RRTMG"

    # Ensemble forecasting parameters
    ensemble_size: int = 800
    perturbation_magnitude: float = 0.5

    # Convergence criteria for adaptive refinement
    convergence_threshold: float = 0.1  # 10% of initial ensemble spread
    max_iterations: int = 3

@dataclass
class ClimateMetrics:
    timestamp: str
    iteration: int
    convergence_rate: float
    energy_conservation_error: float
    max_wind_speed_mps: float
    min_pressure_mb: float
    detected_phenomena: list[str]
    compute_time_seconds: float
    ensemble_spread: float

@dataclass
class SimulationSummary:
    total_iterations: int
    final_resolution_km: float
    avg_convergence_rate: float
    total_compute_time_seconds: float
    hurricanes_detected: int
    heatwaves_detected: int
    converged: bool
    region: str
    output_files: list[File]
    date_range: list[str, str]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

`SimulationParams` defines the core behavior of the simulation, including grid resolution, physics schemes, and ensemble size. The default configuration runs 800 ensemble members, which is sufficient to produce statistically meaningful uncertainty estimates.

> [!NOTE]
> Decreasing the grid spacing via `grid_resolution_km` (for example, from 10 km to 5 km) increases grid resolution and significantly increases memory usage because it introduces more data points and intermediate state. Even with 141 GB of H200 GPU memory, high-resolution or adaptively refined simulations may exceed available VRAM, especially when running large ensembles.
>
> To mitigate this, consider reducing the ensemble size, limiting the refined region, running fewer physics variables, or scaling the simulation across more GPUs so memory is distributed more evenly.

`ClimateMetrics` collects diagnostics at each iteration, such as convergence rate, energy conservation, and detected phenomena. These metrics are streamed to the real-time dashboard so you can monitor how the simulation evolves as it runs.

### Task environments

Different stages need different resources. Flyte's `TaskEnvironment` declares exactly what each task requires:

```
gpu_env = flyte.TaskEnvironment(
    name="climate_modeling_gpu",
    resources=flyte.Resources(
        cpu=5,
        memory="130Gi",
        gpu="H200:1",
    ),
    image=climate_image,
    cache="auto",
)

dask_env = flyte.TaskEnvironment(
    name="climate_modeling_dask",
    plugin_config=Dask(
        scheduler=Scheduler(resources=flyte.Resources(cpu=2, memory="6Gi")),
        workers=WorkerGroup(
            number_of_workers=2,
            resources=flyte.Resources(cpu=2, memory="12Gi"),
        ),
    ),
    image=climate_image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),  # Head node
    cache="auto",
)

cpu_env = flyte.TaskEnvironment(
    name="climate_modeling_cpu",
    resources=flyte.Resources(cpu=8, memory="64Gi"),
    image=climate_image,
    cache="auto",
    secrets=[
        flyte.Secret(key="cds_api_key", as_env_var="ECMWF_DATASTORES_KEY"),
        flyte.Secret(key="cds_api_url", as_env_var="ECMWF_DATASTORES_URL"),
    ],
    depends_on=[gpu_env, dask_env],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

Here’s what each environment is responsible for:

- **`gpu_env`**: Runs the atmospheric simulations on H200 GPUs. The 130 GB of GPU memory is used to hold the ensemble members in VRAM during execution.
- **`dask_env`**: Provides a distributed Dask cluster for preprocessing. A scheduler and multiple workers handle data fusion and transformation in parallel.
- **`cpu_env`**: Handles data ingestion and orchestration. This environment also includes the secrets required to access the ERA5 API.

The `depends_on` setting on `cpu_env` ensures that Flyte builds the GPU and Dask images first. Once those environments are ready, the orchestration task can launch the specialized simulation and preprocessing tasks.

### Data ingestion: multiple sources in parallel

Climate models need data from multiple sources. Each source has different formats, APIs, and failure modes. We handle them as separate Flyte tasks that run concurrently.

**Satellite imagery from NOAA GOES**

```
@cpu_env.task
async def ingest_satellite_data(region: str, date_range: list[str, str]) -> File:
    """Ingest GOES satellite imagery from NOAA's public S3 buckets."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

This task fetches cloud imagery and precipitable water products from NOAA's public S3 buckets. GOES-16 covers the Atlantic; GOES-17 covers the Pacific. The task selects the appropriate satellite based on region, fetches multiple days in parallel using `asyncio.gather`, and combines everything into a single xarray Dataset.

**ERA5 reanalysis from Copernicus**

```
@cpu_env.task
async def ingest_reanalysis_data(region: str, date_range: list[str, str]) -> File:
    """Fetch ERA5 reanalysis from Copernicus Climate Data Store."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

ERA5 provides 3D atmospheric fields such as temperature, wind, humidity at multiple pressure levels from surface to stratosphere. The ECMWF datastores client handles authentication via Flyte secrets. Each day fetches in parallel, then gets concatenated.

**Surface observations from weather stations:**

```
@cpu_env.task
async def ingest_station_data(
    region: str, date_range: list[str, str], max_stations: int = 100
) -> File:
    """Fetch ground observations from NOAA's Integrated Surface Database."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

Ground truth comes from NOAA's Integrated Surface Database. The task filters stations by geographic bounds, fetches hourly observations, and returns a Parquet file for efficient downstream processing.

All three tasks return Flyte `File` objects that hold references to data in blob storage. No data moves until a downstream task actually needs it.

### Preprocessing with Dask

The three data sources need to be combined into a unified atmospheric state. This means:

- Interpolating to a common grid
- Handling missing values
- Merging variables from different sources
- Quality control

This is a perfect fit for Dask to handle lazy evaluation over chunked arrays:

```python
@dask_env.task
async def preprocess_atmospheric_data(
    satellite_data: File,
    reanalysis_data: File,
    station_data: File,
    target_resolution_km: float,
) -> File:
```

This task connects to the Dask cluster provisioned by Flyte, loads the datasets with appropriate chunking, merges satellite and reanalysis grids, fills in missing values, and persists the result. Flyte caches the output, so preprocessing only runs when the inputs change.

### GPU-accelerated atmospheric simulation

Now the core: running atmospheric physics on the GPU. Each ensemble member is an independent forecast with slightly perturbed initial conditions.

```
@gpu_env.task
async def run_atmospheric_simulation(
    input_data: File,
    params: SimulationParams,
    partition_id: int = 0,
    ensemble_start: int | None = None,
    ensemble_end: int | None = None,
) -> tuple[File, ClimateMetrics]:
    """Run GPU-accelerated atmospheric simulation with ensemble forecasting."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

The task accepts a subset of ensemble members (`ensemble_start` to `ensemble_end`). This enables distributing 800 members across multiple GPUs.

The physics step is the computational kernel. It runs advection (wind transport), pressure gradients, Coriolis forces, turbulent diffusion, and moisture condensation:

```
    @torch.compile(mode="reduce-overhead")
    def physics_step(state_tensor, dt_val, dx_val):
        """Compiled atmospheric physics - 3-4x faster with torch.compile."""
        # Advection: transport by wind
        temp_grad_x = torch.roll(state_tensor[:, 0], -1, dims=2) - torch.roll(
            state_tensor[:, 0], 1, dims=2
        )
        temp_grad_y = torch.roll(state_tensor[:, 0], -1, dims=3) - torch.roll(
            state_tensor[:, 0], 1, dims=3
        )
        advection = -(
            state_tensor[:, 3] * temp_grad_x + state_tensor[:, 4] * temp_grad_y
        ) / (2 * dx_val)
        state_tensor[:, 0] = state_tensor[:, 0] + advection * dt_val

        # Pressure gradient with Coriolis
        pressure_grad_x = (
            torch.roll(state_tensor[:, 1], -1, dims=2)
            - torch.roll(state_tensor[:, 1], 1, dims=2)
        ) / (2 * dx_val)
        pressure_grad_y = (
            torch.roll(state_tensor[:, 1], -1, dims=3)
            - torch.roll(state_tensor[:, 1], 1, dims=3)
        ) / (2 * dx_val)

        coriolis_param = 1e-4  # ~45°N latitude
        coriolis_u = coriolis_param * state_tensor[:, 4]
        coriolis_v = -coriolis_param * state_tensor[:, 3]

        state_tensor[:, 3] = (
            state_tensor[:, 3] - pressure_grad_x * dt_val * 0.01 + coriolis_u * dt_val
        )
        state_tensor[:, 4] = (
            state_tensor[:, 4] - pressure_grad_y * dt_val * 0.01 + coriolis_v * dt_val
        )

        # Turbulent diffusion
        diffusion_coeff = 10.0
        laplacian_temp = (
            torch.roll(state_tensor[:, 0], 1, dims=2)
            + torch.roll(state_tensor[:, 0], -1, dims=2)
            + torch.roll(state_tensor[:, 0], 1, dims=3)
            + torch.roll(state_tensor[:, 0], -1, dims=3)
            - 4 * state_tensor[:, 0]
        ) / (dx_val * dx_val)
        state_tensor[:, 0] = (
            state_tensor[:, 0] + diffusion_coeff * laplacian_temp * dt_val
        )

        # Moisture condensation
        sat_vapor_pressure = 611.2 * torch.exp(
            17.67 * state_tensor[:, 0] / (state_tensor[:, 0] + 243.5)
        )
        condensation = torch.clamp(
            state_tensor[:, 2] - sat_vapor_pressure * 0.001, min=0
        )
        state_tensor[:, 2] = state_tensor[:, 2] - condensation * 0.1
        state_tensor[:, 0] = state_tensor[:, 0] + condensation * 2.5e6 / 1005 * dt_val

        return state_tensor
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

`@torch.compile(mode="reduce-overhead")` compiles this function into optimized CUDA kernels. Combined with mixed precision (`torch.cuda.amp.autocast`), this runs 3-4x faster than eager PyTorch.

Every 10 timesteps, the simulation checks for extreme events:

- **Hurricanes**: Wind speed > 33 m/s with low pressure
- **Heatwaves**: Temperature anomalies exceeding thresholds

Detected phenomena get logged to the metrics, which flow to the live dashboard.

### Distributing across multiple GPUs

800 ensemble members is a lot for one GPU, so we distribute them:

```
@cpu_env.task
async def run_distributed_simulation_ensemble(
    preprocessed_data: File, params: SimulationParams, n_gpus: int
) -> tuple[list[File], list[ClimateMetrics]]:
    total_members = params.ensemble_size
    members_per_gpu = total_members // n_gpus

    # Distribute ensemble members across GPUs
    tasks = []
    for gpu_id in range(n_gpus):
        # Calculate ensemble range for this GPU
        ensemble_start = gpu_id * members_per_gpu
        # Last GPU gets any remainder members
        if gpu_id == n_gpus - 1:
            ensemble_end = total_members
        else:
            ensemble_end = ensemble_start + members_per_gpu

        # Launch GPU task with ensemble subset
        gpu_task = run_atmospheric_simulation(
            preprocessed_data,
            params,
            gpu_id,
            ensemble_start=ensemble_start,
            ensemble_end=ensemble_end,
        )
        tasks.append(gpu_task)

    # Execute all GPUs in parallel
    results = await asyncio.gather(*tasks)

    output_files = [r[0] for r in results]
    metrics = [r[1] for r in results]

    return output_files, metrics
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

The task splits the ensemble members evenly across the available GPUs, launches the simulation runs in parallel using `asyncio.gather`, and then aggregates the results. With five GPUs, each GPU runs 160 ensemble members. Flyte takes care of scheduling, so GPU tasks start automatically as soon as resources become available.

### The main workflow

Everything comes together in the orchestration task:

```
@cpu_env.task(report=True)
async def adaptive_climate_modeling_workflow(
    region: str = "atlantic",
    date_range: list[str, str] = ["2024-09-01", "2024-09-10"],
    current_params: SimulationParams = SimulationParams(),
    enable_multi_gpu: bool = True,
    n_gpus: int = 5,
) -> SimulationSummary:
    """Orchestrates multi-source ingestion, GPU simulation, and adaptive refinement."""
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

`report=True` enables Flyte Reports for live monitoring.

```
    # Parallel data ingestion from three sources
    with flyte.group("data-ingestion"):
        satellite_task = ingest_satellite_data(region, date_range)
        reanalysis_task = ingest_reanalysis_data(region, date_range)
        station_task = ingest_station_data(region, date_range)

        satellite_data, reanalysis_data, station_data = await asyncio.gather(
            satellite_task,
            reanalysis_task,
            station_task,
        )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

`flyte.group("data-ingestion")` visually groups the ingestion tasks in the Flyte UI. Inside the group, three tasks launch concurrently. `asyncio.gather` waits for all three to complete before preprocessing begins.

The workflow then enters an iterative loop:

1. Run GPU simulation (single or multi-GPU)
2. Check convergence by comparing forecasts across iterations
3. Detect extreme events
4. If a hurricane is detected and we haven't refined yet, double the grid resolution
5. Stream metrics to the live dashboard
6. Repeat until converged or max iterations reached

Adaptive mesh refinement is the key feature here. When the simulation detects a hurricane forming, it automatically increases resolution to capture the fine-scale dynamics. This is expensive, so we limit it to one refinement per run.

### Running the pipeline

```
if __name__ == "__main__":
    flyte.init_from_config()
    run_multi_gpu = flyte.run(adaptive_climate_modeling_workflow)

    print(f"Run URL: {run_multi_gpu.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate_modeling/simulation.py*

Before running, set up ERA5 API credentials:

```bash
flyte create secret cds_api_key <YOUR_CDS_API_KEY>
flyte create secret cds_api_url https://cds.climate.copernicus.eu/api
```

Then launch:

```bash
flyte create config --endpoint <FLYTE_OR_UNION_ENDPOINT> --project <PROJECT_NAME> --domain <DOMAIN_NAME> --builder remote
uv run simulation.py
```

The default configuration uses the Atlantic region for September 2024, which is hurricane season.

## Key concepts

### Ensemble forecasting

Weather prediction is inherently uncertain. Small errors in the initial conditions grow over time due to chaotic dynamics, which means a single forecast can only ever be one possible outcome.

Ensemble forecasting addresses this uncertainty by:

- Perturbing the initial conditions within known observational error bounds
- Running many independent forecasts
- Computing the ensemble mean as the most likely outcome and the ensemble spread as a measure of uncertainty

### Adaptive mesh refinement

When a hurricane begins to form, coarse spatial grids are not sufficient to resolve critical features like eyewall dynamics. Adaptive mesh refinement allows the simulation to focus compute where it matters most by:

- Increasing grid resolution, for example from 10 km to 5 km
- Reducing the timestep to maintain numerical stability
- Refining only the regions of interest instead of the entire domain

This approach is computationally expensive, but it is essential for producing accurate intensity forecasts.

### Real-time event detection

Rather than analyzing results after a simulation completes, this pipeline detects significant events as the simulation runs.

The system monitors for conditions such as:

- **Hurricanes**: Wind speeds exceeding 33 m/s (Category 1 threshold) combined with central pressure below 980 mb
- **Heatwaves**: Sustained temperature anomalies over a defined period

Detecting these events in real time enables adaptive responses, such as refining the simulation or triggering alerts, and supports earlier warnings for extreme weather.

## Where to go next

This example is intentionally scoped to keep the ideas clear, but there are several natural ways to extend it for more realistic workloads.

To model different ocean basins, change the `region` parameter to values like `"pacific"` or `"indian"`. The ingestion tasks automatically adjust to pull the appropriate satellite coverage for each region.

To run longer forecasts, increase `simulation_hours` in `SimulationParams`. The default of 240 hours, or 10 days, is typical for medium-range forecasting, but you can run longer simulations if you have the compute budget.

Finally, the physics step here is deliberately simplified. Production systems usually incorporate additional components such as radiation schemes, boundary layer parameterizations, and land surface models. These can be added incrementally as separate steps without changing the overall structure of the pipeline.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services ===

# Financial services & fintech

Tutorials for financial research, trading, and other fintech workloads.

### **Financial services & fintech > Financial research agent**

Prep equity briefings for the earnings cycle with grounded You.com Research synthesis and fresh news from the Search API.

### **Financial services & fintech > Fraud detection with Feast**

Train an XGBoost fraud classifier and materialize transaction features in Feast for online scoring.

### **Financial services & fintech > Multi-agent trading simulation**

A multi-agent trading simulation, modeling how agents within a firm might interact, strategize, and make trades collaboratively.

## Subpages
- **Financial services & fintech > Multi-agent trading simulation**
- **Financial services & fintech > Fraud detection with Feast**
- **Financial services & fintech > Financial research agent**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services/trading-agents ===

# Multi-agent trading simulation

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/trading_agents); based on work by [TauricResearch](https://github.com/TauricResearch/TradingAgents).

This example walks you through building a multi-agent trading simulation, modeling how agents within a firm might interact, strategize, and make trades collaboratively.

![Trading agents execution visualization](../../../_static/images/tutorials/trading-agents/execution.png)
_Trading agents execution visualization_

## TL;DR

- You'll build a trading firm made up of agents that analyze, argue, and act, modeled with Python functions.
- You'll use the Flyte SDK to orchestrate this world, giving you visibility, retries, caching, and durability.
- You'll learn how to plug in tools, structure conversations, and track decisions across agents.
- You'll see how agents debate, use context, generate reports, and retain memory via vector DBs.

## What is an agent, anyway?

Agentic workflows are a rising pattern for complex problem-solving with LLMs. Think of agents as:

- An LLM (like GPT-4 or Mistral)
- A loop that keeps them thinking until a goal is met
- A set of optional tools they can call (APIs, search, calculators, etc.)
- Enough tokens to reason about the problem at hand

That's it.

You define tools, bind them to an agent, and let it run, reasoning step-by-step, optionally using those tools, until it finishes.

## What's different here?

We're not building yet another agent framework. You're free to use LangChain, custom code, or whatever setup you like.

What we're giving you is the missing piece: a way to run these workflows **reliably, observably, and at scale, with zero rewrites.**

With Flyte, you get:

- Prompt + tool traceability and full state retention
- Built-in retries, caching, and failure recovery
- A native way to plug in your agents; no magic syntax required

## How it works: step-by-step walkthrough

This simulation is powered by a Flyte task that orchestrates multiple intelligent agents working together to analyze a company's stock and make informed trading decisions.

![Trading agents schema](../../../_static/images/tutorials/trading-agents/schema.png)
_Trading agents schema_

### Entry point

Everything begins with a top-level Flyte task called `main`, which serves as the entry point to the workflow.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.0.0b52",
#     "akshare==1.16.98",
#     "backtrader==1.9.78.123",
#     "boto3==1.39.9",
#     "chainlit==2.5.5",
#     "eodhd==1.0.32",
#     "feedparser==6.0.11",
#     "finnhub-python==2.4.23",
#     "langchain-experimental==0.3.4",
#     "langchain-openai==0.3.23",
#     "pandas==2.3.0",
#     "parsel==1.10.0",
#     "praw==7.8.1",
#     "pytz==2025.2",
#     "questionary==2.1.0",
#     "redis==6.2.0",
#     "requests==2.32.4",
#     "stockstats==0.6.5",
#     "tqdm==4.67.1",
#     "tushare==1.4.21",
#     "typing-extensions==4.14.0",
#     "yfinance==0.2.63",
# ]
# main = "main"
# params = ""
# ///
import asyncio
from copy import deepcopy

import agents
import agents.analysts
from agents.managers import create_research_manager, create_risk_manager
from agents.researchers import create_bear_researcher, create_bull_researcher
from agents.risk_debators import (
    create_neutral_debator,
    create_risky_debator,
    create_safe_debator,
)
from agents.trader import create_trader
from agents.utils.utils import AgentState
from flyte_env import DEEP_THINKING_LLM, QUICK_THINKING_LLM, env, flyte
from langchain_openai import ChatOpenAI
from reflection import (
    reflect_bear_researcher,
    reflect_bull_researcher,
    reflect_research_manager,
    reflect_risk_manager,
    reflect_trader,
)

@env.task
async def process_signal(full_signal: str, QUICK_THINKING_LLM: str) -> str:
    """Process a full trading signal to extract the core decision."""

    messages = [
        {
            "role": "system",
            "content": """You are an efficient assistant designed to analyze paragraphs or
financial reports provided by a group of analysts.
Your task is to extract the investment decision: SELL, BUY, or HOLD.
Provide only the extracted decision (SELL, BUY, or HOLD) as your output,
without adding any additional text or information.""",
        },
        {"role": "human", "content": full_signal},
    ]

    return ChatOpenAI(model=QUICK_THINKING_LLM).invoke(messages).content

async def run_analyst(analyst_name, state, online_tools):
    # Create a copy of the state for isolation
    run_fn = getattr(agents.analysts, f"create_{analyst_name}_analyst")

    # Run the analyst's chain
    result_state = await run_fn(QUICK_THINKING_LLM, state, online_tools)

    # Determine the report key
    report_key = (
        "sentiment_report"
        if analyst_name == "social_media"
        else f"{analyst_name}_report"
    )
    report_value = getattr(result_state, report_key)

    return result_state.messages[1:], report_key, report_value

# {{docs-fragment main}}
@env.task
async def main(
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> tuple[str, AgentState]:
    if not selected_analysts:
        raise ValueError(
            "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social_media."
        )

    state = AgentState(
        messages=[{"role": "human", "content": company_name}],
        company_of_interest=company_name,
        trade_date=str(trade_date),
    )

    # Run all analysts concurrently
    results = await asyncio.gather(
        *[
            run_analyst(analyst, deepcopy(state), online_tools)
            for analyst in selected_analysts
        ]
    )

    # Flatten and append all resulting messages into the shared state
    for messages, report_attr, report in results:
        state.messages.extend(messages)
        setattr(state, report_attr, report)

    # Bull/Bear debate loop
    state = await create_bull_researcher(QUICK_THINKING_LLM, state)  # Start with bull
    while state.investment_debate_state.count < 2 * max_debate_rounds:
        current = state.investment_debate_state.current_response
        if current.startswith("Bull"):
            state = await create_bear_researcher(QUICK_THINKING_LLM, state)
        else:
            state = await create_bull_researcher(QUICK_THINKING_LLM, state)

    state = await create_research_manager(DEEP_THINKING_LLM, state)
    state = await create_trader(QUICK_THINKING_LLM, state)

    # Risk debate loop
    state = await create_risky_debator(QUICK_THINKING_LLM, state)  # Start with risky
    while state.risk_debate_state.count < 3 * max_risk_discuss_rounds:
        speaker = state.risk_debate_state.latest_speaker
        if speaker == "Risky":
            state = await create_safe_debator(QUICK_THINKING_LLM, state)
        elif speaker == "Safe":
            state = await create_neutral_debator(QUICK_THINKING_LLM, state)
        else:
            state = await create_risky_debator(QUICK_THINKING_LLM, state)

    state = await create_risk_manager(DEEP_THINKING_LLM, state)
    decision = await process_signal(state.final_trade_decision, QUICK_THINKING_LLM)

    return decision, state

# {{/docs-fragment main}}

# {{docs-fragment reflect_on_decisions}}
@env.task
async def reflect_and_store(state: AgentState, returns: str) -> str:
    await asyncio.gather(
        reflect_bear_researcher(state, returns),
        reflect_bull_researcher(state, returns),
        reflect_trader(state, returns),
        reflect_risk_manager(state, returns),
        reflect_research_manager(state, returns),
    )

    return "Reflection completed."

# Run the reflection task after the main function
@env.task(cache="disable")
async def reflect_on_decisions(
    returns: str,
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> str:
    _, state = await main(
        selected_analysts,
        max_debate_rounds,
        max_risk_discuss_rounds,
        online_tools,
        company_name,
        trade_date,
    )

    return await reflect_and_store(state, returns)

# {{/docs-fragment reflect_on_decisions}}

# {{docs-fragment execute_main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

    # run = flyte.run(reflect_on_decisions, "+3.2% gain over 5 days")
    # print(run.url)

# {{/docs-fragment execute_main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/main.py*

This task accepts several inputs:

- the list of analysts to run,
- the number of debate and risk discussion rounds,
- a flag to enable online tools,
- the company you're evaluating,
- and the target trading date.

The most interesting parameter here is the list of analysts to run. It determines which analyst agents will be invoked and shapes the overall structure of the simulation. Based on this input, the task dynamically launches agent tasks, running them in parallel.

The `main` task is written as a regular asynchronous Python function wrapped with Flyte's task decorator. No domain-specific language or orchestration glue is needed: just idiomatic Python, optionally using async for better performance. The task environment is configured once and shared across all tasks for consistency.

```
# {{docs-fragment env}}
import flyte

QUICK_THINKING_LLM = "gpt-4o-mini"
DEEP_THINKING_LLM = "o4-mini"

env = flyte.TaskEnvironment(
    name="trading-agents",
    secrets=[
        flyte.Secret(key="finnhub_api_key", as_env_var="FINNHUB_API_KEY"),
        flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY"),
    ],
    image=flyte.Image.from_uv_script("main.py", name="trading-agents", pre=True),
    resources=flyte.Resources(cpu="1"),
    cache="auto",
)

# {{/docs-fragment env}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/flyte_env.py*

### Analyst agents

Each analyst agent comes equipped with a set of tools and a carefully designed prompt tailored to its specific domain. These tools are modular Flyte tasks (for example, downloading financial reports or computing technical indicators) and benefit from Flyte's built-in caching to avoid redundant computation.

```
from datetime import datetime

import pandas as pd
import tools.interface as interface
import yfinance as yf
from flyte_env import env

from flyte.io import File

@env.task
async def get_reddit_news(
    curr_date: str,  # Date you want to get news for in yyyy-mm-dd format
) -> str:
    """
    Retrieve global news from Reddit within a specified time frame.
    Args:
        curr_date (str): Date you want to get news for in yyyy-mm-dd format
    Returns:
        str: A formatted dataframe containing the latest global news
        from Reddit in the specified time frame.
    """

    global_news_result = interface.get_reddit_global_news(curr_date, 7, 5)

    return global_news_result

@env.task
async def get_finnhub_news(
    ticker: str,  # Search query of a company, e.g. 'AAPL, TSM, etc.
    start_date: str,  # Start date in yyyy-mm-dd format
    end_date: str,  # End date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the latest news about a given stock from Finnhub within a date range
    Args:
        ticker (str): Ticker of a company. e.g. AAPL, TSM
        start_date (str): Start date in yyyy-mm-dd format
        end_date (str): End date in yyyy-mm-dd format
    Returns:
        str: A formatted dataframe containing news about the company
        within the date range from start_date to end_date
    """

    end_date_str = end_date

    end_date = datetime.strptime(end_date, "%Y-%m-%d")
    start_date = datetime.strptime(start_date, "%Y-%m-%d")
    look_back_days = (end_date - start_date).days

    finnhub_news_result = interface.get_finnhub_news(
        ticker, end_date_str, look_back_days
    )

    return finnhub_news_result

@env.task
async def get_reddit_stock_info(
    ticker: str,  # Ticker of a company. e.g. AAPL, TSM
    curr_date: str,  # Current date you want to get news for
) -> str:
    """
    Retrieve the latest news about a given stock from Reddit, given the current date.
    Args:
        ticker (str): Ticker of a company. e.g. AAPL, TSM
        curr_date (str): current date in yyyy-mm-dd format to get news for
    Returns:
        str: A formatted dataframe containing the latest news about the company on the given date
    """

    stock_news_results = interface.get_reddit_company_news(ticker, curr_date, 7, 5)

    return stock_news_results

@env.task
async def get_YFin_data(
    symbol: str,  # ticker symbol of the company
    start_date: str,  # Start date in yyyy-mm-dd format
    end_date: str,  # End date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the stock price data for a given ticker symbol from Yahoo Finance.
    Args:
        symbol (str): Ticker symbol of the company, e.g. AAPL, TSM
        start_date (str): Start date in yyyy-mm-dd format
        end_date (str): End date in yyyy-mm-dd format
    Returns:
        str: A formatted dataframe containing the stock price data
        for the specified ticker symbol in the specified date range.
    """

    result_data = interface.get_YFin_data(symbol, start_date, end_date)

    return result_data

@env.task
async def get_YFin_data_online(
    symbol: str,  # ticker symbol of the company
    start_date: str,  # Start date in yyyy-mm-dd format
    end_date: str,  # End date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the stock price data for a given ticker symbol from Yahoo Finance.
    Args:
        symbol (str): Ticker symbol of the company, e.g. AAPL, TSM
        start_date (str): Start date in yyyy-mm-dd format
        end_date (str): End date in yyyy-mm-dd format
    Returns:
        str: A formatted dataframe containing the stock price data
        for the specified ticker symbol in the specified date range.
    """

    result_data = interface.get_YFin_data_online(symbol, start_date, end_date)

    return result_data

@env.task
async def cache_market_data(symbol: str, start_date: str, end_date: str) -> File:
    data_file = f"{symbol}-YFin-data-{start_date}-{end_date}.csv"

    data = yf.download(
        symbol,
        start=start_date,
        end=end_date,
        multi_level_index=False,
        progress=False,
        auto_adjust=True,
    )
    data = data.reset_index()
    data.to_csv(data_file, index=False)

    return await File.from_local(data_file)

@env.task
async def get_stockstats_indicators_report(
    symbol: str,  # ticker symbol of the company
    indicator: str,  # technical indicator to get the analysis and report of
    curr_date: str,  # The current trading date you are trading on, YYYY-mm-dd
    look_back_days: int = 30,  # how many days to look back
) -> str:
    """
    Retrieve stock stats indicators for a given ticker symbol and indicator.
    Args:
        symbol (str): Ticker symbol of the company, e.g. AAPL, TSM
        indicator (str): Technical indicator to get the analysis and report of
        curr_date (str): The current trading date you are trading on, YYYY-mm-dd
        look_back_days (int): How many days to look back, default is 30
    Returns:
        str: A formatted dataframe containing the stock stats indicators
        for the specified ticker symbol and indicator.
    """

    today_date = pd.Timestamp.today()

    end_date = today_date
    start_date = today_date - pd.DateOffset(years=15)
    start_date = start_date.strftime("%Y-%m-%d")
    end_date = end_date.strftime("%Y-%m-%d")

    data_file = await cache_market_data(symbol, start_date, end_date)
    local_data_file = await data_file.download()

    result_stockstats = interface.get_stock_stats_indicators_window(
        symbol, indicator, curr_date, look_back_days, False, local_data_file
    )

    return result_stockstats

# {{docs-fragment get_stockstats_indicators_report_online}}
@env.task
async def get_stockstats_indicators_report_online(
    symbol: str,  # ticker symbol of the company
    indicator: str,  # technical indicator to get the analysis and report of
    curr_date: str,  # The current trading date you are trading on, YYYY-mm-dd"
    look_back_days: int = 30,  # "how many days to look back"
) -> str:
    """
    Retrieve stock stats indicators for a given ticker symbol and indicator.
    Args:
        symbol (str): Ticker symbol of the company, e.g. AAPL, TSM
        indicator (str): Technical indicator to get the analysis and report of
        curr_date (str): The current trading date you are trading on, YYYY-mm-dd
        look_back_days (int): How many days to look back, default is 30
    Returns:
        str: A formatted dataframe containing the stock stats indicators
        for the specified ticker symbol and indicator.
    """

    today_date = pd.Timestamp.today()

    end_date = today_date
    start_date = today_date - pd.DateOffset(years=15)
    start_date = start_date.strftime("%Y-%m-%d")
    end_date = end_date.strftime("%Y-%m-%d")

    data_file = await cache_market_data(symbol, start_date, end_date)
    local_data_file = await data_file.download()

    result_stockstats = interface.get_stock_stats_indicators_window(
        symbol, indicator, curr_date, look_back_days, True, local_data_file
    )

    return result_stockstats

# {{/docs-fragment get_stockstats_indicators_report_online}}

@env.task
async def get_finnhub_company_insider_sentiment(
    ticker: str,  # ticker symbol for the company
    curr_date: str,  # current date of you are trading at, yyyy-mm-dd
) -> str:
    """
    Retrieve insider sentiment information about a company (retrieved
    from public SEC information) for the past 30 days
    Args:
        ticker (str): ticker symbol of the company
        curr_date (str): current date you are trading at, yyyy-mm-dd
    Returns:
        str: a report of the sentiment in the past 30 days starting at curr_date
    """

    data_sentiment = interface.get_finnhub_company_insider_sentiment(
        ticker, curr_date, 30
    )

    return data_sentiment

@env.task
async def get_finnhub_company_insider_transactions(
    ticker: str,  # ticker symbol
    curr_date: str,  # current date you are trading at, yyyy-mm-dd
) -> str:
    """
    Retrieve insider transaction information about a company
    (retrieved from public SEC information) for the past 30 days
    Args:
        ticker (str): ticker symbol of the company
        curr_date (str): current date you are trading at, yyyy-mm-dd
    Returns:
        str: a report of the company's insider transactions/trading information in the past 30 days
    """

    data_trans = interface.get_finnhub_company_insider_transactions(
        ticker, curr_date, 30
    )

    return data_trans

@env.task
async def get_simfin_balance_sheet(
    ticker: str,  # ticker symbol
    freq: str,  # reporting frequency of the company's financial history: annual/quarterly
    curr_date: str,  # current date you are trading at, yyyy-mm-dd
):
    """
    Retrieve the most recent balance sheet of a company
    Args:
        ticker (str): ticker symbol of the company
        freq (str): reporting frequency of the company's financial history: annual / quarterly
        curr_date (str): current date you are trading at, yyyy-mm-dd
    Returns:
        str: a report of the company's most recent balance sheet
    """

    data_balance_sheet = interface.get_simfin_balance_sheet(ticker, freq, curr_date)

    return data_balance_sheet

@env.task
async def get_simfin_cashflow(
    ticker: str,  # ticker symbol
    freq: str,  # reporting frequency of the company's financial history: annual/quarterly
    curr_date: str,  # current date you are trading at, yyyy-mm-dd
) -> str:
    """
    Retrieve the most recent cash flow statement of a company
    Args:
        ticker (str): ticker symbol of the company
        freq (str): reporting frequency of the company's financial history: annual / quarterly
        curr_date (str): current date you are trading at, yyyy-mm-dd
    Returns:
            str: a report of the company's most recent cash flow statement
    """

    data_cashflow = interface.get_simfin_cashflow(ticker, freq, curr_date)

    return data_cashflow

@env.task
async def get_simfin_income_stmt(
    ticker: str,  # ticker symbol
    freq: str,  # reporting frequency of the company's financial history: annual/quarterly
    curr_date: str,  # current date you are trading at, yyyy-mm-dd
) -> str:
    """
    Retrieve the most recent income statement of a company
    Args:
        ticker (str): ticker symbol of the company
        freq (str): reporting frequency of the company's financial history: annual / quarterly
        curr_date (str): current date you are trading at, yyyy-mm-dd
    Returns:
            str: a report of the company's most recent income statement
    """

    data_income_stmt = interface.get_simfin_income_statements(ticker, freq, curr_date)

    return data_income_stmt

@env.task
async def get_google_news(
    query: str,  # Query to search with
    curr_date: str,  # Curr date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the latest news from Google News based on a query and date range.
    Args:
        query (str): Query to search with
        curr_date (str): Current date in yyyy-mm-dd format
        look_back_days (int): How many days to look back
    Returns:
        str: A formatted string containing the latest news from Google News
        based on the query and date range.
    """

    google_news_results = interface.get_google_news(query, curr_date, 7)

    return google_news_results

@env.task
async def get_stock_news_openai(
    ticker: str,  # the company's ticker
    curr_date: str,  # Current date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the latest news about a given stock by using OpenAI's news API.
    Args:
        ticker (str): Ticker of a company. e.g. AAPL, TSM
        curr_date (str): Current date in yyyy-mm-dd format
    Returns:
        str: A formatted string containing the latest news about the company on the given date.
    """

    openai_news_results = interface.get_stock_news_openai(ticker, curr_date)

    return openai_news_results

@env.task
async def get_global_news_openai(
    curr_date: str,  # Current date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the latest macroeconomics news on a given date using OpenAI's macroeconomics news API.
    Args:
        curr_date (str): Current date in yyyy-mm-dd format
    Returns:
        str: A formatted string containing the latest macroeconomic news on the given date.
    """

    openai_news_results = interface.get_global_news_openai(curr_date)

    return openai_news_results

@env.task
async def get_fundamentals_openai(
    ticker: str,  # the company's ticker
    curr_date: str,  # Current date in yyyy-mm-dd format
) -> str:
    """
    Retrieve the latest fundamental information about a given stock
    on a given date by using OpenAI's news API.
    Args:
        ticker (str): Ticker of a company. e.g. AAPL, TSM
        curr_date (str): Current date in yyyy-mm-dd format

    Returns:
        str: A formatted string containing the latest fundamental information
        about the company on the given date.
    """

    openai_fundamentals_results = interface.get_fundamentals_openai(ticker, curr_date)

    return openai_fundamentals_results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/tools/toolkit.py*

When initialized, an analyst enters a structured reasoning loop (via LangChain), where it can call tools, observe outputs, and refine its internal state before generating a final report. These reports are later consumed by downstream agents.

Here's an example of a news analyst that interprets global events and macroeconomic signals. We specify the tools accessible to the analyst, and the LLM selects which ones to use based on context.

```
import asyncio

from agents.utils.utils import AgentState
from flyte_env import env
from langchain_core.messages import ToolMessage, convert_to_openai_messages
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from tools import toolkit

import flyte

MAX_ITERATIONS = 5

# {{docs-fragment agent_helper}}
async def run_chain_with_tools(
    type: str, state: AgentState, llm: str, system_message: str, tool_names: list[str]
) -> AgentState:
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You are a helpful AI assistant, collaborating with other assistants."
                " Use the provided tools to progress towards answering the question."
                " If you are unable to fully answer, that's OK; another assistant with different tools"
                " will help where you left off. Execute what you can to make progress."
                " If you or any other assistant has the FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL** or deliverable,"
                " prefix your response with FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL** so the team knows to stop."
                " You have access to the following tools: {tool_names}.\n{system_message}"
                " For your reference, the current date is {current_date}. The company we want to look at is {ticker}.",
            ),
            MessagesPlaceholder(variable_name="messages"),
        ]
    )

    prompt = prompt.partial(system_message=system_message)
    prompt = prompt.partial(tool_names=", ".join(tool_names))
    prompt = prompt.partial(current_date=state.trade_date)
    prompt = prompt.partial(ticker=state.company_of_interest)

    chain = prompt | ChatOpenAI(model=llm).bind_tools(
        [getattr(toolkit, tool_name).func for tool_name in tool_names]
    )

    iteration = 0
    while iteration < MAX_ITERATIONS:
        result = await chain.ainvoke(state.messages)
        state.messages.append(convert_to_openai_messages(result))

        if not result.tool_calls:
            # Final response — no tools required
            setattr(state, f"{type}_report", result.content or "")
            break

        # Run all tool calls in parallel
        async def run_single_tool(tool_call):
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            tool = getattr(toolkit, tool_name, None)

            if not tool:
                return None

            content = await tool(**tool_args)
            return ToolMessage(
                tool_call_id=tool_call["id"], name=tool_name, content=content
            )

        with flyte.group(f"tool_calls_iteration_{iteration}"):
            tool_messages = await asyncio.gather(
                *[run_single_tool(tc) for tc in result.tool_calls]
            )

        # Add valid tool results to state
        tool_messages = [msg for msg in tool_messages if msg]
        state.messages.extend(convert_to_openai_messages(tool_messages))

        iteration += 1
    else:
        # Reached iteration cap — optionally raise or log
        print(f"Max iterations ({MAX_ITERATIONS}) reached for {type}")

    return state

# {{/docs-fragment agent_helper}}

@env.task
async def create_fundamentals_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [toolkit.get_fundamentals_openai]
    else:
        tools = [
            toolkit.get_finnhub_company_insider_sentiment,
            toolkit.get_finnhub_company_insider_transactions,
            toolkit.get_simfin_balance_sheet,
            toolkit.get_simfin_cashflow,
            toolkit.get_simfin_income_stmt,
        ]

    system_message = (
        "You are a researcher tasked with analyzing fundamental information over the past week about a company. "
        "Please write a comprehensive report of the company's fundamental information such as financial documents, "
        "company profile, basic company financials, company financial history, insider sentiment, and insider "
        "transactions to gain a full view of the company's "
        "fundamental information to inform traders. Make sure to include as much detail as possible. "
        "Do not simply state the trends are mixed, "
        "provide detailed and finegrained analysis and insights that may help traders make decisions. "
        "Make sure to append a Markdown table at the end of the report to organize key points in the report, "
        "organized and easy to read."
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools(
        "fundamentals", state, llm, system_message, tool_names
    )

@env.task
async def create_market_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [
            toolkit.get_YFin_data_online,
            toolkit.get_stockstats_indicators_report_online,
        ]
    else:
        tools = [
            toolkit.get_YFin_data,
            toolkit.get_stockstats_indicators_report,
        ]

    system_message = (
        """You are a trading assistant tasked with analyzing financial markets.
Your role is to select the **most relevant indicators** for a given market condition
or trading strategy from the following list.
The goal is to choose up to **8 indicators** that provide complementary insights without redundancy.
Categories and each category's indicators are:

Moving Averages:
- close_50_sma: 50 SMA: A medium-term trend indicator.
Usage: Identify trend direction and serve as dynamic support/resistance.
Tips: It lags price; combine with faster indicators for timely signals.
- close_200_sma: 200 SMA: A long-term trend benchmark.
Usage: Confirm overall market trend and identify golden/death cross setups.
Tips: It reacts slowly; best for strategic trend confirmation rather than frequent trading entries.
- close_10_ema: 10 EMA: A responsive short-term average.
Usage: Capture quick shifts in momentum and potential entry points.
Tips: Prone to noise in choppy markets; use alongside longer averages for filtering false signals.

MACD Related:
- macd: MACD: Computes momentum via differences of EMAs.
Usage: Look for crossovers and divergence as signals of trend changes.
Tips: Confirm with other indicators in low-volatility or sideways markets.
- macds: MACD Signal: An EMA smoothing of the MACD line.
Usage: Use crossovers with the MACD line to trigger trades.
Tips: Should be part of a broader strategy to avoid false positives.
- macdh: MACD Histogram: Shows the gap between the MACD line and its signal.
Usage: Visualize momentum strength and spot divergence early.
Tips: Can be volatile; complement with additional filters in fast-moving markets.

Momentum Indicators:
- rsi: RSI: Measures momentum to flag overbought/oversold conditions.
Usage: Apply 70/30 thresholds and watch for divergence to signal reversals.
Tips: In strong trends, RSI may remain extreme; always cross-check with trend analysis.

Volatility Indicators:
- boll: Bollinger Middle: A 20 SMA serving as the basis for Bollinger Bands.
Usage: Acts as a dynamic benchmark for price movement.
Tips: Combine with the upper and lower bands to effectively spot breakouts or reversals.
- boll_ub: Bollinger Upper Band: Typically 2 standard deviations above the middle line.
Usage: Signals potential overbought conditions and breakout zones.
Tips: Confirm signals with other tools; prices may ride the band in strong trends.
- boll_lb: Bollinger Lower Band: Typically 2 standard deviations below the middle line.
Usage: Indicates potential oversold conditions.
Tips: Use additional analysis to avoid false reversal signals.
- atr: ATR: Averages true range to measure volatility.
Usage: Set stop-loss levels and adjust position sizes based on current market volatility.
Tips: It's a reactive measure, so use it as part of a broader risk management strategy.

Volume-Based Indicators:
- vwma: VWMA: A moving average weighted by volume.
Usage: Confirm trends by integrating price action with volume data.
Tips: Watch for skewed results from volume spikes; use in combination with other volume analyses.

- Select indicators that provide diverse and complementary information.
Avoid redundancy (e.g., do not select both rsi and stochrsi).
Also briefly explain why they are suitable for the given market context.
When you tool call, please use the exact name of the indicators provided above as they are defined parameters,
otherwise your call will fail.
Please make sure to call get_YFin_data first to retrieve the CSV that is needed to generate indicators.
Write a very detailed and nuanced report of the trends you observe.
Do not simply state the trends are mixed, provide detailed and finegrained analysis
and insights that may help traders make decisions."""
        """ Make sure to append a Markdown table at the end of the report to
        organize key points in the report, organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]
    return await run_chain_with_tools("market", state, llm, system_message, tool_names)

# {{docs-fragment news_analyst}}
@env.task
async def create_news_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [
            toolkit.get_global_news_openai,
            toolkit.get_google_news,
        ]
    else:
        tools = [
            toolkit.get_finnhub_news,
            toolkit.get_reddit_news,
            toolkit.get_google_news,
        ]

    system_message = (
        "You are a news researcher tasked with analyzing recent news and trends over the past week. "
        "Please write a comprehensive report of the current state of the world that is relevant for "
        "trading and macroeconomics. "
        "Look at news from EODHD, and finnhub to be comprehensive. Do not simply state the trends are mixed, "
        "provide detailed and finegrained analysis and insights that may help traders make decisions."
        """ Make sure to append a Markdown table at the end of the report to organize key points in the report,
        organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools("news", state, llm, system_message, tool_names)

# {{/docs-fragment news_analyst}}

@env.task
async def create_social_media_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [toolkit.get_stock_news_openai]
    else:
        tools = [toolkit.get_reddit_stock_info]

    system_message = (
        "You are a social media and company specific news researcher/analyst tasked with analyzing social media posts, "
        "recent company news, and public sentiment for a specific company over the past week. "
        "You will be given a company's name your objective is to write a comprehensive long report "
        "detailing your analysis, insights, and implications for traders and investors on this company's current state "
        "after looking at social media and what people are saying about that company, "
        "analyzing sentiment data of what people feel each day about the company, and looking at recent company news. "
        "Try to look at all sources possible from social media to sentiment to news. Do not simply state the trends "
        "are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions."
        """ Make sure to append a Makrdown table at the end of the report to organize key points in the report,
        organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools(
        "sentiment", state, llm, system_message, tool_names
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/analysts.py*

Each analyst agent uses a helper function to bind tools, iterate through reasoning steps (up to a configurable maximum), and produce an answer. Setting a max iteration count is crucial to prevent runaway loops. As agents reason, their message history is preserved in their internal state and passed along to the next agent in the chain.

```
import asyncio

from agents.utils.utils import AgentState
from flyte_env import env
from langchain_core.messages import ToolMessage, convert_to_openai_messages
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from tools import toolkit

import flyte

MAX_ITERATIONS = 5

# {{docs-fragment agent_helper}}
async def run_chain_with_tools(
    type: str, state: AgentState, llm: str, system_message: str, tool_names: list[str]
) -> AgentState:
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You are a helpful AI assistant, collaborating with other assistants."
                " Use the provided tools to progress towards answering the question."
                " If you are unable to fully answer, that's OK; another assistant with different tools"
                " will help where you left off. Execute what you can to make progress."
                " If you or any other assistant has the FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL** or deliverable,"
                " prefix your response with FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL** so the team knows to stop."
                " You have access to the following tools: {tool_names}.\n{system_message}"
                " For your reference, the current date is {current_date}. The company we want to look at is {ticker}.",
            ),
            MessagesPlaceholder(variable_name="messages"),
        ]
    )

    prompt = prompt.partial(system_message=system_message)
    prompt = prompt.partial(tool_names=", ".join(tool_names))
    prompt = prompt.partial(current_date=state.trade_date)
    prompt = prompt.partial(ticker=state.company_of_interest)

    chain = prompt | ChatOpenAI(model=llm).bind_tools(
        [getattr(toolkit, tool_name).func for tool_name in tool_names]
    )

    iteration = 0
    while iteration < MAX_ITERATIONS:
        result = await chain.ainvoke(state.messages)
        state.messages.append(convert_to_openai_messages(result))

        if not result.tool_calls:
            # Final response — no tools required
            setattr(state, f"{type}_report", result.content or "")
            break

        # Run all tool calls in parallel
        async def run_single_tool(tool_call):
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            tool = getattr(toolkit, tool_name, None)

            if not tool:
                return None

            content = await tool(**tool_args)
            return ToolMessage(
                tool_call_id=tool_call["id"], name=tool_name, content=content
            )

        with flyte.group(f"tool_calls_iteration_{iteration}"):
            tool_messages = await asyncio.gather(
                *[run_single_tool(tc) for tc in result.tool_calls]
            )

        # Add valid tool results to state
        tool_messages = [msg for msg in tool_messages if msg]
        state.messages.extend(convert_to_openai_messages(tool_messages))

        iteration += 1
    else:
        # Reached iteration cap — optionally raise or log
        print(f"Max iterations ({MAX_ITERATIONS}) reached for {type}")

    return state

# {{/docs-fragment agent_helper}}

@env.task
async def create_fundamentals_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [toolkit.get_fundamentals_openai]
    else:
        tools = [
            toolkit.get_finnhub_company_insider_sentiment,
            toolkit.get_finnhub_company_insider_transactions,
            toolkit.get_simfin_balance_sheet,
            toolkit.get_simfin_cashflow,
            toolkit.get_simfin_income_stmt,
        ]

    system_message = (
        "You are a researcher tasked with analyzing fundamental information over the past week about a company. "
        "Please write a comprehensive report of the company's fundamental information such as financial documents, "
        "company profile, basic company financials, company financial history, insider sentiment, and insider "
        "transactions to gain a full view of the company's "
        "fundamental information to inform traders. Make sure to include as much detail as possible. "
        "Do not simply state the trends are mixed, "
        "provide detailed and finegrained analysis and insights that may help traders make decisions. "
        "Make sure to append a Markdown table at the end of the report to organize key points in the report, "
        "organized and easy to read."
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools(
        "fundamentals", state, llm, system_message, tool_names
    )

@env.task
async def create_market_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [
            toolkit.get_YFin_data_online,
            toolkit.get_stockstats_indicators_report_online,
        ]
    else:
        tools = [
            toolkit.get_YFin_data,
            toolkit.get_stockstats_indicators_report,
        ]

    system_message = (
        """You are a trading assistant tasked with analyzing financial markets.
Your role is to select the **most relevant indicators** for a given market condition
or trading strategy from the following list.
The goal is to choose up to **8 indicators** that provide complementary insights without redundancy.
Categories and each category's indicators are:

Moving Averages:
- close_50_sma: 50 SMA: A medium-term trend indicator.
Usage: Identify trend direction and serve as dynamic support/resistance.
Tips: It lags price; combine with faster indicators for timely signals.
- close_200_sma: 200 SMA: A long-term trend benchmark.
Usage: Confirm overall market trend and identify golden/death cross setups.
Tips: It reacts slowly; best for strategic trend confirmation rather than frequent trading entries.
- close_10_ema: 10 EMA: A responsive short-term average.
Usage: Capture quick shifts in momentum and potential entry points.
Tips: Prone to noise in choppy markets; use alongside longer averages for filtering false signals.

MACD Related:
- macd: MACD: Computes momentum via differences of EMAs.
Usage: Look for crossovers and divergence as signals of trend changes.
Tips: Confirm with other indicators in low-volatility or sideways markets.
- macds: MACD Signal: An EMA smoothing of the MACD line.
Usage: Use crossovers with the MACD line to trigger trades.
Tips: Should be part of a broader strategy to avoid false positives.
- macdh: MACD Histogram: Shows the gap between the MACD line and its signal.
Usage: Visualize momentum strength and spot divergence early.
Tips: Can be volatile; complement with additional filters in fast-moving markets.

Momentum Indicators:
- rsi: RSI: Measures momentum to flag overbought/oversold conditions.
Usage: Apply 70/30 thresholds and watch for divergence to signal reversals.
Tips: In strong trends, RSI may remain extreme; always cross-check with trend analysis.

Volatility Indicators:
- boll: Bollinger Middle: A 20 SMA serving as the basis for Bollinger Bands.
Usage: Acts as a dynamic benchmark for price movement.
Tips: Combine with the upper and lower bands to effectively spot breakouts or reversals.
- boll_ub: Bollinger Upper Band: Typically 2 standard deviations above the middle line.
Usage: Signals potential overbought conditions and breakout zones.
Tips: Confirm signals with other tools; prices may ride the band in strong trends.
- boll_lb: Bollinger Lower Band: Typically 2 standard deviations below the middle line.
Usage: Indicates potential oversold conditions.
Tips: Use additional analysis to avoid false reversal signals.
- atr: ATR: Averages true range to measure volatility.
Usage: Set stop-loss levels and adjust position sizes based on current market volatility.
Tips: It's a reactive measure, so use it as part of a broader risk management strategy.

Volume-Based Indicators:
- vwma: VWMA: A moving average weighted by volume.
Usage: Confirm trends by integrating price action with volume data.
Tips: Watch for skewed results from volume spikes; use in combination with other volume analyses.

- Select indicators that provide diverse and complementary information.
Avoid redundancy (e.g., do not select both rsi and stochrsi).
Also briefly explain why they are suitable for the given market context.
When you tool call, please use the exact name of the indicators provided above as they are defined parameters,
otherwise your call will fail.
Please make sure to call get_YFin_data first to retrieve the CSV that is needed to generate indicators.
Write a very detailed and nuanced report of the trends you observe.
Do not simply state the trends are mixed, provide detailed and finegrained analysis
and insights that may help traders make decisions."""
        """ Make sure to append a Markdown table at the end of the report to
        organize key points in the report, organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]
    return await run_chain_with_tools("market", state, llm, system_message, tool_names)

# {{docs-fragment news_analyst}}
@env.task
async def create_news_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [
            toolkit.get_global_news_openai,
            toolkit.get_google_news,
        ]
    else:
        tools = [
            toolkit.get_finnhub_news,
            toolkit.get_reddit_news,
            toolkit.get_google_news,
        ]

    system_message = (
        "You are a news researcher tasked with analyzing recent news and trends over the past week. "
        "Please write a comprehensive report of the current state of the world that is relevant for "
        "trading and macroeconomics. "
        "Look at news from EODHD, and finnhub to be comprehensive. Do not simply state the trends are mixed, "
        "provide detailed and finegrained analysis and insights that may help traders make decisions."
        """ Make sure to append a Markdown table at the end of the report to organize key points in the report,
        organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools("news", state, llm, system_message, tool_names)

# {{/docs-fragment news_analyst}}

@env.task
async def create_social_media_analyst(
    llm: str, state: AgentState, online_tools: bool
) -> AgentState:
    if online_tools:
        tools = [toolkit.get_stock_news_openai]
    else:
        tools = [toolkit.get_reddit_stock_info]

    system_message = (
        "You are a social media and company specific news researcher/analyst tasked with analyzing social media posts, "
        "recent company news, and public sentiment for a specific company over the past week. "
        "You will be given a company's name your objective is to write a comprehensive long report "
        "detailing your analysis, insights, and implications for traders and investors on this company's current state "
        "after looking at social media and what people are saying about that company, "
        "analyzing sentiment data of what people feel each day about the company, and looking at recent company news. "
        "Try to look at all sources possible from social media to sentiment to news. Do not simply state the trends "
        "are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions."
        """ Make sure to append a Makrdown table at the end of the report to organize key points in the report,
        organized and easy to read."""
    )

    tool_names = [tool.func.__name__ for tool in tools]

    return await run_chain_with_tools(
        "sentiment", state, llm, system_message, tool_names
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/analysts.py*

Once all analyst reports are complete, their outputs are collected and passed to the next stage of the workflow.

### Research agents

The research phase consists of two agents: a bullish researcher and a bearish one. They evaluate the company from opposing viewpoints, drawing on the analysts' reports. Unlike analysts, they don't use tools. Their role is to interpret, critique, and develop positions based on the evidence.

```
from agents.utils.utils import AgentState, InvestmentDebateState, memory_init
from flyte_env import env
from langchain_openai import ChatOpenAI

# {{docs-fragment bear_researcher}}
@env.task
async def create_bear_researcher(llm: str, state: AgentState) -> AgentState:
    investment_debate_state = state.investment_debate_state
    history = investment_debate_state.history
    bear_history = investment_debate_state.bear_history

    current_response = investment_debate_state.current_response
    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    memory = await memory_init(name="bear-researcher")

    curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}"
    past_memories = memory.get_memories(curr_situation, n_matches=2)

    past_memory_str = ""
    for rec in past_memories:
        past_memory_str += rec["recommendation"] + "\n\n"

    prompt = f"""You are a Bear Analyst making the case against investing in the stock.
Your goal is to present a well-reasoned argument emphasizing risks, challenges, and negative indicators.
Leverage the provided research and data to highlight potential downsides and counter bullish arguments effectively.

Key points to focus on:

- Risks and Challenges: Highlight factors like market saturation, financial instability,
or macroeconomic threats that could hinder the stock's performance.
- Competitive Weaknesses: Emphasize vulnerabilities such as weaker market positioning, declining innovation,
or threats from competitors.
- Negative Indicators: Use evidence from financial data, market trends, or recent adverse news to support your position.
- Bull Counterpoints: Critically analyze the bull argument with specific data and sound reasoning,
exposing weaknesses or over-optimistic assumptions.
- Engagement: Present your argument in a conversational style, directly engaging with the bull analyst's points
and debating effectively rather than simply listing facts.

Resources available:

Market research report: {market_research_report}
Social media sentiment report: {sentiment_report}
Latest world affairs news: {news_report}
Company fundamentals report: {fundamentals_report}
Conversation history of the debate: {history}
Last bull argument: {current_response}
Reflections from similar situations and lessons learned: {past_memory_str}
Use this information to deliver a compelling bear argument, refute the bull's claims, and engage in a dynamic debate
that demonstrates the risks and weaknesses of investing in the stock.
You must also address reflections and learn from lessons and mistakes you made in the past.
"""

    response = ChatOpenAI(model=llm).invoke(prompt)

    argument = f"Bear Analyst: {response.content}"

    new_investment_debate_state = InvestmentDebateState(
        history=history + "\n" + argument,
        bear_history=bear_history + "\n" + argument,
        bull_history=investment_debate_state.bull_history,
        current_response=argument,
        count=investment_debate_state.count + 1,
    )

    state.investment_debate_state = new_investment_debate_state
    return state

# {{/docs-fragment bear_researcher}}

@env.task
async def create_bull_researcher(llm: str, state: AgentState) -> AgentState:
    investment_debate_state = state.investment_debate_state
    history = investment_debate_state.history
    bull_history = investment_debate_state.bull_history

    current_response = investment_debate_state.current_response
    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    memory = await memory_init(name="bull-researcher")

    curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}"
    past_memories = memory.get_memories(curr_situation, n_matches=2)

    past_memory_str = ""
    for rec in past_memories:
        past_memory_str += rec["recommendation"] + "\n\n"

    prompt = f"""You are a Bull Analyst advocating for investing in the stock.
Your task is to build a strong, evidence-based case emphasizing growth potential, competitive advantages,
and positive market indicators.
Leverage the provided research and data to address concerns and counter bearish arguments effectively.

Key points to focus on:
- Growth Potential: Highlight the company's market opportunities, revenue projections, and scalability.
- Competitive Advantages: Emphasize factors like unique products, strong branding, or dominant market positioning.
- Positive Indicators: Use financial health, industry trends, and recent positive news as evidence.
- Bear Counterpoints: Critically analyze the bear argument with specific data and sound reasoning, addressing
concerns thoroughly and showing why the bull perspective holds stronger merit.
- Engagement: Present your argument in a conversational style, engaging directly with the bear analyst's points
and debating effectively rather than just listing data.

Resources available:
Market research report: {market_research_report}
Social media sentiment report: {sentiment_report}
Latest world affairs news: {news_report}
Company fundamentals report: {fundamentals_report}
Conversation history of the debate: {history}
Last bear argument: {current_response}
Reflections from similar situations and lessons learned: {past_memory_str}
Use this information to deliver a compelling bull argument, refute the bear's concerns, and engage in a dynamic debate
that demonstrates the strengths of the bull position.
You must also address reflections and learn from lessons and mistakes you made in the past.
"""

    response = ChatOpenAI(model=llm).invoke(prompt)

    argument = f"Bull Analyst: {response.content}"

    new_investment_debate_state = InvestmentDebateState(
        history=history + "\n" + argument,
        bull_history=bull_history + "\n" + argument,
        bear_history=investment_debate_state.bear_history,
        current_response=argument,
        count=investment_debate_state.count + 1,
    )

    state.investment_debate_state = new_investment_debate_state
    return state
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/researchers.py*

To aid reasoning, the agents can also retrieve relevant "memories" from a vector database, giving them richer historical context. The number of debate rounds is configurable, and after a few iterations of back-and-forth between the bull and bear, a research manager agent reviews their arguments and makes a final investment decision.

```
from agents.utils.utils import (
    AgentState,
    InvestmentDebateState,
    RiskDebateState,
    memory_init,
)
from flyte_env import env
from langchain_openai import ChatOpenAI

# {{docs-fragment research_manager}}
@env.task
async def create_research_manager(llm: str, state: AgentState) -> AgentState:
    history = state.investment_debate_state.history
    investment_debate_state = state.investment_debate_state
    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    memory = await memory_init(name="research-manager")

    curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}"
    past_memories = memory.get_memories(curr_situation, n_matches=2)

    past_memory_str = ""
    for rec in past_memories:
        past_memory_str += rec["recommendation"] + "\n\n"

    prompt = f"""As the portfolio manager and debate facilitator, your role is to critically evaluate
this round of debate and make a definitive decision:
align with the bear analyst, the bull analyst,
or choose Hold only if it is strongly justified based on the arguments presented.

Summarize the key points from both sides concisely, focusing on the most compelling evidence or reasoning.
Your recommendation—Buy, Sell, or Hold—must be clear and actionable.
Avoid defaulting to Hold simply because both sides have valid points;
commit to a stance grounded in the debate's strongest arguments.

Additionally, develop a detailed investment plan for the trader. This should include:

Your Recommendation: A decisive stance supported by the most convincing arguments.
Rationale: An explanation of why these arguments lead to your conclusion.
Strategic Actions: Concrete steps for implementing the recommendation.
Take into account your past mistakes on similar situations.
Use these insights to refine your decision-making and ensure you are learning and improving.
Present your analysis conversationally, as if speaking naturally, without special formatting.

Here are your past reflections on mistakes:
\"{past_memory_str}\"

Here is the debate:
Debate History:
{history}"""
    response = ChatOpenAI(model=llm).invoke(prompt)

    new_investment_debate_state = InvestmentDebateState(
        judge_decision=response.content,
        history=investment_debate_state.history,
        bear_history=investment_debate_state.bear_history,
        bull_history=investment_debate_state.bull_history,
        current_response=response.content,
        count=investment_debate_state.count,
    )

    state.investment_debate_state = new_investment_debate_state
    state.investment_plan = response.content

    return state

# {{/docs-fragment research_manager}}

@env.task
async def create_risk_manager(llm: str, state: AgentState) -> AgentState:
    history = state.risk_debate_state.history
    risk_debate_state = state.risk_debate_state
    trader_plan = state.investment_plan
    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    memory = await memory_init(name="risk-manager")

    curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}"
    past_memories = memory.get_memories(curr_situation, n_matches=2)

    past_memory_str = ""
    for rec in past_memories:
        past_memory_str += rec["recommendation"] + "\n\n"

    prompt = f"""As the Risk Management Judge and Debate Facilitator,
your goal is to evaluate the debate between three risk analysts—Risky,
Neutral, and Safe/Conservative—and determine the best course of action for the trader.
Your decision must result in a clear recommendation: Buy, Sell, or Hold.
Choose Hold only if strongly justified by specific arguments, not as a fallback when all sides seem valid.
Strive for clarity and decisiveness.

Guidelines for Decision-Making:
1. **Summarize Key Arguments**: Extract the strongest points from each analyst, focusing on relevance to the context.
2. **Provide Rationale**: Support your recommendation with direct quotes and counterarguments from the debate.
3. **Refine the Trader's Plan**: Start with the trader's original plan, **{trader_plan}**,
and adjust it based on the analysts' insights.
4. **Learn from Past Mistakes**: Use lessons from **{past_memory_str}** to address prior misjudgments
and improve the decision you are making now to make sure you don't make a wrong BUY/SELL/HOLD call that loses money.

Deliverables:
- A clear and actionable recommendation: Buy, Sell, or Hold.
- Detailed reasoning anchored in the debate and past reflections.

---

**Analysts Debate History:**
{history}

---

Focus on actionable insights and continuous improvement.
Build on past lessons, critically evaluate all perspectives, and ensure each decision advances better outcomes."""

    response = ChatOpenAI(model=llm).invoke(prompt)

    new_risk_debate_state = RiskDebateState(
        judge_decision=response.content,
        history=risk_debate_state.history,
        risky_history=risk_debate_state.risky_history,
        safe_history=risk_debate_state.safe_history,
        neutral_history=risk_debate_state.neutral_history,
        latest_speaker="Judge",
        current_risky_response=risk_debate_state.current_risky_response,
        current_safe_response=risk_debate_state.current_safe_response,
        current_neutral_response=risk_debate_state.current_neutral_response,
        count=risk_debate_state.count,
    )

    state.risk_debate_state = new_risk_debate_state
    state.final_trade_decision = response.content

    return state
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/managers.py*

### Trading agent

The trader agent consolidates the insights from analysts and researchers to generate a final recommendation. It synthesizes competing signals and produces a conclusion such as _Buy for long-term growth despite short-term volatility_.

```
from agents.utils.utils import AgentState, memory_init
from flyte_env import env
from langchain_core.messages import convert_to_openai_messages
from langchain_openai import ChatOpenAI

# {{docs-fragment trader}}
@env.task
async def create_trader(llm: str, state: AgentState) -> AgentState:
    company_name = state.company_of_interest
    investment_plan = state.investment_plan
    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    memory = await memory_init(name="trader")

    curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}"
    past_memories = memory.get_memories(curr_situation, n_matches=2)

    past_memory_str = ""
    for rec in past_memories:
        past_memory_str += rec["recommendation"] + "\n\n"

    context = {
        "role": "user",
        "content": f"Based on a comprehensive analysis by a team of analysts, "
        f"here is an investment plan tailored for {company_name}. "
        "This plan incorporates insights from current technical market trends, "
        "macroeconomic indicators, and social media sentiment. "
        "Use this plan as a foundation for evaluating your next trading decision.\n\n"
        f"Proposed Investment Plan: {investment_plan}\n\n"
        "Leverage these insights to make an informed and strategic decision.",
    }

    messages = [
        {
            "role": "system",
            "content": f"""You are a trading agent analyzing market data to make investment decisions.
Based on your analysis, provide a specific recommendation to buy, sell, or hold.
End with a firm decision and always conclude your response with 'FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL**'
to confirm your recommendation.
Do not forget to utilize lessons from past decisions to learn from your mistakes.
Here is some reflections from similar situatiosn you traded in and the lessons learned: {past_memory_str}""",
        },
        context,
    ]

    result = ChatOpenAI(model=llm).invoke(messages)

    state.messages.append(convert_to_openai_messages(result))
    state.trader_investment_plan = result.content
    state.sender = "Trader"

    return state

# {{/docs-fragment trader}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/trader.py*

### Risk agents

Risk agents comprise agents with different risk tolerances: a risky debater, a neutral one, and a conservative one. They assess the portfolio through lenses like market volatility, liquidity, and systemic risk. Similar to the bull-bear debate, these agents engage in internal discussion, after which a risk manager makes the final call.

```
from agents.utils.utils import AgentState, RiskDebateState
from flyte_env import env
from langchain_openai import ChatOpenAI

# {{docs-fragment risk_debator}}
@env.task
async def create_risky_debator(llm: str, state: AgentState) -> AgentState:
    risk_debate_state = state.risk_debate_state
    history = risk_debate_state.history
    risky_history = risk_debate_state.risky_history

    current_safe_response = risk_debate_state.current_safe_response
    current_neutral_response = risk_debate_state.current_neutral_response

    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    trader_decision = state.trader_investment_plan

    prompt = f"""As the Risky Risk Analyst, your role is to actively champion high-reward, high-risk opportunities,
emphasizing bold strategies and competitive advantages.
When evaluating the trader's decision or plan, focus intently on the potential upside, growth potential,
and innovative benefits—even when these come with elevated risk.
Use the provided market data and sentiment analysis to strengthen your arguments and challenge the opposing views.
Specifically, respond directly to each point made by the conservative and neutral analysts,
countering with data-driven rebuttals and persuasive reasoning.
Highlight where their caution might miss critical opportunities or where their assumptions may be overly conservative.
Here is the trader's decision:

{trader_decision}

Your task is to create a compelling case for the trader's decision by questioning and critiquing the conservative
and neutral stances to demonstrate why your high-reward perspective offers the best path forward.
Incorporate insights from the following sources into your arguments:

Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report}
Company Fundamentals Report: {fundamentals_report}
Here is the current conversation history: {history}
Here are the last arguments from the conservative analyst: {current_safe_response}
Here are the last arguments from the neutral analyst: {current_neutral_response}.
If there are no responses from the other viewpoints, do not halluncinate and just present your point.

Engage actively by addressing any specific concerns raised, refuting the weaknesses in their logic,
and asserting the benefits of risk-taking to outpace market norms.
Maintain a focus on debating and persuading, not just presenting data.
Challenge each counterpoint to underscore why a high-risk approach is optimal.
Output conversationally as if you are speaking without any special formatting."""

    response = ChatOpenAI(model=llm).invoke(prompt)

    argument = f"Risky Analyst: {response.content}"

    new_risk_debate_state = RiskDebateState(
        history=history + "\n" + argument,
        risky_history=risky_history + "\n" + argument,
        safe_history=risk_debate_state.safe_history,
        neutral_history=risk_debate_state.neutral_history,
        latest_speaker="Risky",
        current_risky_response=argument,
        current_safe_response=current_safe_response,
        current_neutral_response=current_neutral_response,
        count=risk_debate_state.count + 1,
    )

    state.risk_debate_state = new_risk_debate_state
    return state

# {{/docs-fragment risk_debator}}

@env.task
async def create_safe_debator(llm: str, state: AgentState) -> AgentState:
    risk_debate_state = state.risk_debate_state
    history = risk_debate_state.history
    safe_history = risk_debate_state.safe_history

    current_risky_response = risk_debate_state.current_risky_response
    current_neutral_response = risk_debate_state.current_neutral_response

    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    trader_decision = state.trader_investment_plan

    prompt = f"""As the Safe/Conservative Risk Analyst, your primary objective is to protect assets,
minimize volatility, and ensure steady, reliable growth. You prioritize stability, security, and risk mitigation,
carefully assessing potential losses, economic downturns, and market volatility.
When evaluating the trader's decision or plan, critically examine high-risk elements,
pointing out where the decision may expose the firm to undue risk and where more cautious
alternatives could secure long-term gains.
Here is the trader's decision:

{trader_decision}

Your task is to actively counter the arguments of the Risky and Neutral Analysts,
highlighting where their views may overlook potential threats or fail to prioritize sustainability.
Respond directly to their points, drawing from the following data sources
to build a convincing case for a low-risk approach adjustment to the trader's decision:

Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report}
Company Fundamentals Report: {fundamentals_report}
Here is the current conversation history: {history}
Here is the last response from the risky analyst: {current_risky_response}
Here is the last response from the neutral analyst: {current_neutral_response}.
If there are no responses from the other viewpoints, do not halluncinate and just present your point.

Engage by questioning their optimism and emphasizing the potential downsides they may have overlooked.
Address each of their counterpoints to showcase why a conservative stance is ultimately the
safest path for the firm's assets.
Focus on debating and critiquing their arguments to demonstrate the strength of a low-risk strategy
over their approaches.
Output conversationally as if you are speaking without any special formatting."""

    response = ChatOpenAI(model=llm).invoke(prompt)

    argument = f"Safe Analyst: {response.content}"

    new_risk_debate_state = RiskDebateState(
        history=history + "\n" + argument,
        risky_history=risk_debate_state.risky_history,
        safe_history=safe_history + "\n" + argument,
        neutral_history=risk_debate_state.neutral_history,
        latest_speaker="Safe",
        current_risky_response=current_risky_response,
        current_safe_response=argument,
        current_neutral_response=current_neutral_response,
        count=risk_debate_state.count + 1,
    )

    state.risk_debate_state = new_risk_debate_state
    return state

@env.task
async def create_neutral_debator(llm: str, state: AgentState) -> AgentState:
    risk_debate_state = state.risk_debate_state
    history = risk_debate_state.history
    neutral_history = risk_debate_state.neutral_history

    current_risky_response = risk_debate_state.current_risky_response
    current_safe_response = risk_debate_state.current_safe_response

    market_research_report = state.market_report
    sentiment_report = state.sentiment_report
    news_report = state.news_report
    fundamentals_report = state.fundamentals_report

    trader_decision = state.trader_investment_plan

    prompt = f"""As the Neutral Risk Analyst, your role is to provide a balanced perspective,
weighing both the potential benefits and risks of the trader's decision or plan.
You prioritize a well-rounded approach, evaluating the upsides
and downsides while factoring in broader market trends,
potential economic shifts, and diversification strategies.Here is the trader's decision:

{trader_decision}

Your task is to challenge both the Risky and Safe Analysts,
pointing out where each perspective may be overly optimistic or overly cautious.
Use insights from the following data sources to support a moderate, sustainable strategy
to adjust the trader's decision:

Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report}
Company Fundamentals Report: {fundamentals_report}
Here is the current conversation history: {history}
Here is the last response from the risky analyst: {current_risky_response}
Here is the last response from the safe analyst: {current_safe_response}.
If there are no responses from the other viewpoints, do not halluncinate and just present your point.

Engage actively by analyzing both sides critically, addressing weaknesses in the risky
and conservative arguments to advocate for a more balanced approach.
Challenge each of their points to illustrate why a moderate risk strategy might offer the best of both worlds,
providing growth potential while safeguarding against extreme volatility.
Focus on debating rather than simply presenting data, aiming to show that a balanced view can lead to
the most reliable outcomes. Output conversationally as if you are speaking without any special formatting."""

    response = ChatOpenAI(model=llm).invoke(prompt)

    argument = f"Neutral Analyst: {response.content}"

    new_risk_debate_state = RiskDebateState(
        history=history + "\n" + argument,
        risky_history=risk_debate_state.risky_history,
        safe_history=risk_debate_state.safe_history,
        neutral_history=neutral_history + "\n" + argument,
        latest_speaker="Neutral",
        current_risky_response=current_risky_response,
        current_safe_response=current_safe_response,
        current_neutral_response=argument,
        count=risk_debate_state.count + 1,
    )

    state.risk_debate_state = new_risk_debate_state
    return state
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/agents/risk_debators.py*

The outcome of the risk manager, whether to proceed with the trade or not, is considered the final decision of the trading simulation.

You can visualize this full pipeline in the Flyte/Union UI, where every step is logged.
You’ll see input/output metadata for each tool and agent task.
Thanks to Flyte's caching, repeated steps are skipped unless inputs change, saving time and compute resources.

### Retaining agent memory with S3 vectors

To help agents learn from past decisions, we persist their memory in a vector store. In this example, we use an [S3 vector](https://aws.amazon.com/s3/features/vectors/) bucket for their simplicity and tight integration with Flyte and Union, but any vector database can be used.

Note: To use the S3 vector store, make sure your IAM role has the following permissions configured:

```
s3vectors:CreateVectorBucket
s3vectors:CreateIndex
s3vectors:PutVectors
s3vectors:GetIndex
s3vectors:GetVectors
s3vectors:QueryVectors
s3vectors:GetVectorBucket
```

After each trade decision, you can run a `reflect_on_decisions` task. This evaluates whether the final outcome aligned with the agent's recommendation and stores that reflection in the vector store. These stored insights can later be retrieved to provide historical context and improve future decision-making.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.0.0b52",
#     "akshare==1.16.98",
#     "backtrader==1.9.78.123",
#     "boto3==1.39.9",
#     "chainlit==2.5.5",
#     "eodhd==1.0.32",
#     "feedparser==6.0.11",
#     "finnhub-python==2.4.23",
#     "langchain-experimental==0.3.4",
#     "langchain-openai==0.3.23",
#     "pandas==2.3.0",
#     "parsel==1.10.0",
#     "praw==7.8.1",
#     "pytz==2025.2",
#     "questionary==2.1.0",
#     "redis==6.2.0",
#     "requests==2.32.4",
#     "stockstats==0.6.5",
#     "tqdm==4.67.1",
#     "tushare==1.4.21",
#     "typing-extensions==4.14.0",
#     "yfinance==0.2.63",
# ]
# main = "main"
# params = ""
# ///
import asyncio
from copy import deepcopy

import agents
import agents.analysts
from agents.managers import create_research_manager, create_risk_manager
from agents.researchers import create_bear_researcher, create_bull_researcher
from agents.risk_debators import (
    create_neutral_debator,
    create_risky_debator,
    create_safe_debator,
)
from agents.trader import create_trader
from agents.utils.utils import AgentState
from flyte_env import DEEP_THINKING_LLM, QUICK_THINKING_LLM, env, flyte
from langchain_openai import ChatOpenAI
from reflection import (
    reflect_bear_researcher,
    reflect_bull_researcher,
    reflect_research_manager,
    reflect_risk_manager,
    reflect_trader,
)

@env.task
async def process_signal(full_signal: str, QUICK_THINKING_LLM: str) -> str:
    """Process a full trading signal to extract the core decision."""

    messages = [
        {
            "role": "system",
            "content": """You are an efficient assistant designed to analyze paragraphs or
financial reports provided by a group of analysts.
Your task is to extract the investment decision: SELL, BUY, or HOLD.
Provide only the extracted decision (SELL, BUY, or HOLD) as your output,
without adding any additional text or information.""",
        },
        {"role": "human", "content": full_signal},
    ]

    return ChatOpenAI(model=QUICK_THINKING_LLM).invoke(messages).content

async def run_analyst(analyst_name, state, online_tools):
    # Create a copy of the state for isolation
    run_fn = getattr(agents.analysts, f"create_{analyst_name}_analyst")

    # Run the analyst's chain
    result_state = await run_fn(QUICK_THINKING_LLM, state, online_tools)

    # Determine the report key
    report_key = (
        "sentiment_report"
        if analyst_name == "social_media"
        else f"{analyst_name}_report"
    )
    report_value = getattr(result_state, report_key)

    return result_state.messages[1:], report_key, report_value

# {{docs-fragment main}}
@env.task
async def main(
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> tuple[str, AgentState]:
    if not selected_analysts:
        raise ValueError(
            "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social_media."
        )

    state = AgentState(
        messages=[{"role": "human", "content": company_name}],
        company_of_interest=company_name,
        trade_date=str(trade_date),
    )

    # Run all analysts concurrently
    results = await asyncio.gather(
        *[
            run_analyst(analyst, deepcopy(state), online_tools)
            for analyst in selected_analysts
        ]
    )

    # Flatten and append all resulting messages into the shared state
    for messages, report_attr, report in results:
        state.messages.extend(messages)
        setattr(state, report_attr, report)

    # Bull/Bear debate loop
    state = await create_bull_researcher(QUICK_THINKING_LLM, state)  # Start with bull
    while state.investment_debate_state.count < 2 * max_debate_rounds:
        current = state.investment_debate_state.current_response
        if current.startswith("Bull"):
            state = await create_bear_researcher(QUICK_THINKING_LLM, state)
        else:
            state = await create_bull_researcher(QUICK_THINKING_LLM, state)

    state = await create_research_manager(DEEP_THINKING_LLM, state)
    state = await create_trader(QUICK_THINKING_LLM, state)

    # Risk debate loop
    state = await create_risky_debator(QUICK_THINKING_LLM, state)  # Start with risky
    while state.risk_debate_state.count < 3 * max_risk_discuss_rounds:
        speaker = state.risk_debate_state.latest_speaker
        if speaker == "Risky":
            state = await create_safe_debator(QUICK_THINKING_LLM, state)
        elif speaker == "Safe":
            state = await create_neutral_debator(QUICK_THINKING_LLM, state)
        else:
            state = await create_risky_debator(QUICK_THINKING_LLM, state)

    state = await create_risk_manager(DEEP_THINKING_LLM, state)
    decision = await process_signal(state.final_trade_decision, QUICK_THINKING_LLM)

    return decision, state

# {{/docs-fragment main}}

# {{docs-fragment reflect_on_decisions}}
@env.task
async def reflect_and_store(state: AgentState, returns: str) -> str:
    await asyncio.gather(
        reflect_bear_researcher(state, returns),
        reflect_bull_researcher(state, returns),
        reflect_trader(state, returns),
        reflect_risk_manager(state, returns),
        reflect_research_manager(state, returns),
    )

    return "Reflection completed."

# Run the reflection task after the main function
@env.task(cache="disable")
async def reflect_on_decisions(
    returns: str,
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> str:
    _, state = await main(
        selected_analysts,
        max_debate_rounds,
        max_risk_discuss_rounds,
        online_tools,
        company_name,
        trade_date,
    )

    return await reflect_and_store(state, returns)

# {{/docs-fragment reflect_on_decisions}}

# {{docs-fragment execute_main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

    # run = flyte.run(reflect_on_decisions, "+3.2% gain over 5 days")
    # print(run.url)

# {{/docs-fragment execute_main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/main.py*

### Running the simulation

First, set up your OpenAI secret (from [openai.com](https://platform.openai.com/api-keys)) and Finnhub API key (from [finnhub.io](https://finnhub.io/)):

```
flyte create secret openai_api_key <YOUR_OPENAI_API_KEY>
flyte create secret finnhub_api_key <YOUR_FINNHUB_API_KEY>
```

Then [clone the repo](https://github.com/unionai/unionai-examples), navigate to the `tutorials-v2/trading_agents` directory, and run the following commands:

```
flyte create config --endpoint <FLYTE_OR_UNION_ENDPOINT> --project <PROJECT_NAME> --domain <DOMAIN_NAME> --builder remote
uv run main.py
```

If you'd like to run the `reflect_on_decisions` task instead, comment out the `main` function call and uncomment the `reflect_on_decisions` call in the `__main__` block:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.0.0b52",
#     "akshare==1.16.98",
#     "backtrader==1.9.78.123",
#     "boto3==1.39.9",
#     "chainlit==2.5.5",
#     "eodhd==1.0.32",
#     "feedparser==6.0.11",
#     "finnhub-python==2.4.23",
#     "langchain-experimental==0.3.4",
#     "langchain-openai==0.3.23",
#     "pandas==2.3.0",
#     "parsel==1.10.0",
#     "praw==7.8.1",
#     "pytz==2025.2",
#     "questionary==2.1.0",
#     "redis==6.2.0",
#     "requests==2.32.4",
#     "stockstats==0.6.5",
#     "tqdm==4.67.1",
#     "tushare==1.4.21",
#     "typing-extensions==4.14.0",
#     "yfinance==0.2.63",
# ]
# main = "main"
# params = ""
# ///
import asyncio
from copy import deepcopy

import agents
import agents.analysts
from agents.managers import create_research_manager, create_risk_manager
from agents.researchers import create_bear_researcher, create_bull_researcher
from agents.risk_debators import (
    create_neutral_debator,
    create_risky_debator,
    create_safe_debator,
)
from agents.trader import create_trader
from agents.utils.utils import AgentState
from flyte_env import DEEP_THINKING_LLM, QUICK_THINKING_LLM, env, flyte
from langchain_openai import ChatOpenAI
from reflection import (
    reflect_bear_researcher,
    reflect_bull_researcher,
    reflect_research_manager,
    reflect_risk_manager,
    reflect_trader,
)

@env.task
async def process_signal(full_signal: str, QUICK_THINKING_LLM: str) -> str:
    """Process a full trading signal to extract the core decision."""

    messages = [
        {
            "role": "system",
            "content": """You are an efficient assistant designed to analyze paragraphs or
financial reports provided by a group of analysts.
Your task is to extract the investment decision: SELL, BUY, or HOLD.
Provide only the extracted decision (SELL, BUY, or HOLD) as your output,
without adding any additional text or information.""",
        },
        {"role": "human", "content": full_signal},
    ]

    return ChatOpenAI(model=QUICK_THINKING_LLM).invoke(messages).content

async def run_analyst(analyst_name, state, online_tools):
    # Create a copy of the state for isolation
    run_fn = getattr(agents.analysts, f"create_{analyst_name}_analyst")

    # Run the analyst's chain
    result_state = await run_fn(QUICK_THINKING_LLM, state, online_tools)

    # Determine the report key
    report_key = (
        "sentiment_report"
        if analyst_name == "social_media"
        else f"{analyst_name}_report"
    )
    report_value = getattr(result_state, report_key)

    return result_state.messages[1:], report_key, report_value

# {{docs-fragment main}}
@env.task
async def main(
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> tuple[str, AgentState]:
    if not selected_analysts:
        raise ValueError(
            "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social_media."
        )

    state = AgentState(
        messages=[{"role": "human", "content": company_name}],
        company_of_interest=company_name,
        trade_date=str(trade_date),
    )

    # Run all analysts concurrently
    results = await asyncio.gather(
        *[
            run_analyst(analyst, deepcopy(state), online_tools)
            for analyst in selected_analysts
        ]
    )

    # Flatten and append all resulting messages into the shared state
    for messages, report_attr, report in results:
        state.messages.extend(messages)
        setattr(state, report_attr, report)

    # Bull/Bear debate loop
    state = await create_bull_researcher(QUICK_THINKING_LLM, state)  # Start with bull
    while state.investment_debate_state.count < 2 * max_debate_rounds:
        current = state.investment_debate_state.current_response
        if current.startswith("Bull"):
            state = await create_bear_researcher(QUICK_THINKING_LLM, state)
        else:
            state = await create_bull_researcher(QUICK_THINKING_LLM, state)

    state = await create_research_manager(DEEP_THINKING_LLM, state)
    state = await create_trader(QUICK_THINKING_LLM, state)

    # Risk debate loop
    state = await create_risky_debator(QUICK_THINKING_LLM, state)  # Start with risky
    while state.risk_debate_state.count < 3 * max_risk_discuss_rounds:
        speaker = state.risk_debate_state.latest_speaker
        if speaker == "Risky":
            state = await create_safe_debator(QUICK_THINKING_LLM, state)
        elif speaker == "Safe":
            state = await create_neutral_debator(QUICK_THINKING_LLM, state)
        else:
            state = await create_risky_debator(QUICK_THINKING_LLM, state)

    state = await create_risk_manager(DEEP_THINKING_LLM, state)
    decision = await process_signal(state.final_trade_decision, QUICK_THINKING_LLM)

    return decision, state

# {{/docs-fragment main}}

# {{docs-fragment reflect_on_decisions}}
@env.task
async def reflect_and_store(state: AgentState, returns: str) -> str:
    await asyncio.gather(
        reflect_bear_researcher(state, returns),
        reflect_bull_researcher(state, returns),
        reflect_trader(state, returns),
        reflect_risk_manager(state, returns),
        reflect_research_manager(state, returns),
    )

    return "Reflection completed."

# Run the reflection task after the main function
@env.task(cache="disable")
async def reflect_on_decisions(
    returns: str,
    selected_analysts: list[str] = [
        "market",
        "fundamentals",
        "news",
        "social_media",
    ],
    max_debate_rounds: int = 1,
    max_risk_discuss_rounds: int = 1,
    online_tools: bool = True,
    company_name: str = "NVDA",
    trade_date: str = "2024-05-12",
) -> str:
    _, state = await main(
        selected_analysts,
        max_debate_rounds,
        max_risk_discuss_rounds,
        online_tools,
        company_name,
        trade_date,
    )

    return await reflect_and_store(state, returns)

# {{/docs-fragment reflect_on_decisions}}

# {{docs-fragment execute_main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

    # run = flyte.run(reflect_on_decisions, "+3.2% gain over 5 days")
    # print(run.url)

# {{/docs-fragment execute_main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading_agents/main.py*

Then run:

```
uv run main.py
```

## Why Flyte? _(A quick note before you go)_

You might now be wondering: can't I just build all this with Python and LangChain?
Absolutely. But as your project grows, you'll likely run into these challenges:

1. **Observability**: Agent workflows can feel opaque. You send a prompt, get a response, but what happened in between?

    - Were the right tools used?
    - Were correct arguments passed?
    - How did the LLM reason through intermediate steps?
    - Why did it fail?

    Flyte gives you a window into each of these stages.

2. **Multi-agent coordination**: Real-world applications often require multiple agents with distinct roles and responsibilities. In such cases, you'll need:

    - Isolated state per agent,
    - Shared context where needed,
    - And coordination: sequential or parallel.

    Managing this manually gets fragile, fast. Flyte handles it for you.

3. **Scalability**: Agents and tools might need to run in isolated or containerized environments. Whether you're scaling out to more agents or more powerful hardware, Flyte lets you scale without taxing your local machine or racking up unnecessary cloud bills.
4. **Durability & recovery**: LLM-based workflows are often long-running and expensive. If something fails halfway:

    - Do you lose all progress?
    - Replay everything from scratch?

    With Flyte, you get built-in caching, checkpointing, and recovery, so you can resume where you left off.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services/fraud-detection-feast ===

# Fraud detection with Feast

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/fraud_detection_feast).

This tutorial builds a credit-card fraud detection pipeline that combines [Feast](https://feast.dev/) feature store materialization with an XGBoost classifier on the Sparkov simulated transactions dataset. The workflow engineers transaction and user-level features, trains a model, registers features in Feast, and materializes online feature values for low-latency scoring.

Flyte provides:

- **Cached data preparation** for the Kaggle dataset download and feature engineering.
- **Report-backed training** with confusion matrix and ROC-style metrics in the UI.
- **Durable artifacts**: the trained model and Feast repo are returned as `flyte.io.File` and `flyte.io.Dir`.

## Define the task environment

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "feast==0.63.0",
#    "scikit-learn==1.8.0",
#    "xgboost==3.2.0",
#    "joblib",
#    "pandas",
#    "pyarrow",
#    "kagglehub==0.3.12",
# ]
# main = "fraud_detection_pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import shutil
import tempfile
from datetime import datetime, timedelta, timezone

import joblib
import numpy as np
import pandas as pd
import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="fraud-detection-feast", pre=True)

env = flyte.TaskEnvironment(
    name="fraud-detection-feast",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)
# {{/docs-fragment env}}

import report_helpers as rh

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Feature definitions
#
# Transaction features: known at scoring time (from the request)
# User features: pre-computed aggregates stored in Feast
# Derived features: computed at both training and scoring time by
#                   comparing the transaction to the user's profile
# ------------------------------------------------------------------

TXN_FEATURE_COLS = ["amt", "amt_log", "category_encoded", "merch_lat", "merch_long"]

USER_FEATURE_COLS = [
    "txn_count", "mean_amt", "std_amt", "max_amt",
    "home_lat", "home_long", "age",
]

DERIVED_FEATURE_COLS = [
    "amt_zscore", "amt_ratio", "distance_from_home", "hour", "day_of_week",
]

ALL_FEATURE_COLS = TXN_FEATURE_COLS + USER_FEATURE_COLS + DERIVED_FEATURE_COLS

def haversine(lat1, lon1, lat2, lon2):
    """Compute distance in miles between two (lat, lon) points."""
    R = 3959  # Earth radius in miles
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    return 2 * R * np.arcsin(np.sqrt(a))

# ------------------------------------------------------------------
# Task 1: Download dataset and engineer features
# ------------------------------------------------------------------

@env.task(report=True, cache="auto")
async def prepare_data() -> flyte.io.Dir:
    """Download the Sparkov credit card fraud dataset and prepare parquets."""
    import kagglehub

    log.info("Downloading dataset...")
    dataset_path = kagglehub.dataset_download("kartik2112/fraud-detection")
    csv_path = os.path.join(dataset_path, "fraudTrain.csv")
    df = pd.read_csv(csv_path)
    log.info(f"Loaded {len(df):,} transactions ({int(df['is_fraud'].sum()):,} fraudulent)")

    # Sample for workshop speed (stratified to preserve fraud ratio)
    if len(df) > 500_000:
        from sklearn.model_selection import train_test_split
        df, _ = train_test_split(df, train_size=500_000, stratify=df["is_fraud"], random_state=42)
        log.info(f"Sampled to {len(df):,} transactions")

    # ------------------------------------------------------------------
    # Parse timestamps
    # ------------------------------------------------------------------
    df["event_timestamp"] = pd.to_datetime(df["trans_date_trans_time"])
    df["event_timestamp"] = df["event_timestamp"].dt.tz_localize("UTC")
    df["hour"] = df["event_timestamp"].dt.hour
    df["day_of_week"] = df["event_timestamp"].dt.dayofweek

    # ------------------------------------------------------------------
    # Map cc_num → sequential user_id for clean API
    # ------------------------------------------------------------------
    cc_nums = df["cc_num"].unique()
    cc_to_user = {cc: i for i, cc in enumerate(sorted(cc_nums))}
    df["user_id"] = df["cc_num"].map(cc_to_user)

    # ------------------------------------------------------------------
    # Feature engineering
    # ------------------------------------------------------------------
    df["amt_log"] = np.log1p(df["amt"])

    # Label-encode merchant category
    categories = sorted(df["category"].unique())
    cat_to_int = {cat: i for i, cat in enumerate(categories)}
    df["category_encoded"] = df["category"].map(cat_to_int)

    # Compute age from dob
    df["dob"] = pd.to_datetime(df["dob"]).dt.tz_localize("UTC")
    ref_date = df["event_timestamp"].max()
    df["age"] = ((ref_date - df["dob"]).dt.days / 365.25).astype(int)

    # Distance between buyer and merchant
    df["distance"] = haversine(df["lat"], df["long"], df["merch_lat"], df["merch_long"])

    # ------------------------------------------------------------------
    # Build user aggregates
    # ------------------------------------------------------------------
    user_stats = df.groupby("user_id").agg(
        txn_count=("amt", "count"),
        mean_amt=("amt", "mean"),
        std_amt=("amt", "std"),
        max_amt=("amt", "max"),
        home_lat=("lat", "median"),
        home_long=("long", "median"),
        age=("age", "first"),
    ).reset_index()
    user_stats["std_amt"] = user_stats["std_amt"].fillna(0)
    # Use earliest timestamp so Feast point-in-time joins work for all transactions
    earliest_ts = df.groupby("user_id")["event_timestamp"].min().reset_index()
    user_stats = user_stats.merge(earliest_ts, on="user_id")

    # ------------------------------------------------------------------
    # Save to temp directory
    # ------------------------------------------------------------------
    data_dir = tempfile.mkdtemp()

    txn_cols = [
        "user_id", "event_timestamp",
        "amt", "amt_log", "category_encoded", "merch_lat", "merch_long",
        "hour", "day_of_week", "lat", "long", "distance",
        "is_fraud",
    ]
    df[txn_cols].to_parquet(os.path.join(data_dir, "transactions.parquet"), index=False)
    user_stats.to_parquet(os.path.join(data_dir, "user_features.parquet"), index=False)

    # Save category mapping + cc_num mapping for the app
    with open(os.path.join(data_dir, "category_mapping.json"), "w") as f:
        json.dump(cat_to_int, f)
    with open(os.path.join(data_dir, "user_mapping.json"), "w") as f:
        json.dump({str(k): v for k, v in cc_to_user.items()}, f)

    n_fraud = int(df["is_fraud"].sum())
    n_legit = len(df) - n_fraud
    fraud_pct = df["is_fraud"].mean() * 100
    html = (
        '<h2>Data Prepared</h2>'
        + rh.stat_grid([
            (f"{len(df):,}", "Transactions"),
            (f"{n_fraud:,}", "Fraudulent"),
            (f"{fraud_pct:.2f}%", "Fraud Rate"),
            (f"{user_stats['user_id'].nunique():,}", "Users"),
            (f"{len(categories)}", "Categories"),
        ])
        + rh.class_distribution_bar(n_legit, n_fraud)
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    return await flyte.io.Dir.from_local(data_dir)

# ------------------------------------------------------------------
# Task 2: Set up Feast and materialize user profiles to online store
# ------------------------------------------------------------------

@env.task(report=True)
async def materialize_features(data_dir: flyte.io.Dir) -> flyte.io.Dir:
    """Apply Feast definitions and materialize user profiles to SQLite online store."""
    from feast import Entity, FeatureStore, FeatureView, Field, FileSource
    from feast.types import Float64, Int64

    data_path = await data_dir.download()

    # Create a self-contained Feast repo in a temp directory
    feast_dir = tempfile.mkdtemp()

    # Copy parquet into feast dir so the repo is fully self-contained
    shutil.copy2(
        os.path.join(data_path, "user_features.parquet"),
        os.path.join(feast_dir, "user_features.parquet"),
    )

    # Write feature_store.yaml
    yaml_content = (
        "project: fraud_detection\n"
        f"registry: {feast_dir}/registry.db\n"
        "provider: local\n"
        "online_store:\n"
        "  type: sqlite\n"
        f"  path: {feast_dir}/online_store.db\n"
        "offline_store:\n"
        "  type: file\n"
        "entity_key_serialization_version: 3\n"
    )
    yaml_path = os.path.join(feast_dir, "feature_store.yaml")
    with open(yaml_path, "w") as f:
        f.write(yaml_content)

    store = FeatureStore(repo_path=feast_dir)

    # Define entity and feature view
    user = Entity(name="user", join_keys=["user_id"], description="Credit card holder")

    user_source = FileSource(
        path=os.path.join(feast_dir, "user_features.parquet"),
        timestamp_field="event_timestamp",
    )

    user_stats = FeatureView(
        name="user_stats",
        entities=[user],
        ttl=timedelta(days=0),  # No expiry — workshop data has old timestamps
        schema=[
            Field(name="txn_count", dtype=Int64),
            Field(name="mean_amt", dtype=Float64),
            Field(name="std_amt", dtype=Float64),
            Field(name="max_amt", dtype=Float64),
            Field(name="home_lat", dtype=Float64),
            Field(name="home_long", dtype=Float64),
            Field(name="age", dtype=Int64),
        ],
        online=True,
        source=user_source,
    )

    # Apply and materialize
    log.info("Applying Feast definitions...")
    store.apply([user, user_stats])

    log.info("Materializing user profiles to online store...")
    store.materialize(
        start_date=datetime(2018, 1, 1, tzinfo=timezone.utc),
        end_date=datetime.now(timezone.utc),
    )

    # Re-apply with relative paths so the registry is portable across workers
    portable_yaml = (
        "project: fraud_detection\n"
        "registry: registry.db\n"
        "provider: local\n"
        "online_store:\n"
        "  type: sqlite\n"
        "  path: online_store.db\n"
        "offline_store:\n"
        "  type: file\n"
        "entity_key_serialization_version: 3\n"
    )
    with open(yaml_path, "w") as f:
        f.write(portable_yaml)

    # Re-apply with relative source path so get_historical_features works on other workers
    store = FeatureStore(repo_path=feast_dir)
    user_source = FileSource(
        path="user_features.parquet",
        timestamp_field="event_timestamp",
    )
    user_stats = FeatureView(
        name="user_stats",
        entities=[user],
        ttl=timedelta(days=0),
        schema=[
            Field(name="txn_count", dtype=Int64),
            Field(name="mean_amt", dtype=Float64),
            Field(name="std_amt", dtype=Float64),
            Field(name="max_amt", dtype=Float64),
            Field(name="home_lat", dtype=Float64),
            Field(name="home_long", dtype=Float64),
            Field(name="age", dtype=Int64),
        ],
        online=True,
        source=user_source,
    )
    store.apply([user, user_stats])

    features = ["txn_count", "mean_amt", "std_amt", "max_amt", "home_lat", "home_long", "age"]
    html = (
        '<h2>Feature Store Materialized</h2>'
        + rh.stat_grid([
            ("user_stats", "Feature View"),
            (str(len(features)), "Features"),
            ("SQLite", "Online Store"),
        ])
        + '<h3>Materialized Features</h3>'
        '<table>'
        '<tr><th>Feature</th><th>Type</th><th>Description</th></tr>'
        '<tr><td>txn_count</td><td><span class="badge badge-info">Int64</span></td><td>Total transactions</td></tr>'
        '<tr><td>mean_amt</td><td><span class="badge badge-info">Float64</span></td><td>Average transaction amount</td></tr>'
        '<tr><td>std_amt</td><td><span class="badge badge-info">Float64</span></td><td>Std dev of amounts</td></tr>'
        '<tr><td>max_amt</td><td><span class="badge badge-info">Float64</span></td><td>Max transaction amount</td></tr>'
        '<tr><td>home_lat</td><td><span class="badge badge-info">Float64</span></td><td>Home latitude (median)</td></tr>'
        '<tr><td>home_long</td><td><span class="badge badge-info">Float64</span></td><td>Home longitude (median)</td></tr>'
        '<tr><td>age</td><td><span class="badge badge-info">Int64</span></td><td>User age</td></tr>'
        '</table>'
        '<div class="note">User profiles are ready for real-time serving via the scoring app.</div>'
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    return await flyte.io.Dir.from_local(feast_dir)

# ------------------------------------------------------------------
# Task 3: Train XGBoost model
# ------------------------------------------------------------------

@env.task(report=True)
async def train_model(
    data_dir: flyte.io.Dir,
    feast_dir: flyte.io.Dir,
    n_estimators: int = 300,
    max_depth: int = 6,
    learning_rate: float = 0.1,
    min_child_weight: int = 5,
    gamma: float = 1.0,
) -> flyte.io.File:
    """Train an XGBoost classifier using Feast for feature retrieval."""
    from feast import FeatureStore
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import classification_report, roc_auc_score, confusion_matrix
    from xgboost import XGBClassifier

    data_path = await data_dir.download()
    feast_path = await feast_dir.download()
    txn_df = pd.read_parquet(os.path.join(data_path, "transactions.parquet"))

    with open(os.path.join(data_path, "category_mapping.json")) as f:
        category_mapping = json.load(f)

    # Fetch user features from Feast (same path as serving)
    store = FeatureStore(repo_path=feast_path)
    entity_df = txn_df[["user_id", "event_timestamp"]].copy()

    log.info("Fetching user features from Feast (get_historical_features)...")
    training_data = store.get_historical_features(
        entity_df=entity_df,
        features=[
            "user_stats:txn_count",
            "user_stats:mean_amt",
            "user_stats:std_amt",
            "user_stats:max_amt",
            "user_stats:home_lat",
            "user_stats:home_long",
            "user_stats:age",
        ],
    ).to_df()

    # Merge back transaction features (Feast only returns user profile)
    training_data = training_data.merge(
        txn_df[["user_id", "event_timestamp", "amt", "amt_log", "category_encoded",
                "merch_lat", "merch_long", "hour", "day_of_week", "is_fraud"]],
        on=["user_id", "event_timestamp"],
        how="inner",
    )

    # Derived features: compare this transaction to the user's profile
    training_data["amt_zscore"] = (
        (training_data["amt"] - training_data["mean_amt"])
        / training_data["std_amt"].replace(0, 1)
    )
    training_data["amt_ratio"] = (
        training_data["amt"] / training_data["mean_amt"].replace(0, 1)
    )
    training_data["distance_from_home"] = haversine(
        training_data["home_lat"], training_data["home_long"],
        training_data["merch_lat"], training_data["merch_long"],
    )

    training_data = training_data.dropna(subset=ALL_FEATURE_COLS)
    X = training_data[ALL_FEATURE_COLS].values
    y = training_data["is_fraud"].values
    log.info(f"Training on {len(X):,} rows, {int(y.sum()):,} fraud")

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y,
    )

    n_legit = int((y_train == 0).sum())
    n_fraud = int((y_train == 1).sum())
    scale_pos_weight = n_legit / max(n_fraud, 1)

    model = XGBClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        learning_rate=learning_rate,
        scale_pos_weight=scale_pos_weight,
        min_child_weight=min_child_weight,
        gamma=gamma,
        random_state=42,
        eval_metric="logloss",
    )
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_proba)
    cm = confusion_matrix(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=["Legit", "Fraud"])

    log.info(f"AUC-ROC: {auc:.4f}")
    log.info(f"\n{report}")

    # Report
    precision_fraud = cm[1][1] / max(cm[1][1] + cm[0][1], 1) * 100
    recall_fraud = cm[1][1] / max(cm[1][1] + cm[1][0], 1) * 100

    html = (
        '<h2>Model Performance</h2>'
        + rh.stat_grid([
            (f"{auc:.4f}", "AUC-ROC"),
            (f"{len(X_train):,}", "Training Samples"),
            (f"{len(X_test):,}", "Test Samples"),
            (f"{precision_fraud:.1f}%", "Fraud Precision"),
            (f"{recall_fraud:.1f}%", "Fraud Recall"),
        ])
        + rh.confusion_matrix_html(cm)
    )

    # Feature importance bar chart
    importance = model.feature_importances_
    top_idx = np.argsort(importance)[::-1]
    top_labels = [ALL_FEATURE_COLS[i] for i in top_idx]
    top_values = [float(importance[i]) for i in top_idx]
    html += '<h3>Feature Importance</h3>'
    html += f'<div class="card">{rh.horizontal_bar_chart(top_labels, top_values)}</div>'

    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Save model + metadata
    model_path = os.path.join(tempfile.mkdtemp(), "model.joblib")
    joblib.dump({
        "model": model,
        "auc_roc": auc,
        "feature_cols": ALL_FEATURE_COLS,
        "category_mapping": category_mapping,
    }, model_path)

    return await flyte.io.File.from_local(model_path)

# ------------------------------------------------------------------
# Orchestrator: prepare → materialize → train
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def fraud_detection_pipeline(
    n_estimators: int = 300,
    max_depth: int = 6,
    learning_rate: float = 0.1,
    min_child_weight: int = 5,
    gamma: float = 1.0,
) -> tuple[flyte.io.File, flyte.io.Dir]:
    """
    Full fraud detection pipeline:
    1. Download and prepare data
    2. Materialize user profiles to Feast
    3. Train model using Feast for feature retrieval
    Returns model file and Feast artifacts for serving.
    """
    log.info("Starting fraud detection pipeline")
    steps = ["Prepare Data", "Materialize Features", "Train Model", "Done"]

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(0, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    data_dir = await prepare_data()

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(1, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Materialize features first so training can use Feast
    feast_dir = await materialize_features(data_dir)

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(2, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Train model using Feast for user feature retrieval
    model_file = await train_model(
        data_dir,
        feast_dir,
        n_estimators=n_estimators,
        max_depth=max_depth,
        learning_rate=learning_rate,
        min_child_weight=min_child_weight,
        gamma=gamma,
    )

    # Save copies to working directory for local app testing
    model_local = await model_file.download()
    feast_local = await feast_dir.download()
    shutil.copy2(model_local, "model.joblib")
    if os.path.exists("feast_artifacts"):
        shutil.rmtree("feast_artifacts")
    shutil.copytree(feast_local, "feast_artifacts")
    log.info("Saved local copies: model.joblib, feast_artifacts/")

    html = (
        '<h2>Fraud Detection Pipeline</h2>'
        + rh.pipeline_step_indicator(4, steps)
        + '<div class="card">'
        '<div style="font-weight:600;color:#155724;font-size:1.1em;margin-bottom:8px;">Pipeline Complete</div>'
        '<p>Model and feature store artifacts are ready for serving.</p>'
        '<table>'
        '<tr><th>Next Step</th><th>Command</th></tr>'
        '<tr><td>Run locally</td><td><code>python app.py</code></td></tr>'
        '<tr><td>Deploy scoring app</td><td><code>flyte deploy app.py serving_env</code></td></tr>'
        '<tr><td>Deploy dashboard</td><td><code>flyte deploy dashboard.py dashboard_env</code></td></tr>'
        '</table></div>'
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    log.info("Pipeline complete")
    return model_file, feast_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(fraud_detection_pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/fraud_detection_feast/fraud_detection_feast.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "feast==0.63.0",
#    "xgboost==3.2.0",
#    "scikit-learn==1.8.0",
#    "kagglehub==0.3.12",
#    ...
# ]
# ///
```

## Orchestrate the pipeline

The `fraud_detection_pipeline` task downloads data, trains XGBoost, applies Feast feature definitions, and materializes features.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "feast==0.63.0",
#    "scikit-learn==1.8.0",
#    "xgboost==3.2.0",
#    "joblib",
#    "pandas",
#    "pyarrow",
#    "kagglehub==0.3.12",
# ]
# main = "fraud_detection_pipeline"
# params = ""
# ///
import json
import logging
import math
import os
import shutil
import tempfile
from datetime import datetime, timedelta, timezone

import joblib
import numpy as np
import pandas as pd
import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="fraud-detection-feast", pre=True)

env = flyte.TaskEnvironment(
    name="fraud-detection-feast",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)
# {{/docs-fragment env}}

import report_helpers as rh

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Feature definitions
#
# Transaction features: known at scoring time (from the request)
# User features: pre-computed aggregates stored in Feast
# Derived features: computed at both training and scoring time by
#                   comparing the transaction to the user's profile
# ------------------------------------------------------------------

TXN_FEATURE_COLS = ["amt", "amt_log", "category_encoded", "merch_lat", "merch_long"]

USER_FEATURE_COLS = [
    "txn_count", "mean_amt", "std_amt", "max_amt",
    "home_lat", "home_long", "age",
]

DERIVED_FEATURE_COLS = [
    "amt_zscore", "amt_ratio", "distance_from_home", "hour", "day_of_week",
]

ALL_FEATURE_COLS = TXN_FEATURE_COLS + USER_FEATURE_COLS + DERIVED_FEATURE_COLS

def haversine(lat1, lon1, lat2, lon2):
    """Compute distance in miles between two (lat, lon) points."""
    R = 3959  # Earth radius in miles
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    return 2 * R * np.arcsin(np.sqrt(a))

# ------------------------------------------------------------------
# Task 1: Download dataset and engineer features
# ------------------------------------------------------------------

@env.task(report=True, cache="auto")
async def prepare_data() -> flyte.io.Dir:
    """Download the Sparkov credit card fraud dataset and prepare parquets."""
    import kagglehub

    log.info("Downloading dataset...")
    dataset_path = kagglehub.dataset_download("kartik2112/fraud-detection")
    csv_path = os.path.join(dataset_path, "fraudTrain.csv")
    df = pd.read_csv(csv_path)
    log.info(f"Loaded {len(df):,} transactions ({int(df['is_fraud'].sum()):,} fraudulent)")

    # Sample for workshop speed (stratified to preserve fraud ratio)
    if len(df) > 500_000:
        from sklearn.model_selection import train_test_split
        df, _ = train_test_split(df, train_size=500_000, stratify=df["is_fraud"], random_state=42)
        log.info(f"Sampled to {len(df):,} transactions")

    # ------------------------------------------------------------------
    # Parse timestamps
    # ------------------------------------------------------------------
    df["event_timestamp"] = pd.to_datetime(df["trans_date_trans_time"])
    df["event_timestamp"] = df["event_timestamp"].dt.tz_localize("UTC")
    df["hour"] = df["event_timestamp"].dt.hour
    df["day_of_week"] = df["event_timestamp"].dt.dayofweek

    # ------------------------------------------------------------------
    # Map cc_num → sequential user_id for clean API
    # ------------------------------------------------------------------
    cc_nums = df["cc_num"].unique()
    cc_to_user = {cc: i for i, cc in enumerate(sorted(cc_nums))}
    df["user_id"] = df["cc_num"].map(cc_to_user)

    # ------------------------------------------------------------------
    # Feature engineering
    # ------------------------------------------------------------------
    df["amt_log"] = np.log1p(df["amt"])

    # Label-encode merchant category
    categories = sorted(df["category"].unique())
    cat_to_int = {cat: i for i, cat in enumerate(categories)}
    df["category_encoded"] = df["category"].map(cat_to_int)

    # Compute age from dob
    df["dob"] = pd.to_datetime(df["dob"]).dt.tz_localize("UTC")
    ref_date = df["event_timestamp"].max()
    df["age"] = ((ref_date - df["dob"]).dt.days / 365.25).astype(int)

    # Distance between buyer and merchant
    df["distance"] = haversine(df["lat"], df["long"], df["merch_lat"], df["merch_long"])

    # ------------------------------------------------------------------
    # Build user aggregates
    # ------------------------------------------------------------------
    user_stats = df.groupby("user_id").agg(
        txn_count=("amt", "count"),
        mean_amt=("amt", "mean"),
        std_amt=("amt", "std"),
        max_amt=("amt", "max"),
        home_lat=("lat", "median"),
        home_long=("long", "median"),
        age=("age", "first"),
    ).reset_index()
    user_stats["std_amt"] = user_stats["std_amt"].fillna(0)
    # Use earliest timestamp so Feast point-in-time joins work for all transactions
    earliest_ts = df.groupby("user_id")["event_timestamp"].min().reset_index()
    user_stats = user_stats.merge(earliest_ts, on="user_id")

    # ------------------------------------------------------------------
    # Save to temp directory
    # ------------------------------------------------------------------
    data_dir = tempfile.mkdtemp()

    txn_cols = [
        "user_id", "event_timestamp",
        "amt", "amt_log", "category_encoded", "merch_lat", "merch_long",
        "hour", "day_of_week", "lat", "long", "distance",
        "is_fraud",
    ]
    df[txn_cols].to_parquet(os.path.join(data_dir, "transactions.parquet"), index=False)
    user_stats.to_parquet(os.path.join(data_dir, "user_features.parquet"), index=False)

    # Save category mapping + cc_num mapping for the app
    with open(os.path.join(data_dir, "category_mapping.json"), "w") as f:
        json.dump(cat_to_int, f)
    with open(os.path.join(data_dir, "user_mapping.json"), "w") as f:
        json.dump({str(k): v for k, v in cc_to_user.items()}, f)

    n_fraud = int(df["is_fraud"].sum())
    n_legit = len(df) - n_fraud
    fraud_pct = df["is_fraud"].mean() * 100
    html = (
        '<h2>Data Prepared</h2>'
        + rh.stat_grid([
            (f"{len(df):,}", "Transactions"),
            (f"{n_fraud:,}", "Fraudulent"),
            (f"{fraud_pct:.2f}%", "Fraud Rate"),
            (f"{user_stats['user_id'].nunique():,}", "Users"),
            (f"{len(categories)}", "Categories"),
        ])
        + rh.class_distribution_bar(n_legit, n_fraud)
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    return await flyte.io.Dir.from_local(data_dir)

# ------------------------------------------------------------------
# Task 2: Set up Feast and materialize user profiles to online store
# ------------------------------------------------------------------

@env.task(report=True)
async def materialize_features(data_dir: flyte.io.Dir) -> flyte.io.Dir:
    """Apply Feast definitions and materialize user profiles to SQLite online store."""
    from feast import Entity, FeatureStore, FeatureView, Field, FileSource
    from feast.types import Float64, Int64

    data_path = await data_dir.download()

    # Create a self-contained Feast repo in a temp directory
    feast_dir = tempfile.mkdtemp()

    # Copy parquet into feast dir so the repo is fully self-contained
    shutil.copy2(
        os.path.join(data_path, "user_features.parquet"),
        os.path.join(feast_dir, "user_features.parquet"),
    )

    # Write feature_store.yaml
    yaml_content = (
        "project: fraud_detection\n"
        f"registry: {feast_dir}/registry.db\n"
        "provider: local\n"
        "online_store:\n"
        "  type: sqlite\n"
        f"  path: {feast_dir}/online_store.db\n"
        "offline_store:\n"
        "  type: file\n"
        "entity_key_serialization_version: 3\n"
    )
    yaml_path = os.path.join(feast_dir, "feature_store.yaml")
    with open(yaml_path, "w") as f:
        f.write(yaml_content)

    store = FeatureStore(repo_path=feast_dir)

    # Define entity and feature view
    user = Entity(name="user", join_keys=["user_id"], description="Credit card holder")

    user_source = FileSource(
        path=os.path.join(feast_dir, "user_features.parquet"),
        timestamp_field="event_timestamp",
    )

    user_stats = FeatureView(
        name="user_stats",
        entities=[user],
        ttl=timedelta(days=0),  # No expiry — workshop data has old timestamps
        schema=[
            Field(name="txn_count", dtype=Int64),
            Field(name="mean_amt", dtype=Float64),
            Field(name="std_amt", dtype=Float64),
            Field(name="max_amt", dtype=Float64),
            Field(name="home_lat", dtype=Float64),
            Field(name="home_long", dtype=Float64),
            Field(name="age", dtype=Int64),
        ],
        online=True,
        source=user_source,
    )

    # Apply and materialize
    log.info("Applying Feast definitions...")
    store.apply([user, user_stats])

    log.info("Materializing user profiles to online store...")
    store.materialize(
        start_date=datetime(2018, 1, 1, tzinfo=timezone.utc),
        end_date=datetime.now(timezone.utc),
    )

    # Re-apply with relative paths so the registry is portable across workers
    portable_yaml = (
        "project: fraud_detection\n"
        "registry: registry.db\n"
        "provider: local\n"
        "online_store:\n"
        "  type: sqlite\n"
        "  path: online_store.db\n"
        "offline_store:\n"
        "  type: file\n"
        "entity_key_serialization_version: 3\n"
    )
    with open(yaml_path, "w") as f:
        f.write(portable_yaml)

    # Re-apply with relative source path so get_historical_features works on other workers
    store = FeatureStore(repo_path=feast_dir)
    user_source = FileSource(
        path="user_features.parquet",
        timestamp_field="event_timestamp",
    )
    user_stats = FeatureView(
        name="user_stats",
        entities=[user],
        ttl=timedelta(days=0),
        schema=[
            Field(name="txn_count", dtype=Int64),
            Field(name="mean_amt", dtype=Float64),
            Field(name="std_amt", dtype=Float64),
            Field(name="max_amt", dtype=Float64),
            Field(name="home_lat", dtype=Float64),
            Field(name="home_long", dtype=Float64),
            Field(name="age", dtype=Int64),
        ],
        online=True,
        source=user_source,
    )
    store.apply([user, user_stats])

    features = ["txn_count", "mean_amt", "std_amt", "max_amt", "home_lat", "home_long", "age"]
    html = (
        '<h2>Feature Store Materialized</h2>'
        + rh.stat_grid([
            ("user_stats", "Feature View"),
            (str(len(features)), "Features"),
            ("SQLite", "Online Store"),
        ])
        + '<h3>Materialized Features</h3>'
        '<table>'
        '<tr><th>Feature</th><th>Type</th><th>Description</th></tr>'
        '<tr><td>txn_count</td><td><span class="badge badge-info">Int64</span></td><td>Total transactions</td></tr>'
        '<tr><td>mean_amt</td><td><span class="badge badge-info">Float64</span></td><td>Average transaction amount</td></tr>'
        '<tr><td>std_amt</td><td><span class="badge badge-info">Float64</span></td><td>Std dev of amounts</td></tr>'
        '<tr><td>max_amt</td><td><span class="badge badge-info">Float64</span></td><td>Max transaction amount</td></tr>'
        '<tr><td>home_lat</td><td><span class="badge badge-info">Float64</span></td><td>Home latitude (median)</td></tr>'
        '<tr><td>home_long</td><td><span class="badge badge-info">Float64</span></td><td>Home longitude (median)</td></tr>'
        '<tr><td>age</td><td><span class="badge badge-info">Int64</span></td><td>User age</td></tr>'
        '</table>'
        '<div class="note">User profiles are ready for real-time serving via the scoring app.</div>'
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    return await flyte.io.Dir.from_local(feast_dir)

# ------------------------------------------------------------------
# Task 3: Train XGBoost model
# ------------------------------------------------------------------

@env.task(report=True)
async def train_model(
    data_dir: flyte.io.Dir,
    feast_dir: flyte.io.Dir,
    n_estimators: int = 300,
    max_depth: int = 6,
    learning_rate: float = 0.1,
    min_child_weight: int = 5,
    gamma: float = 1.0,
) -> flyte.io.File:
    """Train an XGBoost classifier using Feast for feature retrieval."""
    from feast import FeatureStore
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import classification_report, roc_auc_score, confusion_matrix
    from xgboost import XGBClassifier

    data_path = await data_dir.download()
    feast_path = await feast_dir.download()
    txn_df = pd.read_parquet(os.path.join(data_path, "transactions.parquet"))

    with open(os.path.join(data_path, "category_mapping.json")) as f:
        category_mapping = json.load(f)

    # Fetch user features from Feast (same path as serving)
    store = FeatureStore(repo_path=feast_path)
    entity_df = txn_df[["user_id", "event_timestamp"]].copy()

    log.info("Fetching user features from Feast (get_historical_features)...")
    training_data = store.get_historical_features(
        entity_df=entity_df,
        features=[
            "user_stats:txn_count",
            "user_stats:mean_amt",
            "user_stats:std_amt",
            "user_stats:max_amt",
            "user_stats:home_lat",
            "user_stats:home_long",
            "user_stats:age",
        ],
    ).to_df()

    # Merge back transaction features (Feast only returns user profile)
    training_data = training_data.merge(
        txn_df[["user_id", "event_timestamp", "amt", "amt_log", "category_encoded",
                "merch_lat", "merch_long", "hour", "day_of_week", "is_fraud"]],
        on=["user_id", "event_timestamp"],
        how="inner",
    )

    # Derived features: compare this transaction to the user's profile
    training_data["amt_zscore"] = (
        (training_data["amt"] - training_data["mean_amt"])
        / training_data["std_amt"].replace(0, 1)
    )
    training_data["amt_ratio"] = (
        training_data["amt"] / training_data["mean_amt"].replace(0, 1)
    )
    training_data["distance_from_home"] = haversine(
        training_data["home_lat"], training_data["home_long"],
        training_data["merch_lat"], training_data["merch_long"],
    )

    training_data = training_data.dropna(subset=ALL_FEATURE_COLS)
    X = training_data[ALL_FEATURE_COLS].values
    y = training_data["is_fraud"].values
    log.info(f"Training on {len(X):,} rows, {int(y.sum()):,} fraud")

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y,
    )

    n_legit = int((y_train == 0).sum())
    n_fraud = int((y_train == 1).sum())
    scale_pos_weight = n_legit / max(n_fraud, 1)

    model = XGBClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        learning_rate=learning_rate,
        scale_pos_weight=scale_pos_weight,
        min_child_weight=min_child_weight,
        gamma=gamma,
        random_state=42,
        eval_metric="logloss",
    )
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_proba)
    cm = confusion_matrix(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=["Legit", "Fraud"])

    log.info(f"AUC-ROC: {auc:.4f}")
    log.info(f"\n{report}")

    # Report
    precision_fraud = cm[1][1] / max(cm[1][1] + cm[0][1], 1) * 100
    recall_fraud = cm[1][1] / max(cm[1][1] + cm[1][0], 1) * 100

    html = (
        '<h2>Model Performance</h2>'
        + rh.stat_grid([
            (f"{auc:.4f}", "AUC-ROC"),
            (f"{len(X_train):,}", "Training Samples"),
            (f"{len(X_test):,}", "Test Samples"),
            (f"{precision_fraud:.1f}%", "Fraud Precision"),
            (f"{recall_fraud:.1f}%", "Fraud Recall"),
        ])
        + rh.confusion_matrix_html(cm)
    )

    # Feature importance bar chart
    importance = model.feature_importances_
    top_idx = np.argsort(importance)[::-1]
    top_labels = [ALL_FEATURE_COLS[i] for i in top_idx]
    top_values = [float(importance[i]) for i in top_idx]
    html += '<h3>Feature Importance</h3>'
    html += f'<div class="card">{rh.horizontal_bar_chart(top_labels, top_values)}</div>'

    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Save model + metadata
    model_path = os.path.join(tempfile.mkdtemp(), "model.joblib")
    joblib.dump({
        "model": model,
        "auc_roc": auc,
        "feature_cols": ALL_FEATURE_COLS,
        "category_mapping": category_mapping,
    }, model_path)

    return await flyte.io.File.from_local(model_path)

# ------------------------------------------------------------------
# Orchestrator: prepare → materialize → train
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def fraud_detection_pipeline(
    n_estimators: int = 300,
    max_depth: int = 6,
    learning_rate: float = 0.1,
    min_child_weight: int = 5,
    gamma: float = 1.0,
) -> tuple[flyte.io.File, flyte.io.Dir]:
    """
    Full fraud detection pipeline:
    1. Download and prepare data
    2. Materialize user profiles to Feast
    3. Train model using Feast for feature retrieval
    Returns model file and Feast artifacts for serving.
    """
    log.info("Starting fraud detection pipeline")
    steps = ["Prepare Data", "Materialize Features", "Train Model", "Done"]

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(0, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    data_dir = await prepare_data()

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(1, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Materialize features first so training can use Feast
    feast_dir = await materialize_features(data_dir)

    html = '<h2>Fraud Detection Pipeline</h2>' + rh.pipeline_step_indicator(2, steps)
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    # Train model using Feast for user feature retrieval
    model_file = await train_model(
        data_dir,
        feast_dir,
        n_estimators=n_estimators,
        max_depth=max_depth,
        learning_rate=learning_rate,
        min_child_weight=min_child_weight,
        gamma=gamma,
    )

    # Save copies to working directory for local app testing
    model_local = await model_file.download()
    feast_local = await feast_dir.download()
    shutil.copy2(model_local, "model.joblib")
    if os.path.exists("feast_artifacts"):
        shutil.rmtree("feast_artifacts")
    shutil.copytree(feast_local, "feast_artifacts")
    log.info("Saved local copies: model.joblib, feast_artifacts/")

    html = (
        '<h2>Fraud Detection Pipeline</h2>'
        + rh.pipeline_step_indicator(4, steps)
        + '<div class="card">'
        '<div style="font-weight:600;color:#155724;font-size:1.1em;margin-bottom:8px;">Pipeline Complete</div>'
        '<p>Model and feature store artifacts are ready for serving.</p>'
        '<table>'
        '<tr><th>Next Step</th><th>Command</th></tr>'
        '<tr><td>Run locally</td><td><code>python app.py</code></td></tr>'
        '<tr><td>Deploy scoring app</td><td><code>flyte deploy app.py serving_env</code></td></tr>'
        '<tr><td>Deploy dashboard</td><td><code>flyte deploy dashboard.py dashboard_env</code></td></tr>'
        '</table></div>'
    )
    await flyte.report.replace.aio(rh.wrap(html))
    await flyte.report.flush.aio()

    log.info("Pipeline complete")
    return model_file, feast_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(fraud_detection_pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/fraud_detection_feast/fraud_detection_feast.py*

## Run the workflow

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/fraud_detection_feast):

```
cd v2/tutorials/fraud_detection_feast
uv run --script fraud_detection_feast.py
```

The first run downloads the dataset via `kagglehub` (public dataset, no API key required). Open the run report to review the confusion matrix and feature-importance summary when training completes.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services/financial-research-agent ===

# Financial research agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/financial_research_agent).

This example demonstrates how to build a financial research and earnings-cycle agent on Flyte. For each company, the agent runs grounded, source-cited research and fresh news, then synthesizes an analyst-ready equity briefing.

Financial research benefits from **low-latency, ranked, source-cited results** across both the general web and news streams. The [You.com Research API](https://you.com/docs/research/overview) produces a grounded, citation-backed synthesis, and the [You.com Search API](https://you.com/docs/search/overview) adds a fresh-news layer. [Claude](https://docs.anthropic.com/) via [LiteLLM](https://docs.litellm.ai/) turns that evidence into an analyst-ready briefing. Flyte's `cache="auto"` reuses prior results when runs converge on the same companies.

Flyte provides:

- **Fan-out parallelism** across companies
- **`cache="auto"`** to reuse prior You.com and LLM results across converging runs
- **`@flyte.trace`** on every external call for full prompt → citation lineage
- **Flyte reports** with thesis, risks, watch items, and source citations per company

![Financial research agent report](../../../_static/images/tutorials/financial_research_agent/financial-research-agent.png)

## Setting up the environment

The agent runs in a `TaskEnvironment` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the `uv` script dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE0", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial_research_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

CODE1

## Data types

Each `Briefing` carries a thesis, recent developments, risks, watch items, and a list of `Source` objects from both the Research and Search APIs.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE2", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial_research_agent/main.py*

## You.com Research and Search APIs

The agent uses both You.com APIs in parallel for each company:

- **Research API** (`https://api.you.com/v1/research`): grounded, citation-backed analysis with configurable `research_effort` (`lite`, `standard`, `deep`, `exhaustive`). See the [Research API reference](https://you.com/docs/api-reference/research/v1-research).
- **Search API** (`https://ydc-index.io/v1/search`): fresh news headlines with `freshness` filtering. See the [Search API reference](https://you.com/docs/api-reference/search/v1-search).

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE3", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE4
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE5", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial_research_agent/main.py*

## Research one company

The `research_company` task calls both You.com APIs in parallel, collects sources, and synthesizes a structured briefing.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE6", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial_research_agent/main.py*

## Orchestration

The `financial_research` driver task fans out across all companies and renders a Flyte report with per-company briefings and citations.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "financial_research"
# params = ""
# ///
"""Financial research & earnings-cycle agent.

For each company, runs grounded, source-cited research via the You.com Research
API plus a fresh-news layer via the Search API, then uses Claude to synthesize
an analyst-ready equity briefing that preserves citations. Flyte caching cuts
duplicate spend when runs converge.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="financial-research",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="financial-research", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    domain: str = ""
    snippet: str = ""
    published: str = ""
    favicon: str = ""
    section: str = "research"  # "research", "news", or "web"

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Briefing:
    company: str
    thesis: str
    recent_developments: list[str] = field(default_factory=list)
    risks: list[str] = field(default_factory=list)
    watch_items: list[str] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResearchReport:
    briefings: list[Briefing] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_apis}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_request(method: str, url: str, timeout: float, **kwargs) -> dict:
    """HTTP wrapper with exponential backoff + jitter on 429 rate limits.

    Fanned-out tasks run in separate pods, so we retry on the client side to
    smooth out bursts against the You.com API rate limit.
    """
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    if method == "POST":
        headers["Content-Type"] = "application/json"

    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.request(method, url, headers=headers, **kwargs)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str, freshness: str) -> dict:
    """Grounded, citation-backed research answer."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {"freshness": freshness},
    }
    return await _you_request("POST", YOU_RESEARCH_URL, 300.0, json=body)

@flyte.trace
async def you_news(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[dict]:
    """Fresh news headlines for a company.

    ``boost_domains`` (comma-separated) lifts authoritative financial outlets
    in ranking without restricting results to only those domains, so company
    press releases and niche coverage still surface when relevant.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_request("GET", YOU_SEARCH_URL, 60.0, params=params)

    results = data.get("results", {})
    out: list[dict] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            out.append(
                {
                    "title": item.get("title", ""),
                    "url": url,
                    "domain": _domain(url),
                    "snippet": snippets[0] if snippets else item.get("description", ""),
                    "published": item.get("page_age", "") or "",
                    "favicon": item.get("favicon_url")
                    or _favicon_for(url),
                    "section": section,
                }
            )
    return out
# {{/docs-fragment you_apis}}

# {{docs-fragment llm}}
@flyte.trace
async def synthesize_briefing(company: str, focus: str, research: str, news: str) -> dict:
    """Use Claude to synthesize a structured equity briefing."""
    from litellm import acompletion

    system = (
        "You are an equity research analyst. Using ONLY the grounded research "
        "and news provided, write a concise briefing. Respond ONLY with JSON: "
        '{"thesis": str, "recent_developments": [str], "risks": [str], '
        '"watch_items": [str]}. Keep each list to 3-5 short, specific bullets.'
    )
    user = (
        f"Company: {company}\nFocus: {focus}\n\n"
        f"Grounded research:\n{research}\n\nRecent news:\n{news}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE7", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment research_company}}
# Tier-1 financial outlets that consistently break earnings, M&A, and
# analyst-moving news. boost_domains lifts these in ranking without excluding
# other sources, so company press releases and trade-press coverage still
# surface when relevant.
FINANCE_BOOST_DOMAINS = "reuters.com,bloomberg.com,wsj.com,marketwatch.com,cnbc.com,ft.com"

@env.task(retries=3)
async def research_company(
    company: str,
    focus: str,
    research_effort: str,
    freshness: str,
) -> Briefing:
    """Research one company and synthesize a cited briefing."""
    question = (
        f"Provide a grounded analysis of {company} with respect to: {focus}. "
        f"Cover recent financial performance, strategic moves, competitive "
        f"positioning, and risks."
    )
    research_result, news = await asyncio.gather(
        you_research(question, research_effort, freshness),
        you_news(
            f"{company} earnings news",
            freshness=freshness,
            boost_domains=FINANCE_BOOST_DOMAINS,
        ),
    )

    output = research_result.get("output", {})
    research_text = output.get("content", "")
    if not isinstance(research_text, str):
        research_text = json.dumps(research_text)

    sources: list[Source] = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                domain=_domain(url),
                snippet=str((s.get("snippets") or [""])[0]),
                favicon=_favicon_for(url),
                section="research",
            )
        )
    for n in news:
        sources.append(
            Source(
                title=str(n.get("title", "")),
                url=str(n.get("url", "")),
                domain=str(n.get("domain", "")),
                snippet=str(n.get("snippet", "")),
                published=str(n.get("published", "")),
                favicon=str(n.get("favicon", "")),
                section=str(n.get("section", "web")),
            )
        )
    news_text = "\n".join(
        f"- {n['title']} ({n['published']}) {n['domain']}: {n['snippet'][:120]}"
        for n in news
    )

    parsed = await synthesize_briefing(company, focus, research_text, news_text)

    def _list(key: str) -> list[str]:
        return [str(x) for x in (parsed.get(key) or [])]

    return Briefing(
        company=company,
        thesis=str(parsed.get("thesis", "")),
        recent_developments=_list("recent_developments"),
        risks=_list("risks"),
        watch_items=_list("watch_items"),
        sources=sources,
    )
# {{/docs-fragment research_company}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:18px; margin:0 0 8px; color:#102a43; }
  .rpt .thesis { font-size:14px; line-height:1.5; background:#f7f9fb;
                 border-radius:8px; padding:10px 12px; margin:0 0 14px; }
  .rpt .cols { display:flex; gap:18px; flex-wrap:wrap; }
  .rpt .col { flex:1; min-width:220px; }
  .rpt .col h3 { font-size:12px; text-transform:uppercase; letter-spacing:.04em;
                 color:#627d98; margin:0 0 6px; }
  .rpt .col.risks h3 { color:#c0392b; }
  .rpt ul { margin:0; padding-left:18px; }
  .rpt li { font-size:13px; line-height:1.5; margin:0 0 4px; }
  .rpt .sources { margin-top:14px; border-top:1px solid #f0f2f5; padding-top:10px; }
  .rpt .sources h3 { font-size:12px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.research { background:#5b8def; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com citation (Research or Search source)."""
    if not s.url:
        return ""
    tag_cls = s.section if s.section in ("research", "news") else "web"
    meta_bits = []
    if s.published:
        meta_bits.append(s.published[:10])
    if s.title:
        meta_bits.append(s.title)
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag {tag_cls}'>{s.section}</span>"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: ResearchReport) -> str:
    def _ul(items: list[str]) -> str:
        if not items:
            return "<p class='empty'>None reported.</p>"
        return "<ul>" + "".join(f"<li>{x}</li>" for x in items) + "</ul>"

    cards = []
    for b in report.briefings:
        src = "".join(_cite(s) for s in b.sources[:10])
        cards.append(
            f"<div class='card'><h2>{b.company}</h2>"
            f"<div class='thesis'>{b.thesis or 'No thesis generated.'}</div>"
            f"<div class='cols'>"
            f"<div class='col'><h3>Recent developments</h3>{_ul(b.recent_developments)}</div>"
            f"<div class='col risks'><h3>Risks</h3>{_ul(b.risks)}</div>"
            f"<div class='col'><h3>Watch items</h3>{_ul(b.watch_items)}</div>"
            f"</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(b.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(b.sources) for b in report.briefings)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Financial Research Briefings</h1>
      <p class="sub">Grounded, citation-backed equity briefings — each company
      backed by You.com Research synthesis plus fresh Search news.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.briefings)}</b> companies</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No briefings generated.</p>"}
      <p class="yoube">Research answers from the You.com Research API (grounded
      synthesis with inline citations) plus fresh headlines from the You.com
      Search API (web + auto-classified news with timestamps and snippets).</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def financial_research(
    companies: list[str] = [
        "NVIDIA",
        "Advanced Micro Devices",
        "Microsoft",
        "Alphabet",
        "Amazon",
        "Meta Platforms",
        "Broadcom",
        "Taiwan Semiconductor Manufacturing",
    ],
    focus: str = "Q4 earnings preview and competitive positioning",
    research_effort: str = "standard",
    freshness: str = "month",
) -> ResearchReport:
    """Fan out across companies and aggregate cited equity briefings."""
    with flyte.group("research-companies"):
        briefings = await asyncio.gather(
            *[
                research_company(c, focus, research_effort, freshness)
                for c in companies
            ]
        )

    report = ResearchReport(briefings=list(briefings))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(financial_research)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial_research_agent/main.py*

## Run the agent

### Create secrets

Get a You.com API key from the [You.com platform](https://you.com/platform) (see the [quickstart guide](https://you.com/docs/quickstart)). Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/).

Register both keys as Flyte secrets. The secret key names must match those declared in the `TaskEnvironment`:

```
flyte create secret youdotcom-api-key <YOUR_YOU_API_KEY>
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run locally or remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/financial_research_agent):

```
cd v2/tutorials/financial_research_agent
uv run --script main.py
```

To test locally without Flyte secrets:

```
export YOU_API_KEY=<YOUR_YOU_API_KEY>
export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>

uv run --script main.py
```

When the run completes, open the Flyte report to review equity briefings with thesis, risks, and You.com source citations for each company.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai ===

# Frontier AI

Tutorials for frontier-model pretraining, automated experimentation, and large-scale AI workloads.

### **Frontier AI > Distributed LLM pretraining**

Pretrain large language models at scale with PyTorch Lightning, FSDP, and H200 GPUs, featuring streaming data and real-time metrics.

## Subpages
- **Frontier AI > Distributed LLM pretraining**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/distributed-pretraining ===

# Distributed LLM pretraining

When training large models, infrastructure should not be the hardest part. The real work is in the model architecture, the data, and the hyperparameters. In practice, though, teams often spend weeks just trying to get distributed training to run reliably.

And when it breaks, it usually breaks in familiar ways: out-of-memory crashes, corrupted checkpoints, data loaders that silently fail, or runs that hang with no obvious explanation.

Most distributed training tutorials focus on PyTorch primitives. This one focuses on getting something that actually ships. We go into the technical details, such as how FSDP shards parameters, why gradient clipping behaves differently at scale, and how streaming datasets reduce memory pressure, but always with the goal of building a system that works in production.

Real training jobs need more than a training loop. They need checkpointing, fault tolerance, data streaming, visibility into what’s happening, and the ability to recover from failures. In this tutorial, we build all of that using Flyte, without having to stand up or manage any additional infrastructure.

> [!NOTE]
> Full code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/pretraining/train.py).

## Overview

We're going to pretrain a GPT-2 style language model from scratch. This involves training on raw text data starting from randomly initialized weights, rather than fine-tuning or adapting a pretrained model. This is the same process used to train the original GPT-2, LLaMA, and most other foundation models.

The model learns by predicting the next token. Given "The cat sat on the", it learns to predict "mat". Do this billions of times across terabytes of text, and the model develops surprisingly sophisticated language understanding. That's pretraining.

The challenge is scale. A 30B parameter model doesn't fit on a single GPU. The training dataset, [SlimPajama](https://huggingface.co/datasets/cerebras/SlimPajama-627B) in our case, is 627 billion tokens. Training runs last for days or even weeks. To make this work, you need:

- **Distributed training**: Split the model across multiple GPUs using [FSDP (Fully Sharded Data Parallel)](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html)
- **Data streaming**: Pull training data on-demand instead of downloading terabytes upfront
- **Checkpointing**: Save progress regularly so a failure doesn’t wipe out days of compute
- **Observability**: See what's happening inside a multi-day training run

We’ll build a Flyte pipeline that takes care of all of this, using three tasks with clearly defined responsibilities:

1. **Data preparation**: Tokenizes your dataset and converts it to MDS (MosaicML Data Shard) format for streaming. This Flyte task is cached, so it only needs to be run once and can be reused across runs.
2. **Distributed training**: Runs FSDP across 8 H200 GPUs. Flyte's `Elastic` plugin handles the distributed setup. Checkpoints upload to S3 automatically via Flyte's `File` abstraction.
3. **Real-time reporting**: Streams loss curves and training metrics to Flyte Reports, a live dashboard integrated into the Flyte UI.

Why three separate tasks? Flyte makes this separation efficient:

- **Caching**: The data preparation step runs once. On subsequent runs, Flyte skips it entirely.
- **Resource isolation**: Training uses expensive H200 GPUs only while actively training, while the driver runs on inexpensive CPU instances.
- **Fault boundaries**: If training fails, the data preparation step does not re-run. Training can resume directly from the most recent checkpoint.

## Implementation

Let's walk through the code. We'll start with the infrastructure setup, build the model, then wire everything together into a pipeline.

### Setting up the environment

Every distributed training job needs a consistent environment across all nodes. Flyte handles this with container images:

```
import logging
import math
import os
from pathlib import Path
from typing import Optional

import flyte
import flyte.report
import lightning as L
import numpy as np
import torch
import torch.nn as nn
from flyte.io import Dir, File
from flyteplugins.pytorch.task import Elastic
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The imports tell the story: `flyte` for orchestration, `flyte.report` for live dashboards, `lightning` for training loop management, and `Elastic` from Flyte's PyTorch plugin. This last one is key as it configures PyTorch's distributed launch without you writing any distributed setup code.

```
NUM_NODES = 1
DEVICES_PER_NODE = 8
VOCAB_SIZE = (
    50257  # GPT-2 BPE tokenizer vocabulary size (constant across all model sizes)
)
N_POSITIONS = 2048  # Maximum sequence length (constant across all model sizes)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

These constants define the distributed topology. We're using 1 node with 8 GPUs, but you can scale this up by changing `NUM_NODES`. The vocabulary size (50,257 tokens) and sequence length (2,048 tokens) match GPT-2's [Byte Pair Encoding (BPE) tokenizer](https://huggingface.co/learn/llm-course/en/chapter6/5).

```
image = flyte.Image.from_debian_base(
    name="distributed_training_h200"
).with_pip_packages(
    "transformers==4.57.3",
    "datasets==4.4.1",
    "tokenizers==0.22.1",
    "huggingface-hub==0.34.0",
    "mosaicml-streaming>=0.7.0",
    "pyarrow==22.0.0",
    "flyteplugins-pytorch>=2.0.0b33",
    "torch==2.9.1",
    "lightning==2.5.6",
    "tensorboard==2.20.0",
    "sentencepiece==0.2.1",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

Flyte builds this container automatically when the pipeline is run. All dependencies required for distributed training, including PyTorch, Lightning, the streaming library, and NCCL for GPU communication, are baked in. There's no Dockerfile to maintain and no "works on my machine" debugging.

### Declaring resource requirements

Different parts of the pipeline need different resources. Data tokenization needs CPU and memory. Training needs GPUs. The driver just coordinates. Flyte's `TaskEnvironment` lets you declare exactly what each task needs:

```
data_loading_env = flyte.TaskEnvironment(
    name="data_loading_h200",
    image=image,
    resources=flyte.Resources(cpu=5, memory="28Gi", disk="100Gi"),
    env_vars={
        "HF_DATASETS_CACHE": "/tmp/hf_cache",  # Cache directory for datasets
        "TOKENIZERS_PARALLELISM": "true",  # Enable parallel tokenization
    },
    cache="auto",
)

distributed_llm_training_env = flyte.TaskEnvironment(
    name="distributed_llm_training_h200",
    image=image,
    resources=flyte.Resources(
        cpu=64,
        memory="512Gi",
        gpu=f"H200:{DEVICES_PER_NODE}",
        disk="1Ti",
        shm="16Gi",  # Explicit shared memory for NCCL communication
    ),
    plugin_config=Elastic(nnodes=NUM_NODES, nproc_per_node=DEVICES_PER_NODE),
    env_vars={
        "TORCH_DISTRIBUTED_DEBUG": "INFO",
        "NCCL_DEBUG": "WARN",
    },
    cache="auto",
)

driver_env = flyte.TaskEnvironment(
    name="llm_training_driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    cache="auto",
    depends_on=[data_loading_env, distributed_llm_training_env],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

Let's break down the training environment, since this is where most of the complexity lives:

- **`gpu=f"H200:{DEVICES_PER_NODE}"`**: Flyte provisions exactly 8 H200 GPUs. These have 141GB of memory each, enough to train 30B+ parameter models with FSDP.
- **`shm="16Gi"`**: This allocates explicit shared memory. NCCL (NVIDIA's communication library) uses shared memory for inter-GPU communication on the same node. Without this, you'll see cryptic errors like "NCCL error: unhandled system error", which can be difficult to debug.
- **`Elastic(nnodes=NUM_NODES, nproc_per_node=DEVICES_PER_NODE)`**: This is Flyte's integration with PyTorch's elastic launch. It handles process spawning (one process per GPU), rank assignment (each process knows its ID), and environment setup (master address, world size). This replaces the boilerplate typically written in shell scripts.

The `driver_env` is intentionally lightweight, using 2 CPUs and 4 GB of memory. Its role is limited to orchestrating tasks and passing data between them, so allocating GPUs here would be unnecessary.

### Model configurations

Training a 1.5B model uses different hyperparameters than training a 65B model. Rather than hardcoding values, we define presets:

```
MODEL_CONFIGS = {
    "1.5B": {
        "n_embd": 2048,
        "n_layer": 24,
        "n_head": 16,
        "batch_size": 8,
        "learning_rate": 6e-4,
        "checkpoint_every_n_steps": 10,
        "report_every_n_steps": 5,
        "val_check_interval": 100,
    },  # Good for testing and debugging
    "30B": {
        "n_embd": 6656,
        "n_layer": 48,
        "n_head": 52,
        "batch_size": 1,
        "learning_rate": 1.6e-4,
        "checkpoint_every_n_steps": 7500,
        "report_every_n_steps": 200,
        "val_check_interval": 1000,
    },
    "65B": {
        "n_embd": 8192,
        "n_layer": 80,
        "n_head": 64,
        "batch_size": 1,
        "learning_rate": 1.5e-4,
        "checkpoint_every_n_steps": 10000,
        "report_every_n_steps": 250,
        "val_check_interval": 2000,
    },
}

def get_model_config(model_size: str) -> dict:
    if model_size not in MODEL_CONFIGS:
        available = ", ".join(MODEL_CONFIGS.keys())
        raise ValueError(f"Unknown model size: {model_size}. Available: {available}")

    return MODEL_CONFIGS[model_size]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

A few things to notice:

- **Batch size decreases with model size**: For a fixed GPU memory budget, larger models consume more memory for parameters, optimizer state, and activations, leaving less room for per-GPU batch size. For example, a 1.5B parameter model may fit a batch size of 8 per GPU, while a 65B model may only fit a batch size of 1. This is typically compensated for using gradient accumulation to maintain a larger effective batch size.
- **Learning rate decreases with model size**: Larger models are more sensitive to optimization instability and typically require lower learning rates. The values here follow empirical best practices used in large-scale language model training, informed by work such as the [Chinchilla study](https://arxiv.org/pdf/2203.15556) on compute-optimal scaling.
- **Checkpoint frequency increases with model size**: Checkpointing a 65B model is expensive (the checkpoint is huge). We do it less often but make sure we don't lose too much progress if something fails.

The 1.5B config is good for testing your setup before committing to a serious training run.

### Building the GPT model

Now for the model itself. We're building a GPT-2 style decoder-only transformer from scratch.

First, the configuration class:

```
class GPTConfig:
    """Configuration for GPT model."""

    def __init__(
        self,
        vocab_size: int = VOCAB_SIZE,
        n_positions: int = N_POSITIONS,
        n_embd: int = 2048,
        n_layer: int = 24,
        n_head: int = 16,
        n_inner: Optional[int] = None,
        activation_function: str = "gelu_new",
        dropout: float = 0.1,
        layer_norm_epsilon: float = 1e-5,
    ):
        self.vocab_size = vocab_size
        self.n_positions = n_positions
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_inner = n_inner if n_inner is not None else 4 * n_embd
        self.activation_function = activation_function
        self.dropout = dropout
        self.layer_norm_epsilon = layer_norm_epsilon
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The key architectural parameters:

- **`n_embd`**: The hidden (embedding) dimension. Larger values increase model capacity but also increase memory and compute requirements.
- **`n_layer`**: The number of transformer blocks. Model depth strongly influences expressiveness and performance.
- **`n_head`**: The number of attention heads. Each head can attend to different patterns or relationships in the input.
- **`n_inner`**: The hidden dimension of the feed-forward network (MLP), typically set to 4x the embedding dimension.

Next, we define a single transformer block:

```
class GPTBlock(nn.Module):
    """Transformer block with causal self-attention."""

    def __init__(self, config: GPTConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.attn = nn.MultiheadAttention(
            config.n_embd,
            config.n_head,
            dropout=config.dropout,
            batch_first=True,
        )
        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)

        # Get activation function from config
        ACT_FNS = {
            "gelu": nn.GELU(),
            "gelu_new": nn.GELU(approximate="tanh"),  # GPT-2 uses approximate GELU
            "relu": nn.ReLU(),
            "silu": nn.SiLU(),
            "swish": nn.SiLU(),  # SiLU = Swish
        }
        act_fn = ACT_FNS.get(config.activation_function, nn.GELU())

        self.mlp = nn.Sequential(
            nn.Linear(config.n_embd, config.n_inner),
            act_fn,
            nn.Linear(config.n_inner, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x, causal_mask, key_padding_mask=None):
        x_normed = self.ln_1(x)

        # Self-attention with causal and padding masks
        attn_output, _ = self.attn(
            x_normed,  # query
            x_normed,  # key
            x_normed,  # value
            attn_mask=causal_mask,  # Causal mask: (seq_len, seq_len)
            key_padding_mask=key_padding_mask,  # Padding mask: (batch, seq_len)
            need_weights=False,
        )
        x = x + attn_output

        # MLP with residual
        x = x + self.mlp(self.ln_2(x))
        return x
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

Each block has two sub-layers: causal self-attention and a feed-forward MLP. The causal mask ensures the model can only attend to previous tokens in the sequence, so it can't "cheat" by looking at the answer. This is what makes it *autoregressive*.

The full `GPTModel` class (see the complete code) stacks these blocks and adds token and positional embeddings. One important detail is that the input token embedding matrix is shared with the output projection layer (often called [weight tying](https://mbrenndoerfer.com/writing/weight-tying-shared-embeddings-transformers)). This reduces the number of parameters by roughly 50 million for typical vocabulary sizes and often leads to better generalization and more stable training.

### The Lightning training module

PyTorch Lightning handles the training loop boilerplate. We wrap our model in a `LightningModule` that defines how to train it:

```
class GPTPreTrainingModule(L.LightningModule):
    """PyTorch Lightning module for GPT pre-training."""

    def __init__(
        self,
        vocab_size: int = 50257,
        n_positions: int = 2048,
        n_embd: int = 2048,
        n_layer: int = 24,
        n_head: int = 16,
        learning_rate: float = 6e-4,
        weight_decay: float = 0.1,
        warmup_steps: int = 2000,
        max_steps: int = 100000,
    ):
        super().__init__()
        self.save_hyperparameters()

        config = GPTConfig(
            vocab_size=vocab_size,
            n_positions=n_positions,
            n_embd=n_embd,
            n_layer=n_layer,
            n_head=n_head,
        )
        self.model = GPTModel(config)

    def forward(self, input_ids, attention_mask=None):
        return self.model(input_ids, attention_mask)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The `save_hyperparameters()` call is important because it stores all constructor arguments in the checkpoint. This allows the model to be reloaded later without having to manually reconstruct the original configuration.

The training and validation steps implement standard causal language modeling, where the model is trained to predict the next token given all previous tokens in the sequence.

```
    def training_step(self, batch, _batch_idx):
        # Convert int32 to int64 (long) - MDS stores as int32 but PyTorch expects long
        input_ids = batch["input_ids"].long()
        labels = batch["labels"].long()

        # Get attention mask if present (optional, for padded sequences)
        # attention_mask: 1 = real token, 0 = padding
        # Note: Current data pipeline creates fixed-length sequences without padding,
        # so attention_mask is not present. If using padded sequences, ensure:
        #   - Padded positions in labels are set to -100 (ignored by cross_entropy)
        #   - attention_mask marks real tokens (1) vs padding (0)
        attention_mask = batch.get("attention_mask", None)

        # Forward pass (causal mask is created internally in GPTModel)
        logits = self(input_ids, attention_mask=attention_mask)

        # Shift logits and labels for causal language modeling
        # Before shift: labels[i] = input_ids[i]
        # After shift: predict input_ids[i+1] from input_ids[:i+1]
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        # Calculate loss
        loss = nn.functional.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
        )

        # Log loss
        self.log(
            "train/loss",
            loss,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            sync_dist=True,
        )

        # Calculate and log perplexity only on epoch (exp is costly, less frequent is fine)
        perplexity = torch.exp(torch.clamp(loss, max=20.0))
        self.log(
            "train/perplexity",
            perplexity,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            sync_dist=True,
        )

        return loss

    def validation_step(self, batch, _batch_idx):
        # Convert int32 to int64 (long) - MDS stores as int32 but PyTorch expects long
        input_ids = batch["input_ids"].long()
        labels = batch["labels"].long()

        # Get attention mask if present (optional, for padded sequences)
        attention_mask = batch.get("attention_mask", None)

        # Forward pass (causal mask is created internally in GPTModel)
        logits = self(input_ids, attention_mask=attention_mask)

        # Shift logits and labels
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        # Calculate loss
        loss = nn.functional.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
        )

        # Log loss
        self.log("val/loss", loss, prog_bar=True, sync_dist=True)

        # Calculate and log perplexity (exp is costly, but validation is infrequent so OK)
        perplexity = torch.exp(torch.clamp(loss, max=20.0))
        self.log("val/perplexity", perplexity, prog_bar=True, sync_dist=True)

        return loss
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The model performs a forward pass with a causal (autoregressive) mask created internally, ensuring each token can only attend to earlier positions. To align predictions with targets, the logits and labels are shifted so that the representation at position `i` is used to predict token `i + 1`.

Loss is computed using cross-entropy over the shifted logits and labels. Training loss and perplexity are logged during execution, with metrics synchronized across distributed workers.

The optimizer setup is where a lot of training stability comes from:

```
    def configure_optimizers(self):
        # Separate parameters into weight decay and no weight decay groups
        decay_params = []
        no_decay_params = []

        for param in self.model.parameters():
            if param.requires_grad:
                # 1D parameters (biases, LayerNorm) don't get weight decay
                # 2D+ parameters (weight matrices) get weight decay
                if param.ndim == 1:
                    no_decay_params.append(param)
                else:
                    decay_params.append(param)

        optimizer_grouped_parameters = [
            {"params": decay_params, "weight_decay": self.hparams.weight_decay},
            {"params": no_decay_params, "weight_decay": 0.0},
        ]

        # AdamW optimizer
        optimizer = torch.optim.AdamW(
            optimizer_grouped_parameters,
            lr=self.hparams.learning_rate,
            betas=(0.9, 0.95),
            eps=1e-8,
        )

        # Learning rate scheduler: warmup + cosine decay
        # Warmup: linear increase from 0 to 1.0 over warmup_steps
        # Decay: cosine decay from 1.0 to 0.0 over remaining steps
        def lr_lambda(current_step):
            if current_step < self.hparams.warmup_steps:
                # Linear warmup
                return float(current_step) / float(max(1, self.hparams.warmup_steps))

            # Cosine decay after warmup
            progress = (current_step - self.hparams.warmup_steps) / max(
                1, self.hparams.max_steps - self.hparams.warmup_steps
            )
            # Cosine annealing from 1.0 to 0.0 (returns float, not tensor)
            return 0.5 * (1.0 + math.cos(progress * math.pi))

        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
            },
        }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

Two important choices here:

1. **Separate weight decay groups**: We only apply weight decay to the weight matrices, not to biases or LayerNorm parameters. This follows the original BERT paper and is now standard practice, as regularizing biases and normalization parameters does not improve performance and can be harmful.
2. **Cosine learning rate schedule with warmup**: We start with a low learning rate, ramp up linearly during warmup (helps stabilize early training when gradients are noisy), then decay following a cosine curve. This schedule outperforms constant or step decay for transformer training.

### Checkpointing for fault tolerance

Training a 30B-parameter model for 15,000 steps can take days. Hardware failures and spot instance preemptions are inevitable, which makes checkpointing essential.

```
class S3CheckpointCallback(L.Callback):
    """
    Periodically upload checkpoints to S3 for durability and resumption.

    This ensures checkpoints are safely stored in remote storage even if
    the training job is interrupted or the instance fails.
    """

    def __init__(self, checkpoint_dir: Path, upload_every_n_steps: int):
        super().__init__()
        self.checkpoint_dir = checkpoint_dir
        self.upload_every_n_steps = upload_every_n_steps
        self.last_uploaded_step = -1

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        """Upload checkpoint to S3 every N steps."""
        if trainer.global_rank != 0:
            return  # Only upload from rank 0

        current_step = trainer.global_step

        # Upload every N steps (aligns with ModelCheckpoint's every_n_train_steps)
        if (
            current_step % self.upload_every_n_steps == 0
            and current_step > self.last_uploaded_step
            and current_step > 0
        ):
            try:
                # Find the most recent checkpoint file
                checkpoint_files = list(self.checkpoint_dir.glob("*.ckpt"))
                if not checkpoint_files:
                    print("No checkpoint files found to upload")
                    return

                # Get the latest checkpoint (by modification time)
                latest_checkpoint = max(
                    checkpoint_files, key=lambda p: p.stat().st_mtime
                )

                # Upload the checkpoint file directly to S3 using File.from_local_sync
                checkpoint_file = File.from_local_sync(str(latest_checkpoint))
                print(f"Checkpoint uploaded to S3 at: {checkpoint_file.path}")

                self.last_uploaded_step = current_step
            except Exception as e:
                print(f"Warning: Failed to upload checkpoint to S3: {e}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

This callback runs every `N` training steps and uploads the checkpoint to durable storage. The key line is `File.from_local_sync()` which is a Flyte abstraction for uploading files. There are no blob store credentials to manage and no bucket paths to hardcode. Flyte automatically uses the storage backend configured for your cluster.

The callback only runs on rank 0. In distributed training, all 8 GPUs have identical model states (that's the point of data parallelism). Having all of them upload the same checkpoint would be wasteful and could cause race conditions.

When you restart a failed run, pass the checkpoint via `resume_checkpoint` so training resumes exactly where it left off, including the same step count, optimizer state, and learning rate schedule position.

### Real-time metrics with Flyte reports

Multi-day training runs need observability. Is the loss decreasing? Did training diverge? Is the learning rate schedule behaving correctly? Flyte Reports let you build live dashboards directly in the UI:

```
class FlyteReportingCallback(L.Callback):
    """Custom Lightning callback to report training metrics to Flyte Report."""

    def __init__(self, report_every_n_steps: int = 100):
        super().__init__()
        self.report_every_n_steps = report_every_n_steps
        self.metrics_history = {
            "step": [],
            "train_loss": [],
            "learning_rate": [],
            "val_loss": [],
            "val_perplexity": [],
        }
        self.initialized_report = False
        self.last_logged_step = -1

    def on_train_start(self, trainer, pl_module):
        """Initialize the live dashboard on training start."""
        if trainer.global_rank == 0 and not self.initialized_report:
            self._initialize_report()
            self.initialized_report = True
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The `_initialize_report` method (see complete code) creates an HTML/JavaScript dashboard with interactive charts. The callback then calls `flyte.report.log()` every `N` steps to push new metrics. The charts update in real-time so you can watch your loss curve descend while training runs.

There is no need to deploy Grafana, configure Prometheus, or keep a TensorBoard server running. Using `flyte.report.log()` is sufficient to get live training metrics directly in the Flyte UI.

![Metrics viz](../../../_static/images/tutorials/distributed-llm-pretraining/metrics.png)

### Streaming data at scale

Training datasets are massive. SlimPajama contains 627 billion tokens and spans hundreds of gigabytes even when compressed. Downloading the entire dataset to each training node before starting would take hours and waste storage.

Instead, we convert the data to MDS (MosaicML Data Shard) format and stream it during training:

```
@data_loading_env.task
async def load_and_prepare_streaming_dataset(
    dataset_name: str,
    dataset_config: Optional[str],
    max_length: int,
    train_split: str,
    val_split: Optional[str],
    max_train_samples: Optional[int],
    max_val_samples: Optional[int],
    shard_size_mb: int,
) -> Dir:
    """Tokenize dataset and convert to MDS format for streaming."""
    from datasets import load_dataset
    from streaming import MDSWriter
    from transformers import GPT2TokenizerFast

    output_dir = Path("/tmp/streaming_dataset")
    output_dir.mkdir(parents=True, exist_ok=True)

    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    # MDS schema: what each sample contains
    columns = {
        "input_ids": "ndarray:int32",
        "labels": "ndarray:int32",
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

This task does three things:

1. **Tokenizes the text** using GPT-2's BPE tokenizer
2. **Concatenates documents** into fixed-length sequences (no padding waste)
3. **Writes shards** to storage in a format optimized for streaming

The task returns a Flyte `Dir` object, which is a reference to the output location. It's not the data itself, just a pointer. When the training task receives this `Dir`, it streams shards on-demand rather than downloading everything upfront.

Flyte caches this task automatically. Run the pipeline twice with the same dataset config, and Flyte skips tokenization entirely on the second run. Change the dataset or sequence length, and it re-runs.

### Distributed training with FSDP

Now we get to the core: actually training the model across multiple GPUs. FSDP is what makes this possible for large models.

```
@distributed_llm_training_env.task(report=True)
def train_distributed_llm(
    prepared_dataset: Dir,
    resume_checkpoint: Optional[Dir],
    vocab_size: int,
    n_positions: int,
    n_embd: int,
    n_layer: int,
    n_head: int,
    batch_size: int,
    num_workers: int,
    max_steps: int,
    learning_rate: float,
    weight_decay: float,
    warmup_steps: int,
    use_fsdp: bool,
    checkpoint_upload_steps: int,
    checkpoint_every_n_steps: int,
    report_every_n_steps: int,
    val_check_interval: int,
    grad_accumulation_steps: int = 1,
) -> Optional[Dir]:
    # ... setup code ...
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

Notice `report=True` on the task decorator. It enables Flyte Reports for this specific task.

The training task receives the prepared dataset as a `Dir` and streams data directly from storage:

```
    # StreamingDataset streams shards from the remote Flyte storage on-demand
    # It automatically detects torch.distributed context
    # and shards data across GPUs - each rank gets different data automatically
    train_dataset = StreamingDataset(
        remote=f"{remote_path}/train",  # Remote MDS shard location
        local=str(local_cache / "train"),  # Local cache for downloaded shards
        shuffle=True,  # Shuffle samples
        shuffle_algo="naive",  # Shuffling algorithm
        batch_size=batch_size,  # Used for shuffle buffer sizing
    )

    # Create validation StreamingDataset if it exists
    val_dataset = None
    try:
        val_dataset = StreamingDataset(
            remote=f"{remote_path}/validation",
            local=str(local_cache / "validation"),
            shuffle=False,  # No shuffling for validation
            batch_size=batch_size,
        )
        print(
            f"Validation dataset initialized with streaming from: {remote_path}/validation"
        )
    except Exception as e:
        print(f"No validation dataset found: {e}")

    # Create data loaders
    # StreamingDataset handles distributed sampling internally by detecting
    # torch.distributed.get_rank() and torch.distributed.get_world_size()
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=True,
        drop_last=True,  # Drop incomplete batches for distributed training
        collate_fn=mds_collate_fn,  # Handle read-only arrays
    )

    # Create validation loader if validation dataset exists
    val_loader = None
    if val_dataset is not None:
        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,
            persistent_workers=True,
            drop_last=False,
            collate_fn=mds_collate_fn,
        )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

`prepared_dataset.path` provides the remote storage path for the dataset. MosaicML's `StreamingDataset` automatically shards data across GPUs so that each rank sees different samples, without requiring a manual distributed sampler. The credentials are already in the environment because Flyte set them up.

FSDP is where the memory magic happens. Instead of each GPU holding a full copy of the model (like Distributed Data Parallel (DDP)), FSDP shards the parameters, gradients, and optimizer states across all GPUs. Each GPU only holds 1/8th of the model. When a layer needs to run, FSDP all-gathers the full parameters, runs the computation, then discards them.

```
    # Configure distributed strategy
    if use_fsdp:
        from torch.distributed.fsdp.wrap import ModuleWrapPolicy

        strategy = FSDPStrategy(
            auto_wrap_policy=ModuleWrapPolicy([GPTBlock]),
            activation_checkpointing_policy=None,
            cpu_offload=False,  # H200 has 141GB - no CPU offload needed
            state_dict_type="full",
            sharding_strategy="FULL_SHARD",
            process_group_backend="nccl",
        )
    else:
        from lightning.pytorch.strategies import DDPStrategy

        strategy = DDPStrategy(process_group_backend="nccl")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

We wrap at the `GPTBlock` level because each transformer block becomes an FSDP unit. This balances communication overhead (more units = more all-gathers) against memory savings (smaller units = more granular sharding).

One subtle detail: gradient clipping. With FSDP, gradients are sharded across ranks, so computing a global gradient norm would require an expensive all-reduce operation. Instead of norm-based clipping, we use value-based gradient clipping, which clamps each individual gradient element to a fixed range. This can be done independently on each rank with no coordination overhead and is commonly used for large-scale FSDP training.

```
    # Initialize trainer
    trainer = L.Trainer(
        strategy=strategy,
        accelerator="gpu",
        devices=DEVICES_PER_NODE,
        num_nodes=NUM_NODES,
        # Training configuration
        max_steps=max_steps,
        precision="bf16-mixed",  # BFloat16 for better numerical stability
        # Optimization
        gradient_clip_val=1.0,
        gradient_clip_algorithm=(
            "value" if use_fsdp else "norm"
        ),  # FSDP requires 'value', DDP can use 'norm'
        accumulate_grad_batches=grad_accumulation_steps,
        # Logging and checkpointing
        callbacks=callbacks,
        log_every_n_steps=report_every_n_steps,
        val_check_interval=val_check_interval,
        # Performance
        benchmark=True,
        deterministic=False,
        # Enable gradient checkpointing for memory efficiency
        enable_checkpointing=True,
        use_distributed_sampler=False,  # StreamingDataset handles distributed sampling
    )

    # Train the model (resume from checkpoint if provided)
    trainer.fit(model, train_loader, val_loader, ckpt_path=ckpt_path)

    # Print final results
    if trainer.global_rank == 0:
        if val_loader is not None:
            print(
                f"Final validation loss: {trainer.callback_metrics.get('val/loss', 0.0):.4f}"
            )
            print(
                f"Final validation perplexity: {trainer.callback_metrics.get('val/perplexity', 0.0):.4f}"
            )
        print(f"Checkpoints saved to: {checkpoint_dir}")

        return Dir.from_local_sync(output_dir)

    return None
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The trainer configuration brings together all the pieces we've discussed:

- **`precision="bf16-mixed"`**: BFloat16 mixed precision training. BF16 has the same dynamic range as FP32 (unlike FP16), so you don't need loss scaling. This is the standard choice for modern GPU training.
- **`gradient_clip_val=1.0`**: Clips gradients to prevent exploding gradients during training. Combined with value-based clipping for FSDP compatibility.
- **`accumulate_grad_batches`**: Accumulates gradients over multiple forward passes before updating weights. This lets us hit a larger effective batch size than what fits in GPU memory.
- **`val_check_interval`**: How often to run validation. For long training runs, you don't want to validate every epoch. That would be too infrequent. Instead, validate every `N` training steps.
- **`use_distributed_sampler=False`**: We disable Lightning's built-in distributed sampler because `StreamingDataset` handles data sharding internally. Using both would cause conflicts.
- **`benchmark=True`**: Enables cuDNN autotuning. PyTorch will benchmark different convolution algorithms on the first batch and pick the fastest one for your specific input sizes.

The trainer then calls `fit()` with the model, data loaders, and optionally a checkpoint path to resume from.

### Tying it together

The pipeline task orchestrates everything:

```
@driver_env.task
async def distributed_llm_pipeline(
    model_size: str,
    dataset_name: str = "Salesforce/wikitext",
    dataset_config: str = "wikitext-103-raw-v1",
    max_length: int = 2048,
    max_train_samples: Optional[int] = 10000,
    max_val_samples: Optional[int] = 1000,
    max_steps: int = 100000,
    resume_checkpoint: Optional[Dir] = None,
    checkpoint_upload_steps: int = 1000,
    # Optional overrides (if None, uses model preset defaults)
    batch_size: Optional[int] = None,
    learning_rate: Optional[float] = None,
    use_fsdp: bool = True,
) -> Optional[Dir]:
    # Get model configuration
    model_config = get_model_config(model_size)

    # Use preset values if not overridden
    actual_batch_size = (
        batch_size if batch_size is not None else model_config["batch_size"]
    )
    actual_learning_rate = (
        learning_rate if learning_rate is not None else model_config["learning_rate"]
    )

    # Step 1: Load and prepare streaming dataset
    prepared_dataset = await load_and_prepare_streaming_dataset(
        dataset_name=dataset_name,
        dataset_config=dataset_config,
        max_length=max_length,
        train_split="train",
        val_split="validation",
        max_train_samples=max_train_samples,
        max_val_samples=max_val_samples,
        shard_size_mb=64,  # 64MB shards
    )

    # Step 2: Run distributed training
    if resume_checkpoint is not None:
        print("\nStep 2: Resuming distributed training from checkpoint...")
    else:
        print("\nStep 2: Starting distributed training from scratch...")

    target_global_batch = 256
    world_size = NUM_NODES * DEVICES_PER_NODE
    effective_per_step = world_size * actual_batch_size
    grad_accumulation_steps = max(
        1, math.ceil(target_global_batch / max(1, effective_per_step))
    )

    checkpoint_dir = train_distributed_llm(
        prepared_dataset=prepared_dataset,
        resume_checkpoint=resume_checkpoint,
        vocab_size=VOCAB_SIZE,
        n_positions=N_POSITIONS,
        n_embd=model_config["n_embd"],
        n_layer=model_config["n_layer"],
        n_head=model_config["n_head"],
        batch_size=actual_batch_size,
        num_workers=8,
        max_steps=max_steps,
        learning_rate=actual_learning_rate,
        weight_decay=0.1,
        warmup_steps=500,
        use_fsdp=use_fsdp,
        checkpoint_upload_steps=checkpoint_upload_steps,
        checkpoint_every_n_steps=model_config["checkpoint_every_n_steps"],
        report_every_n_steps=model_config["report_every_n_steps"],
        val_check_interval=model_config["val_check_interval"],
        grad_accumulation_steps=grad_accumulation_steps,
    )

    return checkpoint_dir
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

The flow is straightforward: load the configuration, prepare the data, and run training. Flyte automatically manages the execution graph so data preparation runs first and training waits until it completes. If data preparation is cached from a previous run, training starts immediately.

The gradient accumulation calculation balances two constraints. We want a global batch size of 256 (this affects training dynamics), but each GPU can only fit a small batch. With 8 GPUs and batch size 1 each, we need 32 accumulation steps to hit 256.

## Running the pipeline

With everything defined, running is simple:

```
if __name__ == "__main__":
    flyte.init_from_config()

    run = flyte.run(
        distributed_llm_pipeline,
        model_size="30B",
        dataset_name="cerebras/SlimPajama-627B",
        dataset_config=None,
        max_length=2048,
        max_train_samples=5_000_000,
        max_val_samples=50_000,
        max_steps=15_000,
        use_fsdp=True,
        checkpoint_upload_steps=1000,
    )

    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py*

This configuration is designed for testing and demonstration. Notice `max_train_samples=5_000_000`: that's 5 million samples from a dataset with 627 billion tokens. A tiny fraction, enough to verify everything works without burning through compute.

For a real pretraining run, you would remove this limit by setting `max_train_samples=None`, or increase it significantly. You would also increase `max_steps` to match your compute budget, likely scale to multiple nodes by setting `NUM_NODES=4` or higher, and allocate more resources. The rest of the pipeline remains unchanged.

```bash
flyte create config --endpoint <FLYTE_OR_UNION_ENDPOINT> --project <PROJECT_NAME> --domain <DOMAIN_NAME> --builder remote
uv run train.py
```

When you run this, Flyte:

1. **Builds the container** (cached after first run)
2. **Schedules data prep** on CPU nodes
3. **Waits for data prep** (or skips if cached)
4. **Provisions H200 nodes** and launches distributed training
5. **Streams logs and metrics** to the UI in real-time

Open the Flyte UI to observe the workflow execution. The data preparation task completes first, followed by the training task spinning up. As training begins, the Flyte Reports dashboard starts plotting loss curves. If anything goes wrong, the logs are immediately available in the UI.

![Training Log](../../../_static/images/tutorials/distributed-llm-pretraining/logs.png)

If training fails due to an out-of-memory error, a GPU driver error, or a hardware issue, check the logs, fix the problem, and restart the run with `resume_checkpoint` pointing to the most recent checkpoint. Training resumes from where it left off. Flyte tracks the full execution history, so it is easy to see exactly what happened.

## Going further

If you've run through this tutorial, here's where to go next depending on what you're trying to do:

**You want to train on your own data.** The data prep task accepts any HuggingFace dataset with a `text` column. If your data isn't on HuggingFace, you can modify `load_and_prepare_streaming_dataset` to read from S3, local files, or any other source. The key is getting your data into MDS format. Once it's there, the streaming and sharding just works. For production training, look at SlimPajama, [RedPajama](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T), or [The Pile](https://huggingface.co/datasets/EleutherAI/pile) as starting points.

**You want to scale to more GPUs.** Bump `NUM_NODES` and Flyte handles the rest. The main thing to watch is the effective batch size. As you add more GPUs, you may want to reduce gradient accumulation steps to keep the same global batch size, or increase them if you want to experiment with larger batches.

**Your training keeps failing.** Add `retries=3` to your task decorator for automatic retry on transient failures. This handles spot instance preemption, temporary network issues, and the occasional GPU that decides to stop working. Combined with checkpointing, you get fault-tolerant training that can survive most infrastructure hiccups. For persistent failures, the Flyte UI logs are your friend as they capture stdout/stderr from all ranks.

**You want better visibility into what's happening.** We're actively working on surfacing GPU driver logs (xid/sxid errors), memory utilization breakdowns, and NCCL communication metrics directly in the Flyte UI. If you're hitting issues that the current logs don't explain, reach out. Your feedback helps us prioritize what observability features to build next!


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision ===

# Computer vision

Tutorials for image and vision-language model workloads.

### **Computer vision > Fine-tuning a VLM**

Adapt Qwen2.5-VL to occluded image classification by training a 10K-parameter adapter with multi-node DeepSpeed, automatic recovery, and live training dashboards.

### **Computer vision > RT-DETR object detection**

Fine-tune RT-DETRv2 on a COCO dataset with live training charts, mAP evaluation, and bounding-box demos.

### **Computer vision > Multimodal retrieval evaluation**

Benchmark ColPali, SigLIP, and OCR+BM25 visual document retrieval on ViDoRe with warm GPU containers, dynamic batching, and an interactive report.

## Subpages
- **Computer vision > Fine-tuning a VLM**
- **Computer vision > Multimodal retrieval evaluation**
- **Computer vision > RT-DETR object detection**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/qwen-vl-finetuning ===

# Fine-tuning a VLM

Large vision-language models like Qwen2.5-VL are remarkably capable out of the box. But adapting one to a specialized task raises an immediate question: do you really need to update 3 billion parameters?

Usually, no. The **frozen backbone pattern** is a practical alternative: keep all pretrained weights frozen and train only a small, task-specific adapter inserted before the vision encoder. The adapter learns to transform its input in a way that makes the frozen model perform well on your task without touching the underlying billions of parameters. The result is faster training, lower memory pressure, and a much smaller set of weights to store and version.

This tutorial makes that pattern concrete. We take a partially-occluded image classification task (CIFAR-10 images with random black rectangles covering 22 to 45% of the frame) and train a tiny Conv2d adapter to "see through" the occlusion before the frozen VLM processes it. The adapter has approximately **10,500 trainable parameters**. The backbone has 3 billion.

The machine learning is interesting, but the real focus here is on shipping a production-grade training pipeline:

- **Multi-node distributed training** across 2 nodes × 4 GPUs using PyTorch Elastic and DeepSpeed Stage 2
- **Automatic fault tolerance**: checkpoints upload to object storage after every validation epoch; if training fails, the pipeline returns the last known-good checkpoint instead of crashing
- **Live observability**: a streaming HTML dashboard in the Flyte UI updates in real-time as training runs, no separate monitoring infrastructure required
- **Cached data preparation**: dataset processing runs once and is reused across all reruns
- **Clean task isolation**: each stage runs with exactly the resources it needs, nothing more

> [!NOTE]
> Full code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning).

## Overview

The pipeline has four tasks with clearly defined responsibilities:

1. **Dataset preparation** (`prepare_occlusion_dataset`): Downloads CIFAR-10, applies random occlusions, and writes image manifests as streaming JSONL files to object storage. Runs on CPU and is cached, so it only runs once regardless of how many times you rerun the pipeline with the same config.
2. **Multi-node training** (`train_qwen_adapter_multinode`): Runs PyTorch Lightning with DeepSpeed Stage 2 across 2 nodes × 4 L40s GPUs. Only the adapter trains; the 3B backbone stays frozen.
3. **Evaluation** (`evaluate_qwen_adapter`): Loads the saved adapter, runs inference on validation examples, and produces a predictions report. Runs on a single GPU.
4. **Driver** (`qwen_vl_multinode_deepspeed`): The pipeline entry point. Orchestrates the three tasks above, manages WandB initialization, handles recovery from training failures, and produces a final HTML report in the Flyte UI.

Why this separation? It mirrors how production pipelines should be structured. Data prep is cheap and deterministic so we cache it. Training is expensive and failure-prone so we isolate it with fault tolerance. Evaluation needs different hardware than training. The driver is pure coordination, so it gets minimal resources.

## Implementation

### Setting up the environment

Different tasks need different compute. Flyte's `TaskEnvironment` is how you declare exactly what each task needs.

First, define the container images. Training needs a full CUDA stack with ML libraries, driver compatibility, and DeepSpeed's build tools:

```
gpu_image = (
    flyte.Image.from_base("nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04")
    .clone(name="qwen_vl_multinode_deepspeed", python_version=(3, 13), extendable=True)
    .with_apt_packages("build-essential")
    .with_pip_packages(
        "torch==2.9.1",
        "torchvision==0.24.1",
        "lightning==2.6.1",
        "transformers==4.57.3",
        "deepspeed==0.18.8",
        "datasets==4.4.1",
        "pillow==11.3.0",
        "flyteplugins-pytorch>=2.0.11",
        "flyteplugins-jsonl>=2.0.11",
        "flyteplugins-wandb>=2.0.11",
    )
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/config.py*

`from_base` starts from the official NVIDIA CUDA image, giving you NCCL, cuDNN, and the right driver headers out of the box. `with_apt_packages("build-essential")` is required because DeepSpeed compiles CUDA kernels at first use and without build tools, it silently falls back to slower CPU implementations. The non-GPU image for data preparation and orchestration is much lighter:

```
non_gpu_image = flyte.Image.from_debian_base(
    name="qwen_vl_multinode_deepspeed_non_gpu"
).with_pip_packages(
    "flyteplugins-pytorch>=2.0.11",
    "flyteplugins-jsonl>=2.0.11",
    "flyteplugins-wandb>=2.0.11",
    "lightning==2.6.1",
    "datasets==4.4.1",
    "pillow==11.3.0",
    "torchvision==0.24.1",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/config.py*

With images defined, each task gets its own resource declaration:

```
dataset_env = flyte.TaskEnvironment(
    name="qwen_vl_dataset_prep",
    image=non_gpu_image,
    resources=flyte.Resources(cpu=5, memory="15Gi"),
    cache="auto",
)

training_env = flyte.TaskEnvironment(
    name="qwen_vl_multinode_training",
    image=gpu_image,
    resources=flyte.Resources(
        cpu=42,
        memory="256Gi",
        gpu=f"L40s:{DEVICES_PER_NODE}",
        shm="16Gi",
    ),
    plugin_config=Elastic(nnodes=NUM_NODES, nproc_per_node=DEVICES_PER_NODE),
    secrets=[
        flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")
    ],  # TODO: update with your own secret key
    env_vars={
        "TORCH_DISTRIBUTED_DEBUG": "INFO",
        "NCCL_DEBUG": "WARN",
        "TOKENIZERS_PARALLELISM": "false",
        "CUDA_HOME": "/usr/local/cuda",
        "DS_SKIP_CUDA_CHECK": "1",
    },
)

evaluation_env = flyte.TaskEnvironment(
    name="qwen_vl_adapter_eval",
    image=gpu_image,
    resources=flyte.Resources(cpu=16, memory="64Gi", gpu="L40s:1"),
    cache="auto",
)

driver_env = flyte.TaskEnvironment(
    name="qwen_vl_multinode_driver",
    image=non_gpu_image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[dataset_env, training_env, evaluation_env],
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/config.py*

A few things worth noting here:

- **`Elastic(nnodes=2, nproc_per_node=4)`**: Flyte's integration with PyTorch's elastic launch. It handles process spawning (one process per GPU), rank assignment, and distributed environment setup (master address, world size, rendezvous) without any shell scripting or manual `torchrun` invocations.
- **`shm="16Gi"`**: Shared memory is required for NCCL inter-GPU communication on the same node. Without it, you'll see cryptic errors from the communication library when training starts.
- **`cache="auto"`**: The dataset preparation task is cached by input hash. Running the pipeline twice with the same hyperparameters skips it entirely on the second run.
- **`depends_on`**: The driver task declares that each worker image must finish building before it starts, ensuring containers are ready before the driver begins orchestrating.
- **`secrets`**: The WandB API key is injected from Flyte's secret store as an environment variable. No credentials in code.

All training hyperparameters flow through a single typed dataclass:

```
@dataclass
class Config:
    model_name: str = DEFAULT_MODEL_NAME
    image_size: int = IMAGE_SIZE
    max_train_samples: int = 1024
    max_val_samples: int = 256
    epochs: int = 8
    per_device_batch_size: int = 1
    target_global_batch_size: int = 16
    learning_rate: float = 2e-4
    weight_decay: float = 1e-2
    reconstruction_loss_weight: float = 0.35
    report_every_n_steps: int = 10
    num_workers: int = 4
    max_length: int = 512
    eval_examples: int = 16
    train_occlusion_min: float = 0.22
    train_occlusion_max: float = 0.42
    eval_occlusion_min: float = 0.28
    eval_occlusion_max: float = 0.45
    seed: int = 7

    def to_dict(self) -> dict:
        return asdict(self)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/config.py*

Using a dataclass rather than scattered constants or argparse arguments means the full config is serializable, can be stored in artifact metadata alongside the model checkpoint, and flows cleanly as a typed input between tasks. The `to_dict()` method serializes it for WandB logging.

### Preparing the dataset

The dataset task handles everything: downloading CIFAR-10, generating occlusions, and writing the manifests.

```
@dataset_env.task
async def prepare_occlusion_dataset(config: Config) -> DatasetArtifacts:
    from PIL import Image
    from torchvision.datasets import CIFAR10
    from flyte.io import Dir
    from flyteplugins.jsonl import JsonlFile
    import random

    rng = random.Random(config.seed)

    images_dir = Path("/tmp/qwen_vl_occlusion_images")
    train_images_dir = images_dir / "train" / "images"
    val_images_dir = images_dir / "validation" / "images"
    train_images_dir.mkdir(parents=True, exist_ok=True)
    val_images_dir.mkdir(parents=True, exist_ok=True)

    prompt = (
        "The image may be partially occluded. "
        "Answer with exactly one CIFAR-10 class label: "
        + ", ".join(CLASS_NAMES)
        + ". What is the main object?"
    )

    async def export_split(
        dataset,
        split_name: str,
        limit: int,
        local_image_dir: Path,
        occ_min: float,
        occ_max: float,
    ):
        out = JsonlFile.new_remote(f"{split_name}_manifest.jsonl")
        async with out.writer() as writer:
            for idx in range(limit):
                pil_image, label_idx = dataset[idx]
                resized = pil_image.resize(
                    (config.image_size, config.image_size),
                    resample=Image.Resampling.BICUBIC,
                )
                rel_path = f"{split_name}/images/{split_name}-{idx:05d}.png"
                resized.save(local_image_dir / f"{split_name}-{idx:05d}.png")
                occlusion = build_occlusion_box(
                    width=config.image_size,
                    height=config.image_size,
                    rng=rng,
                    min_fraction=occ_min,
                    max_fraction=occ_max,
                )
                await writer.write(
                    {
                        "image_path": rel_path,
                        "label": CLASS_NAMES[label_idx],
                        "label_index": int(label_idx),
                        "prompt": prompt,
                        "occlusion": occlusion,
                    }
                )
        return out

    train_dataset = CIFAR10(root="/tmp/cifar10", train=True, download=True)
    val_dataset = CIFAR10(root="/tmp/cifar10", train=False, download=True)

    train_manifest = await export_split(
        train_dataset,
        "train",
        config.max_train_samples,
        train_images_dir,
        config.train_occlusion_min,
        config.train_occlusion_max,
    )
    val_manifest = await export_split(
        val_dataset,
        "validation",
        config.max_val_samples,
        val_images_dir,
        config.eval_occlusion_min,
        config.eval_occlusion_max,
    )

    return DatasetArtifacts(
        train_manifest=train_manifest,
        val_manifest=val_manifest,
        images=await Dir.from_local(str(images_dir)),
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/data.py*

Each image gets a randomly-placed black rectangle. The occlusion covers 22 to 42% of the image area during training and 28 to 45% during evaluation. The occlusion is deliberately harder at eval time to test how robust the adapter is. The bounding box coordinates are written into each manifest record alongside the image path and ground-truth label, so the training task can reconstruct the binary occlusion mask as the adapter's fourth input channel.

Two Flyte primitives handle data persistence without any manual storage management:

- **`JsonlFile.new_remote()`** opens a streaming writer that writes directly to remote object storage. The training task reads records back via `jf.iter_records_sync()`, so no local file paths and S3 credentials to manage.
- **`Dir.from_local()`** uploads the local images directory to object storage and returns a typed handle. The training task downloads it to a local path via `Dir.download_sync()`.

Because `cache="auto"` is set on this task, dataset preparation runs once. Subsequent reruns with the same config skip it entirely.

### The adapter

Here's the entire trainable component of the model with `~10,500` parameters:

```
class ResidualOcclusionAdapter(nn.Module):
    def __init__(self, hidden_channels: int = 32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(4, hidden_channels, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Conv2d(hidden_channels, 3, kernel_size=1),
            nn.Tanh(),
        )
        self.gate = nn.Parameter(torch.tensor(0.10))

    def forward(
        self, pixel_values: torch.Tensor, occlusion_mask: torch.Tensor
    ) -> torch.Tensor:
        if pixel_values.ndim != 4:
            raise ValueError(
                "ResidualOcclusionAdapter expects dense image tensors with shape "
                f"(B, C, H, W), but received {tuple(pixel_values.shape)}."
            )
        if occlusion_mask.ndim == 3:
            occlusion_mask = occlusion_mask.unsqueeze(1)
        adapter_input = torch.cat(
            [pixel_values, occlusion_mask.to(pixel_values.dtype)],
            dim=1,
        )
        residual = self.net(adapter_input)
        return pixel_values + torch.tanh(self.gate) * residual
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/model.py*

The adapter takes the occluded image (3 channels) concatenated with the binary occlusion mask (1 channel) as a 4-channel input. It predicts a residual correction through a small convolutional network, then adds that correction back to the original pixels. The learnable `gate` scalar, initialized to `0.10`, controls how strongly the adapter modifies the image. It starts as a near-identity transformation and gradually grows during training as the adapter gains confidence.

The adapter is plugged into Qwen2.5-VL via a Lightning module:

```
class QwenVLAdapterModule(L.LightningModule):
    def __init__(
        self,
        model_name: str,
        learning_rate: float,
        weight_decay: float,
        reconstruction_loss_weight: float,
    ):
        super().__init__()
        from transformers import Qwen2_5_VLForConditionalGeneration

        self.save_hyperparameters()
        self.adapter = ResidualOcclusionAdapter()

        self.backbone = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            attn_implementation="sdpa",
        )
        self.backbone.requires_grad_(False)
        self.backbone.gradient_checkpointing_enable()

        # DeepSpeed checkpoints only persist the trainable adapter weights when
        # `exclude_frozen_parameters=True`. On resume we rebuild the frozen
        # backbone from Hugging Face and load the checkpoint non-strictly.
        self.strict_loading = False

        self.total_params, self.trainable_params = count_parameters(self)
        self.example_input_array = None
        self.vision_patch_size = int(self.backbone.config.vision_config.patch_size)
        self.temporal_patch_size = int(
            getattr(self.backbone.config.vision_config, "temporal_patch_size", 1)
        )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/model.py*

The key line is `self.backbone.requires_grad_(False)`. This freezes all 3 billion backbone parameters which means only the adapter's ~10,500 weights receive gradients. `gradient_checkpointing_enable()` trades compute for memory: instead of keeping the frozen backbone's intermediate activations in GPU memory during the backward pass, they're recomputed on the fly. This is critical when a 3B model is sitting in GPU memory alongside your optimizer state.

`strict_loading = False` handles an important DeepSpeed checkpoint detail. When `exclude_frozen_parameters=True` is set on the strategy, DeepSpeed only saves the adapter weights in checkpoints, not the 3B frozen backbone. On resume, the checkpoint won't contain backbone weights, so loading must be non-strict. The `on_load_checkpoint` hook fills in the missing backbone weights from the freshly-loaded HuggingFace model, combining the best of both worlds: small checkpoints and a fully initialized model.

The training loss combines two objectives:

```
    def _forward_losses(
        self, batch: dict[str, torch.Tensor]
    ) -> dict[str, torch.Tensor]:
        backbone_dtype = next(self.backbone.parameters()).dtype
        if batch["pixel_values"].ndim == 2:
            if "image_grid_thw" not in batch:
                raise ValueError(
                    "Packed Qwen pixel values require `image_grid_thw` to reconstruct "
                    "dense images for the Conv2d adapter."
                )
            grid_thw = batch["image_grid_thw"]
            dense_pixels = packed_pixels_to_dense_images(
                batch["pixel_values"].to(dtype=backbone_dtype),
                grid_thw,
                patch_size=self.vision_patch_size,
                temporal_patch_size=self.temporal_patch_size,
            )
            clean_pixels = packed_pixels_to_dense_images(
                batch["clean_pixel_values"].to(dtype=backbone_dtype),
                grid_thw,
                patch_size=self.vision_patch_size,
                temporal_patch_size=self.temporal_patch_size,
            )
            adapted_dense = self.adapter(dense_pixels, batch["occlusion_mask"])
            adapted_pixels = dense_images_to_packed_pixels(
                adapted_dense,
                grid_thw,
                patch_size=self.vision_patch_size,
                temporal_patch_size=self.temporal_patch_size,
            )
        else:
            clean_pixels = batch["clean_pixel_values"].to(dtype=backbone_dtype)
            adapted_dense = self.adapter(
                batch["pixel_values"].to(dtype=backbone_dtype),
                batch["occlusion_mask"],
            )
            adapted_pixels = adapted_dense

        forward_kwargs = {
            "input_ids": batch["input_ids"],
            "attention_mask": batch["attention_mask"],
            "pixel_values": adapted_pixels,
            "labels": batch["labels"],
        }
        if "image_grid_thw" in batch:
            forward_kwargs["image_grid_thw"] = batch["image_grid_thw"]
        outputs = self.backbone(**forward_kwargs)

        clean_pixels = clean_pixels.to(
            device=adapted_pixels.device, dtype=backbone_dtype
        )
        occlusion_mask = batch["occlusion_mask"].to(
            device=adapted_pixels.device,
            dtype=backbone_dtype,
        )
        if occlusion_mask.ndim == 3:
            occlusion_mask = occlusion_mask.unsqueeze(1)
        if occlusion_mask.shape[-2:] != adapted_dense.shape[-2:]:
            occlusion_mask = F.interpolate(
                occlusion_mask,
                size=adapted_dense.shape[-2:],
                mode="nearest",
            )

        reconstruction_error = (adapted_dense - clean_pixels).abs() * occlusion_mask
        mask_denominator = (occlusion_mask.sum() * adapted_dense.shape[1]).clamp_min(
            1.0
        )

        reconstruction_loss = reconstruction_error.sum() / mask_denominator
        total_loss = (
            outputs.loss + self.hparams.reconstruction_loss_weight * reconstruction_loss
        )

        return {
            "total_loss": total_loss,
            "lm_loss": outputs.loss,
            "reconstruction_loss": reconstruction_loss,
        }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/model.py*

The **language modeling loss** (cross-entropy on the predicted class label tokens) drives the model to produce correct answers. The **reconstruction loss** (mean absolute error between the adapter's output and the clean image, computed only in the occluded region) pushes the adapter to actually restore the missing pixels rather than finding a representation shortcut. Without it, the adapter could overfit the frozen backbone's quirks and produce correct tokens while generating noise in the masked region. The `reconstruction_loss_weight` (default `0.35`) balances these two objectives.

Because Qwen2.5-VL's preprocessor packs image patches into a flat `(num_patches, patch_dim)` tensor, the adapter must unpack this into a spatial `(B, C, H, W)` tensor, apply the convolutions, then repack. The `packed_pixels_to_dense_images` and `dense_images_to_packed_pixels` utilities in `model.py` handle this format conversion transparently.

### Multi-node training with DeepSpeed

The training task is a standard PyTorch Lightning training loop with distributed infrastructure handled by Flyte and DeepSpeed:

CODE6

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

The `@wandb_init` decorator integrates with the `wandb_config` context created in the driver task. It retrieves the initialized WandB run and attaches a `WandbLogger` to the trainer. The `report=True` flag on the task decorator enables Flyte Reports for live dashboard streaming from this task.

![Live Training](../../../_static/images/tutorials/qwen-vl-finetuning/live_training_graph.png)
![Live Training Contd](../../../_static/images/tutorials/qwen-vl-finetuning/losses.png)

DeepSpeed Stage 2 shards optimizer states and gradients across GPUs, reducing per-GPU memory usage significantly. The critical configuration flag here is `exclude_frozen_parameters=True`:

CODE7

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

Without `exclude_frozen_parameters=True`, DeepSpeed would shard and checkpoint the frozen backbone weights too, producing enormous checkpoint files, slow checkpoint saves, and unnecessary communication overhead. With it, only the adapter participates in sharding and checkpointing. The backbone is loaded independently on each worker from HuggingFace.

Gradient accumulation is computed automatically to hit the target global batch size regardless of how many GPUs are actually running:

CODE8

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

With 2 nodes × 4 GPUs × per-device batch size 1, the effective per-step batch is 8. To reach the default target of 16, the trainer accumulates over 2 steps. Change `NUM_NODES` or `per_device_batch_size` and the calculation adjusts automatically.

The trainer brings everything together:

CODE9

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

`precision="bf16-mixed"` uses BFloat16, which matches FP32's dynamic range (unlike FP16), so you don't need loss scaling. This is the standard choice for modern VLM training. `benchmark=True` runs cuDNN autotuning on the first batch to select the fastest kernels for your specific input sizes.

### Fault tolerance and recovery

Multi-node GPU jobs fail. Hardware hiccups, spot instance preemptions, NCCL timeouts, memory spikes, etc. and the question is when, not if. This pipeline handles it with a two-part system.

After every validation epoch, the `RecoveryArtifactCallback` calls `trainer.save_checkpoint()` to write a DeepSpeed checkpoint directory, then uploads all shard files to the recovery URI. Each node's local rank 0 uploads its own shards; global rank 0 uploads the metadata files (`metrics.json`, `summary.json`). A distributed barrier between save and upload ensures all workers finish before training continues.

If training fails, the driver task catches the error and returns the last recovery artifact instead of propagating the failure:

CODE10

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

A failed run still produces useful output: the best checkpoint reached before the failure, along with a partial training report. To resume from that point, pass the recovery artifact as `resume_training_artifacts` on the next run. The training task downloads it, finds the most recent `.ckpt` file, and passes it to `trainer.fit()` as `ckpt_path`. Training picks up at the last saved epoch with optimizer state and metrics history intact.

The recovery URI is constructed from the configurable base path and the run name:

CODE11

This means each run gets its own recovery location, so you can identify exactly which run a checkpoint came from.

### Live observability

`flyte.report` lets you push HTML content directly into the Flyte UI during task execution, with no separate monitoring infrastructure. The `LiveTrainingReportCallback` uses this to stream training metrics in real-time:

CODE12

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/callbacks.py*

`on_train_start` (see the full code) initializes the dashboard with an SVG loss chart and an HTML metrics table. Every `report_every_n_steps` training steps, `_push_update` serializes the latest metrics into a `<script>` block and calls `flyte.report.log()` to append it to the live page. The JavaScript `updateQwenLiveReport()` function then updates the chart polylines and appends a new table row for each step.

For resumed runs, the prior metrics history is seeded into the table on `on_train_start`, so the metrics view is continuous across runs rather than restarting from zero.

![Recovery](../../../_static/images/tutorials/qwen-vl-finetuning/recovery.png)

WandB metrics are logged in parallel by `AdapterMetricsCallback` after each validation epoch, including per-epoch train and validation losses, the LM loss component, the reconstruction loss component, and the current adapter gate value.

![WandB](../../../_static/images/tutorials/qwen-vl-finetuning/wandb.png)

### Evaluation

After training completes, a separate task runs inference on a single GPU:

CODE13

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

The task is async so the driver can `asyncio.gather` the downloads of training artifacts and images in parallel rather than sequentially, a simple speedup that matters when downloading hundreds of megabytes from object storage.

The evaluation task loads the adapter state dict from `adapter_artifact.pt`, rebuilds the frozen backbone fresh from HuggingFace (there's no need to checkpoint 3B weights, only the ~10,500 adapter weights travel with the artifact), and runs greedy decoding on each validation example. The metric is exact-match accuracy between the model's predicted class token and the ground-truth CIFAR-10 label.

### Putting it all together

The driver task is the pipeline entry point that all other tasks flow through:

CODE14

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/tasks.py*

The driver constructs the recovery URI from `checkpoint_base_uri` and the current run name, prepares the dataset (or retrieves it from cache), then executes training inside a `wandb_config` context. The `wandb_config` context manager creates and registers a WandB run; the `@wandb_init` decorator on the training task retrieves it, updates it with the full `Config` dataclass, and attaches a `WandbLogger`. Neither the training task nor the callbacks need any WandB initialization code of their own.

The recovery handler (shown in the previous section) wraps the training call. If training succeeds, evaluation runs next. The driver then downloads both the training and evaluation artifacts concurrently and assembles a final HTML report with training curves, evaluation summary, per-epoch metrics table, and sample prediction cards with the actual occluded images, which is pushed to Flyte Reports.

## Running the tutorial

Before running, update two placeholders in `config.py`:

- `DEFAULT_CHECKPOINT_BASE_URI`: your S3, GCS, or Azure Blob URI for checkpoint storage
- The `wandb_api_key` secret key name to match your cluster's secret store configuration

Then configure and launch:

CODE15

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen_vl_frozen_backbone_finetuning/train.py*

CODE16

When you run this, the pipeline:

1. **Builds containers** once and caches them for subsequent runs
2. **Prepares the dataset**: downloads CIFAR-10, generates occlusions, writes JSONL manifests; cached on subsequent runs with the same config
3. **Launches multi-node training**: provisions 2 × 4 L40s GPUs and starts the Elastic job
4. **Streams metrics to the live dashboard**: the Flyte Reports view starts updating as soon as the first step logs
5. **Runs evaluation**: a single-GPU task loads the adapter and runs inference, computing exact-match accuracy
6. **Generates the final report**: training curves, evaluation summary, and sample prediction cards appear in the Flyte UI

![Final Report](../../../_static/images/tutorials/qwen-vl-finetuning/final_report.png)
![Predictions](../../../_static/images/tutorials/qwen-vl-finetuning/predictions.png)

To resume a failed or interrupted run, uncomment the `resume_training_artifacts` line in `train.py` and point it to the recovery URI from the previous run. Training picks up from the last checkpoint with metrics history intact.

## Going further

**Adapting this to a different task.** The frozen backbone pattern transfers directly. Replace `QwenOcclusionDataset` and `prepare_occlusion_dataset` with your own data, update the prompt template, and adjust the dual loss if a pixel-level reconstruction term doesn't apply to your task. The multi-node Elastic setup, DeepSpeed Stage 2 config, recovery system, and live reporting are completely reusable.

**Using a larger Qwen model.** Change `DEFAULT_MODEL_NAME` to `Qwen/Qwen2.5-VL-7B-Instruct` or a larger variant. You may need to increase `memory` in `training_env` and reduce `per_device_batch_size`. The frozen backbone + adapter pattern becomes more valuable at larger scales where you're always training the same ~10,500-parameter adapter regardless of backbone size.

**Training keeps failing.** Add `retries=3` to the `@training_env.task` decorator. With the recovery callback uploading checkpoints after every validation epoch, Flyte automatically restarts training from the last checkpoint on transient failures. Spot instance preemptions and most hardware hiccups become non-events.

**Scaling to more nodes.** Increase `NUM_NODES` in `config.py`. The Elastic plugin, DeepSpeed strategy, and gradient accumulation calculation all adapt automatically. The recovery system is unchanged as each run still gets its own recovery URI.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/multimodal-retrieval-evaluation ===

# Multimodal retrieval evaluation

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/multimodal-retrieval-evaluation).

This tutorial builds an experiment framework for benchmarking **visual document retrieval** on the [ViDoRe benchmark](https://huggingface.co/vidore). The corpus is a set of PDF page *images* and the queries are plain-text questions. Each retrieval method must find the page that answers a question from the raw image alone.

Three approaches are compared:

- **ColPali-v1.2**: patch-level multi-vector embeddings from a vision-language model (PaliGemma), scored with MaxSim late interaction. No OCR.
- **SigLIP-SO400M**: a single global embedding per page from Google's CLIP successor.
- **OCR + BM25**: a text-only baseline that OCRs each page with [docTR](https://github.com/mindee/doctr) on GPU and ranks with BM25.

Each experiment is one `ExperimentConfig`; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. It's a strong showcase of several Flyte features working together:

- **`ReusePolicy`** keeps warm GPU containers (with the ~7 GB ColPali weights already in VRAM) alive across task calls.
- A process-level **`DynamicBatcher`** aggregates queries from all concurrent search tasks into single GPU batches.
- **`cache="auto"`** so a model's index is built at most once per corpus and shared across experiments.
- **Typed Pydantic inputs/outputs** so every metric is stored alongside the exact config that produced it.

## Define the container image

One image serves every task. `unionai-reuse` provides the actor bridge required by `ReusePolicy`.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

The Python dependencies (ColPali, transformers, docTR, etc.) are declared in the `uv` script header at the top of the file.

## Define the task environments

Each model gets its own GPU environment so their warm-container pools scale independently. The ColPali and SigLIP environments use `ReusePolicy` to keep model weights resident; the driver coordinates orchestration, BM25, evaluation, and reporting.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

## Configuration and data types

An experiment is fully described by an `ExperimentConfig`. Because it's a Pydantic model, Flyte serializes it alongside every output.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

The corpus, queries, retrieval results, and metrics are likewise typed. Page images are stored as `flyte.io.File` handles in blob storage, so tasks read images directly rather than re-fetching over HTTP.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

## Loading, indexing, and search

`load_vidore_pages` downloads a ViDoRe subset and uploads each page image to blob storage (cached, with retries). Indexing tasks (`index_colpali`, `index_siglip`) encode every page into a `.npz` index, and the OCR task (`extract_page_texts`) produces the text baseline. These run on the GPU environments and are cached per corpus.

Search uses the `DynamicBatcher` so queries from all concurrent search-task invocations on a warm container are merged into a single GPU batch:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

> [!NOTE]
> The `DynamicBatcher` implementation lives in the `extras/` package next to the example. Run the script from the example directory so the import resolves.

## Run one experiment

`run_experiment` selects the right index/search path based on the runtime value of `config.model`. Flyte v2's dynamic execution means there's no static DAG to wire up. `flyte.group` wraps each experiment in a named span in the UI.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

## Compare experiments

The driver loads the dataset once, fans out across all configs with `asyncio.gather`, and emits an interactive Chart.js report in the Flyte UI. Experiments sharing a model reuse the cached index, so you only pay GPU time for new work.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

## Run the evaluation

This example has no secrets: datasets and model weights are pulled from public Hugging Face repositories. It does require GPUs, so run it remotely.

The experiment grid is defined in the entry point; adding a model or varying `top_k` is a one-line change:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "colpali-engine>=0.3.1",
#     "transformers>=4.41",
#     "sentencepiece>=0.2",
#     "torch>=2.0",
#     "pillow>=10",
#     "datasets>=2.18",
#     "rank-bm25>=0.2",
#     "numpy>=1.26",
#     "python-doctr[torch]>=0.8",
#     "pydantic>=2.0",
#     "flyte>=2.0.0",
# ]
# ///
"""
Multimodal Retrieval Evaluation Pipeline

This tutorial is an experiment framework for benchmarking visual document
retrieval approaches on the ViDoRe benchmark. Each experiment is defined by
an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and
returns a ranked comparison table with an interactive HTML report.

The corpus is a set of PDF page images; queries are plain-text questions. Each
retrieval method must find the page that answers each question — no text is
provided to the model, only the raw image.

  ColPali-v1.2  — patch-level multi-vector embeddings from a VLM (PaliGemma).
                  No OCR. The model produces one vector per image patch
                  (~1024 per page). MaxSim late-interaction scoring finds the
                  best matching patch for each query token.

  SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP
                  successor. One matrix multiply per query; fast and effective
                  but a single vector cannot localise fine-grained regions.

  OCR + BM25    — text-only baseline. doctr (GPU OCR) extracts text in
                  batches, BM25 matches keywords. Strong on text-dense pages;
                  fails on charts, tables, and figures where content is visual.

"""

import asyncio
import enum
import json
import math
import os
import tempfile
from functools import lru_cache
from io import BytesIO
from itertools import islice

import numpy as np
from PIL import Image as PILImage
from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from extras import DynamicBatcher

import flyte
import flyte.report
from flyte.io import File

# ─────────────────────────────────────────────────────────────────────────────
# Environments
# ─────────────────────────────────────────────────────────────────────────────

# One Docker image for all tasks. The PEP 723 header defines Python deps.
# ca-certificates is required for HTTPS calls to HuggingFace and blob stores.
# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="vidore-eval-v2")
    .with_apt_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0")
    # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy.
    # Without it every reusable container exits with StartError (exit code 128).
    .with_pip_packages("unionai-reuse>=0.1.11")
)
# {{/docs-fragment image}}

# GPU environment for ColPali image encoding and search.
#
# ReusePolicy keeps up to 3 warm GPU containers alive between task calls.
# Without it, every task invocation cold-starts a new container and downloads
# ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model
# weights already loaded into VRAM — is reused for the next task dispatch.
#
#   replicas=1      single warm container — all concurrent shard calls land
#                   here so they share one DynamicBatcher process
#   concurrency=8   up to 8 query-shard tasks run simultaneously on the
#                   container, all feeding the same DynamicBatcher queue
#   idle_ttl=120    keep alive 2 min after the last task finishes
#   scaledown_ttl=60 scale to zero after 1 min of complete inactivity
# {{docs-fragment envs}}
colpali_indexer = flyte.TaskEnvironment(
    name="vidore-colpali-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for SigLIP image encoding and search.
#
# Separate from the ColPali environment so each model's warm containers
# are managed independently — ColPali and SigLIP experiments can scale
# without contending for the same pool of reusable containers.
siglip_indexer = flyte.TaskEnvironment(
    name="vidore-siglip-indexer",
    image=image,
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1),
    reusable=flyte.ReusePolicy(
        replicas=1,
        concurrency=8,
        idle_ttl=120,
        scaledown_ttl=60,
    ),
)

# GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition)
# in batches on GPU — much faster than CPU Tesseract.
# No ReusePolicy needed: the result is cached, so this task runs at most once.
ocr_engine = flyte.TaskEnvironment(
    name="vidore-ocr-engine",
    image=image,
    resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1),
)

# Driver: orchestration, BM25 search, evaluation, and reporting.
# depends_on ensures the shared Docker image is built before all environments
# try to schedule tasks.
driver = flyte.TaskEnvironment(
    name="vidore-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="12Gi"),
    depends_on=[colpali_indexer, siglip_indexer, ocr_engine],
)
# {{/docs-fragment envs}}

# ─────────────────────────────────────────────────────────────────────────────
# Configuration types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment config_types}}
class RetrievalModel(str, enum.Enum):
    """Retrieval backend to evaluate."""

    COLPALI = "colpali-v1.2"  # multi-vector patch embeddings, MaxSim
    SIGLIP = "siglip-so400m"  # single-vector global embedding, cosine sim
    OCR_BM25 = "ocr+bm25"  # text extracted by Tesseract, ranked by BM25

class ExperimentConfig(BaseModel):
    """
    All knobs for one retrieval experiment. Passed as a typed Flyte input.

    Because ExperimentConfig is a Pydantic model, Flyte serialises it
    alongside every task output — so you can always reconstruct which
    config produced which metric without maintaining a separate log.
    """

    name: str  # human-readable label shown in the comparison table
    model: RetrievalModel
    top_k: int = 5  # number of pages to retrieve per query
# {{/docs-fragment config_types}}

# ─────────────────────────────────────────────────────────────────────────────
# Data types
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment data_types}}
class PageQuery(BaseModel):
    """One retrieval query with its ground-truth page."""

    query_id: str
    text: str  # e.g. "What was revenue growth in Q3?"
    relevant_page_id: str  # one correct page per query

class PageDataset(BaseModel):
    """
    A corpus of document page images paired with text queries.

    page_ids:   unique page identifiers (derived from ViDoRe image filenames).
    page_files: the same pages stored in Flyte's blob store as JPEG File
                handles. Tasks read images directly from here; no live HTTP.
    queries:    text questions with ground-truth page IDs for evaluation.
    """

    page_ids: list[str]
    page_files: list[File]
    queries: list[PageQuery]

    class Config:
        arbitrary_types_allowed = True

class RetrievalResult(BaseModel):
    query_id: str
    ranked_page_ids: list[str]  # ordered best → worst

class Metrics(BaseModel):
    recall_at_k: float
    ndcg_at_k: float
    mrr: float
    k: int

class ExperimentResult(BaseModel):
    config: ExperimentConfig
    metrics: Metrics
# {{/docs-fragment data_types}}

class ComparisonReport(BaseModel):
    results: list[ExperimentResult]

    def best_by(self, metric: str = "recall_at_k") -> ExperimentResult:
        return max(self.results, key=lambda r: getattr(r.metrics, metric))

    def summary(self) -> str:
        header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}"
        sep = "─" * len(header)
        rows = [header, sep]
        for r in sorted(self.results, key=lambda x: -x.metrics.recall_at_k):
            rows.append(
                f"{r.config.name:<30} "
                f"{r.config.model.value:<18} "
                f"{r.metrics.recall_at_k:>10.3f} "
                f"{r.metrics.ndcg_at_k:>8.3f} "
                f"{r.metrics.mrr:>7.3f}"
            )
        return "\n".join(rows)

# ─────────────────────────────────────────────────────────────────────────────
# Cached model loaders
# ─────────────────────────────────────────────────────────────────────────────
# These functions are at module level so they are shared across all tasks that
# run on the same warm container (via ReusePolicy). lru_cache(maxsize=1) means
# the model is loaded from disk/HuggingFace exactly once per container process
# and kept in GPU memory for every subsequent task dispatch to that container.

@lru_cache(maxsize=1)
def _colpali_model():
    """Load ColPali-v1.2 into GPU memory and cache the result.

    device_map= is the correct loading pattern for ColPali's PaliGemma
    backbone; it handles weight placement via accelerate. torch.compile is
    skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross-
    invocation batching is the primary GPU utilisation mechanism.
    """
    import torch
    from colpali_engine.models import ColPali, ColPaliProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ColPali.from_pretrained(
        "vidore/colpali-v1.2",
        torch_dtype=torch.bfloat16,
        device_map=device,
    )
    processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
    return model, processor, device

@lru_cache(maxsize=1)
def _siglip_model():
    """Load SigLIP SO400M into GPU memory, compile it, and cache the result.

    torch.compile (mode="reduce-overhead") fuses the vision and text encoder
    transformer layers into optimised CUDA kernels. As with ColPali, the
    compilation overhead is paid once per warm container lifetime.
    """
    import torch
    from transformers import AutoModel, AutoProcessor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained("google/siglip-so400m-patch14-224").to(device)
    if device == "cuda":
        model = torch.compile(model, mode="reduce-overhead")
    processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-224")
    return model, processor, device

@lru_cache(maxsize=1)
def _ocr_model():
    """Load the doctr OCR predictor onto GPU and cache it.

    doctr's ocr_predictor bundles a detection model (DBNet) and a
    recognition model (CRNN/SAR) into a single callable. pretrained=True
    downloads both model weights from doctr's model zoo on first use.
    """
    import torch
    from doctr.models import ocr_predictor

    predictor = ocr_predictor(pretrained=True)
    if torch.cuda.is_available():
        predictor = predictor.cuda()
    return predictor

# ─────────────────────────────────────────────────────────────────────────────
# Search batcher singletons
# ─────────────────────────────────────────────────────────────────────────────
# One DynamicBatcher per model, shared across all concurrent search task
# invocations on the same warm container (concurrency=3). Queries from every
# concurrent caller are aggregated into a single GPU batch, maximizing
# throughput compared to each invocation running its own forward pass.
#
# Initialised lazily on the first search call via double-checked locking and
# lives for the container's lifetime. The process_fn runs GPU work via
# asyncio.to_thread so the aggregation loop can continue collecting queries
# from other callers while the GPU processes the current batch.
#
# File is not hashable so alru_cache cannot be used here; module-level state
# with asyncio.Lock is the correct pattern.
#
# Assumption: index_colpali/index_siglip use cache="auto", so the same corpus
# always produces the same index File across all callers on this container. If
# the index file ever changed between calls, the batcher would silently continue
# using the corpus embeddings loaded from the first call.

_colpali_batcher: DynamicBatcher | None = None
_colpali_batcher_lock = asyncio.Lock()
_siglip_batcher: DynamicBatcher | None = None
_siglip_batcher_lock = asyncio.Lock()

async def _get_colpali_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level ColPali search batcher, creating it on first call."""
    global _colpali_batcher
    if _colpali_batcher is not None:
        return _colpali_batcher
    async with _colpali_batcher_lock:
        if _colpali_batcher is not None:
            return _colpali_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, n_patches, dim)
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _colpali_model()
        corpus_emb = corpus_emb.to(device, dtype=torch.float32)

        async def colpali_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                query_inputs = processor.process_queries([q.text for q in batch])
                query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
                with torch.no_grad():
                    query_embs = model(**query_inputs).float()  # (B, T, D)
                    query_chunk = 8
                    n_pages = corpus_emb.shape[0]
                    all_scores = torch.empty(len(batch), n_pages, device=device)
                    for start in range(0, len(batch), query_chunk):
                        chunk = query_embs[start : start + query_chunk]
                        all_scores[start : start + query_chunk] = (
                            torch.einsum("ctd,pjd->ctpj", chunk, corpus_emb)
                            .max(dim=3).values
                            .sum(dim=1)
                        )
                    sorted_indices = all_scores.argsort(dim=1, descending=True).cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            # Run GPU work in a thread so the event loop — and the batcher's
            # aggregation loop — remain unblocked while the GPU is busy.
            return await asyncio.to_thread(_gpu_work)

        batcher: DynamicBatcher[PageQuery, list[str]] = DynamicBatcher(
            process_fn=colpali_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _colpali_batcher = batcher
    return _colpali_batcher

async def _get_siglip_search_batcher(index_file: File) -> DynamicBatcher:
    """Return the process-level SigLIP search batcher, creating it on first call."""
    global _siglip_batcher
    if _siglip_batcher is not None:
        return _siglip_batcher
    async with _siglip_batcher_lock:
        if _siglip_batcher is not None:
            return _siglip_batcher

        import torch

        data = await _load_npz(index_file)
        corpus_emb = torch.from_numpy(data["embeddings"])  # (n_pages, dim), L2-normalised
        index_page_ids: list[str] = list(data["page_ids"])
        model, processor, device = _siglip_model()
        corpus_emb = corpus_emb.to(device)

        async def siglip_process_fn(batch: list[PageQuery]) -> list[list[str]]:
            def _gpu_work() -> list[list[str]]:
                text_inputs = processor(
                    text=[q.text for q in batch],
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(device)
                with torch.no_grad():
                    text_out = model.text_model(**text_inputs)
                    query_embs = text_out.pooler_output  # (B, dim)
                    query_embs = query_embs / query_embs.norm(dim=-1, keepdim=True)
                    scores_matrix = corpus_emb @ query_embs.T  # (n_pages, B)
                    sorted_indices = scores_matrix.argsort(dim=0, descending=True).T.cpu().tolist()
                return [[index_page_ids[j] for j in ranked] for ranked in sorted_indices]

            return await asyncio.to_thread(_gpu_work)

        batcher = DynamicBatcher(
            process_fn=siglip_process_fn,
            target_batch_cost=128,
            max_batch_size=128,
            batch_timeout_s=0.05,
            default_cost=1,
            prefetch_batches=2,
        )
        await batcher.start()
        _siglip_batcher = batcher
    return _siglip_batcher

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _batches(items: list, batch_size: int):
    """Yield successive fixed-size batches from a list."""
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]

def _load_image_sync(f: File) -> PILImage.Image:
    """Blocking download + decode. Intended to be called from a thread pool."""
    with f.open_sync("rb") as fh:
        data = fh.read()
    return PILImage.open(BytesIO(data)).convert("RGB")

async def _load_image(f: File) -> PILImage.Image:
    """Download and decode a page image in a thread-pool worker.

    asyncio.to_thread runs _load_image_sync in a real OS thread so that
    blocking network I/O can overlap with GPU-bound forward passes when
    images are pre-submitted via loop.run_in_executor before the GPU kernel.
    """
    return await asyncio.to_thread(_load_image_sync, f)

async def _load_npz(index_file: File) -> np.lib.npyio.NpzFile:
    """Download an index File to a local temp path and open with np.load."""
    with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
        async with index_file.open("rb") as fh:
            tmp.write(bytes(await fh.read()))
        return np.load(tmp.name)

def _dcg(relevances: list[int]) -> float:
    return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances))

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — data loading
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(cache="auto", retries=3)
async def load_vidore_pages(subset: str = "docvqa", max_pages: int = 200) -> PageDataset:
    """
    Load a ViDoRe benchmark subset and store page images in Flyte's blob store.

    Supports two dataset formats:

    Legacy (subsampled) — single 'test' split with one row per (query, page)
    pair; fields: image, query, image_filename. streaming=True reads only the
    rows requested via islice — no full-shard download.
    Datasets: vidore/docvqa_test_subsampled, vidore/infovqa_test_subsampled

    V3 — separate corpus / queries / qrels splits following the BEIR retrieval
    benchmark format. corpus contains page images; queries contains question
    text; qrels maps query IDs to relevant corpus page IDs (many-to-many).
    Datasets: vidore/vidore_v3_finance_en  (~2 942 pages, 1 854 queries)

    The first call uploads page images to Flyte's blob store and caches the
    PageDataset; every subsequent call with the same arguments returns the
    cached result instantly. retries=3 guards against transient HuggingFace
    network failures.

    Available subsets: "docvqa", "infovqa", "vidore_v3_finance_en"
    """
    from datasets import load_dataset

    subset_map = {
        "docvqa": "vidore/docvqa_test_subsampled",
        "infovqa": "vidore/infovqa_test_subsampled",
        "vidore_v3_finance_en": "vidore/vidore_v3_finance_en",
    }
    dataset_name = subset_map.get(subset, f"vidore/{subset}_test_subsampled")

    # V3 datasets ship with separate corpus / queries / qrels splits.
    _V3_SUBSETS = {"vidore_v3_finance_en"}

    if subset in _V3_SUBSETS:
        # ── V3 format ─────────────────────────────────────────────────────────
        # corpus / queries / qrels are HuggingFace configs (name=), not splits.
        # corpus uses streaming=True so images are decoded one at a time —
        # loading all 2 942 rows eagerly would hold gigabytes of PIL images in
        # the driver's RAM simultaneously. qrels and queries are text-only and
        # small enough to load fully into memory.
        corpus_ds = load_dataset(dataset_name, name="corpus", split="test", streaming=True)
        qrels_ds = load_dataset(dataset_name, name="qrels", split="test")
        queries_ds = load_dataset(dataset_name, name="queries", split="test")

        # Normalise field names — V3 follows BEIR convention (hyphenated ids).
        def _col(ds, *candidates):
            cols = set(ds.column_names)
            for c in candidates:
                if c in cols:
                    return c
            raise KeyError(f"None of {candidates} found in columns {cols}")

        corpus_id_col = _col(corpus_ds, "corpus-id", "corpus_id", "id", "_id")
        query_id_col = _col(queries_ds, "query-id", "query_id", "id", "_id")
        query_text_col = _col(queries_ds, "query", "text")
        qrel_qid_col = _col(qrels_ds, "query-id", "query_id")
        qrel_cid_col = _col(qrels_ds, "corpus-id", "corpus_id")

        # Slice corpus to max_pages, upload each image to Flyte blob store.
        page_ids: list[str] = []
        page_files: list[File] = []
        corpus_id_to_page_id: dict[str, str] = {}

        for i, row in enumerate(islice(corpus_ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            cid = str(row[corpus_id_col])
            page_id = f"{subset}_{i:04d}"
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                tmp_path = f.name
                img.convert("RGB").save(tmp_path, format="JPEG")
            del img  # free PIL memory before upload
            page_file = await File.from_local(tmp_path)
            os.unlink(tmp_path)
            corpus_id_to_page_id[cid] = page_id
            page_ids.append(page_id)
            page_files.append(page_file)

        # Build query_id → relevant page_id from qrels (first match wins).
        # Only keep relevance judgements whose corpus page is in our slice.
        qrel_map: dict[str, str] = {}
        for row in qrels_ds:
            qid = str(row[qrel_qid_col])
            cid = str(row[qrel_cid_col])
            if cid in corpus_id_to_page_id and qid not in qrel_map:
                qrel_map[qid] = corpus_id_to_page_id[cid]

        # Collect queries that have at least one relevant page in our slice.
        queries: list[PageQuery] = []
        for row in queries_ds:
            qid = str(row[query_id_col])
            if qid not in qrel_map:
                continue
            queries.append(
                PageQuery(
                    query_id=qid,
                    text=str(row[query_text_col]),
                    relevant_page_id=qrel_map[qid],
                )
            )

    else:
        # ── Legacy format ─────────────────────────────────────────────────────
        # Single 'test' split with one row per (query, page) pair.
        ds = load_dataset(dataset_name, split="test", streaming=True)

        page_ids = []
        page_files = []
        queries = []
        seen_pages: dict[str, str] = {}  # image_filename → page_id

        for i, row in enumerate(islice(ds, max_pages)):
            img = row.get("image")
            if not isinstance(img, PILImage.Image):
                continue
            filename: str = row.get("image_filename") or f"page_{i}"
            query_text: str = row.get("query", "")
            if not query_text:
                continue

            # Each unique page is uploaded exactly once; multiple queries may
            # share the same page (same image_filename).
            if filename not in seen_pages:
                page_id = f"{subset}_{len(page_ids):04d}"
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
                    tmp_path = f.name
                    img.convert("RGB").save(tmp_path, format="JPEG")
                del img  # free PIL memory before upload
                page_file = await File.from_local(tmp_path)
                os.unlink(tmp_path)
                seen_pages[filename] = page_id
                page_ids.append(page_id)
                page_files.append(page_file)
            else:
                page_id = seen_pages[filename]

            queries.append(
                PageQuery(
                    query_id=f"q{i:04d}",
                    text=query_text,
                    relevant_page_id=page_id,
                )
            )

    print(f"Loaded {len(page_ids)} unique pages, {len(queries)} queries", flush=True)
    return PageDataset(page_ids=page_ids, page_files=page_files, queries=queries)

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — indexing
# ─────────────────────────────────────────────────────────────────────────────

@colpali_indexer.task(cache="auto", retries=2)
async def index_colpali(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with ColPali-v1.2 and save the multi-vector index.

    ColPali skips OCR entirely. It feeds the raw page image into PaliGemma
    (a vision-language model) and produces one embedding vector per image
    patch — roughly 1,024 patches per page, each of dimension 128.

    _colpali_model() is an lru_cache'd loader. On a cold container, it
    downloads and loads the model once. On a warm container (kept alive by
    ReusePolicy), it returns the already-loaded model instantly from cache —
    no repeated ~7 GB download.

    The index is stored as a .npz file in Flyte's blob store:
      embeddings — float32, shape (n_pages, n_patches, dim)
      page_ids   — matching page ID strings

    cache="auto" + retries=2: the result is stored permanently on success;
    transient failures (e.g. HuggingFace rate limits) are retried twice.
    """
    import torch

    model, processor, device = _colpali_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 4))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor.process_images(images)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            emb = model(**inputs)  # (batch, n_patches, dim)

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"ColPali: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, n_patches, dim)
    out_path = os.path.join(tempfile.gettempdir(), "colpali_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@siglip_indexer.task(cache="auto", retries=2)
async def index_siglip(page_ids: list[str], page_files: list[File]) -> File:
    """
    Encode every page with SigLIP SO400M and save the single-vector index.

    SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss
    instead of softmax — avoiding the normalisation bottleneck that limits
    CLIP's scalability. Produces one global embedding per page.

    _siglip_model() caches the model across warm container reuses.

    The index is stored as a .npz file:
      embeddings — float32, shape (n_pages, dim), L2-normalised
      page_ids   — matching page ID strings
    """
    import torch

    model, processor, device = _siglip_model()

    loop = asyncio.get_running_loop()
    batches = list(_batches(page_files, 8))
    n_batches = len(batches)

    # Submit the first batch to the thread pool before entering the loop so
    # that downloads are already in flight when we first await them.
    prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[0]]

    all_embeddings: list[np.ndarray] = []
    for batch_idx in range(n_batches):
        images = list(await asyncio.gather(*prefetch))

        # Submit next batch downloads immediately — OS threads run these in
        # parallel with the GPU forward pass below.
        if batch_idx + 1 < n_batches:
            prefetch = [loop.run_in_executor(None, _load_image_sync, f) for f in batches[batch_idx + 1]]

        inputs = processor(images=images, return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            outputs = model.vision_model(**inputs)
            emb = outputs.pooler_output  # (batch, dim)
            emb = emb / emb.norm(dim=-1, keepdim=True)  # L2 normalise

        all_embeddings.append(emb.cpu().float().numpy())
        print(f"SigLIP: indexed batch {batch_idx + 1}/{n_batches}", flush=True)

    embeddings = np.concatenate(all_embeddings, axis=0)  # (n_pages, dim)
    out_path = os.path.join(tempfile.gettempdir(), "siglip_index.npz")
    np.savez(out_path, embeddings=embeddings, page_ids=np.array(page_ids))
    return await File.from_local(out_path)

@ocr_engine.task(cache="auto")
async def extract_page_texts(page_files: list[File]) -> list[str]:
    """
    OCR every page with doctr on GPU to produce a text-only baseline.

    doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single
    callable predictor. Pages are downloaded in parallel then fed in batches
    of ocr_batch_size. asyncio.to_thread keeps the event loop unblocked
    during GPU inference.

    Result structure: result.pages[i].blocks[j].lines[k].words[l].value

    Cached: the same corpus is OCR'd at most once across all experiments
    that use the OCR+BM25 backend.
    """
    import gc

    predictor = _ocr_model()

    # Process in batches: download each batch just-in-time so only
    # ocr_batch_size images are in memory at once instead of all 2 000.
    ocr_batch_size = 8
    total = len(page_files)
    texts: list[str] = []
    for start in range(0, total, ocr_batch_size):
        batch_files = page_files[start : start + ocr_batch_size]
        batch_images = list(
            await asyncio.gather(*[asyncio.to_thread(_load_image_sync, f) for f in batch_files])
        )
        batch_np = [np.array(img) for img in batch_images]
        del batch_images
        result = await asyncio.to_thread(predictor, batch_np)
        del batch_np
        for page_output in result.pages:
            texts.append(
                "\n".join(
                    " ".join(word.value for word in line.words)
                    for block in page_output.blocks
                    for line in block.lines
                )
            )
        del result
        gc.collect()
        print(f"OCR: processed {min(start + ocr_batch_size, total)}/{total} pages", flush=True)

    return texts

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — search
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment search_colpali}}
@colpali_indexer.task
async def search_colpali(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher.

    MaxSim score for page p given query q:
        score(q, p) = Σ_{t ∈ query tokens} max_{j ∈ page patches} (q_t · p_j)

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_colpali invocations on the
    same warm container (concurrency=8) into a single GPU batch. This keeps
    the GPU saturated rather than running one small batch per caller.

    The batcher's process_fn runs GPU work in asyncio.to_thread, so the
    aggregation loop stays live while the GPU encodes and scores.
    """
    batcher = await _get_colpali_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]
# {{/docs-fragment search_colpali}}

@siglip_indexer.task
async def search_siglip(
    index_file: File,
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using SigLIP cosine similarity via DynamicBatcher.

    Each query is submitted to the process-level DynamicBatcher, which
    aggregates queries from all concurrent search_siglip invocations on the
    same warm container (concurrency=3) into a single GPU batch.

    SigLIP's single-vector embeddings make full vectorisation safe —
    the scores matrix (n_pages x n_queries) is small enough to materialise
    in one GPU call regardless of batch size.
    """
    batcher = await _get_siglip_search_batcher(index_file)
    futures = await batcher.submit_batch(queries)
    all_ranked: list[list[str]] = list(await asyncio.gather(*futures))

    return [
        RetrievalResult(query_id=q.query_id, ranked_page_ids=ranked[:top_k])
        for q, ranked in zip(queries, all_ranked)
    ]

@driver.task
async def search_bm25(
    page_texts: list[str],
    page_ids: list[str],
    queries: list[PageQuery],
    top_k: int,
) -> list[RetrievalResult]:
    """
    Retrieve pages using BM25 over OCR'd text.

    The standard keyword-based baseline. No GPU required; strong on
    text-dense pages, weak on visual content that Tesseract cannot read.
    """
    tokenized = [text.lower().split() for text in page_texts]
    bm25 = BM25Okapi(tokenized)

    results: list[RetrievalResult] = []
    for q in queries:
        scores = bm25.get_scores(q.text.lower().split())
        ranked = sorted(range(len(page_ids)), key=lambda i: -scores[i])[:top_k]
        results.append(
            RetrievalResult(
                query_id=q.query_id,
                ranked_page_ids=[page_ids[i] for i in ranked],
            )
        )
    return results

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — evaluation
# ─────────────────────────────────────────────────────────────────────────────

@driver.task
async def evaluate(
    results: list[RetrievalResult],
    ground_truth: list[PageQuery],
    k: int,
) -> Metrics:
    """
    Compute Recall@K, NDCG@K, and MRR for a single retrieval model.

    Recall@K  — was the correct page in the top-K results?
    NDCG@K    — normalised discounted cumulative gain; rewards earlier hits.
    MRR       — mean reciprocal rank of the first correct result.

    All three are averaged over all queries. Higher is better.
    """
    gt_map = {q.query_id: q.relevant_page_id for q in ground_truth}
    recall_vals, ndcg_vals, mrr_vals = [], [], []

    for r in results:
        relevant = gt_map.get(r.query_id, "")
        top = r.ranked_page_ids[:k]

        recall_vals.append(1.0 if relevant in top else 0.0)

        rels = [1 if pid == relevant else 0 for pid in top]
        idcg = _dcg([1])  # ideal: correct page at rank 1
        ndcg_vals.append(_dcg(rels) / idcg if idcg > 0 else 0.0)

        rr = 0.0
        for rank, pid in enumerate(r.ranked_page_ids, start=1):
            if pid == relevant:
                rr = 1.0 / rank
                break
        mrr_vals.append(rr)

    return Metrics(
        recall_at_k=float(np.mean(recall_vals)),
        ndcg_at_k=float(np.mean(ndcg_vals)),
        mrr=float(np.mean(mrr_vals)),
        k=k,
    )

# ─────────────────────────────────────────────────────────────────────────────
# Tasks — report
# ─────────────────────────────────────────────────────────────────────────────

@driver.task(report=True)
async def generate_report(report: ComparisonReport) -> None:
    """
    Emit an interactive HTML report visible in the Flyte UI.

    report=True marks this task as a reporting task. Flyte renders the HTML
    returned via flyte.report.replace.aio() directly in the execution detail
    page — no separate dashboard or export step required.

    The report contains:
      - Summary cards: experiment count, best model, best Recall@K.
      - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment.
      - Ranked results table with all three metrics.
    """
    sorted_results = sorted(report.results, key=lambda r: -r.metrics.recall_at_k)
    best = sorted_results[0]

    labels = [r.config.name for r in sorted_results]
    recall_vals = [r.metrics.recall_at_k for r in sorted_results]
    ndcg_vals = [r.metrics.ndcg_at_k for r in sorted_results]
    mrr_vals = [r.metrics.mrr for r in sorted_results]

    table_rows = "".join(
        f"""
        <tr>
          <td>{r.config.name}</td>
          <td>{r.config.model.value}</td>
          <td>{r.metrics.recall_at_k:.3f}</td>
          <td>{r.metrics.ndcg_at_k:.3f}</td>
          <td>{r.metrics.mrr:.3f}</td>
          <td>{r.metrics.k}</td>
        </tr>"""
        for r in sorted_results
    )

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Visual Document Retrieval — Results</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
      background: #f0f2f5; color: #222; padding: 24px;
    }}
    h1 {{ font-size: 1.6em; margin-bottom: 4px; }}
    .subtitle {{ color: #666; margin-bottom: 24px; font-size: 0.95em; }}
    .cards {{
      display: flex; gap: 16px; flex-wrap: wrap; margin-bottom: 28px;
    }}
    .card {{
      background: #fff; border-radius: 10px; padding: 18px 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); min-width: 160px;
    }}
    .card-value {{ font-size: 1.9em; font-weight: 700; color: #4f46e5; }}
    .card-label {{ font-size: 0.8em; color: #888; text-transform: uppercase;
                   letter-spacing: .04em; margin-top: 2px; }}
    .chart-box {{
      background: #fff; border-radius: 10px; padding: 24px;
      box-shadow: 0 1px 4px rgba(0,0,0,.08); margin-bottom: 28px;
    }}
    .chart-box h2 {{ font-size: 1em; margin-bottom: 16px; color: #444; }}
    table {{ width: 100%; border-collapse: collapse; font-size: 0.9em; }}
    th {{
      background: #4f46e5; color: #fff; padding: 10px 14px;
      text-align: left; font-weight: 600;
    }}
    td {{ padding: 9px 14px; border-bottom: 1px solid #eee; }}
    tr:hover td {{ background: #f8f8ff; }}
    tr:first-child td {{ font-weight: 600; }}
  </style>
</head>
<body>
  <h1>Visual Document Retrieval — Experiment Comparison</h1>
  <p class="subtitle">ViDoRe benchmark &middot; {len(report.results)} experiment(s)</p>

  <div class="cards">
    <div class="card">
      <div class="card-value">{len(report.results)}</div>
      <div class="card-label">Experiments</div>
    </div>
    <div class="card">
      <div class="card-value">{best.config.name}</div>
      <div class="card-label">Best by Recall@K</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.recall_at_k:.3f}</div>
      <div class="card-label">Best Recall@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.ndcg_at_k:.3f}</div>
      <div class="card-label">Best NDCG@{best.metrics.k}</div>
    </div>
    <div class="card">
      <div class="card-value">{best.metrics.mrr:.3f}</div>
      <div class="card-label">Best MRR</div>
    </div>
  </div>

  <div class="chart-box">
    <h2>Metrics by Experiment</h2>
    <canvas id="metricsChart" height="100"></canvas>
  </div>

  <div class="chart-box">
    <h2>Ranked Results</h2>
    <table>
      <thead>
        <tr>
          <th>Experiment</th><th>Model</th>
          <th>Recall@K</th><th>NDCG@K</th><th>MRR</th><th>K</th>
        </tr>
      </thead>
      <tbody>{table_rows}</tbody>
    </table>
  </div>

  <script>
    new Chart(document.getElementById('metricsChart'), {{
      type: 'bar',
      data: {{
        labels: {json.dumps(labels)},
        datasets: [
          {{
            label: 'Recall@K',
            data: {json.dumps(recall_vals)},
            backgroundColor: 'rgba(79,70,229,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'NDCG@K',
            data: {json.dumps(ndcg_vals)},
            backgroundColor: 'rgba(16,185,129,0.75)',
            borderRadius: 4
          }},
          {{
            label: 'MRR',
            data: {json.dumps(mrr_vals)},
            backgroundColor: 'rgba(245,158,11,0.75)',
            borderRadius: 4
          }}
        ]
      }},
      options: {{
        responsive: true,
        plugins: {{ legend: {{ position: 'top' }} }},
        scales: {{
          y: {{ beginAtZero: true, max: 1.0,
               title: {{ display: true, text: 'Score' }} }}
        }}
      }}
    }});
  </script>
</body>
</html>"""

    await flyte.report.replace.aio(html)
    await flyte.report.flush.aio()

# ─────────────────────────────────────────────────────────────────────────────
# Experiment orchestration
# ─────────────────────────────────────────────────────────────────────────────

# {{docs-fragment run_experiment}}
@driver.task
async def run_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult:
    """
    End-to-end retrieval pipeline for a single ExperimentConfig.

    Flyte v2's dynamic execution means this driver task can call GPU tasks
    (index_colpali, search_colpali) based on the runtime value of config.model
    — no static DAG wiring required. The if/elif is plain Python; Flyte
    schedules the selected sub-tasks on the appropriate environment.

    Caching: two experiments that share the same model and corpus (e.g. ColPali
    at top_k=5 and top_k=10) will hit the same cached index. GPU work is paid
    at most once per (model, corpus) pair across all experiments.

    Search queries are sharded into chunks of SEARCH_SHARD_SIZE and dispatched
    as concurrent task invocations. All shards land on the single warm container
    (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the
    GPU saturated throughout search rather than processing one large sequential
    batch from a single caller.

    flyte.group wraps each experiment in a named span in the Flyte UI, making
    it easy to compare latencies and drill into individual runs.
    """
    SEARCH_SHARD_SIZE = 256

    with flyte.group(config.name):
        if config.model == RetrievalModel.COLPALI:
            index_file = await index_colpali(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_colpali(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        elif config.model == RetrievalModel.SIGLIP:
            index_file = await index_siglip(dataset.page_ids, dataset.page_files)
            shards = list(_batches(dataset.queries, SEARCH_SHARD_SIZE))
            shard_results = await asyncio.gather(
                *[search_siglip(index_file, shard, config.top_k) for shard in shards]
            )
            results = [r for shard in shard_results for r in shard]

        else:  # RetrievalModel.OCR_BM25
            page_texts = await extract_page_texts(dataset.page_files)
            results = await search_bm25(page_texts, dataset.page_ids, dataset.queries, config.top_k)

        metrics = await evaluate(results, dataset.queries, config.top_k)

    return ExperimentResult(config=config, metrics=metrics)
# {{/docs-fragment run_experiment}}

# {{docs-fragment compare_experiments}}
@driver.task
async def compare_experiments(
    configs: list[ExperimentConfig],
    subset: str = "docvqa",
    max_pages: int = 200,
) -> ComparisonReport:
    """
    Fan out over all experiment configs and return a ranked comparison table.

    The dataset is loaded once and shared across all experiments. Each config
    runs as a concurrent Flyte task via asyncio.gather. Experiments that share
    a model reuse the cached index — you only pay GPU time for new work.

    On completion, generate_report emits an interactive Chart.js HTML report
    visible directly in the Flyte execution detail page.

    Default dataset: vidore_v3_finance_en (~2 942 corpus pages, 1 854 queries)
    with max_pages=2 000 to exercise the GPU pipeline at scale.
    """
    dataset = await load_vidore_pages(subset=subset, max_pages=max_pages)

    # All experiments launch concurrently. Shared cached outputs (same model,
    # same corpus) are served from cache rather than recomputed.
    experiment_coros = [run_experiment(config=cfg, dataset=dataset) for cfg in configs]
    results: list[ExperimentResult] = list(await asyncio.gather(*experiment_coros))

    report = ComparisonReport(results=results)
    print(report.summary())
    best = report.best_by("recall_at_k")
    print(f"\nBest by Recall@{best.metrics.k}: {best.config.name}")

    # Emit the interactive HTML report in the Flyte UI.
    await generate_report(report)

    return report
# {{/docs-fragment compare_experiments}}

# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    flyte.init_from_config()

    # Define the experiment grid. Each ExperimentConfig is one point in the
    # design space. Adding a new model or varying top_k is one line here —
    # no task code changes required.
    #
    # ColPali appears twice with different top_k values. The cache ensures
    # index_colpali runs only once and both experiments share that result.
    # {{docs-fragment grid}}
    configs = [
        ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top_k=5),
        ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top_k=10),
        ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top_k=5),
        ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR_BM25, top_k=5),
    ]
    # {{/docs-fragment grid}}

    run = flyte.with_runcontext().run(
        compare_experiments,
        configs=configs,
        subset="vidore_v3_finance_en",
        max_pages=2000,
    )
    print(f"Run URL: {run.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval_eval.py*

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/multimodal-retrieval-evaluation):

```
cd v2/tutorials/multimodal-retrieval-evaluation
python retrieval_eval.py
```

When the run completes, open the `generate_report` task in the UI to see the summary cards, the grouped Recall@K / NDCG@K / MRR bar chart, and the ranked results table.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/detr-object-detection ===

# RT-DETR object detection

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/detr_object_detection).

This tutorial fine-tunes [RT-DETRv2](https://huggingface.co/PekingU/rtdetr_v2_r18vd) on a custom COCO-format dataset from HuggingFace. The pipeline downloads and splits the data, fine-tunes the detector with live training charts in Flyte reports, evaluates COCO mAP on a validation split, and renders a side-by-side inference demo with ground-truth and predicted bounding boxes.

Flyte highlights:

- **Cached dataset preparation** so re-runs skip the HuggingFace download.
- **Live training reports** with loss curves and optional periodic mAP checkpoints.
- **GPU evaluation and demo tasks** that stream annotated images into the UI.

## Define the task environments

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "torchvision>=0.24.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "huggingface_hub>=0.24.0",
#    "datasets>=3.0.0",
#    "pillow>=10.0.0",
#    "albumentations>=1.4.0",
#    "torchmetrics>=1.4.0",
#    "pycocotools>=2.0.7",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import base64
import io
import json
import logging
import os
import random
import shutil
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="detr-object-detection", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="detr-object-detection-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="detr-object-detection-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Report styling — shared CSS for all task reports
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #16213e; border-bottom: 2px solid #0f3460; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #0f3460; margin-top: 20px; }
  .report .card { background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #e9ecef; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0f3460; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0f3460; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dee2e6; }
  .report tr:nth-child(even) { background: #f8f9fa; }
  .report .highlight { color: #0f3460; font-weight: 700; }
  .report .note { background: #fff3cd; border-left: 4px solid #ffc107; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .img-pair { display: flex; gap: 12px; margin: 16px 0; flex-wrap: wrap; }
  .report .img-pair > div { flex: 1; min-width: 300px; }
  .report .img-pair img { width: 100%; border-radius: 6px; border: 1px solid #dee2e6; }
  .report .img-pair .gt-label { color: #5a7db5; font-weight: 600; }
  .report .img-pair .pred-label { color: #06d6a0; font-weight: 600; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d4edda; color: #155724; }
  .report .badge-info { background: #d1ecf1; color: #0c5460; }
  .report .chart-container { background: #fff; border: 1px solid #dee2e6; border-radius: 8px; padding: 16px; margin: 16px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers — lightweight charts without matplotlib
# ------------------------------------------------------------------

def _make_line_chart(
    data: list[dict],
    x_key: str,
    y_keys: list[str],
    title: str = "",
    x_label: str = "",
    y_label: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    y_max_cap: float | None = None,
    x_range_override: tuple[float, float] | None = None,
    y_display_names: dict[str, str] | None = None,
) -> str:
    """Generate an SVG line chart from a list of dicts.

    Args:
        data: List of dicts, each with x_key and y_keys values.
        x_key: Key for x-axis values.
        y_keys: Keys for y-axis series to plot.
        title: Chart title.
        x_label: X-axis label.
        y_label: Y-axis label.
        colors: Colors for each series (defaults to a built-in palette).
        width: SVG width in pixels.
        height: SVG height in pixels.
        y_max_cap: If set, cap the y-axis at this value (e.g. 1.0 for mAP).
        x_range_override: If set, force the x-axis to this (min, max) range.

    Returns:
        SVG string.
    """

    default_colors = ["#5a7db5", "#0f3460", "#06d6a0", "#ffc107", "#6c757d"]
    colors = colors or default_colors

    # Chart area margins
    ml, mr, mt, mb = 60, 20, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [d[x_key] for d in data] if data else []
    if x_range_override:
        x_min, x_max = x_range_override
    elif x_vals:
        x_min, x_max = min(x_vals), max(x_vals)
    else:
        x_min, x_max = 0, 1
    x_range = x_max - x_min or 1

    # Compute y range across all series
    all_y = []
    for key in y_keys:
        all_y.extend(d[key] for d in data if key in d)
    y_min = min(all_y) if all_y else 0
    y_max = max(all_y) if all_y else 1
    y_pad = (y_max - y_min) * 0.1 or 0.1
    y_min_plot = max(0, y_min - y_pad)
    y_max_plot = y_max + y_pad
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap)
    y_range = y_max_plot - y_min_plot or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min_plot) / y_range * ch

    # Build SVG
    lines = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        # Background
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines (5 horizontal)
    for i in range(6):
        y_tick = y_min_plot + y_range * i / 5
        py = sy(y_tick)
        lines.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        lines.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.3f}</text>'
        )

    # Axes
    lines.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#adb5bd" stroke-width="1.5"/>'
    )
    lines.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#adb5bd" stroke-width="1.5"/>'
    )

    # X-axis ticks
    if x_vals:
        n_x_ticks = min(len(data), 10)
        step = max(1, len(data) // n_x_ticks)
        for i in range(0, len(data), step):
            px = sx(x_vals[i])
            lines.append(
                f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
                f'font-size="11" fill="#6c757d">{x_vals[i]:.0f}</text>'
            )
    else:
        # Empty chart — generate evenly spaced ticks from x range
        for i in range(6):
            x_tick = x_min + x_range * i / 5
            px = sx(x_tick)
            lines.append(
                f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
                f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
            )

    # Plot each series
    if not data:
        # Empty chart placeholder
        lines.append(
            f'<text x="{ml + cw / 2}" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="13" fill="#adb5bd" font-style="italic">Waiting for data...</text>'
        )
    for si, key in enumerate(y_keys):
        color = colors[si % len(colors)]
        points = [(sx(d[x_key]), sy(d[key])) for d in data if key in d]
        if not points:
            continue
        # Draw line if we have 2+ points (dash odd series for visibility)
        if len(points) >= 2:
            path_d = f"M {points[0][0]:.1f},{points[0][1]:.1f}"
            for px, py in points[1:]:
                path_d += f" L {px:.1f},{py:.1f}"
            dash = ' stroke-dasharray="6,3"' if si % 2 == 1 else ""
            lines.append(
                f'<path d="{path_d}" fill="none" stroke="{color}" '
                f'stroke-width="2" stroke-linejoin="round"{dash}/>'
            )
        # Always show dots for sparse data (including single points)
        if len(points) <= 30:
            for px, py in points:
                lines.append(
                    f'<circle cx="{px:.1f}" cy="{py:.1f}" r="3" fill="{color}"/>'
                )

    # Title
    if title:
        lines.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Axis labels
    if x_label:
        lines.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        lines.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    # Legend
    names = y_display_names or {}
    if len(y_keys) > 1:
        lx = ml + 10
        for si, key in enumerate(y_keys):
            color = colors[si % len(colors)]
            ly = mt + 14 + si * 18
            lines.append(
                f'<rect x="{lx}" y="{ly - 6}" width="12" height="12" '
                f'rx="2" fill="{color}"/>'
            )
            label = names.get(key, key)
            lines.append(
                f'<text x="{lx + 16}" y="{ly + 4}" font-size="11" '
                f'fill="#1a1a2e">{label}</text>'
            )

    lines.append("</svg>")
    return "\n".join(lines)

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    y_max_cap: float | None = None,
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels for x-axis.
        series: Dict mapping series name to list of values (same length as labels).
        title: Chart title.
        colors: Colors for each series.
        width: SVG width.
        height: SVG height.
        y_max_cap: If set, cap the y-axis at this value (e.g. 1.0 for mAP).

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#adb5bd", "#0f3460", "#06d6a0", "#5a7db5"]
    colors = colors or default_colors

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    lines_svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        lines_svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        lines_svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.3f}</text>'
        )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            lines_svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            # Value label on top of bar
            lines_svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="10" fill="#1a1a2e">'
                f'{val:.3f}</text>'
            )
        # Group label
        lines_svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 18}" '
            f'text-anchor="middle" font-size="11" fill="#6c757d">{label}</text>'
        )

    # Title
    if title:
        lines_svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Legend
    lx = ml + cw - len(series) * 100
    for si, name in enumerate(series):
        color = colors[si % len(colors)]
        lines_svg.append(
            f'<rect x="{lx + si * 100}" y="{mt + ch + 35}" width="12" '
            f'height="12" rx="2" fill="{color}"/>'
        )
        lines_svg.append(
            f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 46}" font-size="11" '
            f'fill="#1a1a2e">{name}</text>'
        )

    lines_svg.append("</svg>")
    return "\n".join(lines_svg)

# ------------------------------------------------------------------
# Task 1: Prepare dataset — download COCO JSON + images, split train/val
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def prepare_data(
    dataset_repo: str = "sagecodes/union_flyte_swag_object_detection",
    annotations_path: str = "swag/train.json",
    images_subdir: str = "swag/images",
    val_fraction: float = 0.2,
    seed: int = 42,
) -> flyte.io.Dir:
    """Download a COCO-format dataset from HF and split into train/val."""
    from huggingface_hub import snapshot_download

    log.info(f"Downloading dataset: {dataset_repo}")
    local_repo = snapshot_download(
        repo_id=dataset_repo,
        repo_type="dataset",
    )

    ann_file = os.path.join(local_repo, annotations_path)
    img_root = os.path.join(local_repo, images_subdir)

    with open(ann_file) as f:
        coco = json.load(f)

    images = coco["images"]
    annotations = coco["annotations"]
    categories = coco["categories"]

    log.info(
        f"Loaded {len(images)} images, {len(annotations)} annotations, "
        f"{len(categories)} categories"
    )
    log.info(f"Raw category ids: {sorted({c['id'] for c in categories})}")
    log.info(
        f"Raw annotation category_ids (unique): "
        f"{sorted({a['category_id'] for a in annotations})}"
    )

    # Remap category ids to contiguous 0..N-1 — required because HF object
    # detection models size their classifier head to len(id2label) and treat
    # class labels as direct indices into that head. Any gap or 1-indexed id
    # causes an IndexKernel OOB inside the focal-loss scatter.
    #
    # Build the remap from the UNION of ids declared in `categories` and ids
    # actually used in `annotations` — some datasets have orphaned annotations
    # referencing categories that aren't declared (this one does).
    declared_ids = {c["id"] for c in categories}
    used_ids = {a["category_id"] for a in annotations}
    orphans = used_ids - declared_ids
    if orphans:
        log.warning(
            f"Annotations reference undeclared category ids {sorted(orphans)} — "
            f"adding stub categories."
        )

    all_cat_ids = sorted(declared_ids | used_ids)
    id_remap = {old: new for new, old in enumerate(all_cat_ids)}
    existing_names = {c["id"]: c["name"] for c in categories}
    categories = [
        {"id": id_remap[old], "name": existing_names.get(old, f"category_{old}")}
        for old in all_cat_ids
    ]
    annotations = [
        {**a, "category_id": id_remap[a["category_id"]]} for a in annotations
    ]
    log.info(f"Remapped category ids: {id_remap}")
    log.info(f"Final categories: {categories}")

    # Split by image id
    rng = random.Random(seed)
    img_ids = [im["id"] for im in images]
    rng.shuffle(img_ids)
    n_val = max(1, int(len(img_ids) * val_fraction))
    val_ids = set(img_ids[:n_val])
    train_ids = set(img_ids[n_val:])

    def filter_coco(keep_ids: set) -> dict:
        return {
            "info": coco.get("info", {}),
            "categories": categories,
            "images": [im for im in images if im["id"] in keep_ids],
            "annotations": [a for a in annotations if a["image_id"] in keep_ids],
        }

    train_coco = filter_coco(train_ids)
    val_coco = filter_coco(val_ids)

    log.info(
        f"Split: {len(train_coco['images'])} train / {len(val_coco['images'])} val images"
    )

    # Pack output dir: images/ + train.json + val.json
    out_dir = tempfile.mkdtemp(prefix="coco_split_")
    out_img = os.path.join(out_dir, "images")
    shutil.copytree(img_root, out_img)

    with open(os.path.join(out_dir, "train.json"), "w") as f:
        json.dump(train_coco, f)
    with open(os.path.join(out_dir, "val.json"), "w") as f:
        json.dump(val_coco, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Helpers — torch Dataset wrapping COCO JSON
# ------------------------------------------------------------------

def _build_torch_dataset(coco_path: str, images_root: str, augment: bool):
    """Build a torch Dataset that yields {image, target} for the HF image processor."""
    import albumentations as A
    import numpy as np
    from PIL import Image
    from torch.utils.data import Dataset

    with open(coco_path) as f:
        coco = json.load(f)

    images_by_id = {im["id"]: im for im in coco["images"]}
    anns_by_image: dict[int, list] = {}
    for a in coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    image_ids = list(images_by_id.keys())

    # NOTE: we deliberately don't resize here — the HF image processor handles
    # resize+pad. Augmentation only.
    if augment:
        transform = A.Compose(
            [
                A.HorizontalFlip(p=0.5),
                A.VerticalFlip(p=0.1),
                A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
                A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=30, val_shift_limit=20, p=0.4),
                A.Rotate(limit=15, border_mode=0, p=0.4),
                A.RandomScale(scale_limit=0.2, p=0.4),
                A.GaussianBlur(blur_limit=(3, 5), p=0.2),
                A.GaussNoise(p=0.2),
            ],
            bbox_params=A.BboxParams(
                format="coco",
                label_fields=["category"],
                min_area=4,
                min_visibility=0.1,
                clip=True,
            ),
        )
    else:
        transform = A.Compose(
            [A.NoOp()],
            bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True),
        )

    class CocoDataset(Dataset):
        def __len__(self) -> int:
            return len(image_ids)

        def __getitem__(self, idx: int):
            img_id = image_ids[idx]
            meta = images_by_id[img_id]
            img_path = os.path.join(images_root, os.path.basename(meta["file_name"]))
            if not os.path.exists(img_path):
                img_path = os.path.join(images_root, meta["file_name"])
            image = np.array(Image.open(img_path).convert("RGB"))

            anns = anns_by_image.get(img_id, [])
            bboxes = [a["bbox"] for a in anns]
            categories = [a["category_id"] for a in anns]

            out = transform(image=image, bboxes=bboxes, category=categories)
            image_t = out["image"]
            bboxes_t = out["bboxes"]
            categories_t = out["category"]

            target_anns = []
            for bb, cat in zip(bboxes_t, categories_t):
                x, y, w, h = bb
                target_anns.append(
                    {
                        "image_id": img_id,
                        "category_id": int(cat),
                        "bbox": [float(x), float(y), float(w), float(h)],
                        "area": float(w * h),
                        "iscrowd": 0,
                    }
                )

            return {
                "image": image_t,
                "target": {"image_id": img_id, "annotations": target_anns},
            }

    return CocoDataset(), coco["categories"]

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    epochs: int = 30,
    lr: float = 5e-5,
    batch_size: int = 4,
    weight_decay: float = 1e-4,
    eval_every_n_epochs: int | None = None,
) -> flyte.io.Dir:
    """Fine-tune RT-DETR (or any HuggingFace object-detection model) on COCO data."""
    import torch
    from transformers import (
        AutoImageProcessor,
        AutoModelForObjectDetection,
        Trainer,
        TrainerCallback,
        TrainingArguments,
    )

    log.info(f"Training: model={model_name}")
    await flyte.report.replace.aio(_wrap_report(
        f"<h2>Loading model...</h2><p>{model_name}</p>"
        f"<p>Preparing dataset and initializing weights...</p>"
    ), do_flush=True)

    # -- Load data --
    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    train_json = os.path.join(data_path, "train.json")

    with open(train_json) as f:
        categories = json.load(f)["categories"]
    id2label = {c["id"]: c["name"] for c in categories}
    label2id = {v: k for k, v in id2label.items()}

    train_ds, _ = _build_torch_dataset(train_json, images_root, augment=True)
    log.info(f"Train examples: {len(train_ds)} | Categories: {id2label}")

    # -- Optionally load val set for periodic mAP evaluation --
    val_json = os.path.join(data_path, "val.json")
    val_images = None
    val_targets = None
    if eval_every_n_epochs and os.path.exists(val_json):
        import torch as _torch
        from PIL import Image

        with open(val_json) as f:
            val_coco = json.load(f)
        images_by_id = {im["id"]: im for im in val_coco["images"]}
        anns_by_image: dict[int, list] = {}
        for a in val_coco["annotations"]:
            anns_by_image.setdefault(a["image_id"], []).append(a)
        val_images = []
        val_targets = []
        for img_id, meta in images_by_id.items():
            path = os.path.join(images_root, os.path.basename(meta["file_name"]))
            if not os.path.exists(path):
                path = os.path.join(images_root, meta["file_name"])
            val_images.append(Image.open(path).convert("RGB"))
            boxes_xyxy = []
            labels = []
            for a in anns_by_image.get(img_id, []):
                x, y, w, h = a["bbox"]
                boxes_xyxy.append([x, y, x + w, y + h])
                labels.append(a["category_id"])
            val_targets.append({
                "boxes": _torch.tensor(boxes_xyxy, dtype=_torch.float32).reshape(-1, 4),
                "labels": _torch.tensor(labels, dtype=_torch.long),
            })
        log.info(f"Val examples for periodic eval: {len(val_images)}")

    # -- Processor + model --
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = AutoModelForObjectDetection.from_pretrained(
        model_name,
        num_labels=len(id2label),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True,
    )

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log.info(
        f"Parameters: {trainable_params:,} / {total_params:,} "
        f"({trainable_params / total_params * 100:.1f}%)"
    )

    # -- Collator — runs the image processor on each batch --
    def collate_fn(batch):
        images = [b["image"] for b in batch]
        targets = [b["target"] for b in batch]
        enc = processor(images=images, annotations=targets, return_tensors="pt")
        return {"pixel_values": enc["pixel_values"], "labels": enc["labels"]}

    # -- Sanity check: peek at one batch and verify class_labels fit --
    sample = collate_fn([train_ds[i] for i in range(min(2, len(train_ds)))])
    all_labels = []
    for lbl in sample["labels"]:
        all_labels.extend(lbl["class_labels"].tolist())
    log.info(
        f"Sanity check — class_labels in first batch: {sorted(set(all_labels))} | "
        f"model num_labels: {model.config.num_labels} | "
        f"id2label: {model.config.id2label}"
    )
    if all_labels and max(all_labels) >= model.config.num_labels:
        raise ValueError(
            f"class_label {max(all_labels)} out of range for num_labels="
            f"{model.config.num_labels}. Check category id remapping in prepare_data."
        )

    # -- Collect training metrics and update the report chart live.
    # trainer.train() runs in a background thread (via asyncio.to_thread),
    # so the asyncio event loop stays free. We use run_coroutine_threadsafe
    # to push report updates from the callback thread onto that loop.
    training_log: list[dict] = []
    eval_log: list[dict] = []  # periodic mAP checkpoints (epoch, map, map_50)
    loop = asyncio.get_running_loop()

    cat_badges = " ".join(
        f'<span class="badge badge-info">{name}</span>'
        for name in id2label.values()
    )

    def _build_training_report(max_steps: int) -> str:
        """Build the live training report HTML from current training_log."""
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{len(train_ds)}</div><div class="label">Train Examples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>
          <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
        </div>
        <p>Categories: {cat_badges}</p>
        """

        charts_html = ""
        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs} |
              Loss: <span class="highlight">{current['loss']:.4f}</span>
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_chart = _make_line_chart(
                data=training_log,
                x_key="epoch",
                y_keys=["loss"],
                title="Training Loss",
                x_label="Epoch",
                y_label="Loss",
                colors=["#5a7db5"],
            )
            charts_html += f'<div class="chart-container">{loss_chart}</div>'

            if eval_every_n_epochs:
                # Match x-axis to the loss chart: start at 0, end at current epoch
                current_max_epoch = current["epoch"]
                map_chart = _make_line_chart(
                    data=eval_log,
                    x_key="epoch",
                    y_keys=["map", "map_50"],
                    title="Validation mAP (periodic)",
                    x_label="Epoch",
                    y_label="mAP",
                    colors=["#0f3460", "#06d6a0"],
                    y_max_cap=1.0,
                    y_display_names={"map": "mAP (0.50:0.95)", "map_50": "mAP@50"},
                    x_range_override=(0, current_max_epoch),
                )
                charts_html += f'<div class="chart-container">{map_chart}</div>'

            if "lr" in training_log[0]:
                lr_chart = _make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["lr"],
                    title="Learning Rate Schedule",
                    x_label="Epoch",
                    y_label="LR",
                    colors=["#0f3460"],
                )
                charts_html += f'<div class="chart-container">{lr_chart}</div>'

        return _wrap_report(stats_html + charts_html)

    class MetricsCallback(TrainerCallback):
        def __init__(self):
            self._last_eval_epoch = 0

        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs or "loss" not in logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
                "loss": round(logs["loss"], 4),
            }
            if "learning_rate" in logs:
                entry["lr"] = logs["learning_rate"]
            if "grad_norm" in logs:
                entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
            training_log.append(entry)
            log.info(
                f"step={state.global_step}/{state.max_steps} "
                f"epoch={entry['epoch']:.2f} "
                f"loss={entry['loss']:.4f}"
            )

            # Push a live report update onto the asyncio event loop.
            # do_flush=True dispatches the update to the UI immediately.
            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

        def on_epoch_end(self, args, state, control, model=None, **kwargs):
            if not eval_every_n_epochs or val_images is None:
                return
            current_epoch = round(state.epoch)
            if current_epoch % eval_every_n_epochs != 0:
                return
            if current_epoch == self._last_eval_epoch:
                return
            self._last_eval_epoch = current_epoch

            log.info(f"Running periodic mAP eval at epoch {current_epoch}...")
            from torchmetrics.detection.mean_ap import MeanAveragePrecision

            device = next(model.parameters()).device
            # _run_inference sets model.eval(); restore train mode after.
            preds = _run_inference(model, processor, val_images, device, threshold=0.3)
            model.train()

            formatted = [
                {"boxes": p["boxes"], "scores": p["scores"], "labels": p["labels"]}
                for p in preds
            ]
            metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
            metric.update(formatted, val_targets)
            result = metric.compute()
            map_val = round(result["map"].item(), 4)
            map_50 = round(result["map_50"].item(), 4)

            eval_log.append({
                "epoch": current_epoch,
                "map": map_val,
                "map_50": map_50,
            })
            log.info(f"Epoch {current_epoch} — mAP: {map_val:.4f}, mAP@50: {map_50:.4f}")

            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=lr,
        weight_decay=weight_decay,
        logging_steps=5,
        save_strategy="no",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        warmup_ratio=0.1,
        remove_unused_columns=False,
        dataloader_num_workers=2,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        data_collator=collate_fn,
        callbacks=[MetricsCallback()],
    )

    log.info("Starting training...")
    # Run the sync HF training loop in a thread so the asyncio event loop
    # stays free for Flyte's syncify bridge.
    await asyncio.to_thread(trainer.train)
    log.info("Training complete.")

    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
    trainer.save_model(save_dir)
    processor.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Build final training report --
    stats_html = f"""
    <h2>Training Complete</h2>
    <h3>{model_name}</h3>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(train_ds)}</div><div class="label">Train Examples</div></div>
      <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
      <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
      <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
      <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>
      <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
    </div>
    <p>Categories: {cat_badges}</p>
    """

    charts_html = ""
    if training_log:
        final_loss = training_log[-1]["loss"]
        min_loss = min(d["loss"] for d in training_log)
        initial_loss = training_log[0]["loss"]
        total_steps = training_log[-1]["step"]

        charts_html += f"""
        <div class="card">
          <b>Training Summary:</b>
          Initial loss: {initial_loss:.4f} |
          Final loss: <span class="highlight">{final_loss:.4f}</span> |
          Min loss: {min_loss:.4f} |
          Total steps: {total_steps}
        </div>
        """

        epoch_range = (0, epochs)

        loss_chart = _make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["loss"],
            title="Training Loss",
            x_label="Epoch",
            y_label="Loss",
            colors=["#5a7db5"],
            x_range_override=epoch_range,
        )
        charts_html += f'<div class="chart-container">{loss_chart}</div>'

    if eval_log:
        map_chart = _make_line_chart(
            data=eval_log,
            x_key="epoch",
            y_keys=["map", "map_50"],
            title="Validation mAP (periodic)",
            x_label="Epoch",
            y_label="mAP",
            colors=["#0f3460", "#06d6a0"],
            y_max_cap=1.0,
            x_range_override=(0, epochs),
            y_display_names={"map": "mAP (0.50:0.95)", "map_50": "mAP@50"},
        )
        charts_html += f'<div class="chart-container">{map_chart}</div>'

    if training_log and "lr" in training_log[0]:
        lr_chart = _make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["lr"],
            title="Learning Rate Schedule",
            x_label="Epoch",
            y_label="LR",
            colors=["#0f3460"],
            x_range_override=(0, epochs),
        )
        charts_html += f'<div class="chart-container">{lr_chart}</div>'

    await flyte.report.replace.aio(_wrap_report(stats_html + charts_html), do_flush=True)

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Inference helpers
# ------------------------------------------------------------------

def _run_inference(model, processor, images, device, threshold: float = 0.3):
    """Run object detection on a list of PIL images. Returns list of dicts."""
    import torch

    results = []
    model.eval()
    for img in images:
        inputs = processor(images=img, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
        target_size = torch.tensor([img.size[::-1]], device=device)  # (h, w)
        post = processor.post_process_object_detection(
            outputs, target_sizes=target_size, threshold=threshold
        )[0]
        results.append(
            {
                "scores": post["scores"].cpu(),
                "labels": post["labels"].cpu(),
                "boxes": post["boxes"].cpu(),  # xyxy in original image coords
            }
        )
    return results

def _draw_boxes(image, boxes, labels, scores, id2label, color: str = "lime"):
    """Draw bounding boxes on a PIL image. Returns a new PIL image."""
    from PIL import ImageDraw, ImageFont

    img = image.copy()
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=max(14, img.width // 60))
    except Exception:
        font = ImageFont.load_default()

    for box, label, score in zip(boxes.tolist(), labels.tolist(), scores.tolist()):
        x0, y0, x1, y1 = box
        width = max(2, img.width // 400)
        draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
        name = id2label.get(int(label), str(int(label)))
        caption = f"{name} {score:.2f}"
        text_bg = draw.textbbox((x0, y0), caption, font=font)
        draw.rectangle(text_bg, fill=color)
        draw.text((x0, y0), caption, fill="black", font=font)
    return img

def _img_to_data_uri(img, max_dim: int = 800) -> str:
    """PIL image → base64 data URI, downscaled for the report."""
    w, h = img.size
    if max(w, h) > max_dim:
        scale = max_dim / max(w, h)
        img = img.resize((int(w * scale), int(h * scale)))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()

# ------------------------------------------------------------------
# Task 3: Evaluate — COCO mAP on fine-tuned model
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    threshold: float = 0.5,
) -> str:
    """Compute COCO mAP for the fine-tuned model on the val split."""
    import torch
    from PIL import Image
    from torchmetrics.detection.mean_ap import MeanAveragePrecision
    from transformers import AutoImageProcessor, AutoModelForObjectDetection

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(_wrap_report(
        "<h2>Evaluation</h2><p>Loading val split and scoring model...</p>"
    ), do_flush=True)

    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    val_json = os.path.join(data_path, "val.json")

    with open(val_json) as f:
        val_coco = json.load(f)

    images_by_id = {im["id"]: im for im in val_coco["images"]}
    anns_by_image: dict[int, list] = {}
    for a in val_coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    pil_images = []
    targets = []
    for img_id, meta in images_by_id.items():
        path = os.path.join(images_root, os.path.basename(meta["file_name"]))
        if not os.path.exists(path):
            path = os.path.join(images_root, meta["file_name"])
        pil_images.append(Image.open(path).convert("RGB"))
        boxes_xyxy = []
        labels = []
        for a in anns_by_image.get(img_id, []):
            x, y, w, h = a["bbox"]
            boxes_xyxy.append([x, y, x + w, y + h])
            labels.append(a["category_id"])
        targets.append(
            {
                "boxes": torch.tensor(boxes_xyxy, dtype=torch.float32).reshape(-1, 4),
                "labels": torch.tensor(labels, dtype=torch.long),
            }
        )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    ft_path = await finetuned_dir.download()
    log.info(f"Scoring fine-tuned model: {ft_path}")
    processor = AutoImageProcessor.from_pretrained(ft_path)
    model = AutoModelForObjectDetection.from_pretrained(ft_path).to(device)
    preds = _run_inference(model, processor, pil_images, device, threshold=threshold)

    formatted_preds = [
        {"boxes": p["boxes"], "scores": p["scores"], "labels": p["labels"]}
        for p in preds
    ]

    metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
    metric.update(formatted_preds, targets)

    def to_python(v):
        if hasattr(v, "numel"):
            return v.item() if v.numel() == 1 else v.tolist()
        return v

    ft_metrics = {k: to_python(v) for k, v in metric.compute().items()}
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    log.info(f"Fine-tuned mAP: {ft_metrics.get('map', 0):.3f}")

    metric_keys = ["map", "map_50", "map_75", "mar_10"]
    metric_display = {
        "map": "mAP",
        "map_50": "mAP@50",
        "map_75": "mAP@75",
        "mar_10": "mAR@10",
    }

    rows = []
    for key in metric_keys:
        ft_val = ft_metrics.get(key, 0)
        rows.append(
            f"<tr><td><b>{metric_display.get(key, key)}</b></td>"
            f"<td class='highlight'>{ft_val:.3f}</td></tr>"
        )
    table = (
        "<table><tr><th>Metric</th><th>Score</th></tr>"
        + "".join(rows)
        + "</table>"
    )

    bar_chart = _make_bar_chart(
        labels=[metric_display.get(k, k) for k in metric_keys],
        series={"Fine-tuned": [ft_metrics.get(k, 0) for k in metric_keys]},
        title="COCO Evaluation Metrics",
        colors=["#0f3460"],
        y_max_cap=1.0,
    )

    ft_map = ft_metrics.get("map", 0)
    ft_map50 = ft_metrics.get("map_50", 0)

    eval_html = f"""
    <h2>Evaluation — COCO mAP</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(pil_images)}</div><div class="label">Val Images</div></div>
      <div class="stat"><div class="value">{threshold}</div><div class="label">Threshold</div></div>
      <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
      <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
    </div>
    <div class="chart-container">{bar_chart}</div>
    {table}
    <div class="note">
      <b>mAP</b> (mean Average Precision) measures how accurately the model
      detects objects — balancing whether predictions are correct (precision)
      and whether all objects are found (recall). The @50 and @75 variants
      require IoU overlaps of 50% and 75% between predicted and ground-truth
      boxes. <b>mAR</b> (mean Average Recall) measures how many ground-truth
      objects the model finds, with @1 and @10 limiting detections to 1 or 10
      per image.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(eval_html), do_flush=True)

    return json.dumps(
        {
            "finetuned": {k: round(v, 4) for k, v in ft_metrics.items() if isinstance(v, (int, float))},
            "num_val_images": len(pil_images),
        }
    )

# ------------------------------------------------------------------
# Task 4: Inference demo — render bboxes on val images
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def inference_demo(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    threshold: float = 0.5,
    max_images: int = 8,
    metrics_json: str = "{}",
) -> str:
    """Run the fine-tuned model on val images, render bboxes, embed in the report."""
    import torch
    from PIL import Image
    from torchmetrics.detection.mean_ap import MeanAveragePrecision
    from transformers import AutoImageProcessor, AutoModelForObjectDetection

    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    val_json = os.path.join(data_path, "val.json")

    with open(val_json) as f:
        val_coco = json.load(f)

    id2label = {c["id"]: c["name"] for c in val_coco["categories"]}
    metas = val_coco["images"][:max_images]
    anns_by_image: dict[int, list] = {}
    for a in val_coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    pil_images = []
    gt_per_image = []
    for meta in metas:
        path = os.path.join(images_root, os.path.basename(meta["file_name"]))
        if not os.path.exists(path):
            path = os.path.join(images_root, meta["file_name"])
        pil_images.append(Image.open(path).convert("RGB"))

        boxes_xyxy = []
        labels = []
        for a in anns_by_image.get(meta["id"], []):
            x, y, w, h = a["bbox"]
            boxes_xyxy.append([x, y, x + w, y + h])
            labels.append(a["category_id"])
        gt_per_image.append(
            {
                "boxes": torch.tensor(boxes_xyxy, dtype=torch.float32).reshape(-1, 4),
                "labels": torch.tensor(labels, dtype=torch.long),
                "scores": torch.ones(len(labels)),
            }
        )

    ft_path = await finetuned_dir.download()
    processor = AutoImageProcessor.from_pretrained(ft_path)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModelForObjectDetection.from_pretrained(ft_path).to(device)

    preds = _run_inference(model, processor, pil_images, device, threshold=threshold)

    html_blocks = []
    total_gt = 0
    total_pred = 0
    for i, (img, pred, gt) in enumerate(zip(pil_images, preds, gt_per_image)):
        n_gt = len(gt["labels"])
        n_pred = len(pred["labels"])
        total_gt += n_gt
        total_pred += n_pred

        # Per-image mAP
        metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
        metric.update(
            [{"boxes": pred["boxes"], "scores": pred["scores"], "labels": pred["labels"]}],
            [{"boxes": gt["boxes"], "labels": gt["labels"]}],
        )
        img_metrics = metric.compute()
        img_map = img_metrics["map"].item()
        img_map_badge = (
            f'<span class="badge badge-success">mAP {img_map:.2f}</span>'
            if img_map >= 0.5
            else f'<span class="badge badge-info">mAP {img_map:.2f}</span>'
        )

        pred_img = _draw_boxes(
            img, pred["boxes"], pred["labels"], pred["scores"],
            id2label, color="lime",
        )
        gt_img = _draw_boxes(
            img, gt["boxes"], gt["labels"], gt["scores"],
            id2label, color="dodgerblue",
        )
        html_blocks.append(f"""
        <div class="card">
          <b>Image {i + 1}</b> {img_map_badge}
          <div class="img-pair">
            <div>
              <p><span class="gt-label">Ground Truth</span>
                 <span class="badge badge-info">{n_gt} boxes</span></p>
              <img src="{_img_to_data_uri(gt_img)}" />
            </div>
            <div>
              <p><span class="pred-label">Predictions</span>
                 <span class="badge badge-success">{n_pred} boxes</span>
                 (threshold={threshold})</p>
              <img src="{_img_to_data_uri(pred_img)}" />
            </div>
          </div>
        </div>""")

    # Parse metrics if provided (from evaluate task)
    metrics = json.loads(metrics_json)
    ft_metrics = metrics.get("finetuned", {})
    ft_map = ft_metrics.get("map", None)
    ft_map50 = ft_metrics.get("map_50", None)

    metrics_stats = ""
    if ft_map is not None:
        metrics_stats = f"""
        <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
        <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
        """

    demo_html = f"""
    <h2>Inference Demo</h2>
    <h3>Fine-tuned RT-DETR on validation images</h3>
    <div class="stat-grid">
      {metrics_stats}
      <div class="stat"><div class="value">{len(pil_images)}</div><div class="label">Images Shown</div></div>
      <div class="stat"><div class="value">{total_gt}</div><div class="label">Ground Truth Boxes</div></div>
      <div class="stat"><div class="value">{total_pred}</div><div class="label">Predicted Boxes</div></div>
      <div class="stat"><div class="value">{threshold}</div><div class="label">Confidence Threshold</div></div>
    </div>
    <p><span class="gt-label">Blue = ground truth</span> |
       <span class="pred-label">Green = predictions</span></p>
    {"".join(html_blocks)}
    """

    await flyte.report.replace.aio(_wrap_report(demo_html), do_flush=True)

    return json.dumps(
        {
            "num_images": len(pil_images),
            "predictions_per_image": [len(p["labels"]) for p in preds],
        }
    )

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "PekingU/rtdetr_v2_r18vd",
    dataset_repo: str = "sagecodes/union_flyte_swag_object_detection",
    annotations_path: str = "swag/train.json",
    images_subdir: str = "swag/images",
    epochs: int = 30,
    lr: float = 5e-5,
    batch_size: int = 4,
    val_fraction: float = 0.2,
    threshold: float = 0.5,
    demo_images: int = 8,
    eval_every_n_epochs: int | None = None,
) -> tuple[flyte.io.Dir, str]:
    """
    End-to-end RT-DETRv2 fine-tuning pipeline.

    Returns the fine-tuned model directory and a JSON summary.

    1. Download COCO dataset from HuggingFace and split train/val
    2. Fine-tune RT-DETRv2 on the train split
    3. Evaluate: COCO mAP comparison (base vs fine-tuned)
    4. Inference demo: render bounding boxes on val images
    """
    log.info(f"Pipeline: {model_name} | dataset={dataset_repo}")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = ["Preparing Data", "Fine-tuning", "Evaluating", "Inference Demo"]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#06d6a0;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#e94560;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>RT-DETRv2 Object Detection Pipeline</h2>
        <p><b>Model:</b> {model_name} | <b>Dataset:</b> {dataset_repo}</p>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Downloading and splitting dataset...")),
        do_flush=True,
    )

    data_dir = await prepare_data(
        dataset_repo=dataset_repo,
        annotations_path=annotations_path,
        images_subdir=images_subdir,
        val_fraction=val_fraction,
    )

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Fine-tuning model...")),
        do_flush=True,
    )

    finetuned_dir = await train(
        model_name, data_dir, epochs, lr, batch_size,
        eval_every_n_epochs=eval_every_n_epochs,
    )

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Running COCO mAP evaluation...")),
        do_flush=True,
    )

    metrics_json = await evaluate(finetuned_dir, data_dir, threshold)
    metrics = json.loads(metrics_json)

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Rendering bounding box demo...")),
        do_flush=True,
    )

    demo_json = await inference_demo(
        finetuned_dir, data_dir, threshold, demo_images,
        metrics_json=metrics_json,
    )

    ft_map = metrics["finetuned"].get("map", 0)
    ft_map50 = metrics["finetuned"].get("map_50", 0)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <h3>{model_name}</h3>
    <div class="stat-grid">
      <div class="stat"><div class="value">{metrics['num_val_images']}</div><div class="label">Val Images</div></div>
      <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
      <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
    </div>
    <div class="card">
      <b>Configuration:</b> {epochs} epochs | LR {lr} | Batch size {batch_size} |
      Val fraction {val_fraction} | Threshold {threshold}
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)

    log.info(f"Pipeline complete. Fine-tuned mAP: {ft_map:.3f}")
    return finetuned_dir, json.dumps({"metrics": metrics, "demo": json.loads(demo_json)})

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
CODE0
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "transformers>=4.49.0",
#    "albumentations>=1.4.0",
#    "torchmetrics>=1.4.0",
#    ...
# ]
# ///
```

## Orchestrate the pipeline

The `pipeline` task prepares data, fine-tunes RT-DETR, evaluates mAP, and renders an inference demo.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.9.0",
#    "torchvision>=0.24.0",
#    "transformers>=4.49.0",
#    "accelerate>=0.34.0",
#    "huggingface_hub>=0.24.0",
#    "datasets>=3.0.0",
#    "pillow>=10.0.0",
#    "albumentations>=1.4.0",
#    "torchmetrics>=1.4.0",
#    "pycocotools>=2.0.7",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import base64
import io
import json
import logging
import os
import random
import shutil
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="detr-object-detection", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="detr-object-detection-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
)

cpu_env = flyte.TaskEnvironment(
    name="detr-object-detection-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    depends_on=[gpu_env],
)
# {{/docs-fragment env}}

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Report styling — shared CSS for all task reports
# ------------------------------------------------------------------

REPORT_CSS = """
<style>
  .report { font-family: system-ui, -apple-system, sans-serif; max-width: 960px; margin: 0 auto; color: #1a1a2e; }
  .report h2 { color: #16213e; border-bottom: 2px solid #0f3460; padding-bottom: 8px; margin-top: 24px; }
  .report h3 { color: #0f3460; margin-top: 20px; }
  .report .card { background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 8px; padding: 16px; margin: 12px 0; }
  .report .stat-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; margin: 12px 0; }
  .report .stat { background: #fff; border: 1px solid #e9ecef; border-radius: 6px; padding: 12px; text-align: center; }
  .report .stat .value { font-size: 1.5em; font-weight: 700; color: #0f3460; }
  .report .stat .label { font-size: 0.85em; color: #6c757d; margin-top: 4px; }
  .report table { border-collapse: collapse; width: 100%; margin: 12px 0; }
  .report th { background: #0f3460; color: #fff; padding: 10px 14px; text-align: left; font-weight: 600; }
  .report td { padding: 8px 14px; border-bottom: 1px solid #dee2e6; }
  .report tr:nth-child(even) { background: #f8f9fa; }
  .report .highlight { color: #0f3460; font-weight: 700; }
  .report .note { background: #fff3cd; border-left: 4px solid #ffc107; padding: 10px 14px; border-radius: 4px; margin: 12px 0; font-size: 0.9em; }
  .report .img-pair { display: flex; gap: 12px; margin: 16px 0; flex-wrap: wrap; }
  .report .img-pair > div { flex: 1; min-width: 300px; }
  .report .img-pair img { width: 100%; border-radius: 6px; border: 1px solid #dee2e6; }
  .report .img-pair .gt-label { color: #5a7db5; font-weight: 600; }
  .report .img-pair .pred-label { color: #06d6a0; font-weight: 600; }
  .report .badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: 0.8em; font-weight: 600; }
  .report .badge-success { background: #d4edda; color: #155724; }
  .report .badge-info { background: #d1ecf1; color: #0c5460; }
  .report .chart-container { background: #fff; border: 1px solid #dee2e6; border-radius: 8px; padding: 16px; margin: 16px 0; }
</style>
"""

def _wrap_report(html: str) -> str:
    """Wrap HTML content with report styling."""
    return f'{REPORT_CSS}<div class="report">{html}</div>'

# ------------------------------------------------------------------
# SVG chart helpers — lightweight charts without matplotlib
# ------------------------------------------------------------------

def _make_line_chart(
    data: list[dict],
    x_key: str,
    y_keys: list[str],
    title: str = "",
    x_label: str = "",
    y_label: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    y_max_cap: float | None = None,
    x_range_override: tuple[float, float] | None = None,
    y_display_names: dict[str, str] | None = None,
) -> str:
    """Generate an SVG line chart from a list of dicts.

    Args:
        data: List of dicts, each with x_key and y_keys values.
        x_key: Key for x-axis values.
        y_keys: Keys for y-axis series to plot.
        title: Chart title.
        x_label: X-axis label.
        y_label: Y-axis label.
        colors: Colors for each series (defaults to a built-in palette).
        width: SVG width in pixels.
        height: SVG height in pixels.
        y_max_cap: If set, cap the y-axis at this value (e.g. 1.0 for mAP).
        x_range_override: If set, force the x-axis to this (min, max) range.

    Returns:
        SVG string.
    """

    default_colors = ["#5a7db5", "#0f3460", "#06d6a0", "#ffc107", "#6c757d"]
    colors = colors or default_colors

    # Chart area margins
    ml, mr, mt, mb = 60, 20, 40, 50
    cw = width - ml - mr
    ch = height - mt - mb

    x_vals = [d[x_key] for d in data] if data else []
    if x_range_override:
        x_min, x_max = x_range_override
    elif x_vals:
        x_min, x_max = min(x_vals), max(x_vals)
    else:
        x_min, x_max = 0, 1
    x_range = x_max - x_min or 1

    # Compute y range across all series
    all_y = []
    for key in y_keys:
        all_y.extend(d[key] for d in data if key in d)
    y_min = min(all_y) if all_y else 0
    y_max = max(all_y) if all_y else 1
    y_pad = (y_max - y_min) * 0.1 or 0.1
    y_min_plot = max(0, y_min - y_pad)
    y_max_plot = y_max + y_pad
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap)
    y_range = y_max_plot - y_min_plot or 1

    def sx(v):
        return ml + (v - x_min) / x_range * cw

    def sy(v):
        return mt + ch - (v - y_min_plot) / y_range * ch

    # Build SVG
    lines = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        # Background
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines (5 horizontal)
    for i in range(6):
        y_tick = y_min_plot + y_range * i / 5
        py = sy(y_tick)
        lines.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        lines.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.3f}</text>'
        )

    # Axes
    lines.append(
        f'<line x1="{ml}" y1="{mt}" x2="{ml}" y2="{mt + ch}" '
        f'stroke="#adb5bd" stroke-width="1.5"/>'
    )
    lines.append(
        f'<line x1="{ml}" y1="{mt + ch}" x2="{ml + cw}" y2="{mt + ch}" '
        f'stroke="#adb5bd" stroke-width="1.5"/>'
    )

    # X-axis ticks
    if x_vals:
        n_x_ticks = min(len(data), 10)
        step = max(1, len(data) // n_x_ticks)
        for i in range(0, len(data), step):
            px = sx(x_vals[i])
            lines.append(
                f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
                f'font-size="11" fill="#6c757d">{x_vals[i]:.0f}</text>'
            )
    else:
        # Empty chart — generate evenly spaced ticks from x range
        for i in range(6):
            x_tick = x_min + x_range * i / 5
            px = sx(x_tick)
            lines.append(
                f'<text x="{px:.1f}" y="{mt + ch + 20}" text-anchor="middle" '
                f'font-size="11" fill="#6c757d">{x_tick:.0f}</text>'
            )

    # Plot each series
    if not data:
        # Empty chart placeholder
        lines.append(
            f'<text x="{ml + cw / 2}" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="13" fill="#adb5bd" font-style="italic">Waiting for data...</text>'
        )
    for si, key in enumerate(y_keys):
        color = colors[si % len(colors)]
        points = [(sx(d[x_key]), sy(d[key])) for d in data if key in d]
        if not points:
            continue
        # Draw line if we have 2+ points (dash odd series for visibility)
        if len(points) >= 2:
            path_d = f"M {points[0][0]:.1f},{points[0][1]:.1f}"
            for px, py in points[1:]:
                path_d += f" L {px:.1f},{py:.1f}"
            dash = ' stroke-dasharray="6,3"' if si % 2 == 1 else ""
            lines.append(
                f'<path d="{path_d}" fill="none" stroke="{color}" '
                f'stroke-width="2" stroke-linejoin="round"{dash}/>'
            )
        # Always show dots for sparse data (including single points)
        if len(points) <= 30:
            for px, py in points:
                lines.append(
                    f'<circle cx="{px:.1f}" cy="{py:.1f}" r="3" fill="{color}"/>'
                )

    # Title
    if title:
        lines.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Axis labels
    if x_label:
        lines.append(
            f'<text x="{ml + cw / 2}" y="{height - 6}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d">{x_label}</text>'
        )
    if y_label:
        lines.append(
            f'<text x="14" y="{mt + ch / 2}" text-anchor="middle" '
            f'font-size="12" fill="#6c757d" '
            f'transform="rotate(-90, 14, {mt + ch / 2})">{y_label}</text>'
        )

    # Legend
    names = y_display_names or {}
    if len(y_keys) > 1:
        lx = ml + 10
        for si, key in enumerate(y_keys):
            color = colors[si % len(colors)]
            ly = mt + 14 + si * 18
            lines.append(
                f'<rect x="{lx}" y="{ly - 6}" width="12" height="12" '
                f'rx="2" fill="{color}"/>'
            )
            label = names.get(key, key)
            lines.append(
                f'<text x="{lx + 16}" y="{ly + 4}" font-size="11" '
                f'fill="#1a1a2e">{label}</text>'
            )

    lines.append("</svg>")
    return "\n".join(lines)

def _make_bar_chart(
    labels: list[str],
    series: dict[str, list[float]],
    title: str = "",
    colors: list[str] | None = None,
    width: int = 700,
    height: int = 300,
    y_max_cap: float | None = None,
) -> str:
    """Generate an SVG grouped bar chart.

    Args:
        labels: Category labels for x-axis.
        series: Dict mapping series name to list of values (same length as labels).
        title: Chart title.
        colors: Colors for each series.
        width: SVG width.
        height: SVG height.
        y_max_cap: If set, cap the y-axis at this value (e.g. 1.0 for mAP).

    Returns:
        SVG string.
    """
    if not labels:
        return ""

    default_colors = ["#adb5bd", "#0f3460", "#06d6a0", "#5a7db5"]
    colors = colors or default_colors

    ml, mr, mt, mb = 60, 20, 40, 60
    cw = width - ml - mr
    ch = height - mt - mb

    all_vals = [v for vals in series.values() for v in vals]
    y_max = max(all_vals) if all_vals else 1
    y_max_plot = y_max * 1.15 or 1
    if y_max_cap is not None:
        y_max_plot = min(y_max_plot, y_max_cap) or y_max_cap

    n_groups = len(labels)
    n_series = len(series)
    group_width = cw / n_groups
    bar_width = group_width * 0.7 / max(n_series, 1)
    gap = group_width * 0.15

    def sy(v):
        return mt + ch - (v / y_max_plot) * ch

    lines_svg = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;max-width:{width}px;height:auto;">',
        f'<rect width="{width}" height="{height}" fill="#fff" rx="6"/>',
    ]

    # Grid lines
    for i in range(6):
        y_tick = y_max_plot * i / 5
        py = sy(y_tick)
        lines_svg.append(
            f'<line x1="{ml}" y1="{py:.1f}" x2="{ml + cw}" y2="{py:.1f}" '
            f'stroke="#e9ecef" stroke-width="1"/>'
        )
        lines_svg.append(
            f'<text x="{ml - 8}" y="{py + 4:.1f}" text-anchor="end" '
            f'font-size="11" fill="#6c757d">{y_tick:.3f}</text>'
        )

    # Bars
    for gi, label in enumerate(labels):
        gx = ml + gi * group_width + gap
        for si, (name, vals) in enumerate(series.items()):
            color = colors[si % len(colors)]
            bx = gx + si * bar_width
            val = vals[gi]
            by = sy(val)
            bh = mt + ch - by
            lines_svg.append(
                f'<rect x="{bx:.1f}" y="{by:.1f}" width="{bar_width - 1:.1f}" '
                f'height="{bh:.1f}" fill="{color}" rx="2"/>'
            )
            # Value label on top of bar
            lines_svg.append(
                f'<text x="{bx + bar_width / 2:.1f}" y="{by - 4:.1f}" '
                f'text-anchor="middle" font-size="10" fill="#1a1a2e">'
                f'{val:.3f}</text>'
            )
        # Group label
        lines_svg.append(
            f'<text x="{gx + n_series * bar_width / 2:.1f}" y="{mt + ch + 18}" '
            f'text-anchor="middle" font-size="11" fill="#6c757d">{label}</text>'
        )

    # Title
    if title:
        lines_svg.append(
            f'<text x="{width / 2}" y="22" text-anchor="middle" '
            f'font-size="14" font-weight="600" fill="#1a1a2e">{title}</text>'
        )

    # Legend
    lx = ml + cw - len(series) * 100
    for si, name in enumerate(series):
        color = colors[si % len(colors)]
        lines_svg.append(
            f'<rect x="{lx + si * 100}" y="{mt + ch + 35}" width="12" '
            f'height="12" rx="2" fill="{color}"/>'
        )
        lines_svg.append(
            f'<text x="{lx + si * 100 + 16}" y="{mt + ch + 46}" font-size="11" '
            f'fill="#1a1a2e">{name}</text>'
        )

    lines_svg.append("</svg>")
    return "\n".join(lines_svg)

# ------------------------------------------------------------------
# Task 1: Prepare dataset — download COCO JSON + images, split train/val
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def prepare_data(
    dataset_repo: str = "sagecodes/union_flyte_swag_object_detection",
    annotations_path: str = "swag/train.json",
    images_subdir: str = "swag/images",
    val_fraction: float = 0.2,
    seed: int = 42,
) -> flyte.io.Dir:
    """Download a COCO-format dataset from HF and split into train/val."""
    from huggingface_hub import snapshot_download

    log.info(f"Downloading dataset: {dataset_repo}")
    local_repo = snapshot_download(
        repo_id=dataset_repo,
        repo_type="dataset",
    )

    ann_file = os.path.join(local_repo, annotations_path)
    img_root = os.path.join(local_repo, images_subdir)

    with open(ann_file) as f:
        coco = json.load(f)

    images = coco["images"]
    annotations = coco["annotations"]
    categories = coco["categories"]

    log.info(
        f"Loaded {len(images)} images, {len(annotations)} annotations, "
        f"{len(categories)} categories"
    )
    log.info(f"Raw category ids: {sorted({c['id'] for c in categories})}")
    log.info(
        f"Raw annotation category_ids (unique): "
        f"{sorted({a['category_id'] for a in annotations})}"
    )

    # Remap category ids to contiguous 0..N-1 — required because HF object
    # detection models size their classifier head to len(id2label) and treat
    # class labels as direct indices into that head. Any gap or 1-indexed id
    # causes an IndexKernel OOB inside the focal-loss scatter.
    #
    # Build the remap from the UNION of ids declared in `categories` and ids
    # actually used in `annotations` — some datasets have orphaned annotations
    # referencing categories that aren't declared (this one does).
    declared_ids = {c["id"] for c in categories}
    used_ids = {a["category_id"] for a in annotations}
    orphans = used_ids - declared_ids
    if orphans:
        log.warning(
            f"Annotations reference undeclared category ids {sorted(orphans)} — "
            f"adding stub categories."
        )

    all_cat_ids = sorted(declared_ids | used_ids)
    id_remap = {old: new for new, old in enumerate(all_cat_ids)}
    existing_names = {c["id"]: c["name"] for c in categories}
    categories = [
        {"id": id_remap[old], "name": existing_names.get(old, f"category_{old}")}
        for old in all_cat_ids
    ]
    annotations = [
        {**a, "category_id": id_remap[a["category_id"]]} for a in annotations
    ]
    log.info(f"Remapped category ids: {id_remap}")
    log.info(f"Final categories: {categories}")

    # Split by image id
    rng = random.Random(seed)
    img_ids = [im["id"] for im in images]
    rng.shuffle(img_ids)
    n_val = max(1, int(len(img_ids) * val_fraction))
    val_ids = set(img_ids[:n_val])
    train_ids = set(img_ids[n_val:])

    def filter_coco(keep_ids: set) -> dict:
        return {
            "info": coco.get("info", {}),
            "categories": categories,
            "images": [im for im in images if im["id"] in keep_ids],
            "annotations": [a for a in annotations if a["image_id"] in keep_ids],
        }

    train_coco = filter_coco(train_ids)
    val_coco = filter_coco(val_ids)

    log.info(
        f"Split: {len(train_coco['images'])} train / {len(val_coco['images'])} val images"
    )

    # Pack output dir: images/ + train.json + val.json
    out_dir = tempfile.mkdtemp(prefix="coco_split_")
    out_img = os.path.join(out_dir, "images")
    shutil.copytree(img_root, out_img)

    with open(os.path.join(out_dir, "train.json"), "w") as f:
        json.dump(train_coco, f)
    with open(os.path.join(out_dir, "val.json"), "w") as f:
        json.dump(val_coco, f)

    return await flyte.io.Dir.from_local(out_dir)

# ------------------------------------------------------------------
# Helpers — torch Dataset wrapping COCO JSON
# ------------------------------------------------------------------

def _build_torch_dataset(coco_path: str, images_root: str, augment: bool):
    """Build a torch Dataset that yields {image, target} for the HF image processor."""
    import albumentations as A
    import numpy as np
    from PIL import Image
    from torch.utils.data import Dataset

    with open(coco_path) as f:
        coco = json.load(f)

    images_by_id = {im["id"]: im for im in coco["images"]}
    anns_by_image: dict[int, list] = {}
    for a in coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    image_ids = list(images_by_id.keys())

    # NOTE: we deliberately don't resize here — the HF image processor handles
    # resize+pad. Augmentation only.
    if augment:
        transform = A.Compose(
            [
                A.HorizontalFlip(p=0.5),
                A.VerticalFlip(p=0.1),
                A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
                A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=30, val_shift_limit=20, p=0.4),
                A.Rotate(limit=15, border_mode=0, p=0.4),
                A.RandomScale(scale_limit=0.2, p=0.4),
                A.GaussianBlur(blur_limit=(3, 5), p=0.2),
                A.GaussNoise(p=0.2),
            ],
            bbox_params=A.BboxParams(
                format="coco",
                label_fields=["category"],
                min_area=4,
                min_visibility=0.1,
                clip=True,
            ),
        )
    else:
        transform = A.Compose(
            [A.NoOp()],
            bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True),
        )

    class CocoDataset(Dataset):
        def __len__(self) -> int:
            return len(image_ids)

        def __getitem__(self, idx: int):
            img_id = image_ids[idx]
            meta = images_by_id[img_id]
            img_path = os.path.join(images_root, os.path.basename(meta["file_name"]))
            if not os.path.exists(img_path):
                img_path = os.path.join(images_root, meta["file_name"])
            image = np.array(Image.open(img_path).convert("RGB"))

            anns = anns_by_image.get(img_id, [])
            bboxes = [a["bbox"] for a in anns]
            categories = [a["category_id"] for a in anns]

            out = transform(image=image, bboxes=bboxes, category=categories)
            image_t = out["image"]
            bboxes_t = out["bboxes"]
            categories_t = out["category"]

            target_anns = []
            for bb, cat in zip(bboxes_t, categories_t):
                x, y, w, h = bb
                target_anns.append(
                    {
                        "image_id": img_id,
                        "category_id": int(cat),
                        "bbox": [float(x), float(y), float(w), float(h)],
                        "area": float(w * h),
                        "iscrowd": 0,
                    }
                )

            return {
                "image": image_t,
                "target": {"image_id": img_id, "annotations": target_anns},
            }

    return CocoDataset(), coco["categories"]

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    epochs: int = 30,
    lr: float = 5e-5,
    batch_size: int = 4,
    weight_decay: float = 1e-4,
    eval_every_n_epochs: int | None = None,
) -> flyte.io.Dir:
    """Fine-tune RT-DETR (or any HuggingFace object-detection model) on COCO data."""
    import torch
    from transformers import (
        AutoImageProcessor,
        AutoModelForObjectDetection,
        Trainer,
        TrainerCallback,
        TrainingArguments,
    )

    log.info(f"Training: model={model_name}")
    await flyte.report.replace.aio(_wrap_report(
        f"<h2>Loading model...</h2><p>{model_name}</p>"
        f"<p>Preparing dataset and initializing weights...</p>"
    ), do_flush=True)

    # -- Load data --
    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    train_json = os.path.join(data_path, "train.json")

    with open(train_json) as f:
        categories = json.load(f)["categories"]
    id2label = {c["id"]: c["name"] for c in categories}
    label2id = {v: k for k, v in id2label.items()}

    train_ds, _ = _build_torch_dataset(train_json, images_root, augment=True)
    log.info(f"Train examples: {len(train_ds)} | Categories: {id2label}")

    # -- Optionally load val set for periodic mAP evaluation --
    val_json = os.path.join(data_path, "val.json")
    val_images = None
    val_targets = None
    if eval_every_n_epochs and os.path.exists(val_json):
        import torch as _torch
        from PIL import Image

        with open(val_json) as f:
            val_coco = json.load(f)
        images_by_id = {im["id"]: im for im in val_coco["images"]}
        anns_by_image: dict[int, list] = {}
        for a in val_coco["annotations"]:
            anns_by_image.setdefault(a["image_id"], []).append(a)
        val_images = []
        val_targets = []
        for img_id, meta in images_by_id.items():
            path = os.path.join(images_root, os.path.basename(meta["file_name"]))
            if not os.path.exists(path):
                path = os.path.join(images_root, meta["file_name"])
            val_images.append(Image.open(path).convert("RGB"))
            boxes_xyxy = []
            labels = []
            for a in anns_by_image.get(img_id, []):
                x, y, w, h = a["bbox"]
                boxes_xyxy.append([x, y, x + w, y + h])
                labels.append(a["category_id"])
            val_targets.append({
                "boxes": _torch.tensor(boxes_xyxy, dtype=_torch.float32).reshape(-1, 4),
                "labels": _torch.tensor(labels, dtype=_torch.long),
            })
        log.info(f"Val examples for periodic eval: {len(val_images)}")

    # -- Processor + model --
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = AutoModelForObjectDetection.from_pretrained(
        model_name,
        num_labels=len(id2label),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True,
    )

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log.info(
        f"Parameters: {trainable_params:,} / {total_params:,} "
        f"({trainable_params / total_params * 100:.1f}%)"
    )

    # -- Collator — runs the image processor on each batch --
    def collate_fn(batch):
        images = [b["image"] for b in batch]
        targets = [b["target"] for b in batch]
        enc = processor(images=images, annotations=targets, return_tensors="pt")
        return {"pixel_values": enc["pixel_values"], "labels": enc["labels"]}

    # -- Sanity check: peek at one batch and verify class_labels fit --
    sample = collate_fn([train_ds[i] for i in range(min(2, len(train_ds)))])
    all_labels = []
    for lbl in sample["labels"]:
        all_labels.extend(lbl["class_labels"].tolist())
    log.info(
        f"Sanity check — class_labels in first batch: {sorted(set(all_labels))} | "
        f"model num_labels: {model.config.num_labels} | "
        f"id2label: {model.config.id2label}"
    )
    if all_labels and max(all_labels) >= model.config.num_labels:
        raise ValueError(
            f"class_label {max(all_labels)} out of range for num_labels="
            f"{model.config.num_labels}. Check category id remapping in prepare_data."
        )

    # -- Collect training metrics and update the report chart live.
    # trainer.train() runs in a background thread (via asyncio.to_thread),
    # so the asyncio event loop stays free. We use run_coroutine_threadsafe
    # to push report updates from the callback thread onto that loop.
    training_log: list[dict] = []
    eval_log: list[dict] = []  # periodic mAP checkpoints (epoch, map, map_50)
    loop = asyncio.get_running_loop()

    cat_badges = " ".join(
        f'<span class="badge badge-info">{name}</span>'
        for name in id2label.values()
    )

    def _build_training_report(max_steps: int) -> str:
        """Build the live training report HTML from current training_log."""
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{len(train_ds)}</div><div class="label">Train Examples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>
          <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
        </div>
        <p>Categories: {cat_badges}</p>
        """

        charts_html = ""
        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs} |
              Loss: <span class="highlight">{current['loss']:.4f}</span>
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_chart = _make_line_chart(
                data=training_log,
                x_key="epoch",
                y_keys=["loss"],
                title="Training Loss",
                x_label="Epoch",
                y_label="Loss",
                colors=["#5a7db5"],
            )
            charts_html += f'<div class="chart-container">{loss_chart}</div>'

            if eval_every_n_epochs:
                # Match x-axis to the loss chart: start at 0, end at current epoch
                current_max_epoch = current["epoch"]
                map_chart = _make_line_chart(
                    data=eval_log,
                    x_key="epoch",
                    y_keys=["map", "map_50"],
                    title="Validation mAP (periodic)",
                    x_label="Epoch",
                    y_label="mAP",
                    colors=["#0f3460", "#06d6a0"],
                    y_max_cap=1.0,
                    y_display_names={"map": "mAP (0.50:0.95)", "map_50": "mAP@50"},
                    x_range_override=(0, current_max_epoch),
                )
                charts_html += f'<div class="chart-container">{map_chart}</div>'

            if "lr" in training_log[0]:
                lr_chart = _make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["lr"],
                    title="Learning Rate Schedule",
                    x_label="Epoch",
                    y_label="LR",
                    colors=["#0f3460"],
                )
                charts_html += f'<div class="chart-container">{lr_chart}</div>'

        return _wrap_report(stats_html + charts_html)

    class MetricsCallback(TrainerCallback):
        def __init__(self):
            self._last_eval_epoch = 0

        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs or "loss" not in logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
                "loss": round(logs["loss"], 4),
            }
            if "learning_rate" in logs:
                entry["lr"] = logs["learning_rate"]
            if "grad_norm" in logs:
                entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
            training_log.append(entry)
            log.info(
                f"step={state.global_step}/{state.max_steps} "
                f"epoch={entry['epoch']:.2f} "
                f"loss={entry['loss']:.4f}"
            )

            # Push a live report update onto the asyncio event loop.
            # do_flush=True dispatches the update to the UI immediately.
            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

        def on_epoch_end(self, args, state, control, model=None, **kwargs):
            if not eval_every_n_epochs or val_images is None:
                return
            current_epoch = round(state.epoch)
            if current_epoch % eval_every_n_epochs != 0:
                return
            if current_epoch == self._last_eval_epoch:
                return
            self._last_eval_epoch = current_epoch

            log.info(f"Running periodic mAP eval at epoch {current_epoch}...")
            from torchmetrics.detection.mean_ap import MeanAveragePrecision

            device = next(model.parameters()).device
            # _run_inference sets model.eval(); restore train mode after.
            preds = _run_inference(model, processor, val_images, device, threshold=0.3)
            model.train()

            formatted = [
                {"boxes": p["boxes"], "scores": p["scores"], "labels": p["labels"]}
                for p in preds
            ]
            metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
            metric.update(formatted, val_targets)
            result = metric.compute()
            map_val = round(result["map"].item(), 4)
            map_50 = round(result["map_50"].item(), 4)

            eval_log.append({
                "epoch": current_epoch,
                "map": map_val,
                "map_50": map_50,
            })
            log.info(f"Epoch {current_epoch} — mAP: {map_val:.4f}, mAP@50: {map_50:.4f}")

            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=lr,
        weight_decay=weight_decay,
        logging_steps=5,
        save_strategy="no",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        warmup_ratio=0.1,
        remove_unused_columns=False,
        dataloader_num_workers=2,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        data_collator=collate_fn,
        callbacks=[MetricsCallback()],
    )

    log.info("Starting training...")
    # Run the sync HF training loop in a thread so the asyncio event loop
    # stays free for Flyte's syncify bridge.
    await asyncio.to_thread(trainer.train)
    log.info("Training complete.")

    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
    trainer.save_model(save_dir)
    processor.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Build final training report --
    stats_html = f"""
    <h2>Training Complete</h2>
    <h3>{model_name}</h3>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(train_ds)}</div><div class="label">Train Examples</div></div>
      <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
      <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
      <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
      <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>
      <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
    </div>
    <p>Categories: {cat_badges}</p>
    """

    charts_html = ""
    if training_log:
        final_loss = training_log[-1]["loss"]
        min_loss = min(d["loss"] for d in training_log)
        initial_loss = training_log[0]["loss"]
        total_steps = training_log[-1]["step"]

        charts_html += f"""
        <div class="card">
          <b>Training Summary:</b>
          Initial loss: {initial_loss:.4f} |
          Final loss: <span class="highlight">{final_loss:.4f}</span> |
          Min loss: {min_loss:.4f} |
          Total steps: {total_steps}
        </div>
        """

        epoch_range = (0, epochs)

        loss_chart = _make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["loss"],
            title="Training Loss",
            x_label="Epoch",
            y_label="Loss",
            colors=["#5a7db5"],
            x_range_override=epoch_range,
        )
        charts_html += f'<div class="chart-container">{loss_chart}</div>'

    if eval_log:
        map_chart = _make_line_chart(
            data=eval_log,
            x_key="epoch",
            y_keys=["map", "map_50"],
            title="Validation mAP (periodic)",
            x_label="Epoch",
            y_label="mAP",
            colors=["#0f3460", "#06d6a0"],
            y_max_cap=1.0,
            x_range_override=(0, epochs),
            y_display_names={"map": "mAP (0.50:0.95)", "map_50": "mAP@50"},
        )
        charts_html += f'<div class="chart-container">{map_chart}</div>'

    if training_log and "lr" in training_log[0]:
        lr_chart = _make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["lr"],
            title="Learning Rate Schedule",
            x_label="Epoch",
            y_label="LR",
            colors=["#0f3460"],
            x_range_override=(0, epochs),
        )
        charts_html += f'<div class="chart-container">{lr_chart}</div>'

    await flyte.report.replace.aio(_wrap_report(stats_html + charts_html), do_flush=True)

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Inference helpers
# ------------------------------------------------------------------

def _run_inference(model, processor, images, device, threshold: float = 0.3):
    """Run object detection on a list of PIL images. Returns list of dicts."""
    import torch

    results = []
    model.eval()
    for img in images:
        inputs = processor(images=img, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
        target_size = torch.tensor([img.size[::-1]], device=device)  # (h, w)
        post = processor.post_process_object_detection(
            outputs, target_sizes=target_size, threshold=threshold
        )[0]
        results.append(
            {
                "scores": post["scores"].cpu(),
                "labels": post["labels"].cpu(),
                "boxes": post["boxes"].cpu(),  # xyxy in original image coords
            }
        )
    return results

def _draw_boxes(image, boxes, labels, scores, id2label, color: str = "lime"):
    """Draw bounding boxes on a PIL image. Returns a new PIL image."""
    from PIL import ImageDraw, ImageFont

    img = image.copy()
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=max(14, img.width // 60))
    except Exception:
        font = ImageFont.load_default()

    for box, label, score in zip(boxes.tolist(), labels.tolist(), scores.tolist()):
        x0, y0, x1, y1 = box
        width = max(2, img.width // 400)
        draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
        name = id2label.get(int(label), str(int(label)))
        caption = f"{name} {score:.2f}"
        text_bg = draw.textbbox((x0, y0), caption, font=font)
        draw.rectangle(text_bg, fill=color)
        draw.text((x0, y0), caption, fill="black", font=font)
    return img

def _img_to_data_uri(img, max_dim: int = 800) -> str:
    """PIL image → base64 data URI, downscaled for the report."""
    w, h = img.size
    if max(w, h) > max_dim:
        scale = max_dim / max(w, h)
        img = img.resize((int(w * scale), int(h * scale)))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()

# ------------------------------------------------------------------
# Task 3: Evaluate — COCO mAP on fine-tuned model
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    threshold: float = 0.5,
) -> str:
    """Compute COCO mAP for the fine-tuned model on the val split."""
    import torch
    from PIL import Image
    from torchmetrics.detection.mean_ap import MeanAveragePrecision
    from transformers import AutoImageProcessor, AutoModelForObjectDetection

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(_wrap_report(
        "<h2>Evaluation</h2><p>Loading val split and scoring model...</p>"
    ), do_flush=True)

    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    val_json = os.path.join(data_path, "val.json")

    with open(val_json) as f:
        val_coco = json.load(f)

    images_by_id = {im["id"]: im for im in val_coco["images"]}
    anns_by_image: dict[int, list] = {}
    for a in val_coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    pil_images = []
    targets = []
    for img_id, meta in images_by_id.items():
        path = os.path.join(images_root, os.path.basename(meta["file_name"]))
        if not os.path.exists(path):
            path = os.path.join(images_root, meta["file_name"])
        pil_images.append(Image.open(path).convert("RGB"))
        boxes_xyxy = []
        labels = []
        for a in anns_by_image.get(img_id, []):
            x, y, w, h = a["bbox"]
            boxes_xyxy.append([x, y, x + w, y + h])
            labels.append(a["category_id"])
        targets.append(
            {
                "boxes": torch.tensor(boxes_xyxy, dtype=torch.float32).reshape(-1, 4),
                "labels": torch.tensor(labels, dtype=torch.long),
            }
        )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    ft_path = await finetuned_dir.download()
    log.info(f"Scoring fine-tuned model: {ft_path}")
    processor = AutoImageProcessor.from_pretrained(ft_path)
    model = AutoModelForObjectDetection.from_pretrained(ft_path).to(device)
    preds = _run_inference(model, processor, pil_images, device, threshold=threshold)

    formatted_preds = [
        {"boxes": p["boxes"], "scores": p["scores"], "labels": p["labels"]}
        for p in preds
    ]

    metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
    metric.update(formatted_preds, targets)

    def to_python(v):
        if hasattr(v, "numel"):
            return v.item() if v.numel() == 1 else v.tolist()
        return v

    ft_metrics = {k: to_python(v) for k, v in metric.compute().items()}
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    log.info(f"Fine-tuned mAP: {ft_metrics.get('map', 0):.3f}")

    metric_keys = ["map", "map_50", "map_75", "mar_10"]
    metric_display = {
        "map": "mAP",
        "map_50": "mAP@50",
        "map_75": "mAP@75",
        "mar_10": "mAR@10",
    }

    rows = []
    for key in metric_keys:
        ft_val = ft_metrics.get(key, 0)
        rows.append(
            f"<tr><td><b>{metric_display.get(key, key)}</b></td>"
            f"<td class='highlight'>{ft_val:.3f}</td></tr>"
        )
    table = (
        "<table><tr><th>Metric</th><th>Score</th></tr>"
        + "".join(rows)
        + "</table>"
    )

    bar_chart = _make_bar_chart(
        labels=[metric_display.get(k, k) for k in metric_keys],
        series={"Fine-tuned": [ft_metrics.get(k, 0) for k in metric_keys]},
        title="COCO Evaluation Metrics",
        colors=["#0f3460"],
        y_max_cap=1.0,
    )

    ft_map = ft_metrics.get("map", 0)
    ft_map50 = ft_metrics.get("map_50", 0)

    eval_html = f"""
    <h2>Evaluation — COCO mAP</h2>
    <div class="stat-grid">
      <div class="stat"><div class="value">{len(pil_images)}</div><div class="label">Val Images</div></div>
      <div class="stat"><div class="value">{threshold}</div><div class="label">Threshold</div></div>
      <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
      <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
    </div>
    <div class="chart-container">{bar_chart}</div>
    {table}
    <div class="note">
      <b>mAP</b> (mean Average Precision) measures how accurately the model
      detects objects — balancing whether predictions are correct (precision)
      and whether all objects are found (recall). The @50 and @75 variants
      require IoU overlaps of 50% and 75% between predicted and ground-truth
      boxes. <b>mAR</b> (mean Average Recall) measures how many ground-truth
      objects the model finds, with @1 and @10 limiting detections to 1 or 10
      per image.
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(eval_html), do_flush=True)

    return json.dumps(
        {
            "finetuned": {k: round(v, 4) for k, v in ft_metrics.items() if isinstance(v, (int, float))},
            "num_val_images": len(pil_images),
        }
    )

# ------------------------------------------------------------------
# Task 4: Inference demo — render bboxes on val images
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def inference_demo(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    threshold: float = 0.5,
    max_images: int = 8,
    metrics_json: str = "{}",
) -> str:
    """Run the fine-tuned model on val images, render bboxes, embed in the report."""
    import torch
    from PIL import Image
    from torchmetrics.detection.mean_ap import MeanAveragePrecision
    from transformers import AutoImageProcessor, AutoModelForObjectDetection

    data_path = await data_dir.download()
    images_root = os.path.join(data_path, "images")
    val_json = os.path.join(data_path, "val.json")

    with open(val_json) as f:
        val_coco = json.load(f)

    id2label = {c["id"]: c["name"] for c in val_coco["categories"]}
    metas = val_coco["images"][:max_images]
    anns_by_image: dict[int, list] = {}
    for a in val_coco["annotations"]:
        anns_by_image.setdefault(a["image_id"], []).append(a)

    pil_images = []
    gt_per_image = []
    for meta in metas:
        path = os.path.join(images_root, os.path.basename(meta["file_name"]))
        if not os.path.exists(path):
            path = os.path.join(images_root, meta["file_name"])
        pil_images.append(Image.open(path).convert("RGB"))

        boxes_xyxy = []
        labels = []
        for a in anns_by_image.get(meta["id"], []):
            x, y, w, h = a["bbox"]
            boxes_xyxy.append([x, y, x + w, y + h])
            labels.append(a["category_id"])
        gt_per_image.append(
            {
                "boxes": torch.tensor(boxes_xyxy, dtype=torch.float32).reshape(-1, 4),
                "labels": torch.tensor(labels, dtype=torch.long),
                "scores": torch.ones(len(labels)),
            }
        )

    ft_path = await finetuned_dir.download()
    processor = AutoImageProcessor.from_pretrained(ft_path)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModelForObjectDetection.from_pretrained(ft_path).to(device)

    preds = _run_inference(model, processor, pil_images, device, threshold=threshold)

    html_blocks = []
    total_gt = 0
    total_pred = 0
    for i, (img, pred, gt) in enumerate(zip(pil_images, preds, gt_per_image)):
        n_gt = len(gt["labels"])
        n_pred = len(pred["labels"])
        total_gt += n_gt
        total_pred += n_pred

        # Per-image mAP
        metric = MeanAveragePrecision(box_format="xyxy", iou_type="bbox")
        metric.update(
            [{"boxes": pred["boxes"], "scores": pred["scores"], "labels": pred["labels"]}],
            [{"boxes": gt["boxes"], "labels": gt["labels"]}],
        )
        img_metrics = metric.compute()
        img_map = img_metrics["map"].item()
        img_map_badge = (
            f'<span class="badge badge-success">mAP {img_map:.2f}</span>'
            if img_map >= 0.5
            else f'<span class="badge badge-info">mAP {img_map:.2f}</span>'
        )

        pred_img = _draw_boxes(
            img, pred["boxes"], pred["labels"], pred["scores"],
            id2label, color="lime",
        )
        gt_img = _draw_boxes(
            img, gt["boxes"], gt["labels"], gt["scores"],
            id2label, color="dodgerblue",
        )
        html_blocks.append(f"""
        <div class="card">
          <b>Image {i + 1}</b> {img_map_badge}
          <div class="img-pair">
            <div>
              <p><span class="gt-label">Ground Truth</span>
                 <span class="badge badge-info">{n_gt} boxes</span></p>
              <img src="{_img_to_data_uri(gt_img)}" />
            </div>
            <div>
              <p><span class="pred-label">Predictions</span>
                 <span class="badge badge-success">{n_pred} boxes</span>
                 (threshold={threshold})</p>
              <img src="{_img_to_data_uri(pred_img)}" />
            </div>
          </div>
        </div>""")

    # Parse metrics if provided (from evaluate task)
    metrics = json.loads(metrics_json)
    ft_metrics = metrics.get("finetuned", {})
    ft_map = ft_metrics.get("map", None)
    ft_map50 = ft_metrics.get("map_50", None)

    metrics_stats = ""
    if ft_map is not None:
        metrics_stats = f"""
        <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
        <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
        """

    demo_html = f"""
    <h2>Inference Demo</h2>
    <h3>Fine-tuned RT-DETR on validation images</h3>
    <div class="stat-grid">
      {metrics_stats}
      <div class="stat"><div class="value">{len(pil_images)}</div><div class="label">Images Shown</div></div>
      <div class="stat"><div class="value">{total_gt}</div><div class="label">Ground Truth Boxes</div></div>
      <div class="stat"><div class="value">{total_pred}</div><div class="label">Predicted Boxes</div></div>
      <div class="stat"><div class="value">{threshold}</div><div class="label">Confidence Threshold</div></div>
    </div>
    <p><span class="gt-label">Blue = ground truth</span> |
       <span class="pred-label">Green = predictions</span></p>
    {"".join(html_blocks)}
    """

    await flyte.report.replace.aio(_wrap_report(demo_html), do_flush=True)

    return json.dumps(
        {
            "num_images": len(pil_images),
            "predictions_per_image": [len(p["labels"]) for p in preds],
        }
    )

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "PekingU/rtdetr_v2_r18vd",
    dataset_repo: str = "sagecodes/union_flyte_swag_object_detection",
    annotations_path: str = "swag/train.json",
    images_subdir: str = "swag/images",
    epochs: int = 30,
    lr: float = 5e-5,
    batch_size: int = 4,
    val_fraction: float = 0.2,
    threshold: float = 0.5,
    demo_images: int = 8,
    eval_every_n_epochs: int | None = None,
) -> tuple[flyte.io.Dir, str]:
    """
    End-to-end RT-DETRv2 fine-tuning pipeline.

    Returns the fine-tuned model directory and a JSON summary.

    1. Download COCO dataset from HuggingFace and split train/val
    2. Fine-tune RT-DETRv2 on the train split
    3. Evaluate: COCO mAP comparison (base vs fine-tuned)
    4. Inference demo: render bounding boxes on val images
    """
    log.info(f"Pipeline: {model_name} | dataset={dataset_repo}")

    def _pipeline_progress(step: int, label: str) -> str:
        steps = ["Preparing Data", "Fine-tuning", "Evaluating", "Inference Demo"]
        dots = ""
        for i, s in enumerate(steps):
            if i + 1 < step:
                icon = '<span style="color:#06d6a0;">&#10003;</span>'
            elif i + 1 == step:
                icon = '<span style="color:#e94560;">&#9679;</span>'
            else:
                icon = '<span style="color:#adb5bd;">&#9675;</span>'
            dots += f"<span style='margin:0 8px;'>{icon} {s}</span>"
        return f"""
        <h2>RT-DETRv2 Object Detection Pipeline</h2>
        <p><b>Model:</b> {model_name} | <b>Dataset:</b> {dataset_repo}</p>
        <div class="card" style="text-align:center;">{dots}</div>
        <p>{label}</p>
        """

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(1, "Downloading and splitting dataset...")),
        do_flush=True,
    )

    data_dir = await prepare_data(
        dataset_repo=dataset_repo,
        annotations_path=annotations_path,
        images_subdir=images_subdir,
        val_fraction=val_fraction,
    )

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(2, "Fine-tuning model...")),
        do_flush=True,
    )

    finetuned_dir = await train(
        model_name, data_dir, epochs, lr, batch_size,
        eval_every_n_epochs=eval_every_n_epochs,
    )

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(3, "Running COCO mAP evaluation...")),
        do_flush=True,
    )

    metrics_json = await evaluate(finetuned_dir, data_dir, threshold)
    metrics = json.loads(metrics_json)

    await flyte.report.replace.aio(
        _wrap_report(_pipeline_progress(4, "Rendering bounding box demo...")),
        do_flush=True,
    )

    demo_json = await inference_demo(
        finetuned_dir, data_dir, threshold, demo_images,
        metrics_json=metrics_json,
    )

    ft_map = metrics["finetuned"].get("map", 0)
    ft_map50 = metrics["finetuned"].get("map_50", 0)

    final_html = f"""
    <h2>Pipeline Complete</h2>
    <h3>{model_name}</h3>
    <div class="stat-grid">
      <div class="stat"><div class="value">{metrics['num_val_images']}</div><div class="label">Val Images</div></div>
      <div class="stat"><div class="value highlight">{ft_map:.3f}</div><div class="label">mAP</div></div>
      <div class="stat"><div class="value highlight">{ft_map50:.3f}</div><div class="label">mAP@50</div></div>
    </div>
    <div class="card">
      <b>Configuration:</b> {epochs} epochs | LR {lr} | Batch size {batch_size} |
      Val fraction {val_fraction} | Threshold {threshold}
    </div>
    """

    await flyte.report.replace.aio(_wrap_report(final_html), do_flush=True)

    log.info(f"Pipeline complete. Fine-tuned mAP: {ft_map:.3f}")
    return finetuned_dir, json.dumps({"metrics": metrics, "demo": json.loads(demo_json)})

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
CODE1
cd v2/tutorials/detr_object_detection
uv run --script detr_object_detection.py
CODE2
flyte run detr_object_detection.py pipeline --epochs 1 --batch_size 2
```

This workflow needs a GPU. Check the **train**, **evaluate**, and **inference_demo** task reports for charts and annotated images.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents ===

# Agents

Tutorials for building agentic workflows and autonomous LLM-powered systems.

### **Agents > Autoresearch agent**

Run an autonomous research loop that drives Claude Code in a GPU container to run experiments, then commits results and opens a pull request.

### **Agents > Parallelized autoresearch agent**

Scale autoresearch with a code-mode MLE agent that batches train.py edits and runs sandbox experiments in parallel via flyte.map.

### **Agents > Code mode analytics agent**

Chat with a dataset in the browser: Claude writes a Python program that runs in the Monty sandbox, with the heavy DuckDB query dispatched as a durable Flyte task.

### **Agents > AutoSec researcher agent**

Fan out vulnerability analysis across C targets, hypothesize exploits with an LLM agent, and validate PoCs in an isolated sandbox.

### **Agents > Coding agent**

Securely execute and iterate on LLM-generated code using a code agent with error reflection and retry logic.

### **Agents > Competitive intelligence agent**

Fan out across competitors, extract source-cited market deltas with the You.com Search API, and build a knowledge-graph-ready intelligence table.

### **Agents > Compliance monitoring agent**

Monitor trusted regulatory sources with the You.com Research API and route citation-precise findings to the right team.

### **Agents > Deep research**

Build an agentic workflow for deep research with multi-step reasoning and evaluation.

### **Agents > LangGraph research agent**

Combine LangGraph control flow with Flyte tasks for multi-topic web research with quality-check loops.

### **Agents > Field data enrichment agent**

Enrich geo-tagged operational events with real-world public context using the You.com Search API with country and freshness targeting.

### **Agents > MLE bot: an autonomous ML engineer**

An autonomous ML agent that designs, runs, and iterates on experiments using Flyte's durable sandbox for safe LLM-generated code execution.

### **Agents > Support resolution agent**

Ground support tickets in fresh public sources via the You.com Research API and draft cited, customer-ready replies for human review.

## Subpages
- **Agents > Autoresearch agent**
- **Agents > Parallelized autoresearch agent**
- **Agents > AutoSec researcher agent**
- **Agents > Coding agent**
- **Agents > Competitive intelligence agent**
- **Agents > Deep research**
- **Agents > LangGraph research agent**
- **Agents > MLE bot: an autonomous ML engineer**
- **Agents > Compliance monitoring agent**
- **Agents > Field data enrichment agent**
- **Agents > Support resolution agent**
- **Agents > Code mode analytics agent**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/autoresearch ===

# Autoresearch agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autoresearch).

This tutorial wraps an autonomous AI research loop in a single Flyte task. The task spins up a GPU container, installs the [Claude Code](https://docs.anthropic.com/en/docs/claude-code/overview) CLI, clones a research repository, and points Claude Code at a `program.md` brief. The agent runs experiments to improve a model, writes results to disk, and the task then commits the changes and opens a pull request, with a progress plot rendered both in the PR and in the Flyte UI.

It's an example of using Flyte as durable infrastructure for long-running, autonomous agent work:

- **A GPU `TaskEnvironment`** with the API-key and GitHub secrets the agent needs.
- **`report=True`** to stream a progress plot into the Flyte UI.
- **A reconnecting `run.wait()`** loop in the driver so a dropped client connection doesn't lose track of a multi-hour run.

> [!WARNING]
> This example drives a coding agent that executes arbitrary code and pushes commits to a GitHub repository. Run it against a repository you control, and review the constants described below before launching.

## Define the container image

The image is kept in its own `_image.py` module so edits to the agent logic in `run.py` don't invalidate the image cache. Node.js and the Claude Code CLI are installed at run time (see below) to keep the image small.

```
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "flyte>=2.0.0b22",
#     "PyGithub>=2.5.0",
#     "matplotlib>=3.7.0",
#     "pandas>=2.0.0",
# ]
# ///
#
# Stable image definition — kept separate from run.py so edits to run.py
# don't invalidate the image cache. Only touch this file when the image itself needs to change.

import flyte

# {{docs-fragment image}}
image = (
    flyte.Image.from_uv_script(__file__, name="autoresearch-agent", pre=True)
    .with_apt_packages("git")
)
# {{/docs-fragment image}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/_image.py*

## Define the task environment

The task needs a GPU, a generous disk for the cloned repo and model weights, and two secrets: a GitHub token (to clone and push) and an Anthropic API key (for Claude Code).

```
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "flyte>=2.0.0b22",
#     "PyGithub>=2.5.0",
#     "matplotlib>=3.7.0",
# ]
# ///

"""
AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment.

This agent:
1. Starts a GPU-enabled container
2. Installs Claude Code CLI
3. Clones the autoresearch repository
4. Points Claude Code at program.md as the prompt and lets it run
5. Commits the result (CSV + code changes in train/) and creates a PR
"""

import os
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from github import Auth, Github

import flyte
import flyte.report
from _image import image as autoresearch_image

GITHUB_USERNAME = "parnianz"
GITHUB_EMAIL = "parnianzargham@gmail.com"
AUTORESEARCH_REPO_URL = "https://github.com/unionai-oss/autoresearch.git"
AUTORESEARCH_REPO_FULL_NAME = "unionai-oss/autoresearch"

# {{docs-fragment env}}
autoresearch_env = flyte.TaskEnvironment(
    name="autoresearch-agent",
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    secrets=[
        flyte.Secret(key="github_token", as_env_var="GITHUB_TOKEN"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=autoresearch_image,
)
# {{/docs-fragment env}}

# {{docs-fragment result}}
@dataclass
class AutoResearchResult:
    """Result of the autoresearch run."""

    pr_url: str
    pr_number: int
    branch_name: str
    files_changed: list[str]
    success: bool
    error_message: Optional[str] = None
# {{/docs-fragment result}}

def clone_repository(repo_url: str, work_dir: Path, github_token: str) -> Path:
    """Clone the autoresearch repository with authentication."""
    repo_name = repo_url.rstrip("/").split("/")[-1].replace(".git", "")
    repo_path = work_dir / repo_name

    # Inject token into HTTPS URL for authentication
    authenticated_url = repo_url.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )

    if repo_path.exists():
        subprocess.run(["git", "pull"], cwd=repo_path, check=True)
    else:
        subprocess.run(["git", "clone", authenticated_url, str(repo_path)], check=True)

    return repo_path

# {{docs-fragment task}}
@autoresearch_env.task(report=True)
async def run_autoresearch() -> AutoResearchResult:
    """
    Run the autoresearch workflow end-to-end.

    Steps:
    - Clone https://github.com/unionai-oss/autoresearch
    - Configure git identity
    - Create a new branch
    - Run Claude Code CLI with program.md as the prompt
    - Commit results (CSV + train/ changes)
    - Push and open a PR against the autoresearch repo
    """
    github_token = os.environ["GITHUB_TOKEN"]
    anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]

    # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) ---
    import tarfile
    import urllib.request as _urllib

    subprocess.run(["apt-get", "update", "-y"], check=False)
    subprocess.run(["apt-get", "install", "-y", "git"], check=False)

    node_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz"
    node_tar = Path("/tmp/node.tar.gz")
    print(f"Downloading Node.js from {node_url}...", flush=True)
    _urllib.urlretrieve(node_url, node_tar)
    size_mb = node_tar.stat().st_size / 1024 / 1024
    print(f"Downloaded {size_mb:.1f} MB to {node_tar}", flush=True)
    if size_mb < 1:
        raise RuntimeError(f"Node.js download appears empty/corrupt ({size_mb:.2f} MB) — network may be restricted")
    node_dir = Path("/tmp/node")
    node_dir.mkdir(exist_ok=True)
    print("Extracting Node.js...", flush=True)
    with tarfile.open(node_tar, "r:gz") as tar:
        members = [m for m in tar.getmembers() if m.name.split("/", 1)[-1]]
        for m in members:
            m.name = m.name.split("/", 1)[-1]
        tar.extractall(str(node_dir), members=[m for m in members if m.name])

    # Add node/npm to PATH for this process and all subprocesses
    node_bin = str(node_dir / "bin")
    os.environ["PATH"] = node_bin + ":" + os.environ.get("PATH", "")
    print(f"Node version: {subprocess.run(['node', '--version'], capture_output=True, text=True).stdout.strip()}", flush=True)

    npm_prefix = "/tmp/npm-global"
    Path(npm_prefix).mkdir(exist_ok=True)
    subprocess.run(["npm", "install", "-g", "--prefix", npm_prefix, "@anthropic-ai/claude-code"], check=True)
    os.environ["PATH"] = str(Path(npm_prefix) / "bin") + ":" + os.environ["PATH"]
    print("Node.js + Claude Code installed.", flush=True)

    # --- Clone repo ---
    work_dir = Path("/tmp/autoresearch_workspace")
    work_dir.mkdir(exist_ok=True, parents=True)
    repo_path = clone_repository(AUTORESEARCH_REPO_URL, work_dir, github_token)

    # --- Git identity ---
    subprocess.run(
        ["git", "config", "--global", "user.email", GITHUB_EMAIL], check=True
    )
    subprocess.run(
        ["git", "config", "--global", "user.name", GITHUB_USERNAME], check=True
    )

    # --- Create branch ---
    import time as _time
    branch_name = f"autoresearch/claude-run-{int(_time.time())}"
    try:
        subprocess.run(
            ["git", "checkout", "-b", branch_name],
            cwd=repo_path,
            check=True,
        )
    except subprocess.CalledProcessError:
        subprocess.run(
            ["git", "checkout", branch_name],
            cwd=repo_path,
            check=True,
        )

    # --- Read program.md to use as the Claude Code prompt ---
    program_md = repo_path / "program.md"
    if not program_md.exists():
        raise FileNotFoundError(
            f"program.md not found in {repo_path}. "
            "Make sure the autoresearch repo has a program.md at its root."
        )

    program_md_content = program_md.read_text()
    print(f"Loaded prompt from program.md ({len(program_md_content)} chars)")
    # {{/docs-fragment task}}

    # Install repo dependencies before handing off to Claude
    for pip_cmd in [
        ["pip", "install", "-e", "."],
        ["pip", "install", "-r", "requirements.txt"],
    ]:
        req_file = repo_path / pip_cmd[-1] if pip_cmd[-1].startswith("req") else None
        if req_file is None or req_file.exists():
            dep_result = subprocess.run(
                pip_cmd, cwd=repo_path, capture_output=True, text=True
            )
            print(f"{' '.join(pip_cmd)}:\n{dep_result.stdout}", flush=True)
            if dep_result.returncode != 0:
                print(f"(non-fatal) {dep_result.stderr}", flush=True)

    # Wrap the program.md content with explicit instructions to write outputs to disk
    prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files.

Here are your instructions from program.md:

{program_md_content}

LOGGING INSTRUCTIONS (follow exactly):
- Before you start any training, print this exact line: [AUTORESEARCH] Training started
- Before training, print what change you are testing: [AUTORESEARCH] Change: <one line description of the code change being tested>
- When training finishes, print this exact line: [AUTORESEARCH] Training finished
- After training, print the key metric value: [AUTORESEARCH] Metric: <metric name>=<value>
- When writing results to CSV, print this exact line: [AUTORESEARCH] Writing results to CSV

IMPORTANT: After completing the above instructions, make sure you have:
1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar)
2. Saved all code changes you made to the train/ directory (or wherever the training code lives)
3. All files must be written to the current working directory so they appear in git status
If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk."""

    # --- Pre-flight: verify claude is installed and API key is reachable ---
    version_check = subprocess.run(
        ["claude", "--version"], capture_output=True, text=True
    )
    print(f"claude version: {version_check.stdout.strip()} | stderr: {version_check.stderr.strip()}", flush=True)
    if version_check.returncode != 0:
        raise RuntimeError(f"claude CLI not found or broken: {version_check.stderr}")

    # --- Disable Claude Code sandbox ---
    # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container
    # which fails silently and causes file writes to go to an ephemeral space instead of
    # the real working directory. Disabling it makes writes land in the actual filesystem.
    claude_config_dir = Path("/root/.claude")
    claude_config_dir.mkdir(parents=True, exist_ok=True)
    settings = claude_config_dir / "settings.json"
    import json as _json
    existing = _json.loads(settings.read_text()) if settings.exists() else {}
    existing["sandbox"] = False
    settings.write_text(_json.dumps(existing, indent=2))
    print(f"Wrote Claude Code settings: {settings.read_text()}", flush=True)

    # --- Run Claude Code CLI ---
    # Matches swe_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode
    cmd = [
        "claude",
        "--dangerously-skip-permissions",
        "--max-turns", "100",
        "--model", "claude-haiku-4-5-20251001",
        prompt,
    ]

    print(f"Running: {shlex.join(cmd[:3])} <prompt>", flush=True)

    claude_env = {
        **os.environ,
        "ANTHROPIC_API_KEY": anthropic_api_key,
        "CLAUDE_SKIP_PERMISSIONS": "true",
        "CI": "true",  # Enables non-interactive mode (no TTY required)
    }

    # Stream output line by line so logs appear in real time instead of buffering until done
    proc = subprocess.Popen(
        cmd,
        cwd=repo_path,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # merge stderr into stdout stream
        text=True,
        env=claude_env,
    )

    stdout_lines = []
    for line in proc.stdout:
        line = line.rstrip("\n")
        print(line, flush=True)
        stdout_lines.append(line)

    proc.wait()
    full_output = "\n".join(stdout_lines)
    print(f"Claude Code exit code: {proc.returncode}", flush=True)

    if proc.returncode != 0:
        raise RuntimeError(
            f"Claude Code CLI exited with code {proc.returncode}\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Collect changed files ---
    git_status = subprocess.run(
        ["git", "status", "--porcelain"],
        cwd=repo_path,
        capture_output=True,
        text=True,
        check=True,
    )

    print(f"Git status:\n{git_status.stdout}", flush=True)

    files_changed = []
    for line in git_status.stdout.strip().splitlines():
        if line:
            # git status --porcelain: first two chars are XY status flags
            file_path = line[3:].strip()
            files_changed.append(file_path)

    # Also list all files in repo dir for debugging
    all_files = subprocess.run(
        ["find", ".", "-type", "f", "-not", "-path", "./.git/*"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"All files in repo:\n{all_files.stdout}", flush=True)

    if not files_changed:
        raise RuntimeError(
            "Claude Code ran successfully but produced no file changes.\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Commit ---
    subprocess.run(["git", "add", "."], cwd=repo_path, check=True)
    subprocess.run(["git", "add", "-f", "results.tsv"], cwd=repo_path, check=False)
    subprocess.run(["git", "add", "-f", "results/"], cwd=repo_path, check=False)
    commit_message = (
        "feat: autoresearch run via Claude Code\n\n"
        "Added research results (CSV) and updated train/ code changes.\n"
        "Generated by the autoresearch Flyte agent."
    )
    subprocess.run(
        ["git", "commit", "-m", commit_message],
        cwd=repo_path,
        check=True,
    )

    # --- Push ---
    print(f"GitHub token present: {bool(github_token)}, length: {len(github_token) if github_token else 0}", flush=True)
    authenticated_url = AUTORESEARCH_REPO_URL.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )
    subprocess.run(
        ["git", "remote", "set-url", "origin", authenticated_url],
        cwd=repo_path,
        check=True,
    )
    push_result = subprocess.run(
        ["git", "push", "-u", "origin", branch_name, "--force"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"Push stdout: {push_result.stdout}", flush=True)
    print(f"Push stderr: {push_result.stderr}", flush=True)
    if push_result.returncode != 0:
        raise RuntimeError(f"git push failed (exit {push_result.returncode}):\n{push_result.stderr}")

    # --- Create PR via PyGithub ---
    auth = Auth.Token(github_token)
    gh = Github(auth=auth)
    repo = gh.get_repo(AUTORESEARCH_REPO_FULL_NAME)

    csv_files = [f for f in files_changed if f.endswith(".csv")]
    train_files = [f for f in files_changed if "train" in f]

    pr_body = f"""## AutoResearch Run

This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI.

### What changed
- **Result CSV files**: {', '.join(f'`{f}`' for f in csv_files) or 'none detected'}
- **Train code changes**: {', '.join(f'`{f}`' for f in train_files) or 'none detected'}

### All changed files
{chr(10).join(f'- `{f}`' for f in files_changed)}

---
🤖 Generated by [autoresearch Flyte agent](https://github.com/unionai-oss/autoresearch)
"""

    existing_prs = list(repo.get_pulls(state="open", head=f"unionai-oss:{branch_name}"))
    if existing_prs:
        pr = existing_prs[0]
        print(f"PR already exists: {pr.html_url}", flush=True)
    else:
        pr = repo.create_pull(
            title="feat: autoresearch results + train changes",
            body=pr_body,
            head=branch_name,
            base="master",
        )
        print(f"PR created: {pr.html_url}", flush=True)

    # --- Generate progress plot from results.tsv ---
    plot_path = repo_path / "progress.png"
    results_tsv = repo_path / "results.tsv"
    if results_tsv.exists():
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import pandas as pd

        df = pd.read_csv(str(results_tsv), sep="\t")
        df["val_bpb"] = pd.to_numeric(df["val_bpb"], errors="coerce")
        df["memory_gb"] = pd.to_numeric(df["memory_gb"], errors="coerce")
        df["status"] = df["status"].str.strip().str.upper()

        # Filter out crashes for plotting
        valid = df[df["status"] != "CRASH"].copy()
        valid = valid.reset_index(drop=True)

        if len(valid) > 0 and valid["val_bpb"].notna().any():
            baseline_bpb = valid.loc[0, "val_bpb"]
            best = valid["val_bpb"].min()

            # Only plot points at or below baseline (the interesting region)
            below = valid[valid["val_bpb"] <= baseline_bpb + 0.0005]

            fig, ax = plt.subplots(figsize=(16, 8))

            # Plot discarded as faint background dots
            disc = below[below["status"] == "DISCARD"]
            ax.scatter(disc.index, disc["val_bpb"],
                       c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded")

            # Plot kept experiments as prominent green dots
            kept_v = below[below["status"] == "KEEP"]
            ax.scatter(kept_v.index, kept_v["val_bpb"],
                       c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5)

            # Running minimum step line
            kept_mask = valid["status"] == "KEEP"
            kept_idx = valid.index[kept_mask]
            kept_bpb = valid.loc[kept_mask, "val_bpb"]
            running_min = kept_bpb.cummin()
            ax.step(kept_idx, running_min, where="post", color="#27ae60",
                    linewidth=2, alpha=0.7, zorder=3, label="Running best")

            # Label each kept experiment with its description
            for idx, bpb in zip(kept_idx, kept_bpb):
                desc = str(valid.loc[idx, "description"]).strip()
                if len(desc) > 45:
                    desc = desc[:42] + "..."
                ax.annotate(desc, (idx, bpb),
                            textcoords="offset points",
                            xytext=(6, 6), fontsize=8.0,
                            color="#1a7a3a", alpha=0.9,
                            rotation=30, ha="left", va="bottom")

            n_total = len(df)
            n_kept = len(df[df["status"] == "KEEP"])
            ax.set_xlabel("Experiment #", fontsize=12)
            ax.set_ylabel("Validation BPB (lower is better)", fontsize=12)
            ax.set_title(f"Autoresearch Progress: {n_total} Experiments, {n_kept} Kept Improvements", fontsize=14)
            ax.legend(loc="upper right", fontsize=9)
            ax.grid(True, alpha=0.2)

            margin = (baseline_bpb - best) * 0.15
            ax.set_ylim(best - margin, baseline_bpb + margin)

            plt.tight_layout()
            plt.savefig(str(plot_path), dpi=150, bbox_inches="tight")
            plt.close(fig)
            print(f"Saved plot to {plot_path}", flush=True)

            # Upload plot to PR as a comment with base64 inline image
            import base64
            img_b64 = base64.b64encode(plot_path.read_bytes()).decode()
            pr_comment = (
                "## Autoresearch Progress\n\n"
                f"![Autoresearch Progress](data:image/png;base64,{img_b64})"
            )
            pr.create_issue_comment(pr_comment)
            print("Posted plot as PR comment.", flush=True)

            # Force-add plot to git and amend commit
            subprocess.run(["git", "add", "-f", str(plot_path)], cwd=repo_path, check=False)
            subprocess.run(
                ["git", "commit", "--amend", "--no-edit"],
                cwd=repo_path, check=False,
            )
            subprocess.run(
                ["git", "push", "-u", "origin", branch_name, "--force"],
                cwd=repo_path, check=False,
            )

            # Show plot in Flyte UI via report
            await flyte.report.replace.aio(
                f"<h2>Autoresearch Progress</h2>"
                f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
                f'<p><a href="{pr.html_url}">View PR</a></p>'
            )
            await flyte.report.flush.aio()
        else:
            print("results.tsv found but no valid val_bpb rows — skipping plot.", flush=True)
    else:
        print("results.tsv not found — skipping plot.", flush=True)

    return AutoResearchResult(
        pr_url=pr.html_url,
        pr_number=pr.number,
        branch_name=branch_name,
        files_changed=files_changed,
        success=True,
    )

# {{docs-fragment main}}
if __name__ == "__main__":
    import time

    flyte.init_from_config()

    run = flyte.with_runcontext(mode="remote").run(run_autoresearch)

    print(f"AutoResearch run started: {run.url}")
    print("Waiting for completion...")

    while True:
        try:
            run.wait()
            break
        except Exception as e:
            print(f"Connection dropped ({e}), reconnecting in 30s...")
            time.sleep(30)

    print(f"Done! See run at: {run.url}")
# {{/docs-fragment main}}
CODE1
GITHUB_USERNAME = "<YOUR_GITHUB_USERNAME>"
GITHUB_EMAIL = "you@example.com"
AUTORESEARCH_REPO_URL = "https://github.com/<YOUR_ORG>/<YOUR_REPO>.git"
AUTORESEARCH_REPO_FULL_NAME = "<YOUR_ORG>/<YOUR_REPO>"
CODE2
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "flyte>=2.0.0b22",
#     "PyGithub>=2.5.0",
#     "matplotlib>=3.7.0",
# ]
# ///

"""
AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment.

This agent:
1. Starts a GPU-enabled container
2. Installs Claude Code CLI
3. Clones the autoresearch repository
4. Points Claude Code at program.md as the prompt and lets it run
5. Commits the result (CSV + code changes in train/) and creates a PR
"""

import os
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from github import Auth, Github

import flyte
import flyte.report
from _image import image as autoresearch_image

GITHUB_USERNAME = "parnianz"
GITHUB_EMAIL = "parnianzargham@gmail.com"
AUTORESEARCH_REPO_URL = "https://github.com/unionai-oss/autoresearch.git"
AUTORESEARCH_REPO_FULL_NAME = "unionai-oss/autoresearch"

# {{docs-fragment env}}
autoresearch_env = flyte.TaskEnvironment(
    name="autoresearch-agent",
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    secrets=[
        flyte.Secret(key="github_token", as_env_var="GITHUB_TOKEN"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=autoresearch_image,
)
# {{/docs-fragment env}}

# {{docs-fragment result}}
@dataclass
class AutoResearchResult:
    """Result of the autoresearch run."""

    pr_url: str
    pr_number: int
    branch_name: str
    files_changed: list[str]
    success: bool
    error_message: Optional[str] = None
# {{/docs-fragment result}}

def clone_repository(repo_url: str, work_dir: Path, github_token: str) -> Path:
    """Clone the autoresearch repository with authentication."""
    repo_name = repo_url.rstrip("/").split("/")[-1].replace(".git", "")
    repo_path = work_dir / repo_name

    # Inject token into HTTPS URL for authentication
    authenticated_url = repo_url.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )

    if repo_path.exists():
        subprocess.run(["git", "pull"], cwd=repo_path, check=True)
    else:
        subprocess.run(["git", "clone", authenticated_url, str(repo_path)], check=True)

    return repo_path

# {{docs-fragment task}}
@autoresearch_env.task(report=True)
async def run_autoresearch() -> AutoResearchResult:
    """
    Run the autoresearch workflow end-to-end.

    Steps:
    - Clone https://github.com/unionai-oss/autoresearch
    - Configure git identity
    - Create a new branch
    - Run Claude Code CLI with program.md as the prompt
    - Commit results (CSV + train/ changes)
    - Push and open a PR against the autoresearch repo
    """
    github_token = os.environ["GITHUB_TOKEN"]
    anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]

    # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) ---
    import tarfile
    import urllib.request as _urllib

    subprocess.run(["apt-get", "update", "-y"], check=False)
    subprocess.run(["apt-get", "install", "-y", "git"], check=False)

    node_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz"
    node_tar = Path("/tmp/node.tar.gz")
    print(f"Downloading Node.js from {node_url}...", flush=True)
    _urllib.urlretrieve(node_url, node_tar)
    size_mb = node_tar.stat().st_size / 1024 / 1024
    print(f"Downloaded {size_mb:.1f} MB to {node_tar}", flush=True)
    if size_mb < 1:
        raise RuntimeError(f"Node.js download appears empty/corrupt ({size_mb:.2f} MB) — network may be restricted")
    node_dir = Path("/tmp/node")
    node_dir.mkdir(exist_ok=True)
    print("Extracting Node.js...", flush=True)
    with tarfile.open(node_tar, "r:gz") as tar:
        members = [m for m in tar.getmembers() if m.name.split("/", 1)[-1]]
        for m in members:
            m.name = m.name.split("/", 1)[-1]
        tar.extractall(str(node_dir), members=[m for m in members if m.name])

    # Add node/npm to PATH for this process and all subprocesses
    node_bin = str(node_dir / "bin")
    os.environ["PATH"] = node_bin + ":" + os.environ.get("PATH", "")
    print(f"Node version: {subprocess.run(['node', '--version'], capture_output=True, text=True).stdout.strip()}", flush=True)

    npm_prefix = "/tmp/npm-global"
    Path(npm_prefix).mkdir(exist_ok=True)
    subprocess.run(["npm", "install", "-g", "--prefix", npm_prefix, "@anthropic-ai/claude-code"], check=True)
    os.environ["PATH"] = str(Path(npm_prefix) / "bin") + ":" + os.environ["PATH"]
    print("Node.js + Claude Code installed.", flush=True)

    # --- Clone repo ---
    work_dir = Path("/tmp/autoresearch_workspace")
    work_dir.mkdir(exist_ok=True, parents=True)
    repo_path = clone_repository(AUTORESEARCH_REPO_URL, work_dir, github_token)

    # --- Git identity ---
    subprocess.run(
        ["git", "config", "--global", "user.email", GITHUB_EMAIL], check=True
    )
    subprocess.run(
        ["git", "config", "--global", "user.name", GITHUB_USERNAME], check=True
    )

    # --- Create branch ---
    import time as _time
    branch_name = f"autoresearch/claude-run-{int(_time.time())}"
    try:
        subprocess.run(
            ["git", "checkout", "-b", branch_name],
            cwd=repo_path,
            check=True,
        )
    except subprocess.CalledProcessError:
        subprocess.run(
            ["git", "checkout", branch_name],
            cwd=repo_path,
            check=True,
        )

    # --- Read program.md to use as the Claude Code prompt ---
    program_md = repo_path / "program.md"
    if not program_md.exists():
        raise FileNotFoundError(
            f"program.md not found in {repo_path}. "
            "Make sure the autoresearch repo has a program.md at its root."
        )

    program_md_content = program_md.read_text()
    print(f"Loaded prompt from program.md ({len(program_md_content)} chars)")
    # {{/docs-fragment task}}

    # Install repo dependencies before handing off to Claude
    for pip_cmd in [
        ["pip", "install", "-e", "."],
        ["pip", "install", "-r", "requirements.txt"],
    ]:
        req_file = repo_path / pip_cmd[-1] if pip_cmd[-1].startswith("req") else None
        if req_file is None or req_file.exists():
            dep_result = subprocess.run(
                pip_cmd, cwd=repo_path, capture_output=True, text=True
            )
            print(f"{' '.join(pip_cmd)}:\n{dep_result.stdout}", flush=True)
            if dep_result.returncode != 0:
                print(f"(non-fatal) {dep_result.stderr}", flush=True)

    # Wrap the program.md content with explicit instructions to write outputs to disk
    prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files.

Here are your instructions from program.md:

{program_md_content}

LOGGING INSTRUCTIONS (follow exactly):
- Before you start any training, print this exact line: [AUTORESEARCH] Training started
- Before training, print what change you are testing: [AUTORESEARCH] Change: <one line description of the code change being tested>
- When training finishes, print this exact line: [AUTORESEARCH] Training finished
- After training, print the key metric value: [AUTORESEARCH] Metric: <metric name>=<value>
- When writing results to CSV, print this exact line: [AUTORESEARCH] Writing results to CSV

IMPORTANT: After completing the above instructions, make sure you have:
1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar)
2. Saved all code changes you made to the train/ directory (or wherever the training code lives)
3. All files must be written to the current working directory so they appear in git status
If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk."""

    # --- Pre-flight: verify claude is installed and API key is reachable ---
    version_check = subprocess.run(
        ["claude", "--version"], capture_output=True, text=True
    )
    print(f"claude version: {version_check.stdout.strip()} | stderr: {version_check.stderr.strip()}", flush=True)
    if version_check.returncode != 0:
        raise RuntimeError(f"claude CLI not found or broken: {version_check.stderr}")

    # --- Disable Claude Code sandbox ---
    # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container
    # which fails silently and causes file writes to go to an ephemeral space instead of
    # the real working directory. Disabling it makes writes land in the actual filesystem.
    claude_config_dir = Path("/root/.claude")
    claude_config_dir.mkdir(parents=True, exist_ok=True)
    settings = claude_config_dir / "settings.json"
    import json as _json
    existing = _json.loads(settings.read_text()) if settings.exists() else {}
    existing["sandbox"] = False
    settings.write_text(_json.dumps(existing, indent=2))
    print(f"Wrote Claude Code settings: {settings.read_text()}", flush=True)

    # --- Run Claude Code CLI ---
    # Matches swe_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode
    cmd = [
        "claude",
        "--dangerously-skip-permissions",
        "--max-turns", "100",
        "--model", "claude-haiku-4-5-20251001",
        prompt,
    ]

    print(f"Running: {shlex.join(cmd[:3])} <prompt>", flush=True)

    claude_env = {
        **os.environ,
        "ANTHROPIC_API_KEY": anthropic_api_key,
        "CLAUDE_SKIP_PERMISSIONS": "true",
        "CI": "true",  # Enables non-interactive mode (no TTY required)
    }

    # Stream output line by line so logs appear in real time instead of buffering until done
    proc = subprocess.Popen(
        cmd,
        cwd=repo_path,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # merge stderr into stdout stream
        text=True,
        env=claude_env,
    )

    stdout_lines = []
    for line in proc.stdout:
        line = line.rstrip("\n")
        print(line, flush=True)
        stdout_lines.append(line)

    proc.wait()
    full_output = "\n".join(stdout_lines)
    print(f"Claude Code exit code: {proc.returncode}", flush=True)

    if proc.returncode != 0:
        raise RuntimeError(
            f"Claude Code CLI exited with code {proc.returncode}\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Collect changed files ---
    git_status = subprocess.run(
        ["git", "status", "--porcelain"],
        cwd=repo_path,
        capture_output=True,
        text=True,
        check=True,
    )

    print(f"Git status:\n{git_status.stdout}", flush=True)

    files_changed = []
    for line in git_status.stdout.strip().splitlines():
        if line:
            # git status --porcelain: first two chars are XY status flags
            file_path = line[3:].strip()
            files_changed.append(file_path)

    # Also list all files in repo dir for debugging
    all_files = subprocess.run(
        ["find", ".", "-type", "f", "-not", "-path", "./.git/*"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"All files in repo:\n{all_files.stdout}", flush=True)

    if not files_changed:
        raise RuntimeError(
            "Claude Code ran successfully but produced no file changes.\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Commit ---
    subprocess.run(["git", "add", "."], cwd=repo_path, check=True)
    subprocess.run(["git", "add", "-f", "results.tsv"], cwd=repo_path, check=False)
    subprocess.run(["git", "add", "-f", "results/"], cwd=repo_path, check=False)
    commit_message = (
        "feat: autoresearch run via Claude Code\n\n"
        "Added research results (CSV) and updated train/ code changes.\n"
        "Generated by the autoresearch Flyte agent."
    )
    subprocess.run(
        ["git", "commit", "-m", commit_message],
        cwd=repo_path,
        check=True,
    )

    # --- Push ---
    print(f"GitHub token present: {bool(github_token)}, length: {len(github_token) if github_token else 0}", flush=True)
    authenticated_url = AUTORESEARCH_REPO_URL.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )
    subprocess.run(
        ["git", "remote", "set-url", "origin", authenticated_url],
        cwd=repo_path,
        check=True,
    )
    push_result = subprocess.run(
        ["git", "push", "-u", "origin", branch_name, "--force"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"Push stdout: {push_result.stdout}", flush=True)
    print(f"Push stderr: {push_result.stderr}", flush=True)
    if push_result.returncode != 0:
        raise RuntimeError(f"git push failed (exit {push_result.returncode}):\n{push_result.stderr}")

    # --- Create PR via PyGithub ---
    auth = Auth.Token(github_token)
    gh = Github(auth=auth)
    repo = gh.get_repo(AUTORESEARCH_REPO_FULL_NAME)

    csv_files = [f for f in files_changed if f.endswith(".csv")]
    train_files = [f for f in files_changed if "train" in f]

    pr_body = f"""## AutoResearch Run

This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI.

### What changed
- **Result CSV files**: {', '.join(f'`{f}`' for f in csv_files) or 'none detected'}
- **Train code changes**: {', '.join(f'`{f}`' for f in train_files) or 'none detected'}

### All changed files
{chr(10).join(f'- `{f}`' for f in files_changed)}

---
🤖 Generated by [autoresearch Flyte agent](https://github.com/unionai-oss/autoresearch)
"""

    existing_prs = list(repo.get_pulls(state="open", head=f"unionai-oss:{branch_name}"))
    if existing_prs:
        pr = existing_prs[0]
        print(f"PR already exists: {pr.html_url}", flush=True)
    else:
        pr = repo.create_pull(
            title="feat: autoresearch results + train changes",
            body=pr_body,
            head=branch_name,
            base="master",
        )
        print(f"PR created: {pr.html_url}", flush=True)

    # --- Generate progress plot from results.tsv ---
    plot_path = repo_path / "progress.png"
    results_tsv = repo_path / "results.tsv"
    if results_tsv.exists():
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import pandas as pd

        df = pd.read_csv(str(results_tsv), sep="\t")
        df["val_bpb"] = pd.to_numeric(df["val_bpb"], errors="coerce")
        df["memory_gb"] = pd.to_numeric(df["memory_gb"], errors="coerce")
        df["status"] = df["status"].str.strip().str.upper()

        # Filter out crashes for plotting
        valid = df[df["status"] != "CRASH"].copy()
        valid = valid.reset_index(drop=True)

        if len(valid) > 0 and valid["val_bpb"].notna().any():
            baseline_bpb = valid.loc[0, "val_bpb"]
            best = valid["val_bpb"].min()

            # Only plot points at or below baseline (the interesting region)
            below = valid[valid["val_bpb"] <= baseline_bpb + 0.0005]

            fig, ax = plt.subplots(figsize=(16, 8))

            # Plot discarded as faint background dots
            disc = below[below["status"] == "DISCARD"]
            ax.scatter(disc.index, disc["val_bpb"],
                       c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded")

            # Plot kept experiments as prominent green dots
            kept_v = below[below["status"] == "KEEP"]
            ax.scatter(kept_v.index, kept_v["val_bpb"],
                       c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5)

            # Running minimum step line
            kept_mask = valid["status"] == "KEEP"
            kept_idx = valid.index[kept_mask]
            kept_bpb = valid.loc[kept_mask, "val_bpb"]
            running_min = kept_bpb.cummin()
            ax.step(kept_idx, running_min, where="post", color="#27ae60",
                    linewidth=2, alpha=0.7, zorder=3, label="Running best")

            # Label each kept experiment with its description
            for idx, bpb in zip(kept_idx, kept_bpb):
                desc = str(valid.loc[idx, "description"]).strip()
                if len(desc) > 45:
                    desc = desc[:42] + "..."
                ax.annotate(desc, (idx, bpb),
                            textcoords="offset points",
                            xytext=(6, 6), fontsize=8.0,
                            color="#1a7a3a", alpha=0.9,
                            rotation=30, ha="left", va="bottom")

            n_total = len(df)
            n_kept = len(df[df["status"] == "KEEP"])
            ax.set_xlabel("Experiment #", fontsize=12)
            ax.set_ylabel("Validation BPB (lower is better)", fontsize=12)
            ax.set_title(f"Autoresearch Progress: {n_total} Experiments, {n_kept} Kept Improvements", fontsize=14)
            ax.legend(loc="upper right", fontsize=9)
            ax.grid(True, alpha=0.2)

            margin = (baseline_bpb - best) * 0.15
            ax.set_ylim(best - margin, baseline_bpb + margin)

            plt.tight_layout()
            plt.savefig(str(plot_path), dpi=150, bbox_inches="tight")
            plt.close(fig)
            print(f"Saved plot to {plot_path}", flush=True)

            # Upload plot to PR as a comment with base64 inline image
            import base64
            img_b64 = base64.b64encode(plot_path.read_bytes()).decode()
            pr_comment = (
                "## Autoresearch Progress\n\n"
                f"![Autoresearch Progress](data:image/png;base64,{img_b64})"
            )
            pr.create_issue_comment(pr_comment)
            print("Posted plot as PR comment.", flush=True)

            # Force-add plot to git and amend commit
            subprocess.run(["git", "add", "-f", str(plot_path)], cwd=repo_path, check=False)
            subprocess.run(
                ["git", "commit", "--amend", "--no-edit"],
                cwd=repo_path, check=False,
            )
            subprocess.run(
                ["git", "push", "-u", "origin", branch_name, "--force"],
                cwd=repo_path, check=False,
            )

            # Show plot in Flyte UI via report
            await flyte.report.replace.aio(
                f"<h2>Autoresearch Progress</h2>"
                f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
                f'<p><a href="{pr.html_url}">View PR</a></p>'
            )
            await flyte.report.flush.aio()
        else:
            print("results.tsv found but no valid val_bpb rows — skipping plot.", flush=True)
    else:
        print("results.tsv not found — skipping plot.", flush=True)

    return AutoResearchResult(
        pr_url=pr.html_url,
        pr_number=pr.number,
        branch_name=branch_name,
        files_changed=files_changed,
        success=True,
    )

# {{docs-fragment main}}
if __name__ == "__main__":
    import time

    flyte.init_from_config()

    run = flyte.with_runcontext(mode="remote").run(run_autoresearch)

    print(f"AutoResearch run started: {run.url}")
    print("Waiting for completion...")

    while True:
        try:
            run.wait()
            break
        except Exception as e:
            print(f"Connection dropped ({e}), reconnecting in 30s...")
            time.sleep(30)

    print(f"Done! See run at: {run.url}")
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py*

## The autoresearch task

The task is a long, sequential procedure. It starts by installing Node.js and Claude Code at run time, cloning the repo, configuring git, creating a branch, and loading `program.md` as the prompt:

```
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "flyte>=2.0.0b22",
#     "PyGithub>=2.5.0",
#     "matplotlib>=3.7.0",
# ]
# ///

"""
AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment.

This agent:
1. Starts a GPU-enabled container
2. Installs Claude Code CLI
3. Clones the autoresearch repository
4. Points Claude Code at program.md as the prompt and lets it run
5. Commits the result (CSV + code changes in train/) and creates a PR
"""

import os
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from github import Auth, Github

import flyte
import flyte.report
from _image import image as autoresearch_image

GITHUB_USERNAME = "parnianz"
GITHUB_EMAIL = "parnianzargham@gmail.com"
AUTORESEARCH_REPO_URL = "https://github.com/unionai-oss/autoresearch.git"
AUTORESEARCH_REPO_FULL_NAME = "unionai-oss/autoresearch"

# {{docs-fragment env}}
autoresearch_env = flyte.TaskEnvironment(
    name="autoresearch-agent",
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    secrets=[
        flyte.Secret(key="github_token", as_env_var="GITHUB_TOKEN"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=autoresearch_image,
)
# {{/docs-fragment env}}

# {{docs-fragment result}}
@dataclass
class AutoResearchResult:
    """Result of the autoresearch run."""

    pr_url: str
    pr_number: int
    branch_name: str
    files_changed: list[str]
    success: bool
    error_message: Optional[str] = None
# {{/docs-fragment result}}

def clone_repository(repo_url: str, work_dir: Path, github_token: str) -> Path:
    """Clone the autoresearch repository with authentication."""
    repo_name = repo_url.rstrip("/").split("/")[-1].replace(".git", "")
    repo_path = work_dir / repo_name

    # Inject token into HTTPS URL for authentication
    authenticated_url = repo_url.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )

    if repo_path.exists():
        subprocess.run(["git", "pull"], cwd=repo_path, check=True)
    else:
        subprocess.run(["git", "clone", authenticated_url, str(repo_path)], check=True)

    return repo_path

# {{docs-fragment task}}
@autoresearch_env.task(report=True)
async def run_autoresearch() -> AutoResearchResult:
    """
    Run the autoresearch workflow end-to-end.

    Steps:
    - Clone https://github.com/unionai-oss/autoresearch
    - Configure git identity
    - Create a new branch
    - Run Claude Code CLI with program.md as the prompt
    - Commit results (CSV + train/ changes)
    - Push and open a PR against the autoresearch repo
    """
    github_token = os.environ["GITHUB_TOKEN"]
    anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]

    # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) ---
    import tarfile
    import urllib.request as _urllib

    subprocess.run(["apt-get", "update", "-y"], check=False)
    subprocess.run(["apt-get", "install", "-y", "git"], check=False)

    node_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz"
    node_tar = Path("/tmp/node.tar.gz")
    print(f"Downloading Node.js from {node_url}...", flush=True)
    _urllib.urlretrieve(node_url, node_tar)
    size_mb = node_tar.stat().st_size / 1024 / 1024
    print(f"Downloaded {size_mb:.1f} MB to {node_tar}", flush=True)
    if size_mb < 1:
        raise RuntimeError(f"Node.js download appears empty/corrupt ({size_mb:.2f} MB) — network may be restricted")
    node_dir = Path("/tmp/node")
    node_dir.mkdir(exist_ok=True)
    print("Extracting Node.js...", flush=True)
    with tarfile.open(node_tar, "r:gz") as tar:
        members = [m for m in tar.getmembers() if m.name.split("/", 1)[-1]]
        for m in members:
            m.name = m.name.split("/", 1)[-1]
        tar.extractall(str(node_dir), members=[m for m in members if m.name])

    # Add node/npm to PATH for this process and all subprocesses
    node_bin = str(node_dir / "bin")
    os.environ["PATH"] = node_bin + ":" + os.environ.get("PATH", "")
    print(f"Node version: {subprocess.run(['node', '--version'], capture_output=True, text=True).stdout.strip()}", flush=True)

    npm_prefix = "/tmp/npm-global"
    Path(npm_prefix).mkdir(exist_ok=True)
    subprocess.run(["npm", "install", "-g", "--prefix", npm_prefix, "@anthropic-ai/claude-code"], check=True)
    os.environ["PATH"] = str(Path(npm_prefix) / "bin") + ":" + os.environ["PATH"]
    print("Node.js + Claude Code installed.", flush=True)

    # --- Clone repo ---
    work_dir = Path("/tmp/autoresearch_workspace")
    work_dir.mkdir(exist_ok=True, parents=True)
    repo_path = clone_repository(AUTORESEARCH_REPO_URL, work_dir, github_token)

    # --- Git identity ---
    subprocess.run(
        ["git", "config", "--global", "user.email", GITHUB_EMAIL], check=True
    )
    subprocess.run(
        ["git", "config", "--global", "user.name", GITHUB_USERNAME], check=True
    )

    # --- Create branch ---
    import time as _time
    branch_name = f"autoresearch/claude-run-{int(_time.time())}"
    try:
        subprocess.run(
            ["git", "checkout", "-b", branch_name],
            cwd=repo_path,
            check=True,
        )
    except subprocess.CalledProcessError:
        subprocess.run(
            ["git", "checkout", branch_name],
            cwd=repo_path,
            check=True,
        )

    # --- Read program.md to use as the Claude Code prompt ---
    program_md = repo_path / "program.md"
    if not program_md.exists():
        raise FileNotFoundError(
            f"program.md not found in {repo_path}. "
            "Make sure the autoresearch repo has a program.md at its root."
        )

    program_md_content = program_md.read_text()
    print(f"Loaded prompt from program.md ({len(program_md_content)} chars)")
    # {{/docs-fragment task}}

    # Install repo dependencies before handing off to Claude
    for pip_cmd in [
        ["pip", "install", "-e", "."],
        ["pip", "install", "-r", "requirements.txt"],
    ]:
        req_file = repo_path / pip_cmd[-1] if pip_cmd[-1].startswith("req") else None
        if req_file is None or req_file.exists():
            dep_result = subprocess.run(
                pip_cmd, cwd=repo_path, capture_output=True, text=True
            )
            print(f"{' '.join(pip_cmd)}:\n{dep_result.stdout}", flush=True)
            if dep_result.returncode != 0:
                print(f"(non-fatal) {dep_result.stderr}", flush=True)

    # Wrap the program.md content with explicit instructions to write outputs to disk
    prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files.

Here are your instructions from program.md:

{program_md_content}

LOGGING INSTRUCTIONS (follow exactly):
- Before you start any training, print this exact line: [AUTORESEARCH] Training started
- Before training, print what change you are testing: [AUTORESEARCH] Change: <one line description of the code change being tested>
- When training finishes, print this exact line: [AUTORESEARCH] Training finished
- After training, print the key metric value: [AUTORESEARCH] Metric: <metric name>=<value>
- When writing results to CSV, print this exact line: [AUTORESEARCH] Writing results to CSV

IMPORTANT: After completing the above instructions, make sure you have:
1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar)
2. Saved all code changes you made to the train/ directory (or wherever the training code lives)
3. All files must be written to the current working directory so they appear in git status
If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk."""

    # --- Pre-flight: verify claude is installed and API key is reachable ---
    version_check = subprocess.run(
        ["claude", "--version"], capture_output=True, text=True
    )
    print(f"claude version: {version_check.stdout.strip()} | stderr: {version_check.stderr.strip()}", flush=True)
    if version_check.returncode != 0:
        raise RuntimeError(f"claude CLI not found or broken: {version_check.stderr}")

    # --- Disable Claude Code sandbox ---
    # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container
    # which fails silently and causes file writes to go to an ephemeral space instead of
    # the real working directory. Disabling it makes writes land in the actual filesystem.
    claude_config_dir = Path("/root/.claude")
    claude_config_dir.mkdir(parents=True, exist_ok=True)
    settings = claude_config_dir / "settings.json"
    import json as _json
    existing = _json.loads(settings.read_text()) if settings.exists() else {}
    existing["sandbox"] = False
    settings.write_text(_json.dumps(existing, indent=2))
    print(f"Wrote Claude Code settings: {settings.read_text()}", flush=True)

    # --- Run Claude Code CLI ---
    # Matches swe_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode
    cmd = [
        "claude",
        "--dangerously-skip-permissions",
        "--max-turns", "100",
        "--model", "claude-haiku-4-5-20251001",
        prompt,
    ]

    print(f"Running: {shlex.join(cmd[:3])} <prompt>", flush=True)

    claude_env = {
        **os.environ,
        "ANTHROPIC_API_KEY": anthropic_api_key,
        "CLAUDE_SKIP_PERMISSIONS": "true",
        "CI": "true",  # Enables non-interactive mode (no TTY required)
    }

    # Stream output line by line so logs appear in real time instead of buffering until done
    proc = subprocess.Popen(
        cmd,
        cwd=repo_path,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # merge stderr into stdout stream
        text=True,
        env=claude_env,
    )

    stdout_lines = []
    for line in proc.stdout:
        line = line.rstrip("\n")
        print(line, flush=True)
        stdout_lines.append(line)

    proc.wait()
    full_output = "\n".join(stdout_lines)
    print(f"Claude Code exit code: {proc.returncode}", flush=True)

    if proc.returncode != 0:
        raise RuntimeError(
            f"Claude Code CLI exited with code {proc.returncode}\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Collect changed files ---
    git_status = subprocess.run(
        ["git", "status", "--porcelain"],
        cwd=repo_path,
        capture_output=True,
        text=True,
        check=True,
    )

    print(f"Git status:\n{git_status.stdout}", flush=True)

    files_changed = []
    for line in git_status.stdout.strip().splitlines():
        if line:
            # git status --porcelain: first two chars are XY status flags
            file_path = line[3:].strip()
            files_changed.append(file_path)

    # Also list all files in repo dir for debugging
    all_files = subprocess.run(
        ["find", ".", "-type", "f", "-not", "-path", "./.git/*"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"All files in repo:\n{all_files.stdout}", flush=True)

    if not files_changed:
        raise RuntimeError(
            "Claude Code ran successfully but produced no file changes.\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Commit ---
    subprocess.run(["git", "add", "."], cwd=repo_path, check=True)
    subprocess.run(["git", "add", "-f", "results.tsv"], cwd=repo_path, check=False)
    subprocess.run(["git", "add", "-f", "results/"], cwd=repo_path, check=False)
    commit_message = (
        "feat: autoresearch run via Claude Code\n\n"
        "Added research results (CSV) and updated train/ code changes.\n"
        "Generated by the autoresearch Flyte agent."
    )
    subprocess.run(
        ["git", "commit", "-m", commit_message],
        cwd=repo_path,
        check=True,
    )

    # --- Push ---
    print(f"GitHub token present: {bool(github_token)}, length: {len(github_token) if github_token else 0}", flush=True)
    authenticated_url = AUTORESEARCH_REPO_URL.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )
    subprocess.run(
        ["git", "remote", "set-url", "origin", authenticated_url],
        cwd=repo_path,
        check=True,
    )
    push_result = subprocess.run(
        ["git", "push", "-u", "origin", branch_name, "--force"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"Push stdout: {push_result.stdout}", flush=True)
    print(f"Push stderr: {push_result.stderr}", flush=True)
    if push_result.returncode != 0:
        raise RuntimeError(f"git push failed (exit {push_result.returncode}):\n{push_result.stderr}")

    # --- Create PR via PyGithub ---
    auth = Auth.Token(github_token)
    gh = Github(auth=auth)
    repo = gh.get_repo(AUTORESEARCH_REPO_FULL_NAME)

    csv_files = [f for f in files_changed if f.endswith(".csv")]
    train_files = [f for f in files_changed if "train" in f]

    pr_body = f"""## AutoResearch Run

This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI.

### What changed
- **Result CSV files**: {', '.join(f'`{f}`' for f in csv_files) or 'none detected'}
- **Train code changes**: {', '.join(f'`{f}`' for f in train_files) or 'none detected'}

### All changed files
{chr(10).join(f'- `{f}`' for f in files_changed)}

---
🤖 Generated by [autoresearch Flyte agent](https://github.com/unionai-oss/autoresearch)
"""

    existing_prs = list(repo.get_pulls(state="open", head=f"unionai-oss:{branch_name}"))
    if existing_prs:
        pr = existing_prs[0]
        print(f"PR already exists: {pr.html_url}", flush=True)
    else:
        pr = repo.create_pull(
            title="feat: autoresearch results + train changes",
            body=pr_body,
            head=branch_name,
            base="master",
        )
        print(f"PR created: {pr.html_url}", flush=True)

    # --- Generate progress plot from results.tsv ---
    plot_path = repo_path / "progress.png"
    results_tsv = repo_path / "results.tsv"
    if results_tsv.exists():
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import pandas as pd

        df = pd.read_csv(str(results_tsv), sep="\t")
        df["val_bpb"] = pd.to_numeric(df["val_bpb"], errors="coerce")
        df["memory_gb"] = pd.to_numeric(df["memory_gb"], errors="coerce")
        df["status"] = df["status"].str.strip().str.upper()

        # Filter out crashes for plotting
        valid = df[df["status"] != "CRASH"].copy()
        valid = valid.reset_index(drop=True)

        if len(valid) > 0 and valid["val_bpb"].notna().any():
            baseline_bpb = valid.loc[0, "val_bpb"]
            best = valid["val_bpb"].min()

            # Only plot points at or below baseline (the interesting region)
            below = valid[valid["val_bpb"] <= baseline_bpb + 0.0005]

            fig, ax = plt.subplots(figsize=(16, 8))

            # Plot discarded as faint background dots
            disc = below[below["status"] == "DISCARD"]
            ax.scatter(disc.index, disc["val_bpb"],
                       c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded")

            # Plot kept experiments as prominent green dots
            kept_v = below[below["status"] == "KEEP"]
            ax.scatter(kept_v.index, kept_v["val_bpb"],
                       c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5)

            # Running minimum step line
            kept_mask = valid["status"] == "KEEP"
            kept_idx = valid.index[kept_mask]
            kept_bpb = valid.loc[kept_mask, "val_bpb"]
            running_min = kept_bpb.cummin()
            ax.step(kept_idx, running_min, where="post", color="#27ae60",
                    linewidth=2, alpha=0.7, zorder=3, label="Running best")

            # Label each kept experiment with its description
            for idx, bpb in zip(kept_idx, kept_bpb):
                desc = str(valid.loc[idx, "description"]).strip()
                if len(desc) > 45:
                    desc = desc[:42] + "..."
                ax.annotate(desc, (idx, bpb),
                            textcoords="offset points",
                            xytext=(6, 6), fontsize=8.0,
                            color="#1a7a3a", alpha=0.9,
                            rotation=30, ha="left", va="bottom")

            n_total = len(df)
            n_kept = len(df[df["status"] == "KEEP"])
            ax.set_xlabel("Experiment #", fontsize=12)
            ax.set_ylabel("Validation BPB (lower is better)", fontsize=12)
            ax.set_title(f"Autoresearch Progress: {n_total} Experiments, {n_kept} Kept Improvements", fontsize=14)
            ax.legend(loc="upper right", fontsize=9)
            ax.grid(True, alpha=0.2)

            margin = (baseline_bpb - best) * 0.15
            ax.set_ylim(best - margin, baseline_bpb + margin)

            plt.tight_layout()
            plt.savefig(str(plot_path), dpi=150, bbox_inches="tight")
            plt.close(fig)
            print(f"Saved plot to {plot_path}", flush=True)

            # Upload plot to PR as a comment with base64 inline image
            import base64
            img_b64 = base64.b64encode(plot_path.read_bytes()).decode()
            pr_comment = (
                "## Autoresearch Progress\n\n"
                f"![Autoresearch Progress](data:image/png;base64,{img_b64})"
            )
            pr.create_issue_comment(pr_comment)
            print("Posted plot as PR comment.", flush=True)

            # Force-add plot to git and amend commit
            subprocess.run(["git", "add", "-f", str(plot_path)], cwd=repo_path, check=False)
            subprocess.run(
                ["git", "commit", "--amend", "--no-edit"],
                cwd=repo_path, check=False,
            )
            subprocess.run(
                ["git", "push", "-u", "origin", branch_name, "--force"],
                cwd=repo_path, check=False,
            )

            # Show plot in Flyte UI via report
            await flyte.report.replace.aio(
                f"<h2>Autoresearch Progress</h2>"
                f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
                f'<p><a href="{pr.html_url}">View PR</a></p>'
            )
            await flyte.report.flush.aio()
        else:
            print("results.tsv found but no valid val_bpb rows — skipping plot.", flush=True)
    else:
        print("results.tsv not found — skipping plot.", flush=True)

    return AutoResearchResult(
        pr_url=pr.html_url,
        pr_number=pr.number,
        branch_name=branch_name,
        files_changed=files_changed,
        success=True,
    )

# {{docs-fragment main}}
if __name__ == "__main__":
    import time

    flyte.init_from_config()

    run = flyte.with_runcontext(mode="remote").run(run_autoresearch)

    print(f"AutoResearch run started: {run.url}")
    print("Waiting for completion...")

    while True:
        try:
            run.wait()
            break
        except Exception as e:
            print(f"Connection dropped ({e}), reconnecting in 30s...")
            time.sleep(30)

    print(f"Done! See run at: {run.url}")
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py*

From there the task:

1. Wraps the `program.md` brief with explicit logging and "write outputs to disk" instructions.
2. Disables the Claude Code sandbox (it conflicts with the Flyte pod's container) and runs the CLI non-interactively, streaming its output to the Flyte logs in real time.
3. Collects the files the agent changed via `git status`, commits them, and force-pushes the branch.
4. Opens (or reuses) a pull request with [PyGithub](https://pygithub.readthedocs.io/).
5. If the agent produced a `results.tsv`, renders a progress plot of validation bits-per-byte, attaches it to the PR, and streams it into the Flyte UI:

```
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "flyte>=2.0.0b22",
#     "PyGithub>=2.5.0",
#     "matplotlib>=3.7.0",
# ]
# ///

"""
AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment.

This agent:
1. Starts a GPU-enabled container
2. Installs Claude Code CLI
3. Clones the autoresearch repository
4. Points Claude Code at program.md as the prompt and lets it run
5. Commits the result (CSV + code changes in train/) and creates a PR
"""

import os
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from github import Auth, Github

import flyte
import flyte.report
from _image import image as autoresearch_image

GITHUB_USERNAME = "parnianz"
GITHUB_EMAIL = "parnianzargham@gmail.com"
AUTORESEARCH_REPO_URL = "https://github.com/unionai-oss/autoresearch.git"
AUTORESEARCH_REPO_FULL_NAME = "unionai-oss/autoresearch"

# {{docs-fragment env}}
autoresearch_env = flyte.TaskEnvironment(
    name="autoresearch-agent",
    resources=flyte.Resources(
        cpu=8,
        memory="32Gi",
        gpu="T4:1",
        disk="100Gi",
    ),
    secrets=[
        flyte.Secret(key="github_token", as_env_var="GITHUB_TOKEN"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=autoresearch_image,
)
# {{/docs-fragment env}}

# {{docs-fragment result}}
@dataclass
class AutoResearchResult:
    """Result of the autoresearch run."""

    pr_url: str
    pr_number: int
    branch_name: str
    files_changed: list[str]
    success: bool
    error_message: Optional[str] = None
# {{/docs-fragment result}}

def clone_repository(repo_url: str, work_dir: Path, github_token: str) -> Path:
    """Clone the autoresearch repository with authentication."""
    repo_name = repo_url.rstrip("/").split("/")[-1].replace(".git", "")
    repo_path = work_dir / repo_name

    # Inject token into HTTPS URL for authentication
    authenticated_url = repo_url.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )

    if repo_path.exists():
        subprocess.run(["git", "pull"], cwd=repo_path, check=True)
    else:
        subprocess.run(["git", "clone", authenticated_url, str(repo_path)], check=True)

    return repo_path

# {{docs-fragment task}}
@autoresearch_env.task(report=True)
async def run_autoresearch() -> AutoResearchResult:
    """
    Run the autoresearch workflow end-to-end.

    Steps:
    - Clone https://github.com/unionai-oss/autoresearch
    - Configure git identity
    - Create a new branch
    - Run Claude Code CLI with program.md as the prompt
    - Commit results (CSV + train/ changes)
    - Push and open a PR against the autoresearch repo
    """
    github_token = os.environ["GITHUB_TOKEN"]
    anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]

    # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) ---
    import tarfile
    import urllib.request as _urllib

    subprocess.run(["apt-get", "update", "-y"], check=False)
    subprocess.run(["apt-get", "install", "-y", "git"], check=False)

    node_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz"
    node_tar = Path("/tmp/node.tar.gz")
    print(f"Downloading Node.js from {node_url}...", flush=True)
    _urllib.urlretrieve(node_url, node_tar)
    size_mb = node_tar.stat().st_size / 1024 / 1024
    print(f"Downloaded {size_mb:.1f} MB to {node_tar}", flush=True)
    if size_mb < 1:
        raise RuntimeError(f"Node.js download appears empty/corrupt ({size_mb:.2f} MB) — network may be restricted")
    node_dir = Path("/tmp/node")
    node_dir.mkdir(exist_ok=True)
    print("Extracting Node.js...", flush=True)
    with tarfile.open(node_tar, "r:gz") as tar:
        members = [m for m in tar.getmembers() if m.name.split("/", 1)[-1]]
        for m in members:
            m.name = m.name.split("/", 1)[-1]
        tar.extractall(str(node_dir), members=[m for m in members if m.name])

    # Add node/npm to PATH for this process and all subprocesses
    node_bin = str(node_dir / "bin")
    os.environ["PATH"] = node_bin + ":" + os.environ.get("PATH", "")
    print(f"Node version: {subprocess.run(['node', '--version'], capture_output=True, text=True).stdout.strip()}", flush=True)

    npm_prefix = "/tmp/npm-global"
    Path(npm_prefix).mkdir(exist_ok=True)
    subprocess.run(["npm", "install", "-g", "--prefix", npm_prefix, "@anthropic-ai/claude-code"], check=True)
    os.environ["PATH"] = str(Path(npm_prefix) / "bin") + ":" + os.environ["PATH"]
    print("Node.js + Claude Code installed.", flush=True)

    # --- Clone repo ---
    work_dir = Path("/tmp/autoresearch_workspace")
    work_dir.mkdir(exist_ok=True, parents=True)
    repo_path = clone_repository(AUTORESEARCH_REPO_URL, work_dir, github_token)

    # --- Git identity ---
    subprocess.run(
        ["git", "config", "--global", "user.email", GITHUB_EMAIL], check=True
    )
    subprocess.run(
        ["git", "config", "--global", "user.name", GITHUB_USERNAME], check=True
    )

    # --- Create branch ---
    import time as _time
    branch_name = f"autoresearch/claude-run-{int(_time.time())}"
    try:
        subprocess.run(
            ["git", "checkout", "-b", branch_name],
            cwd=repo_path,
            check=True,
        )
    except subprocess.CalledProcessError:
        subprocess.run(
            ["git", "checkout", branch_name],
            cwd=repo_path,
            check=True,
        )

    # --- Read program.md to use as the Claude Code prompt ---
    program_md = repo_path / "program.md"
    if not program_md.exists():
        raise FileNotFoundError(
            f"program.md not found in {repo_path}. "
            "Make sure the autoresearch repo has a program.md at its root."
        )

    program_md_content = program_md.read_text()
    print(f"Loaded prompt from program.md ({len(program_md_content)} chars)")
    # {{/docs-fragment task}}

    # Install repo dependencies before handing off to Claude
    for pip_cmd in [
        ["pip", "install", "-e", "."],
        ["pip", "install", "-r", "requirements.txt"],
    ]:
        req_file = repo_path / pip_cmd[-1] if pip_cmd[-1].startswith("req") else None
        if req_file is None or req_file.exists():
            dep_result = subprocess.run(
                pip_cmd, cwd=repo_path, capture_output=True, text=True
            )
            print(f"{' '.join(pip_cmd)}:\n{dep_result.stdout}", flush=True)
            if dep_result.returncode != 0:
                print(f"(non-fatal) {dep_result.stderr}", flush=True)

    # Wrap the program.md content with explicit instructions to write outputs to disk
    prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files.

Here are your instructions from program.md:

{program_md_content}

LOGGING INSTRUCTIONS (follow exactly):
- Before you start any training, print this exact line: [AUTORESEARCH] Training started
- Before training, print what change you are testing: [AUTORESEARCH] Change: <one line description of the code change being tested>
- When training finishes, print this exact line: [AUTORESEARCH] Training finished
- After training, print the key metric value: [AUTORESEARCH] Metric: <metric name>=<value>
- When writing results to CSV, print this exact line: [AUTORESEARCH] Writing results to CSV

IMPORTANT: After completing the above instructions, make sure you have:
1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar)
2. Saved all code changes you made to the train/ directory (or wherever the training code lives)
3. All files must be written to the current working directory so they appear in git status
If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk."""

    # --- Pre-flight: verify claude is installed and API key is reachable ---
    version_check = subprocess.run(
        ["claude", "--version"], capture_output=True, text=True
    )
    print(f"claude version: {version_check.stdout.strip()} | stderr: {version_check.stderr.strip()}", flush=True)
    if version_check.returncode != 0:
        raise RuntimeError(f"claude CLI not found or broken: {version_check.stderr}")

    # --- Disable Claude Code sandbox ---
    # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container
    # which fails silently and causes file writes to go to an ephemeral space instead of
    # the real working directory. Disabling it makes writes land in the actual filesystem.
    claude_config_dir = Path("/root/.claude")
    claude_config_dir.mkdir(parents=True, exist_ok=True)
    settings = claude_config_dir / "settings.json"
    import json as _json
    existing = _json.loads(settings.read_text()) if settings.exists() else {}
    existing["sandbox"] = False
    settings.write_text(_json.dumps(existing, indent=2))
    print(f"Wrote Claude Code settings: {settings.read_text()}", flush=True)

    # --- Run Claude Code CLI ---
    # Matches swe_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode
    cmd = [
        "claude",
        "--dangerously-skip-permissions",
        "--max-turns", "100",
        "--model", "claude-haiku-4-5-20251001",
        prompt,
    ]

    print(f"Running: {shlex.join(cmd[:3])} <prompt>", flush=True)

    claude_env = {
        **os.environ,
        "ANTHROPIC_API_KEY": anthropic_api_key,
        "CLAUDE_SKIP_PERMISSIONS": "true",
        "CI": "true",  # Enables non-interactive mode (no TTY required)
    }

    # Stream output line by line so logs appear in real time instead of buffering until done
    proc = subprocess.Popen(
        cmd,
        cwd=repo_path,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # merge stderr into stdout stream
        text=True,
        env=claude_env,
    )

    stdout_lines = []
    for line in proc.stdout:
        line = line.rstrip("\n")
        print(line, flush=True)
        stdout_lines.append(line)

    proc.wait()
    full_output = "\n".join(stdout_lines)
    print(f"Claude Code exit code: {proc.returncode}", flush=True)

    if proc.returncode != 0:
        raise RuntimeError(
            f"Claude Code CLI exited with code {proc.returncode}\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Collect changed files ---
    git_status = subprocess.run(
        ["git", "status", "--porcelain"],
        cwd=repo_path,
        capture_output=True,
        text=True,
        check=True,
    )

    print(f"Git status:\n{git_status.stdout}", flush=True)

    files_changed = []
    for line in git_status.stdout.strip().splitlines():
        if line:
            # git status --porcelain: first two chars are XY status flags
            file_path = line[3:].strip()
            files_changed.append(file_path)

    # Also list all files in repo dir for debugging
    all_files = subprocess.run(
        ["find", ".", "-type", "f", "-not", "-path", "./.git/*"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"All files in repo:\n{all_files.stdout}", flush=True)

    if not files_changed:
        raise RuntimeError(
            "Claude Code ran successfully but produced no file changes.\n"
            f"output: {full_output[-2000:]}"
        )

    # --- Commit ---
    subprocess.run(["git", "add", "."], cwd=repo_path, check=True)
    subprocess.run(["git", "add", "-f", "results.tsv"], cwd=repo_path, check=False)
    subprocess.run(["git", "add", "-f", "results/"], cwd=repo_path, check=False)
    commit_message = (
        "feat: autoresearch run via Claude Code\n\n"
        "Added research results (CSV) and updated train/ code changes.\n"
        "Generated by the autoresearch Flyte agent."
    )
    subprocess.run(
        ["git", "commit", "-m", commit_message],
        cwd=repo_path,
        check=True,
    )

    # --- Push ---
    print(f"GitHub token present: {bool(github_token)}, length: {len(github_token) if github_token else 0}", flush=True)
    authenticated_url = AUTORESEARCH_REPO_URL.replace(
        "https://", f"https://{GITHUB_USERNAME}:{github_token}@"
    )
    subprocess.run(
        ["git", "remote", "set-url", "origin", authenticated_url],
        cwd=repo_path,
        check=True,
    )
    push_result = subprocess.run(
        ["git", "push", "-u", "origin", branch_name, "--force"],
        cwd=repo_path,
        capture_output=True,
        text=True,
    )
    print(f"Push stdout: {push_result.stdout}", flush=True)
    print(f"Push stderr: {push_result.stderr}", flush=True)
    if push_result.returncode != 0:
        raise RuntimeError(f"git push failed (exit {push_result.returncode}):\n{push_result.stderr}")

    # --- Create PR via PyGithub ---
    auth = Auth.Token(github_token)
    gh = Github(auth=auth)
    repo = gh.get_repo(AUTORESEARCH_REPO_FULL_NAME)

    csv_files = [f for f in files_changed if f.endswith(".csv")]
    train_files = [f for f in files_changed if "train" in f]

    pr_body = f"""## AutoResearch Run

This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI.

### What changed
- **Result CSV files**: {', '.join(f'`{f}`' for f in csv_files) or 'none detected'}
- **Train code changes**: {', '.join(f'`{f}`' for f in train_files) or 'none detected'}

### All changed files
{chr(10).join(f'- `{f}`' for f in files_changed)}

---
🤖 Generated by [autoresearch Flyte agent](https://github.com/unionai-oss/autoresearch)
"""

    existing_prs = list(repo.get_pulls(state="open", head=f"unionai-oss:{branch_name}"))
    if existing_prs:
        pr = existing_prs[0]
        print(f"PR already exists: {pr.html_url}", flush=True)
    else:
        pr = repo.create_pull(
            title="feat: autoresearch results + train changes",
            body=pr_body,
            head=branch_name,
            base="master",
        )
        print(f"PR created: {pr.html_url}", flush=True)

    # --- Generate progress plot from results.tsv ---
    plot_path = repo_path / "progress.png"
    results_tsv = repo_path / "results.tsv"
    if results_tsv.exists():
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import pandas as pd

        df = pd.read_csv(str(results_tsv), sep="\t")
        df["val_bpb"] = pd.to_numeric(df["val_bpb"], errors="coerce")
        df["memory_gb"] = pd.to_numeric(df["memory_gb"], errors="coerce")
        df["status"] = df["status"].str.strip().str.upper()

        # Filter out crashes for plotting
        valid = df[df["status"] != "CRASH"].copy()
        valid = valid.reset_index(drop=True)

        if len(valid) > 0 and valid["val_bpb"].notna().any():
            baseline_bpb = valid.loc[0, "val_bpb"]
            best = valid["val_bpb"].min()

            # Only plot points at or below baseline (the interesting region)
            below = valid[valid["val_bpb"] <= baseline_bpb + 0.0005]

            fig, ax = plt.subplots(figsize=(16, 8))

            # Plot discarded as faint background dots
            disc = below[below["status"] == "DISCARD"]
            ax.scatter(disc.index, disc["val_bpb"],
                       c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded")

            # Plot kept experiments as prominent green dots
            kept_v = below[below["status"] == "KEEP"]
            ax.scatter(kept_v.index, kept_v["val_bpb"],
                       c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5)

            # Running minimum step line
            kept_mask = valid["status"] == "KEEP"
            kept_idx = valid.index[kept_mask]
            kept_bpb = valid.loc[kept_mask, "val_bpb"]
            running_min = kept_bpb.cummin()
            ax.step(kept_idx, running_min, where="post", color="#27ae60",
                    linewidth=2, alpha=0.7, zorder=3, label="Running best")

            # Label each kept experiment with its description
            for idx, bpb in zip(kept_idx, kept_bpb):
                desc = str(valid.loc[idx, "description"]).strip()
                if len(desc) > 45:
                    desc = desc[:42] + "..."
                ax.annotate(desc, (idx, bpb),
                            textcoords="offset points",
                            xytext=(6, 6), fontsize=8.0,
                            color="#1a7a3a", alpha=0.9,
                            rotation=30, ha="left", va="bottom")

            n_total = len(df)
            n_kept = len(df[df["status"] == "KEEP"])
            ax.set_xlabel("Experiment #", fontsize=12)
            ax.set_ylabel("Validation BPB (lower is better)", fontsize=12)
            ax.set_title(f"Autoresearch Progress: {n_total} Experiments, {n_kept} Kept Improvements", fontsize=14)
            ax.legend(loc="upper right", fontsize=9)
            ax.grid(True, alpha=0.2)

            margin = (baseline_bpb - best) * 0.15
            ax.set_ylim(best - margin, baseline_bpb + margin)

            plt.tight_layout()
            plt.savefig(str(plot_path), dpi=150, bbox_inches="tight")
            plt.close(fig)
            print(f"Saved plot to {plot_path}", flush=True)

            # Upload plot to PR as a comment with base64 inline image
            import base64
            img_b64 = base64.b64encode(plot_path.read_bytes()).decode()
            pr_comment = (
                "## Autoresearch Progress\n\n"
                f"![Autoresearch Progress](data:image/png;base64,{img_b64})"
            )
            pr.create_issue_comment(pr_comment)
            print("Posted plot as PR comment.", flush=True)

            # Force-add plot to git and amend commit
            subprocess.run(["git", "add", "-f", str(plot_path)], cwd=repo_path, check=False)
            subprocess.run(
                ["git", "commit", "--amend", "--no-edit"],
                cwd=repo_path, check=False,
            )
            subprocess.run(
                ["git", "push", "-u", "origin", branch_name, "--force"],
                cwd=repo_path, check=False,
            )

            # Show plot in Flyte UI via report
            await flyte.report.replace.aio(
                f"<h2>Autoresearch Progress</h2>"
                f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
                f'<p><a href="{pr.html_url}">View PR</a></p>'
            )
            await flyte.report.flush.aio()
        else:
            print("results.tsv found but no valid val_bpb rows — skipping plot.", flush=True)
    else:
        print("results.tsv not found — skipping plot.", flush=True)

    return AutoResearchResult(
        pr_url=pr.html_url,
        pr_number=pr.number,
        branch_name=branch_name,
        files_changed=files_changed,
        success=True,
    )

# {{docs-fragment main}}
if __name__ == "__main__":
    import time

    flyte.init_from_config()

    run = flyte.with_runcontext(mode="remote").run(run_autoresearch)

    print(f"AutoResearch run started: {run.url}")
    print("Waiting for completion...")

    while True:
        try:
            run.wait()
            break
        except Exception as e:
            print(f"Connection dropped ({e}), reconnecting in 30s...")
            time.sleep(30)

    print(f"Done! See run at: {run.url}")
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py*

The entry point submits the task in `remote` mode and reconnects automatically if the client connection drops during the long run.

## Run the agent

### Create secrets

Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/) and a [GitHub personal access token](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens) with permission to push and open PRs on the target repository.

Register both as Flyte secrets. The key names must match those declared in the `TaskEnvironment`:

CODE3

See [Secrets](../../../user-guide/tasks/task-configuration/secrets/page.md) for scoping and file-based secrets.

### Prepare the research repository

The target repository must contain a `program.md` at its root describing the research task for the agent. Point `AUTORESEARCH_REPO_URL` / `AUTORESEARCH_REPO_FULL_NAME` (and the git identity constants) at a repo you control.

### Run remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autoresearch):

CODE4

This task runs remotely (it needs a GPU and network access). Follow the printed run URL to watch the agent's logs stream in, and open the run's report panel to see the progress plot once results are available. When the task finishes, the returned `AutoResearchResult` contains the pull request URL.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/parallelized-autoresearch-agent ===

# Parallelized autoresearch agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/parallelized_autoresearch).

This tutorial extends the [Autoresearch agent](../autoresearch/_index) pattern with a code-mode MLE agent that plans **batches** of training experiments, saves distinct `train.py` edits, and runs them **in parallel** via `flyte.map`. It follows the [karpathy/autoresearch](https://github.com/karpathy/autoresearch) loop (minimize validation bits-per-byte on a TinyGPT variant) but orchestrates fan-out batches with durable Flyte tasks and [unionai-sandbox](../../../user-guide/agents/sandboxing/_index) execution.

Compared to the single-threaded Claude Code autoresearch tutorial, this agent:

- Edits full `train.py` source (upstream karpathy style) instead of calling a remote coding CLI
- Uses **`code_mode=True`** so the LLM writes Python plans that call batch tools such as `run_experiment_batch`
- Persists a **leaderboard**, code-edit history, and batch plans in `MemoryStore`
- **Right-sizes each experiment** with an LLM via a `@tool` **`call_handler`**, then retries on Flyte or sandbox OOM by bumping memory

Each experiment has different compute needs (wider models, larger batch sizes, longer training loops). A single static `flyte.Resources` on the task would either waste cluster memory or OOM on the heavy configs. Instead, this example uses the same **Agents > Build an agent > Flyte-native agents** as the Flyte SDK self-correcting agent: before every run, a sizing LLM reads the tool name, docstring, and call arguments and returns a JSON resource spec; the handler applies it with `tool_fn.target.override(resources=...).aio(**kwargs)` and retries with more memory when needed.

## Define the task environments

The example uses three environments (bundle preparation, sandbox experiments, and the agent driver) sharing a Debian-based image with PyTorch and sandbox tooling.

```
"""Shared Flyte environments and climbmix dataset bundle tasks."""

from __future__ import annotations

import os
import tempfile
from dataclasses import dataclass
from pathlib import Path

import flyte
from flyte.io import Dir

from autoresearch_types import DatasetProfile
from autoresearch_types import DEFAULT_NUM_SHARDS

TRAIN_PIP_PACKAGES = ["torch", "numpy", "pyarrow", "requests", "tiktoken", "rustbpe"]

_TUTORIAL_DIR = Path(__file__).parent
_INCLUDE = [str(p) for p in sorted(_TUTORIAL_DIR.glob("*.py"))]

image = flyte.Image.from_debian_base(name="mle-autoresearch").with_pip_packages(
    "litellm",
    "httpx",
    "pydantic-monty",
    "unionai-sandbox[flyte]",
    *TRAIN_PIP_PACKAGES,
)

bundle_env = flyte.TaskEnvironment(
    name="autoresearch-bundle",
    resources=flyte.Resources(cpu=4, memory="8Gi"),
    image=image,
    include=_INCLUDE,
)

experiment_env = flyte.TaskEnvironment(
    name="autoresearch-experiment",
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    image=image,
    include=_INCLUDE,
    secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
)

# {{docs-fragment env}}
agent_env = flyte.TaskEnvironment(
    name="autoresearch-agent",
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    image=image,
    include=_INCLUDE,
    secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
    depends_on=[experiment_env, bundle_env],
)
# {{/docs-fragment env}}

@dataclass
class AutoresearchBundle:
    data_dir: Dir
    tokenizer_dir: Dir

@bundle_env.task(cache="auto")
async def build_bundle(num_shards: int = DEFAULT_NUM_SHARDS, download_workers: int = 4) -> AutoresearchBundle:
    """Download climbmix shards + train the BPE tokenizer; cache the result."""
    import prepare

    cache = tempfile.mkdtemp(prefix="autoresearch-cache-")
    os.environ["AUTORESEARCH_CACHE"] = cache
    prepare.download_data(num_shards, download_workers=download_workers)
    prepare.train_tokenizer()
    data_dir = await Dir.from_local(prepare.data_dir())
    tokenizer_dir = await Dir.from_local(prepare.tokenizer_dir())
    return AutoresearchBundle(data_dir=data_dir, tokenizer_dir=tokenizer_dir)

@bundle_env.task(cache="auto")
async def profile_bundle(bundle: AutoresearchBundle) -> DatasetProfile:
    """Summarize the prepared bundle for the agent's context."""
    import prepare

    data_dir = await bundle.data_dir.download()
    tokenizer_dir = await bundle.tokenizer_dir.download()
    parquet_files = sorted(p.name for p in Path(data_dir).glob("*.parquet"))
    data_bytes = sum(p.stat().st_size for p in Path(data_dir).glob("**/*") if p.is_file())
    tok_bytes = sum(p.stat().st_size for p in Path(tokenizer_dir).glob("**/*") if p.is_file())
    return DatasetProfile(
        n_parquet_files=len(parquet_files),
        parquet_files=parquet_files,
        vocab_size=prepare.VOCAB_SIZE,
        data_bytes=data_bytes,
        tokenizer_bytes=tok_bytes,
    )

async def materialize_cache(bundle: AutoresearchBundle) -> str:
    """Download the bundle into an AUTORESEARCH_CACHE-shaped scratch dir."""
    cache = tempfile.mkdtemp(prefix="autoresearch-run-")
    os.environ["AUTORESEARCH_CACHE"] = cache
    await bundle.data_dir.download(local_path=os.path.join(cache, "data"))
    await bundle.tokenizer_dir.download(local_path=os.path.join(cache, "tokenizer"))
    return cache
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/parallelized_autoresearch/bundle.py*

Supporting modules (`train.py`, `prepare.py`, `tools.py`, and `ui.py`) live alongside the entry point in the example directory.

## Right-size experiments with `call_handler`

The right-sizing logic lives in `tools.py`. `execute_with_right_sizing` asks the LLM for a resource estimate, runs the underlying `@env.task` with `override(resources=...)`, and loops on `flyte.errors.OOMError` or a sandbox-reported OOM flag until the run succeeds or retries are exhausted:

```
"""Agent tools, sandbox execution, and memory helpers for parallelized autoresearch."""

from __future__ import annotations

import asyncio
import dataclasses
import hashlib
import json
import re
import textwrap
import xml.etree.ElementTree as ET
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import flyte
import flyte.errors
from flyte.ai.agents import LLMCallable, LLMMessage, MemoryStore, ToolFn, tool
from flyte.ai.agents._llm import _default_call_llm

from autoresearch_types import (
    CONFIG_ONLY_EDIT_LIMIT,
    DEFAULT_NUM_SHARDS,
    DatasetProfile,
    ExperimentConfig,
    HypothesisEntry,
    MAX_DEVICE_BATCH_SIZE,
    MAX_MAX_STEPS,
    MAX_N_EMBD,
    MAX_N_HEAD,
    MAX_N_LAYER,
)
from bundle import agent_env, build_bundle, bundle_env, profile_bundle

MEMORY_KEY_FANOUT = "parallelized-autoresearch"

MAX_LLM_RETRIES = 5
INITIAL_BACKOFF_SEC = 2.0

async def call_llm(
    model: str,
    system: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]] | None,
) -> LLMMessage:
    """Call litellm via the Flyte default callback, retrying transient provider errors."""
    import litellm

    backoff = INITIAL_BACKOFF_SEC
    last_exc: Exception | None = None
    for attempt in range(MAX_LLM_RETRIES):
        try:
            return await _default_call_llm(model, system, messages, tools)
        except litellm.InternalServerError as exc:
            last_exc = exc
            if attempt >= MAX_LLM_RETRIES - 1:
                break
            flyte.logger.warning(
                "LLM InternalServerError (attempt %d/%d); retrying in %.1fs: %s",
                attempt + 1,
                MAX_LLM_RETRIES,
                backoff,
                exc,
            )
            await asyncio.sleep(backoff)
            backoff *= 2
    assert last_exc is not None
    raise last_exc

RESOURCE_FLOOR = flyte.Resources(cpu=2, memory="2Gi")
RESOURCE_CEILING = flyte.Resources(cpu=16, memory="32Gi")
_MEM_RE = re.compile(r"^\s*([0-9]*\.?[0-9]+)\s*([A-Za-z]+)?\s*$")

def _memory_to_mib(memory: str | None) -> int:
    if not memory:
        return 2048
    match = _MEM_RE.match(memory)
    if not match:
        return 2048
    value = float(match.group(1))
    unit = (match.group(2) or "Mi").lower()
    if unit in ("gi", "g", "gb"):
        return int(value * 1024)
    if unit in ("mi", "m", "mb"):
        return int(value)
    if unit in ("ki", "k", "kb"):
        return max(1, int(value // 1024))
    return int(value)

def _mib_to_memory(mib: int) -> str:
    if mib >= 1024 and mib % 1024 == 0:
        return f"{mib // 1024}Gi"
    return f"{mib}Mi"

def _cap_resources(resources: flyte.Resources) -> flyte.Resources:
    floor_cpu = int(RESOURCE_FLOOR.cpu or 2)
    ceil_cpu = int(RESOURCE_CEILING.cpu or 16)
    cpu = int(resources.cpu or floor_cpu)
    cpu = max(floor_cpu, min(ceil_cpu, cpu))

    floor_mib = _memory_to_mib(
        RESOURCE_FLOOR.memory if isinstance(RESOURCE_FLOOR.memory, str) else "2Gi"
    )
    ceil_mib = _memory_to_mib(
        RESOURCE_CEILING.memory if isinstance(RESOURCE_CEILING.memory, str) else "32Gi"
    )
    mem_mib = _memory_to_mib(resources.memory if isinstance(resources.memory, str) else None)
    mem_mib = max(floor_mib, min(ceil_mib, mem_mib))
    return flyte.Resources(cpu=cpu, memory=_mib_to_memory(mem_mib))

def _ensure_oom_increase(resources: flyte.Resources, previous: flyte.Resources) -> flyte.Resources:
    """If memory did not grow after OOM, bump deterministically up to the ceiling."""
    prev_mib = _memory_to_mib(previous.memory if isinstance(previous.memory, str) else None)
    new_mib = _memory_to_mib(resources.memory if isinstance(resources.memory, str) else None)
    if new_mib <= prev_mib:
        ceil_mib = _memory_to_mib(
            RESOURCE_CEILING.memory if isinstance(RESOURCE_CEILING.memory, str) else "32Gi"
        )
        new_mib = min(ceil_mib, max(prev_mib * 2, prev_mib + 2048))
        resources = dataclasses.replace(resources, memory=_mib_to_memory(new_mib))
    prev_cpu = int(previous.cpu or RESOURCE_FLOOR.cpu or 2)
    new_cpu = int(resources.cpu or prev_cpu)
    if new_cpu < prev_cpu:
        resources = dataclasses.replace(resources, cpu=prev_cpu)
    return _cap_resources(resources)

def bump_memory(resources: flyte.Resources) -> flyte.Resources:
    """Deterministic memory bump after OOM."""
    return _ensure_oom_increase(resources, resources)

MAX_OOM_RETRIES = 3

RESOURCE_SIZING_SYSTEM_PROMPT = """\
You are a Kubernetes capacity planner for Flyte autoresearch sandbox training runs. \
Given a task's name, its docstring, and the concrete arguments it is about to be \
called with, estimate the *minimum sensible* compute it needs to finish without \
being OOM-killed, while not wildly over-provisioning.

Reason about the work implied by the arguments:
- TinyGPT training is memory-bound: scale with model width/depth (n_layer, n_embd, \
n_head), device_batch_size, and sequence length (512 in this workshop).
- Larger models and batch sizes need more RAM; CPU helps dataloader throughput but \
memory is usually the bottleneck.
- Sandbox runs are capped at a short time_budget_sec wall clock — prefer enough \
memory to survive peak activation usage over extra CPU.

Respond with ONLY a JSON object (no prose, no code fences) with any of these keys:
  - "cpu":    a number of cores, e.g. 2, 4, 8
  - "memory": a Kubernetes memory string, e.g. "4Gi", "16Gi"
  - "disk":   a Kubernetes disk string, e.g. "10Gi" (omit unless large I/O)
Omit a key to accept the default. Do not include any other keys. No GPUs are \
available on this cluster.

Example response: {"cpu": 4, "memory": "8Gi"}
"""

_ALLOWED_RESOURCE_KEYS = ("cpu", "memory", "disk", "shm")
_JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)

def _extract_json(text: str | None) -> dict[str, Any]:
    """Best-effort extraction of a single JSON object from an LLM reply."""
    if not text:
        return {}
    match = _JSON_OBJECT_RE.search(text)
    if not match:
        return {}
    try:
        parsed = json.loads(match.group(0))
    except json.JSONDecodeError:
        return {}
    return parsed if isinstance(parsed, dict) else {}

def _resources_from_spec(spec: dict[str, Any], floor: flyte.Resources) -> flyte.Resources:
    """Merge an LLM-produced spec onto the floor, keeping only known keys."""
    kwargs: dict[str, Any] = {
        "cpu": floor.cpu,
        "memory": floor.memory,
        "gpu": floor.gpu,
        "disk": floor.disk,
        "shm": floor.shm,
    }
    for key in _ALLOWED_RESOURCE_KEYS:
        value = spec.get(key)
        if value in (None, "", "null"):
            continue
        kwargs[key] = value
    try:
        return _cap_resources(flyte.Resources(**kwargs))
    except Exception as exc:  # pragma: no cover - defensive against bad model output
        flyte.logger.warning("Invalid resource spec %s (%s); falling back to floor.", spec, exc)
        return floor

async def estimate_resources(
    call_llm: LLMCallable,
    model: str,
    tool_name: str,
    description: str,
    args: dict[str, Any],
) -> flyte.Resources:
    """Ask the LLM to size the compute for a single tool call."""
    user = json.dumps({"tool": tool_name, "description": description, "arguments": args}, default=str)
    try:
        reply = await call_llm(
            model,
            RESOURCE_SIZING_SYSTEM_PROMPT,
            [{"role": "user", "content": user}],
            None,
        )
        spec = _extract_json(reply.content)
    except Exception as exc:  # pragma: no cover - never let sizing break the tool
        flyte.logger.warning("Resource right-sizing LLM call failed (%s); using floor.", exc)
        spec = {}
    resources = _resources_from_spec(spec, RESOURCE_FLOOR)
    flyte.logger.info("right-size %s %s -> %s", tool_name, args, resources)
    return resources

# {{docs-fragment right_size}}
async def execute_with_right_sizing(
    call_llm: LLMCallable,
    target_task: Any,
    *,
    model: str,
    tool_name: str,
    description: str,
    max_oom_retries: int = MAX_OOM_RETRIES,
    **kwargs: Any,
) -> dict:
    """LLM-size *target_task*, run it, and retry with more memory on OOM."""
    resources = await estimate_resources(call_llm, model, tool_name, description, kwargs)
    attempt = 0
    while True:
        try:
            with flyte.group(f"{tool_name}-attempt-{attempt + 1}"):
                result = await target_task.override(resources=resources).aio(**kwargs)
        except flyte.errors.OOMError:
            if attempt >= max_oom_retries:
                flyte.logger.error("%s Flyte OOM after %d retries; giving up.", tool_name, attempt)
                raise
            resources = bump_memory(resources)
            attempt += 1
            flyte.logger.warning(
                "%s Flyte OOM; retrying with memory=%s",
                tool_name,
                resources.memory,
            )
            continue

        if isinstance(result, dict):
            result["resources"] = f"cpu={resources.cpu}, mem={resources.memory}"
            result["oom_retries"] = attempt

        if isinstance(result, dict) and result.get("oom"):
            if attempt >= max_oom_retries:
                return result
            resources = bump_memory(resources)
            attempt += 1
            flyte.logger.warning(
                "%s sandbox OOM; retrying with memory=%s",
                tool_name,
                resources.memory,
            )
            continue

        return result

def right_sizing_handler(*, max_oom_retries: int = MAX_OOM_RETRIES):
    """Build a ``@tool`` ``call_handler`` that right-sizes and self-heals on OOM."""

    async def handle(call_llm: LLMCallable, tool_fn: ToolFn, **kwargs: Any) -> Any:
        return await execute_with_right_sizing(
            call_llm,
            tool_fn.target,
            model=tool_fn.model,
            tool_name=tool_fn.name,
            description=tool_fn.description,
            max_oom_retries=max_oom_retries,
            **kwargs,
        )

    return handle

right_size = right_sizing_handler(max_oom_retries=MAX_OOM_RETRIES)
# {{/docs-fragment right_size}}

def _find_leaderboard_entry(entries: list[dict[str, Any]], title: str) -> dict[str, Any] | None:
    title_lower = title.strip().lower()
    for entry in entries:
        if str(entry.get("title", "")).strip().lower() == title_lower:
            return entry
    for entry in entries:
        if title_lower in str(entry.get("title", "")).strip().lower():
            return entry
    return None

@tool
@agent_env.task(retries=3)
async def search_arxiv(query: str, max_results: int = 4) -> str:
    """Search arXiv for recent papers relevant to the next experiment.

    Use this to gather external context on architectures, optimizers, or
    evaluation metrics before proposing a new TinyGPT configuration.

    Args:
        query: Free-text search query, e.g. ``small language model depth width``.
        max_results: Maximum number of papers to return (default 4).

    Returns:
        A markdown-ish bullet list of titles and short summaries, or a note
        if the search failed or returned nothing.
    """
    import httpx

    if not (query and query.strip()):
        return "(empty query; skip literature search)"

    url = "https://export.arxiv.org/api/query"
    params = {"search_query": f"all:{query}", "start": 0, "max_results": max_results}
    try:
        async with httpx.AsyncClient(timeout=30, follow_redirects=True) as client:
            resp = await client.get(url, params=params)
            resp.raise_for_status()
        root = ET.fromstring(resp.text)
        ns = {"atom": "http://www.w3.org/2005/Atom"}
        lines: list[str] = []
        for entry in root.findall("atom:entry", ns)[:max_results]:
            title_el = entry.find("atom:title", ns)
            title = " ".join((title_el.text or "").split())
            summary_el = entry.find("atom:summary", ns)
            summary = " ".join((summary_el.text or "").split())[:400]
            lines.append(f"- {title}\n  {summary}")
        return "\n".join(lines) if lines else "(no arXiv results; proceed without external context)"
    except (httpx.TimeoutException, httpx.ConnectError, httpx.NetworkError) as exc:
        return f"(literature search failed: {exc})"
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code >= 500:
            return f"(literature search failed: {exc})"
        raise

@tool
@bundle_env.task(cache="auto")
async def inspect_dataset(num_shards: int = DEFAULT_NUM_SHARDS) -> dict:
    """Inspect the prepared climbmix corpus and BPE tokenizer bundle.

    Call this at the start of a research session to understand what data you
    are training on before spending experiment budget.

    Args:
        num_shards: Number of climbmix parquet shards to include in the bundle.

    Returns:
        A dict with shard/file metadata, vocab size, byte counts, and fixed
        training constants (``max_seq_len``, ``val_metric``).
    """
    import prepare

    bundle = await build_bundle(num_shards=num_shards)
    profile: DatasetProfile = await profile_bundle(bundle)
    return {
        **dataclasses.asdict(profile),
        "max_seq_len": prepare.MAX_SEQ_LEN,
        "val_metric": "val_bpb (lower is better)",
        "corpus": "karpathy/climbmix-400b-shuffle",
    }

@tool
@agent_env.task
async def record_hypothesis(
    title: str,
    hypothesis: str,
    expected_effect: str,
    memory_key: str = MEMORY_KEY_FANOUT,
) -> dict:
    """Record a structured hypothesis before running an experiment.

    Persists to the agent's keyed memory so later runs can see what you
    expected and whether it panned out.

    Args:
        title: Experiment title this hypothesis applies to.
        hypothesis: What you are trying and why.
        expected_effect: How you expect val_bpb to move (e.g. ``decrease ~5%``).
        memory_key: Memory namespace (use the key from your directive).

    Returns:
        The recorded hypothesis entry.
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    prior: list[dict[str, Any]] = await memory.read_json.aio("memory/hypotheses.json", default=[])
    entry = HypothesisEntry(
        title=title,
        hypothesis=hypothesis,
        expected_effect=expected_effect,
        recorded_at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
    )
    prior.append(dataclasses.asdict(entry))
    await memory.write_json.aio(
        "memory/hypotheses.json",
        prior,
        actor="mle-autoresearch-agent",
        reason=f"hypothesis for {title}",
    )
    await memory.save.aio()
    return dataclasses.asdict(entry)

@tool
@agent_env.task
async def get_leaderboard(memory_key: str = MEMORY_KEY_FANOUT) -> dict:
    """Return the persisted experiment leaderboard from agent memory.

    Use this to recall prior runs across sessions. Experiments from the
    *current* session also appear in your tool-call transcript.

    Args:
        memory_key: Memory namespace (use the key from your directive).

    Returns:
        A dict with ``entries`` (list) and ``best`` (entry or null).
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    entries: list[dict[str, Any]] = await memory.read_json.aio("memory/leaderboard.json", default=[])
    best: dict[str, Any] | None = None
    best_val = float("inf")
    for entry in entries:
        val = entry.get("val_bpb")
        if val is not None and float(val) < best_val:
            best_val = float(val)
            best = entry
    best_f = best_val if best_val != float("inf") else None
    enriched: list[dict[str, Any]] = []
    for entry in entries:
        val = entry.get("val_bpb")
        val_f = float(val) if val is not None else None
        enriched.append(
            {
                **entry,
                "beat_best": val_f is not None and best_f is not None and val_f <= best_f,
                "delta_vs_best": (val_f - best_f) if val_f is not None and best_f is not None else None,
            }
        )
    return {
        "entries": enriched,
        "best": best,
        "best_val_bpb": best_f,
        "count": len(enriched),
    }

@tool
@agent_env.task
async def compare_experiments(
    title_a: str,
    title_b: str,
    memory_key: str = MEMORY_KEY_FANOUT,
) -> dict:
    """Compare two prior experiments side-by-side.

    Looks up both titles in the persisted leaderboard. For experiments run in
    the current session that are not yet persisted, use the values from your
    recent ``run_experiment`` tool results instead.

    Args:
        title_a: Title of the first experiment.
        title_b: Title of the second experiment.
        memory_key: Memory namespace (use the key from your directive).

    Returns:
        A dict with ``a``, ``b``, and ``delta_val_bpb`` (a minus b; negative
        means a is better).
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    entries: list[dict[str, Any]] = await memory.read_json.aio("memory/leaderboard.json", default=[])
    a = _find_leaderboard_entry(entries, title_a)
    b = _find_leaderboard_entry(entries, title_b)
    missing = [t for t, e in ((title_a, a), (title_b, b)) if e is None]
    delta: float | None = None
    if a is not None and b is not None and a.get("val_bpb") is not None and b.get("val_bpb") is not None:
        delta = float(a["val_bpb"]) - float(b["val_bpb"])
    return {
        "a": a,
        "b": b,
        "delta_val_bpb": delta,
        "missing": missing,
        "note": (
            "Some titles were not found in persisted memory; check recent run_experiment "
            "tool results in your transcript for the current session."
            if missing
            else None
        ),
    }

_CONFIG_FIELDS = {f.name for f in dataclasses.fields(ExperimentConfig)} - {"title"}
_RUN_TRAINING_DOC = re.compile(
    r"(def run_training\(config: ExperimentConfig\)[^:]*:\n(?:    \"\"\"[\s\S]*?\"\"\"\n))"
)

def normalize_train_py(text: str) -> str:
    return text.replace("\r\n", "\n").strip()

def baseline_train_py() -> str:
    """Return the repo baseline ``train.py`` (single source of truth for diffs)."""
    import train

    assert train.__file__ is not None
    return Path(train.__file__).read_text()

def filter_config_overrides(overrides: dict[str, Any] | None) -> dict[str, Any]:
    if not overrides:
        return {}
    filtered = {k: v for k, v in overrides.items() if k in _CONFIG_FIELDS}
    if "n_layer" in filtered:
        filtered["n_layer"] = max(1, min(int(filtered["n_layer"]), MAX_N_LAYER))
    if "n_head" in filtered:
        filtered["n_head"] = max(1, min(int(filtered["n_head"]), MAX_N_HEAD))
    if "n_embd" in filtered:
        filtered["n_embd"] = max(1, min(int(filtered["n_embd"]), MAX_N_EMBD))
    if "device_batch_size" in filtered:
        filtered["device_batch_size"] = max(1, min(int(filtered["device_batch_size"]), MAX_DEVICE_BATCH_SIZE))
    if "max_steps" in filtered:
        filtered["max_steps"] = max(1, min(int(filtered["max_steps"]), MAX_MAX_STEPS))
    if "n_embd" in filtered and "n_head" in filtered and int(filtered["n_embd"]) % int(filtered["n_head"]) != 0:
        head = int(filtered["n_head"])
        filtered["n_embd"] = (int(filtered["n_embd"]) // head) * head
    return filtered

def is_config_only_edit(train_py: str, overrides: dict[str, Any] | None) -> bool:
    """True when *train_py* differs from baseline only via ``config_overrides`` injection."""
    baseline = baseline_train_py()
    filtered = filter_config_overrides(overrides)
    if not filtered:
        return normalize_train_py(train_py) == normalize_train_py(baseline)
    expected = build_train_py_with_config_overrides(baseline, filtered)
    return normalize_train_py(train_py) == normalize_train_py(expected)

def experiment_config_signature(train_py: str, overrides: dict[str, Any] | None) -> str:
    """Stable hash of effective train code + config overrides for duplicate detection."""
    filtered = filter_config_overrides(overrides)
    payload = {
        "train_py": normalize_train_py(train_py),
        "overrides": sorted(filtered.items()),
    }
    return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode()).hexdigest()[:16]

async def check_duplicate_config(
    memory_key: str,
    title: str,
    train_py: str,
    overrides: dict[str, Any] | None,
) -> dict[str, Any] | None:
    """Return duplicate metadata if this config was already run under another title."""
    sig = experiment_config_signature(train_py, overrides)
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    sigs: dict[str, str] = await memory.read_json.aio("memory/config_signatures.json", default={})
    prior_title = sigs.get(sig)
    title_key = title.strip().lower()
    if prior_title and prior_title.strip().lower() != title_key:
        return {"duplicate_of": prior_title, "config_signature": sig}
    return None

async def register_config_signature(
    memory_key: str,
    title: str,
    train_py: str,
    overrides: dict[str, Any] | None,
    *,
    actor: str = "mle-autoresearch-code-agent",
) -> str:
    """Record the config signature for *title* after a successful edit or run."""
    sig = experiment_config_signature(train_py, overrides)
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    sigs: dict[str, str] = await memory.read_json.aio("memory/config_signatures.json", default={})
    sigs[sig] = title
    await memory.write_json.aio(
        "memory/config_signatures.json",
        sigs,
        actor=actor,
        reason=f"config signature for {title}",
    )
    await memory.save.aio()
    return sig

def build_train_py_with_config_overrides(
    base_code: str,
    overrides: dict[str, Any],
) -> str:
    """Inject ``dataclasses.replace(config, ...)`` at the top of ``run_training``."""
    filtered = filter_config_overrides(overrides)
    if not filtered:
        return base_code

    parts = [f"{k}={v!r}" for k, v in sorted(filtered.items())]
    injection = f"    import dataclasses\n    config = dataclasses.replace(config, {', '.join(parts)})\n"
    match = _RUN_TRAINING_DOC.search(base_code)
    if match:
        insert_at = match.end()
        return base_code[:insert_at] + injection + base_code[insert_at:]
    return base_code

async def load_config_overrides(memory_key: str, title: str) -> dict[str, Any]:
    """Load persisted ``ExperimentConfig`` overrides for an experiment title."""
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    slug = slugify(title)
    stored = await memory.read_json.aio(f"memory/config/{slug}.json", default={})
    if stored:
        return filter_config_overrides(stored)

    index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    title_lower = title.strip().lower()
    for entry in index:
        if str(entry.get("title", "")).strip().lower() == title_lower:
            slug = str(entry.get("slug", slug))
            stored = await memory.read_json.aio(f"memory/config/{slug}.json", default={})
            if stored:
                return filter_config_overrides(stored)
            return filter_config_overrides(entry.get("config_overrides") or {})
    return {}

def slugify(title: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", title.lower()).strip("-")
    return slug[:80] or "experiment"

async def load_train_code(memory_key: str, title: str) -> str:
    """Load edited ``train.py`` for *title*, falling back to the repo baseline."""
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    slug = slugify(title)
    saved = await memory.read_text.aio(f"memory/code/{slug}.py", default="")
    if saved.strip():
        return saved

    index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    title_lower = title.strip().lower()
    for entry in index:
        if str(entry.get("title", "")).strip().lower() == title_lower:
            slug = entry.get("slug", slug)
            saved = await memory.read_text.aio(f"memory/code/{slug}.py", default="")
            if saved.strip():
                return saved

    return baseline_train_py()

async def _global_best_val_bpb(memory: MemoryStore, *, exclude_title: str | None = None) -> float:
    """Lowest val_bpb recorded in memory (optionally excluding one title)."""
    exclude = (exclude_title or "").strip().lower()
    leaderboard: list[dict[str, Any]] = await memory.read_json.aio("memory/leaderboard.json", default=[])
    promising: list[dict[str, Any]] = await memory.read_json.aio("memory/promising_code.json", default=[])
    vals: list[float] = []
    for row in leaderboard + promising:
        if exclude and str(row.get("title", "")).strip().lower() == exclude:
            continue
        val = row.get("val_bpb")
        if val is not None:
            vals.append(float(val))
    return min(vals, default=float("inf"))

async def _update_promising_code(
    memory_key: str,
    *,
    title: str,
    slug: str,
    val_bpb: float,
    change_summary: str,
) -> None:
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    promising: list[dict[str, Any]] = await memory.read_json.aio("memory/promising_code.json", default=[])
    prior_best = await _global_best_val_bpb(memory, exclude_title=title)
    kept = val_bpb < prior_best
    promising.append(
        {
            "title": title,
            "slug": slug,
            "val_bpb": val_bpb,
            "kept": kept,
            "change_summary": change_summary,
            "recorded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        }
    )
    await memory.write_json.aio(
        "memory/promising_code.json",
        promising,
        actor="mle-autoresearch-code-agent",
        reason=f"promising code after {title} val_bpb={val_bpb}",
    )
    await memory.save.aio()

async def _resolve_train_py_for_edit(
    memory_key: str,
    spec: dict[str, Any],
) -> tuple[str, dict[str, Any], str | None]:
    """Build the effective ``train.py`` source and overrides for one edit spec."""
    train_py = spec.get("train_py", "")
    if not isinstance(train_py, str):
        train_py = ""
    config_overrides = filter_config_overrides(
        spec.get("config_overrides") or spec.get("config") or {}
    )
    parent_title = spec.get("parent_title")
    parent_title = str(parent_title).strip() if parent_title else None

    baseline = baseline_train_py()
    if config_overrides:
        base_code = await load_train_code(memory_key, parent_title) if parent_title else baseline
        if not train_py.strip() or normalize_train_py(train_py) == normalize_train_py(baseline):
            train_py = build_train_py_with_config_overrides(base_code, config_overrides)
        elif parent_title and normalize_train_py(train_py) == normalize_train_py(base_code):
            train_py = build_train_py_with_config_overrides(base_code, config_overrides)

    return train_py, config_overrides, parent_title

async def _persist_train_edits(
    memory_key: str,
    edits: list[dict[str, Any]],
    *,
    actor: str = "mle-autoresearch-code-agent",
) -> dict[str, Any]:
    """Save one or more ``train.py`` edits in a single memory transaction."""
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    saved: list[dict[str, Any]] = []
    errors: list[dict[str, Any]] = []
    now = datetime.now(timezone.utc).isoformat(timespec="seconds")

    for spec in edits:
        title = str(spec.get("title", "")).strip()
        change_summary = str(spec.get("change_summary", ""))
        if not title:
            errors.append({"title": title or "(missing)", "saved": False, "error": "title is required"})
            continue

        train_py, config_overrides, parent_title = await _resolve_train_py_for_edit(memory_key, spec)
        if not train_py.strip():
            errors.append(
                {
                    "title": title,
                    "saved": False,
                    "error": "train_py or config_overrides is required",
                }
            )
            continue
        if is_config_only_edit(train_py, config_overrides) and len(index) >= CONFIG_ONLY_EDIT_LIMIT:
            errors.append(
                {
                    "title": title,
                    "saved": False,
                    "error": (
                        f"Batch 2+ requires substantive train.py edits (LR schedule, optimizer, "
                        f"weight decay, grad clip, etc.), not config_overrides alone. "
                        f"You already have {len(index)} saved edit(s)."
                    ),
                }
            )
            continue
        if normalize_train_py(train_py) == normalize_train_py(baseline_train_py()) and not config_overrides:
            errors.append(
                {
                    "title": title,
                    "saved": False,
                    "error": (
                        "train.py matches baseline with no config_overrides; "
                        "pass config_overrides={n_layer: 6, ...} or edit run_training"
                    ),
                }
            )
            continue
        if "def run_training" not in train_py:
            errors.append(
                {
                    "title": title,
                    "saved": False,
                    "error": "train_py must define run_training(config) like the baseline train.py",
                }
            )
            continue

        slug = slugify(title)
        await memory.write_text.aio(
            f"memory/code/{slug}.py",
            train_py,
            actor=actor,
            reason=f"edit train.py for {title}",
        )
        if config_overrides:
            await memory.write_json.aio(
                f"memory/config/{slug}.json",
                config_overrides,
                actor=actor,
                reason=f"config overrides for {title}",
            )
        index.append(
            {
                "title": title,
                "slug": slug,
                "change_summary": change_summary,
                "lines": len(train_py.splitlines()),
                "edited_at": now,
                "config_overrides": config_overrides,
                "parent_title": parent_title,
            }
        )
        saved.append(
            {
                "saved": True,
                "title": title,
                "slug": slug,
                "lines": len(train_py.splitlines()),
                "change_summary": change_summary,
                "train_py": train_py,
                "config_overrides": config_overrides,
                "parent_title": parent_title,
                "memory_path": f"memory/code/{slug}.py",
            }
        )

    if saved:
        await memory.write_json.aio(
            "memory/code_index.json",
            index,
            actor=actor,
            reason=f"code index update ({len(saved)} edit(s))",
        )
        await memory.save.aio()

    return {
        "count": len(saved),
        "titles": [row["title"] for row in saved],
        "edits": saved,
        "errors": errors,
    }

@tool
@agent_env.task
async def get_baseline_train_code() -> dict:
    """Return the baseline ``train.py`` from the repo (the karpathy/autoresearch recipe).

    Use this once at the start to understand the starting point before editing.

    Returns:
        A dict with ``title``, ``train_py`` (full source), and ``lines``.
    """
    code = baseline_train_py()
    return {"title": "baseline", "train_py": code, "lines": len(code.splitlines())}

@tool
@agent_env.task
async def edit_train_code(
    title: str,
    train_py: str = "",
    change_summary: str = "",
    memory_key: str = MEMORY_KEY_FANOUT,
    config_overrides: dict[str, Any] | None = None,
    parent_title: str | None = None,
) -> dict:
    """Save an edited ``train.py`` for this experiment to agent memory.

    The code must keep a ``run_training(config: ExperimentConfig) -> ExperimentResult``
    entry point (same as the baseline). Only edit architecture, optimizer, and
    training-loop knobs inside the file.

    Alternatively pass ``config_overrides`` (e.g. ``{"n_layer": 6, "learning_rate": 1e-4}``)
    instead of a full ``train_py`` rewrite — the platform injects
    ``dataclasses.replace(config, ...)`` into ``run_training`` for you.

    Args:
        title: Short human-readable experiment name (used as the memory key slug).
        train_py: Full Python source for the edited training script (optional if
            ``config_overrides`` is set).
        change_summary: One-line description of what you changed and why.
        memory_key: Memory namespace from your directive.
        config_overrides: Optional ``ExperimentConfig`` field overrides.
        parent_title: Optional prior experiment to fork before applying overrides.

    Returns:
        Metadata about the saved edit, including the full ``train_py`` source
        (visible in the Flyte task output UI).
    """
    result = await _persist_train_edits(
        memory_key,
        [
            {
                "title": title,
                "train_py": train_py,
                "change_summary": change_summary,
                "config_overrides": config_overrides,
                "parent_title": parent_title,
            }
        ],
    )
    if result["edits"]:
        return result["edits"][0]
    err = result["errors"][0] if result["errors"] else {"saved": False, "error": "unknown error"}
    return err

@tool
@agent_env.task
async def edit_train_code_batch(
    edits: list[dict[str, Any]],
    memory_key: str = MEMORY_KEY_FANOUT,
) -> dict:
    """Save multiple edited ``train.py`` files in one atomic memory write.

    Use this when preparing a parallel experiment batch — avoids sequential
    ``edit_train_code`` calls and race conditions on ``memory/code_index.json``.

    Each item in ``edits`` must include ``title`` and ``change_summary``, plus either
    ``train_py`` (full source) or ``config_overrides`` (e.g. ``{"n_layer": 6}``).
    Optional ``parent_title`` forks a prior experiment before applying overrides.
    Every ``train_py`` must keep the ``run_training(config)`` entry point.

    Args:
        edits: List of edit specs, e.g.
            ``[{"title": "deeper-6L", "config_overrides": {"n_layer": 6}, "change_summary": "..."}]``.
        memory_key: Memory namespace from your directive.

    Returns:
        A dict with ``count``, ``titles``, ``edits`` (each includes ``train_py``),
        and ``errors`` (rejected).
    """
    if not edits:
        return {"count": 0, "titles": [], "edits": [], "errors": [{"error": "edits list is empty"}]}
    return await _persist_train_edits(
        memory_key,
        edits,
        actor="parallelized-autoresearch",
    )

@tool
@agent_env.task
async def read_train_code(title: str, memory_key: str = MEMORY_KEY_FANOUT) -> dict:
    """Read a previously saved ``train.py`` edit from memory (or the baseline).

    Args:
        title: Experiment title whose code you want to inspect.
        memory_key: Memory namespace from your directive.

    Returns:
        A dict with ``title``, ``train_py``, and ``lines``.
    """
    code = await load_train_code(memory_key, title)
    return {"title": title, "train_py": code, "lines": len(code.splitlines())}

@tool
@agent_env.task
async def get_promising_code(memory_key: str = MEMORY_KEY_FANOUT) -> dict:
    """Return promising ``train.py`` edits, the current best, and deltas vs best.

    Each entry records ``val_bpb`` after a successful run. Use ``read_train_code``
    with the best entry's title to inspect its source. Prefer ``get_code_edit_history``
    for the full cross-session table of edits, results, and regressions.

    Args:
        memory_key: Memory namespace from your directive.

    Returns:
        A dict with ``entries``, ``best``, ``best_val_bpb``, and ``count``.
    """
    history = await load_research_history(memory_key)
    best_val = history.get("best_val_bpb")
    entries: list[dict[str, Any]] = []
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    promising: list[dict[str, Any]] = await memory.read_json.aio("memory/promising_code.json", default=[])
    for row in promising:
        val = row.get("val_bpb")
        val_f = float(val) if val is not None else None
        entries.append(
            {
                **row,
                "beat_best": val_f is not None and best_val is not None and val_f <= best_val,
                "delta_vs_best": (val_f - best_val) if val_f is not None and best_val is not None else None,
            }
        )
    best: dict[str, Any] | None = None
    if history.get("best_title"):
        best_key = str(history["best_title"]).strip().lower()
        for entry in reversed(entries):
            if str(entry.get("title", "")).strip().lower() == best_key:
                best = entry
                break
    return {
        "entries": entries,
        "best": best,
        "best_val_bpb": best_val,
        "best_title": history.get("best_title"),
        "count": len(entries),
    }

@tool
@agent_env.task
async def get_code_edit_history(memory_key: str = MEMORY_KEY_FANOUT) -> dict:
    """Return all prior code edits, run results, and whether each beat the current best.

    Call this at the start of a session when ``memory_key`` already has experiments.
    Shows every saved ``train.py`` edit, its ``change_summary``, ``val_bpb`` (if run),
    ``delta_vs_best`` (negative means better), ``outcome`` (``new_best`` / ``regression`` /
    ``failed`` / ``not_run``), and linked hypotheses.

    Args:
        memory_key: Memory namespace from your directive.

    Returns:
        A dict with ``best_val_bpb``, ``best_title``, ``trials``, and summary counts.
    """
    return await load_research_history(memory_key)

async def load_saved_code_edits(memory_key: str) -> list[dict[str, Any]]:
    """Load all saved ``train.py`` edits from memory for reporting."""
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    promising: list[dict[str, Any]] = await memory.read_json.aio("memory/promising_code.json", default=[])
    val_by_title = {
        str(row.get("title", "")).strip().lower(): row.get("val_bpb")
        for row in promising
        if row.get("val_bpb") is not None
    }
    kept_titles = {
        str(row.get("title", "")).strip().lower()
        for row in promising
        if row.get("kept")
    }

    baseline = baseline_train_py()
    edits: list[dict[str, Any]] = []
    for entry in index:
        slug = str(entry.get("slug", slugify(str(entry.get("title", "")))))
        train_py = await memory.read_text.aio(f"memory/code/{slug}.py", default="")
        title = str(entry.get("title", ""))
        title_key = title.strip().lower()
        config_overrides = filter_config_overrides(entry.get("config_overrides") or {})
        if not config_overrides:
            config_overrides = filter_config_overrides(
                await memory.read_json.aio(f"memory/config/{slug}.json", default={})
            )
        if config_overrides and normalize_train_py(train_py) == normalize_train_py(baseline):
            parent_title = entry.get("parent_title")
            base_code = (
                await load_train_code(memory_key, str(parent_title))
                if parent_title
                else baseline
            )
            train_py = build_train_py_with_config_overrides(base_code, config_overrides)
        edits.append(
            {
                **entry,
                "slug": slug,
                "train_py": train_py,
                "config_overrides": config_overrides,
                "memory_path": f"memory/code/{slug}.py",
                "val_bpb": val_by_title.get(title_key),
                "kept": title_key in kept_titles,
            }
        )
    return edits

async def record_experiment_result(
    memory_key: str,
    result: dict[str, Any],
    *,
    actor: str = "mle-autoresearch-code-agent",
) -> None:
    """Upsert one experiment outcome into ``memory/leaderboard.json``."""
    title = str(result.get("title", "")).strip()
    if not title:
        return
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    leaderboard: list[dict[str, Any]] = await memory.read_json.aio("memory/leaderboard.json", default=[])
    row: dict[str, Any] = {
        "title": title,
        "success": bool(result.get("success")),
        "val_bpb": float(result["val_bpb"]) if result.get("val_bpb") is not None else None,
        "model_name": result.get("model_name"),
        "n_params": result.get("n_params"),
        "steps": int(result["steps"]) if result.get("steps") is not None else None,
        "resources": result.get("resources"),
        "oom_retries": int(result.get("oom_retries", 0)),
    }
    if not result.get("success"):
        err = result.get("error") or result.get("stderr") or "failed"
        row["error"] = str(err)[:200]

    title_key = title.lower()
    replaced = False
    for idx, existing in enumerate(leaderboard):
        if str(existing.get("title", "")).strip().lower() == title_key:
            leaderboard[idx] = row
            replaced = True
            break
    if not replaced:
        leaderboard.append(row)

    await memory.write_json.aio(
        "memory/leaderboard.json",
        leaderboard,
        actor=actor,
        reason=f"experiment result for {title}",
    )
    await memory.save.aio()

async def record_promising_run(
    memory_key: str,
    title: str,
    result: dict[str, Any],
    change_summary: str = "",
) -> None:
    """Persist a successful run's code to the promising-code ledger."""
    if not result.get("success") or result.get("val_bpb") is None:
        return
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    code_index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    summary = change_summary
    slug = slugify(title)
    for entry in reversed(code_index):
        if str(entry.get("title", "")).strip().lower() == title.strip().lower():
            summary = summary or str(entry.get("change_summary", ""))
            slug = str(entry.get("slug", slug))
            break
    await _update_promising_code(
        memory_key,
        title=title,
        slug=slug,
        val_bpb=float(result["val_bpb"]),
        change_summary=summary or "successful run",
    )

@tool
@agent_env.task
async def record_batch_plan(
    batch_id: str,
    experiments: list[dict[str, Any]],
    memory_key: str = MEMORY_KEY_FANOUT,
) -> dict:
    """Persist a batch of planned experiments before editing or running them.

    Each experiment dict should include at least ``title`` and ``hypothesis``.
    Optional keys: ``expected_effect``, ``change_summary``, ``parent_title``.

    Args:
        batch_id: Short identifier for this batch (e.g. ``batch-1-depth-sweep``).
        experiments: Planned experiment specs for parallel execution.
        memory_key: Memory namespace from your directive.

    Returns:
        The saved batch record with ``batch_id``, ``count``, and ``experiments``.
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    batches: list[dict[str, Any]] = await memory.read_json.aio("memory/batches.json", default=[])
    record = {
        "batch_id": batch_id,
        "experiments": experiments,
        "count": len(experiments),
        "status": "planned",
        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }
    batches.append(record)
    await memory.write_json.aio(
        "memory/batches.json",
        batches,
        actor="parallelized-autoresearch",
        reason=f"batch plan {batch_id}",
    )
    await memory.save.aio()
    return record

@tool
@agent_env.task
async def get_batch_plan(batch_id: str, memory_key: str = MEMORY_KEY_FANOUT) -> dict:
    """Load a previously recorded batch plan by ``batch_id``.

    Args:
        batch_id: Identifier passed to ``record_batch_plan``.
        memory_key: Memory namespace from your directive.

    Returns:
        The batch record, or ``{"found": False}`` if missing.
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    batches: list[dict[str, Any]] = await memory.read_json.aio("memory/batches.json", default=[])
    batch_id_lower = batch_id.strip().lower()
    for batch in reversed(batches):
        if str(batch.get("batch_id", "")).strip().lower() == batch_id_lower:
            return {"found": True, **batch}
    return {"found": False, "batch_id": batch_id}

@tool
@agent_env.task
async def record_batch_hypotheses(
    experiments: list[dict[str, Any]],
    memory_key: str = MEMORY_KEY_FANOUT,
) -> dict:
    """Record hypotheses for every experiment in a batch (before ``run_experiment_batch``).

    Each item needs ``title``, ``hypothesis``, and ``expected_effect``.

    Args:
        experiments: List of hypothesis dicts (one per planned experiment title).
        memory_key: Memory namespace from your directive.

    Returns:
        A dict with ``recorded`` count and the appended entries.
    """
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    prior: list[dict[str, Any]] = await memory.read_json.aio("memory/hypotheses.json", default=[])
    recorded: list[dict[str, Any]] = []
    for spec in experiments:
        entry = HypothesisEntry(
            title=str(spec.get("title", "")),
            hypothesis=str(spec.get("hypothesis", "")),
            expected_effect=str(spec.get("expected_effect", "")),
            recorded_at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
        )
        row = dataclasses.asdict(entry)
        prior.append(row)
        recorded.append(row)
    await memory.write_json.aio(
        "memory/hypotheses.json",
        prior,
        actor="parallelized-autoresearch",
        reason=f"batch hypotheses ({len(recorded)} experiments)",
    )
    await memory.save.aio()
    return {"recorded": len(recorded), "entries": recorded}

def evaluate_batch_results_impl(
    results: list[dict[str, Any]],
    batch_id: str = "",
) -> dict[str, Any]:
    """Rank and summarize the outcome of a parallel experiment batch."""
    successes: list[dict[str, Any]] = []
    failures: list[dict[str, Any]] = []
    for result in results:
        if not isinstance(result, dict):
            failures.append({"title": "?", "error": str(result)})
            continue
        if result.get("success") and result.get("val_bpb") is not None:
            successes.append(result)
        else:
            failures.append(
                {
                    "title": result.get("title", "?"),
                    "error": result.get("error") or (result.get("stderr") or "")[:200],
                    "oom": result.get("oom", False),
                }
            )

    ranked = sorted(successes, key=lambda r: float(r["val_bpb"]))
    best = ranked[0] if ranked else None
    return {
        "batch_id": batch_id or None,
        "total": len(results),
        "n_success": len(successes),
        "n_failed": len(failures),
        "ranked": [
            {
                "title": r.get("title"),
                "val_bpb": r.get("val_bpb"),
                "model_name": r.get("model_name"),
                "steps": r.get("steps"),
                "resources": r.get("resources"),
                "oom_retries": r.get("oom_retries", 0),
            }
            for r in ranked
        ],
        "best": best,
        "failures": failures,
    }

@tool
@agent_env.task
async def evaluate_batch_results(
    results: list[dict[str, Any]],
    batch_id: str = "",
) -> dict:
    """Rank and summarize the outcome of a parallel experiment batch.

    Use after ``run_experiment_batch`` or ``flyte_map("run_experiment", ...)``.
    Lower ``val_bpb`` is better.

    Args:
        results: List of ``run_experiment`` result dicts (same order as titles).
        batch_id: Optional batch label for the summary.

    Returns:
        A dict with ``successes``, ``failures``, ``ranked``, ``best``, and ``batch_id``.
    """
    return evaluate_batch_results_impl(results, batch_id=batch_id)

async def persist_run_results_to_leaderboard(
    memory_key: str,
    results: list[dict[str, Any]],
    *,
    actor: str = "parallelized-autoresearch",
) -> int:
    """Persist run results (success or failure) to ``memory/leaderboard.json``."""
    added = 0
    for result in results:
        if not isinstance(result, dict) or not result.get("title"):
            continue
        await record_experiment_result(memory_key, result, actor=actor)
        added += 1
    return added

async def run_experiment_batch_impl(
    run_experiment_task: Any,
    titles: list[str],
    *,
    time_budget_sec: int = 45,
    memory_key: str = MEMORY_KEY_FANOUT,
    concurrency: int = 4,
    group_name: str | None = None,
) -> dict[str, Any]:
    """Fan out ``run_experiment`` across *titles* via ``flyte.map``."""
    if not titles:
        return {"batch_size": 0, "results": [], "titles": []}

    n = len(titles)
    budgets = [time_budget_sec] * n
    keys = [memory_key] * n
    map_kwargs: dict[str, Any] = {"concurrency": concurrency, "return_exceptions": True}
    if group_name:
        map_kwargs["group_name"] = group_name

    results: list[Any] = []
    async for item in flyte.map.aio(run_experiment_task, titles, budgets, keys, **map_kwargs):
        if isinstance(item, BaseException):
            results.append({"success": False, "title": "?", "error": str(item)})
        else:
            results.append(item)

    return {
        "batch_size": n,
        "titles": titles,
        "results": results,
        "concurrency": concurrency,
        "group_name": group_name,
    }

OOM_MARKERS = (
    "out of memory",
    "oom",
    "cannot allocate memory",
    "can't allocate memory",
    "unable to allocate",
    "memoryerror",
    "killed",
    "signal 9",
    "std::bad_alloc",
    "defaultcpuallocator",
    "bad_alloc",
)

def is_oom(stderr: str, returncode: int | None, *, stdout: str = "") -> bool:
    """Detect OOM from sandbox stderr / exit code (137 = SIGKILL/OOM-kill)."""
    if returncode in (137, -9):
        return True
    text = f"{stderr}\n{stdout}".lower()
    return any(marker in text for marker in OOM_MARKERS)

def parse_metrics(stdout: str) -> dict[str, Any] | None:
    """Parse the ``AUTORESEARCH_METRICS=`` line emitted by the driver script."""
    for line in stdout.splitlines():
        if line.startswith("AUTORESEARCH_METRICS="):
            return json.loads(line.split("=", 1)[1])
    return None

def write_driver_script(title: str, time_budget_sec: int, eval_tokens: int) -> str:
    """Return a small driver that imports the agent-edited ``train.py`` and prints metrics."""
    return textwrap.dedent(
        f'''
        import json
        import os
        import sys

        workdir = os.path.dirname(os.path.abspath(__file__))
        os.chdir(workdir)
        os.environ["AUTORESEARCH_CACHE"] = workdir
        sys.path.insert(0, workdir)
        os.environ.setdefault("AUTORESEARCH_EVAL_TOKENS", "{eval_tokens}")

        from autoresearch_types import ExperimentConfig
        import train

        overrides = {{}}
        overrides_path = os.path.join(workdir, "config_overrides.json")
        if os.path.exists(overrides_path):
            with open(overrides_path) as f:
                overrides = json.load(f)

        config = ExperimentConfig(title={title!r}, time_budget_sec={time_budget_sec})
        if overrides:
            import dataclasses
            config = dataclasses.replace(config, **overrides)
        result = train.run_training(config)
        payload = {{
            "title": result.title,
            "val_bpb": round(result.val_bpb, 6),
            "model_name": result.model_name,
            "n_params": result.n_params,
            "steps": result.steps,
            "device": result.device,
            "notes": result.notes,
        }}
        print("AUTORESEARCH_METRICS=" + json.dumps(payload))
        '''
    ).strip()

def stage_sandbox_files(
    work_dir: str,
    train_py: str,
    *,
    title: str,
    time_budget_sec: int,
    eval_tokens: int | None = None,
    config_overrides: dict[str, Any] | None = None,
) -> Path:
    """Copy support modules + edited train code into the sandbox work directory."""
    import autoresearch_types
    import prepare

    if eval_tokens is None:
        eval_tokens = 32 * prepare.MAX_SEQ_LEN
    root = Path(work_dir)
    root.mkdir(parents=True, exist_ok=True)
    (root / "train.py").write_text(train_py)
    if config_overrides:
        (root / "config_overrides.json").write_text(json.dumps(config_overrides))
    (root / "prepare.py").write_text(Path(prepare.__file__).read_text())
    (root / "autoresearch_types.py").write_text(Path(autoresearch_types.__file__).read_text())
    driver = write_driver_script(title, time_budget_sec, eval_tokens)
    driver_path = root / "driver.py"
    driver_path.write_text(driver)
    return driver_path

async def run_train_in_sandbox(
    work_dir: str,
    train_py: str,
    *,
    title: str,
    time_budget_sec: int,
    config_overrides: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Execute ``train.py`` via ``async with sb.on_device.session(backend='userns')``."""
    from union import sandbox as sb

    driver_path = stage_sandbox_files(
        work_dir,
        train_py,
        title=title,
        time_budget_sec=time_budget_sec,
        config_overrides=config_overrides,
    )
    timeout_s = max(time_budget_sec + 180, 300)

    try:
        async with sb.on_device.session(backend="userns", host_work_dir=work_dir) as sbx:
            proc = await sbx.run(
                f"python {driver_path}",
                stdout=True,
                stderr=True,
                network_mode="blocked",
                timeout_s=timeout_s,
            )
            stdout, stderr = await proc.communicate_text()
    except Exception as exc:
        err_text = str(exc)
        oom = is_oom(err_text, None)
        return {
            "success": False,
            "oom": oom,
            "title": title,
            "exit_code": None,
            "stdout_tail": "",
            "stderr": err_text,
            "error": (
                "Training run was OOM-killed; the platform will retry with more memory."
                if oom
                else f"Sandbox execution failed: {err_text}"
            ),
        }

    metrics = parse_metrics(stdout or "")
    oom = is_oom(stderr or "", proc.returncode, stdout=stdout or "")

    if metrics is not None and proc.returncode == 0:
        return {
            "success": True,
            "oom": False,
            **metrics,
            "exit_code": proc.returncode,
            "stderr_tail": (stderr or "")[-800:],
        }

    return {
        "success": False,
        "oom": oom,
        "title": title,
        "exit_code": proc.returncode,
        "stdout_tail": (stdout or "")[-1500:],
        "stderr": stderr or "",
        "error": (
            "Training run was OOM-killed; the platform will retry with more memory."
            if oom
            else f"Training failed (exit {proc.returncode}). See stderr for details."
        ),
    }

def _title_key(title: str) -> str:
    return str(title or "").strip().lower()

def _best_from_entries(entries: list[dict[str, Any]]) -> tuple[float | None, str | None]:
    best_val: float | None = None
    best_title: str | None = None
    for row in entries:
        val = row.get("val_bpb")
        if val is None:
            continue
        fval = float(val)
        if best_val is None or fval < best_val:
            best_val = fval
            best_title = str(row.get("title", ""))
    return best_val, best_title

def _latest_by_title(rows: list[dict[str, Any]], *, title_field: str = "title") -> dict[str, dict[str, Any]]:
    out: dict[str, dict[str, Any]] = {}
    for row in rows:
        key = _title_key(str(row.get(title_field, "")))
        if key:
            out[key] = row
    return out

def _outcome_label(
    *,
    val_bpb: float | None,
    success: bool | None,
    best_val: float | None,
) -> str:
    if success is False or (val_bpb is None and success is not True):
        return "failed"
    if val_bpb is None:
        return "not_run"
    if best_val is None:
        return "ran"
    delta = float(val_bpb) - best_val
    if delta <= 0:
        return "new_best"
    return "regression"

def _vs_best_text(val_bpb: float | None, best_val: float | None) -> str:
    if val_bpb is None or best_val is None:
        return "—"
    delta = float(val_bpb) - best_val
    if abs(delta) < 1e-12:
        return "0 (ties best)"
    sign = "+" if delta > 0 else ""
    quality = "worse" if delta > 0 else "better"
    return f"{sign}{delta:.6g} ({quality})"

async def load_research_history(memory_key: str) -> dict[str, Any]:
    """Merge saved edits, run results, and outcomes for cross-session agent context."""
    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    code_index: list[dict[str, Any]] = await memory.read_json.aio("memory/code_index.json", default=[])
    leaderboard: list[dict[str, Any]] = await memory.read_json.aio("memory/leaderboard.json", default=[])
    promising: list[dict[str, Any]] = await memory.read_json.aio("memory/promising_code.json", default=[])
    hypotheses: list[dict[str, Any]] = await memory.read_json.aio("memory/hypotheses.json", default=[])

    lb_by_title = _latest_by_title(leaderboard)
    prom_by_title = _latest_by_title(promising)
    hyp_by_title = _latest_by_title(hypotheses)

    best_val, best_title = _best_from_entries(leaderboard)
    if best_val is None:
        best_val, best_title = _best_from_entries(promising)

    trials: list[dict[str, Any]] = []
    seen: set[str] = set()

    for edit in code_index:
        title = str(edit.get("title", ""))
        key = _title_key(title)
        if not key:
            continue
        seen.add(key)
        lb = lb_by_title.get(key, {})
        prom = prom_by_title.get(key, {})
        hyp = hyp_by_title.get(key, {})

        val = lb.get("val_bpb")
        if val is None:
            val = prom.get("val_bpb")
        val_f = float(val) if val is not None else None

        success = lb.get("success")
        if success is None and val_f is not None:
            success = True
        if success is None and lb.get("error"):
            success = False

        beat_best = val_f is not None and best_val is not None and val_f <= best_val
        trials.append(
            {
                "title": title,
                "change_summary": edit.get("change_summary") or prom.get("change_summary") or "",
                "edited_at": edit.get("edited_at"),
                "val_bpb": val_f,
                "model_name": lb.get("model_name"),
                "success": success,
                "error": lb.get("error"),
                "hypothesis": hyp.get("hypothesis"),
                "expected_effect": hyp.get("expected_effect"),
                "beat_best": beat_best,
                "delta_vs_best": (val_f - best_val) if val_f is not None and best_val is not None else None,
                "vs_best": _vs_best_text(val_f, best_val),
                "outcome": _outcome_label(val_bpb=val_f, success=success, best_val=best_val),
                "kept": bool(prom.get("kept")),
            }
        )

    for key, lb in lb_by_title.items():
        if key in seen:
            continue
        val = lb.get("val_bpb")
        val_f = float(val) if val is not None else None
        success = lb.get("success")
        if success is None and val_f is not None:
            success = True
        trials.append(
            {
                "title": lb.get("title", key),
                "change_summary": "",
                "edited_at": None,
                "val_bpb": val_f,
                "model_name": lb.get("model_name"),
                "success": success,
                "error": lb.get("error"),
                "hypothesis": hyp_by_title.get(key, {}).get("hypothesis"),
                "expected_effect": hyp_by_title.get(key, {}).get("expected_effect"),
                "beat_best": val_f is not None and best_val is not None and val_f <= best_val,
                "delta_vs_best": (val_f - best_val) if val_f is not None and best_val is not None else None,
                "vs_best": _vs_best_text(val_f, best_val),
                "outcome": _outcome_label(val_bpb=val_f, success=success, best_val=best_val),
                "kept": bool(prom_by_title.get(key, {}).get("kept")),
            }
        )

    trials.sort(key=lambda t: (t.get("edited_at") or "", t.get("title", "")))

    return {
        "memory_key": memory_key,
        "best_val_bpb": best_val,
        "best_title": best_title,
        "trials": trials,
        "count_edits": len(code_index),
        "count_runs": sum(1 for t in trials if t.get("val_bpb") is not None or t.get("success") is False),
        "count_regressions": sum(1 for t in trials if t.get("outcome") == "regression"),
        "count_new_best": sum(1 for t in trials if t.get("outcome") == "new_best"),
    }

def format_research_history_for_directive(history: dict[str, Any], *, max_rows: int = 20) -> str:
    """Render prior edits/results as a compact block for the run directive."""
    trials: list[dict[str, Any]] = history.get("trials") or []
    if not trials:
        return ""

    best_val = history.get("best_val_bpb")
    best_title = history.get("best_title")
    header = "\n\n## Prior research (from memory — continue, do not repeat)\n"
    if best_val is not None:
        header += f"Current best: **val_bpb={best_val:.6g}** ({best_title}). Lower is better.\n"
    else:
        header += "No successful runs recorded yet.\n"

    header += (
        "Call ``get_code_edit_history()`` at the start to refresh this table. "
        "Use ``read_train_code`` on the best title to fork winners.\n\n"
    )

    lines = [
        "| Title | Change | val_bpb | vs best | Outcome |",
        "| --- | --- | --- | --- | --- |",
    ]
    for trial in trials[-max_rows:]:
        title = str(trial.get("title", ""))
        change = str(trial.get("change_summary", ""))[:72]
        val = trial.get("val_bpb")
        val_s = f"{float(val):.6g}" if val is not None else ("failed" if trial.get("success") is False else "—")
        lines.append(
            f"| {title} | {change} | {val_s} | {trial.get('vs_best', '—')} | {trial.get('outcome', '—')} |"
        )

    omitted = len(trials) - max_rows
    footer = ""
    if omitted > 0:
        footer = f"\n({omitted} older trial(s) omitted — use get_code_edit_history for the full list.)\n"

    return header + "\n".join(lines) + footer
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/parallelized_autoresearch/tools.py*

`right_size` is the pre-built handler passed to `@tool(call_handler=...)`. The agent does not need a back-reference to the `Agent` instance: the harness passes `call_llm` and `tool_fn.model` into the handler on each invocation.

The experiment task stacks `@tool(call_handler=tools.right_size)` on `@experiment_env.task`. The task body only loads edited code and runs sandbox training; sizing and OOM recovery happen in the handler:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.5",
#    "litellm",
#    "httpx",
#    "pydantic-monty",
#    "unionai-sandbox[flyte]",
#    "torch",
#    "numpy",
#    "pyarrow",
#    "requests",
#    "tiktoken",
#    "rustbpe",
# ]
# main = "parallelized_autoresearch"
# params = "--n-experiments 6 --batch-size 3 --num-shards 1"
# ///
"""Parallelized autoresearch agent — code-mode MLE agent with batched sandbox experiments."""

from __future__ import annotations

import dataclasses
from typing import Any

import flyte
import flyte.report
from flyte.ai.agents import Agent, MemoryStore, agent_progress_cb, tool

from autoresearch_types import AutoresearchOutput, DEFAULT_MAX_STEPS, DEFAULT_NUM_SHARDS, MAX_DEVICE_BATCH_SIZE, MAX_N_EMBD, MAX_N_HEAD, MAX_N_LAYER
from bundle import agent_env, build_bundle, experiment_env, materialize_cache, profile_bundle
import tools
import ui

MODEL = "claude-sonnet-4-6"

# {{docs-fragment run_experiment}}
@tool(call_handler=tools.right_size)
@experiment_env.task
async def run_experiment(
    title: str,
    time_budget_sec: int = 45,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
) -> dict:
    """Train using agent-edited ``train.py`` with LLM right-sizing and OOM self-healing."""
    train_py = await tools.load_train_code(memory_key, title)
    config_overrides = await tools.load_config_overrides(memory_key, title)
    duplicate = await tools.check_duplicate_config(memory_key, title, train_py, config_overrides)
    if duplicate:
        result = {
            "success": False,
            "title": title,
            "error": (
                f"Duplicate config of '{duplicate['duplicate_of']}' "
                f"(signature {duplicate['config_signature']}); change train.py or overrides."
            ),
            "duplicate_of": duplicate["duplicate_of"],
        }
        await tools.record_experiment_result(
            memory_key,
            result,
            actor="parallelized-autoresearch",
        )
        return result
    bundle = await build_bundle()
    cache_dir = await materialize_cache(bundle)
    result = await tools.run_train_in_sandbox(
        cache_dir,
        train_py,
        title=title,
        time_budget_sec=time_budget_sec,
        config_overrides=config_overrides or None,
    )
    if result.get("success"):
        await tools.record_promising_run(memory_key, title, result)
        await tools.register_config_signature(
            memory_key,
            title,
            train_py,
            config_overrides,
            actor="parallelized-autoresearch",
        )
    await tools.record_experiment_result(
        memory_key,
        result,
        actor="parallelized-autoresearch",
    )
    return result

# ``flyte.map`` invokes ``run_experiment.aio`` directly (not through the agent
# registry), so bind the LLM callback and model here for ``call_handler`` right-sizing.
run_experiment = dataclasses.replace(
    run_experiment,
    call_llm=tools.call_llm,
    model=MODEL,
)
# {{/docs-fragment run_experiment}}

@tool
@agent_env.task
async def run_experiment_batch(
    titles: list[str],
    time_budget_sec: int = 45,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
    concurrency: int = 4,
    batch_id: str = "",
) -> dict:
    """Run multiple ``run_experiment`` calls in parallel via ``flyte.map``.

    Prefer this over hand-rolling ``flyte_map`` when you already have a list of
    experiment titles with saved ``train.py`` edits.

    Args:
        titles: Experiment titles whose code was saved with ``edit_train_code_batch``.
        time_budget_sec: Wall-clock budget passed to each run.
        memory_key: Memory namespace from your directive.
        concurrency: Max parallel sandbox runs (default 4).
        batch_id: Optional label attached to the returned batch metadata.

    Returns:
        A dict with ``batch_size``, ``titles``, ``results``, and ``evaluation``
        (from :func:`evaluate_batch_results`).
    """
    group = batch_id or f"batch-{len(titles)}"
    payload = await tools.run_experiment_batch_impl(
        run_experiment,
        titles,
        time_budget_sec=time_budget_sec,
        memory_key=memory_key,
        concurrency=concurrency,
        group_name=group,
    )
    payload["evaluation"] = tools.evaluate_batch_results_impl(payload["results"], batch_id=batch_id)
    await tools.persist_run_results_to_leaderboard(memory_key, payload["results"])
    return payload

INSTRUCTIONS = f"""\
You are a senior ML-engineer agent running karpathy/autoresearch-style research by
**editing train.py** and **batching parallel experiments**. Your goal: MINIMIZE
val_bpb (LOWER is better).

You operate in CODE MODE. Each turn, write ONE CODE1 block that calls the
available functions, OR reply in plain text when finished. The last expression in
your code block is returned as the observation.

Core tools:
- get_code_edit_history — **call first on resumed sessions**: prior edits, val_bpb, vs-best deltas
- get_baseline_train_code, edit_train_code_batch, read_train_code, get_promising_code
- inspect_dataset, search_arxiv
- get_leaderboard, compare_experiments

Saving edits (required for visible diffs and distinct runs):
- **Batch 1 only:** you may use ``config_overrides`` for a quick architecture/LR sweep via
  ``edit_train_code_batch(edits=[{{"title": "...", "config_overrides": {{"n_layer": 6}}, "change_summary": "..."}}])``.
- **Batch 2 and later:** every edit must include a **substantive ``train_py`` change**
  (learning-rate schedule, optimizer/weight_decay, grad clipping, warmup, etc.).
  ``config_overrides`` alone is **rejected** after the first batch — fork with
  ``parent_title`` and edit the training loop in ``train_py``.
- ``config_overrides`` fields: ``n_layer``, ``n_head``, ``n_embd``, ``dropout``,
  ``device_batch_size``, ``learning_rate``, ``time_budget_sec``, ``max_steps``.
- To fork a winner: set ``parent_title`` to the best title, then edit ``train_py``.
- Do **not** save baseline ``train.py`` without overrides — the platform rejects identical edits.
- Duplicate configs (same effective train.py + overrides) are rejected at run time.

Training budget (fair comparison across architectures):
- Default **max_steps={DEFAULT_MAX_STEPS}** with **time_budget_sec=45** as a safety cap.
  All models train for the same step count unless they hit the wall-clock limit.
- Check ``steps`` in batch results — if a run stopped early on time, the model may be too large.

Batch / fan-out tools:
- record_batch_plan(batch_id, experiments) — persist a multi-experiment plan
- get_batch_plan(batch_id) — reload a plan
- record_batch_hypotheses(experiments) — write hypotheses for every title in a batch
- edit_train_code_batch(edits) — save all ``train.py`` edits in one memory transaction
- run_experiment_batch(titles, concurrency=...) — parallel sandbox runs (LLM right-sized; OOM-healed)
- evaluate_batch_results(results, batch_id=...) — rank successes vs failures

Typical batch loop (aim for **≤8 code turns** before your plain-text summary):
0. If prior research exists in your directive, ``get_code_edit_history()`` then
   ``read_train_code(best_title)`` before planning new batches.
1. Turn 1: ``get_baseline_train_code()`` + ``inspect_dataset()``.
2. Turn 2: ``record_batch_plan`` then ``edit_train_code_batch(edits=[...])`` for the whole batch.
3. Turn 3: ``record_batch_hypotheses`` + ``run_experiment_batch(titles, concurrency=...)``.
4. Turn 4+: fork winners into the next batch with **train.py** edits, or reply in plain text when done.

Batch diversity (required):
- Every title in a batch must test a **distinct hypothesis** — no duplicate configs or renames.
- **Spread axes across the batch**: e.g. one edit tweaks depth/width, another changes the
  **training loop** (cosine LR, AdamW betas, weight decay), another regularization or batch size.
- Avoid LR micro-sweeps (±30% of the current best LR) after batch 1 — those rarely beat a plateau.
- Vary **one or two knobs per edit**; state the change in ``change_summary`` and
  ``record_batch_hypotheses``.
- Use ``evaluate_batch_results`` to see **which axis** helped, then explore under-tested axes.

Plateau rule (required):
- If **3 consecutive batches** fail to beat the global best val_bpb by more than **0.01**,
  stop hyperparameter micro-sweeps. Switch to **training-loop code edits** in ``train.py``
  (scheduler, optimizer, regularization, data/loss changes).

Rules:
- Use ``edit_train_code_batch`` for all code saves (including a single title: ``edits=[{{...}}]``).
- Every edit must keep ``run_training(config: ExperimentConfig) -> ExperimentResult``.
- Do NOT size compute — each run is LLM right-sized and retried automatically on OOM.
- Workshop limits: n_layer<={MAX_N_LAYER}, n_embd<={MAX_N_EMBD}, n_head<={MAX_N_HEAD},
  device_batch_size<={MAX_DEVICE_BATCH_SIZE}, seq_len=512.
- Monty sandbox: no imports, no dict mutation, no augmented assignment (`+=`).
- **Always finish with plain text (no code block)** once you have results to report.
"""

DEFAULT_MAX_TURNS = 50

def build_fanout_agent(*, max_turns: int = DEFAULT_MAX_TURNS) -> Agent:
    """Construct the fan-out agent (``code_mode=True``) with a configurable turn budget."""
    return Agent(
        name="parallelized-autoresearch",
        instructions=INSTRUCTIONS,
        model=MODEL,
        tools=[
            tools.search_arxiv,
            tools.inspect_dataset,
            tools.get_baseline_train_code,
            tools.get_code_edit_history,
            tools.edit_train_code_batch,
            tools.read_train_code,
            tools.get_promising_code,
            tools.get_leaderboard,
            tools.compare_experiments,
            tools.record_batch_plan,
            tools.get_batch_plan,
            tools.record_batch_hypotheses,
            run_experiment_batch,
            tools.evaluate_batch_results,
        ],
        max_turns=max_turns,
        call_llm=tools.call_llm,
        code_mode=True,
    )

# {{docs-fragment agent}}
@agent_env.task(report=True)
async def parallelized_autoresearch(
    n_experiments: int = 6,
    num_shards: int = DEFAULT_NUM_SHARDS,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
    batch_size: int = 3,
    max_turns: int = DEFAULT_MAX_TURNS,
) -> AutoresearchOutput:
    """Drive the fan-out code-edit MLE agent with sandbox batch execution."""
    bundle = await build_bundle(num_shards=num_shards)
    profile = await profile_bundle(bundle)

    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    persisted = await memory.read_json.aio("memory/leaderboard.json", default=[])
    promising = await memory.read_json.aio("memory/promising_code.json", default=[])
    history = await tools.load_research_history(memory_key)
    flyte.logger.info(
        "Fan-out agent restored %d messages, %d experiments, %d promising edits, best val_bpb=%s.",
        len(memory.messages),
        len(persisted),
        len(promising),
        history.get("best_val_bpb"),
    )

    events: list[dict[str, Any]] = []

    async def on_event(ev) -> None:
        events.append({"type": ev.type, "data": ev.data})
        if ev.type in ("tool_start", "tool_end", "tool_error", "turn_start", "agent_end"):
            tab = flyte.report.get_tab("Activity")
            tab.replace(ui.render_activity_log(events))
            await flyte.report.flush.aio()
        if ev.type == "tool_end" and ev.data.get("tool") in (
            "edit_train_code_batch",
            "<sandbox>",
        ):
            edits = await tools.load_saved_code_edits(memory_key)
            if edits:
                flyte.report.get_tab("Code edits").replace(ui.render_code_edits_panel(edits))
                await flyte.report.flush.aio()

    directive_text = ui.directive_code_edit_fanout(
        n_experiments,
        profile,
        memory_key,
        batch_size=batch_size,
        history=history,
    )

    token = agent_progress_cb.set(on_event)
    run_agent = build_fanout_agent(max_turns=max_turns)
    try:
        result = await run_agent.run.aio(directive_text, memory=memory)
    finally:
        agent_progress_cb.reset(token)

    leaderboard, best = ui.parse_leaderboard(
        memory.messages,
        promising_fallback=promising,
    )
    leaderboard_dicts = [dataclasses.asdict(e) for e in leaderboard]
    code_edits = await tools.load_saved_code_edits(memory_key)

    tab_lb = flyte.report.get_tab("Leaderboard")
    tab_lb.replace(ui.render_leaderboard(leaderboard, best))

    flyte.report.get_tab("Code edits").replace(
        ui.render_code_edits_panel(code_edits, best_title=best.title if best else None)
    )

    await memory.write_json.aio(
        "memory/leaderboard.json",
        leaderboard_dicts,
        actor="parallelized-autoresearch",
        reason=f"leaderboard after {len(leaderboard)} experiments",
    )
    await memory.save.aio()
    audit = await memory.audit_tail(20)
    hypotheses = await memory.read_json.aio("memory/hypotheses.json", default=[])
    promising = await memory.read_json.aio("memory/promising_code.json", default=[])

    tab_mem = flyte.report.get_tab("Memory")
    tab_mem.replace(
        ui.render_memory_panel(
            memory_key,
            len(memory.messages),
            leaderboard_dicts,
            audit,
            hypotheses,
            persisted_promising=promising,
            code_edits=code_edits,
        )
    )

    summary_body = result.summary or result.error or ""
    if result.error and leaderboard:
        best_line = f" Best val_bpb so far: {best.val_bpb} ({best.title})." if best and best.val_bpb else ""
        summary_body = f"{result.error}{best_line}"

    await flyte.report.replace.aio(
        ui.render_summary(
            directive_text,
            leaderboard,
            best,
            summary_body,
            code_edits=code_edits,
        )
    )
    await flyte.report.flush.aio()

    return AutoresearchOutput(
        directive=directive_text,
        dataset_profile=profile,
        best=best,
        leaderboard=leaderboard,
        summary=summary_body,
        memory_key=memory_key,
        total_experiments=len(leaderboard),
    )

# {{/docs-fragment agent}}

# {{docs-fragment main}}
if __name__ == "__main__":
    import argparse
    import asyncio
    import os

    parser = argparse.ArgumentParser(description="Parallelized autoresearch agent (CODE MODE)")
    parser.add_argument("--n-experiments", type=int, default=6)
    parser.add_argument("--batch-size", type=int, default=3)
    parser.add_argument("--max-turns", type=int, default=DEFAULT_MAX_TURNS)
    parser.add_argument("--num-shards", type=int, default=DEFAULT_NUM_SHARDS)
    parser.add_argument("--memory-key", default=tools.MEMORY_KEY_FANOUT)
    parser.add_argument(
        "--config",
        default=os.environ.get("FLYTE_CONFIG", os.path.expanduser("~/.flyte/config.yaml")),
    )
    args = parser.parse_args()

    flyte.init_from_config(args.config, image_builder="remote")

    async def main() -> None:
        run = await flyte.with_runcontext(copy_style="all").run.aio(
            parallelized_autoresearch,
            n_experiments=args.n_experiments,
            num_shards=args.num_shards,
            memory_key=args.memory_key,
            batch_size=args.batch_size,
            max_turns=args.max_turns,
        )
        print(f"View run at: {run.url}")

    asyncio.run(main())
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/parallelized_autoresearch/parallelized_autoresearch.py*

Batch fan-out calls `flyte.map.aio(run_experiment, ...)` from `run_experiment_batch`. That path invokes `run_experiment.aio()` directly (**not** through the agent registry) so the example binds `call_llm` and `model` on the tool after construction (see the `dataclasses.replace` block above). With Flyte SDK ≥ 2.5.5, `AgentTool.aio` routes through `call_handler`, so every mapped experiment gets LLM right-sizing even when the agent only exposes `run_experiment_batch` in code mode.

## The fan-out agent task

The driver task `parallelized_autoresearch` restores prior memory (default key `parallelized-autoresearch`), streams Activity / Leaderboard / Code edits / Memory report tabs, and runs the code-mode agent loop. The agent tool registry is trimmed to the batch workflow: `run_experiment` is internal to `run_experiment_batch`, not a sandbox function the LLM calls directly.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.5.5",
#    "litellm",
#    "httpx",
#    "pydantic-monty",
#    "unionai-sandbox[flyte]",
#    "torch",
#    "numpy",
#    "pyarrow",
#    "requests",
#    "tiktoken",
#    "rustbpe",
# ]
# main = "parallelized_autoresearch"
# params = "--n-experiments 6 --batch-size 3 --num-shards 1"
# ///
"""Parallelized autoresearch agent — code-mode MLE agent with batched sandbox experiments."""

from __future__ import annotations

import dataclasses
from typing import Any

import flyte
import flyte.report
from flyte.ai.agents import Agent, MemoryStore, agent_progress_cb, tool

from autoresearch_types import AutoresearchOutput, DEFAULT_MAX_STEPS, DEFAULT_NUM_SHARDS, MAX_DEVICE_BATCH_SIZE, MAX_N_EMBD, MAX_N_HEAD, MAX_N_LAYER
from bundle import agent_env, build_bundle, experiment_env, materialize_cache, profile_bundle
import tools
import ui

MODEL = "claude-sonnet-4-6"

# {{docs-fragment run_experiment}}
@tool(call_handler=tools.right_size)
@experiment_env.task
async def run_experiment(
    title: str,
    time_budget_sec: int = 45,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
) -> dict:
    """Train using agent-edited ``train.py`` with LLM right-sizing and OOM self-healing."""
    train_py = await tools.load_train_code(memory_key, title)
    config_overrides = await tools.load_config_overrides(memory_key, title)
    duplicate = await tools.check_duplicate_config(memory_key, title, train_py, config_overrides)
    if duplicate:
        result = {
            "success": False,
            "title": title,
            "error": (
                f"Duplicate config of '{duplicate['duplicate_of']}' "
                f"(signature {duplicate['config_signature']}); change train.py or overrides."
            ),
            "duplicate_of": duplicate["duplicate_of"],
        }
        await tools.record_experiment_result(
            memory_key,
            result,
            actor="parallelized-autoresearch",
        )
        return result
    bundle = await build_bundle()
    cache_dir = await materialize_cache(bundle)
    result = await tools.run_train_in_sandbox(
        cache_dir,
        train_py,
        title=title,
        time_budget_sec=time_budget_sec,
        config_overrides=config_overrides or None,
    )
    if result.get("success"):
        await tools.record_promising_run(memory_key, title, result)
        await tools.register_config_signature(
            memory_key,
            title,
            train_py,
            config_overrides,
            actor="parallelized-autoresearch",
        )
    await tools.record_experiment_result(
        memory_key,
        result,
        actor="parallelized-autoresearch",
    )
    return result

# ``flyte.map`` invokes ``run_experiment.aio`` directly (not through the agent
# registry), so bind the LLM callback and model here for ``call_handler`` right-sizing.
run_experiment = dataclasses.replace(
    run_experiment,
    call_llm=tools.call_llm,
    model=MODEL,
)
# {{/docs-fragment run_experiment}}

@tool
@agent_env.task
async def run_experiment_batch(
    titles: list[str],
    time_budget_sec: int = 45,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
    concurrency: int = 4,
    batch_id: str = "",
) -> dict:
    """Run multiple ``run_experiment`` calls in parallel via ``flyte.map``.

    Prefer this over hand-rolling ``flyte_map`` when you already have a list of
    experiment titles with saved ``train.py`` edits.

    Args:
        titles: Experiment titles whose code was saved with ``edit_train_code_batch``.
        time_budget_sec: Wall-clock budget passed to each run.
        memory_key: Memory namespace from your directive.
        concurrency: Max parallel sandbox runs (default 4).
        batch_id: Optional label attached to the returned batch metadata.

    Returns:
        A dict with ``batch_size``, ``titles``, ``results``, and ``evaluation``
        (from :func:`evaluate_batch_results`).
    """
    group = batch_id or f"batch-{len(titles)}"
    payload = await tools.run_experiment_batch_impl(
        run_experiment,
        titles,
        time_budget_sec=time_budget_sec,
        memory_key=memory_key,
        concurrency=concurrency,
        group_name=group,
    )
    payload["evaluation"] = tools.evaluate_batch_results_impl(payload["results"], batch_id=batch_id)
    await tools.persist_run_results_to_leaderboard(memory_key, payload["results"])
    return payload

INSTRUCTIONS = f"""\
You are a senior ML-engineer agent running karpathy/autoresearch-style research by
**editing train.py** and **batching parallel experiments**. Your goal: MINIMIZE
val_bpb (LOWER is better).

You operate in CODE MODE. Each turn, write ONE CODE2 block that calls the
available functions, OR reply in plain text when finished. The last expression in
your code block is returned as the observation.

Core tools:
- get_code_edit_history — **call first on resumed sessions**: prior edits, val_bpb, vs-best deltas
- get_baseline_train_code, edit_train_code_batch, read_train_code, get_promising_code
- inspect_dataset, search_arxiv
- get_leaderboard, compare_experiments

Saving edits (required for visible diffs and distinct runs):
- **Batch 1 only:** you may use ``config_overrides`` for a quick architecture/LR sweep via
  ``edit_train_code_batch(edits=[{{"title": "...", "config_overrides": {{"n_layer": 6}}, "change_summary": "..."}}])``.
- **Batch 2 and later:** every edit must include a **substantive ``train_py`` change**
  (learning-rate schedule, optimizer/weight_decay, grad clipping, warmup, etc.).
  ``config_overrides`` alone is **rejected** after the first batch — fork with
  ``parent_title`` and edit the training loop in ``train_py``.
- ``config_overrides`` fields: ``n_layer``, ``n_head``, ``n_embd``, ``dropout``,
  ``device_batch_size``, ``learning_rate``, ``time_budget_sec``, ``max_steps``.
- To fork a winner: set ``parent_title`` to the best title, then edit ``train_py``.
- Do **not** save baseline ``train.py`` without overrides — the platform rejects identical edits.
- Duplicate configs (same effective train.py + overrides) are rejected at run time.

Training budget (fair comparison across architectures):
- Default **max_steps={DEFAULT_MAX_STEPS}** with **time_budget_sec=45** as a safety cap.
  All models train for the same step count unless they hit the wall-clock limit.
- Check ``steps`` in batch results — if a run stopped early on time, the model may be too large.

Batch / fan-out tools:
- record_batch_plan(batch_id, experiments) — persist a multi-experiment plan
- get_batch_plan(batch_id) — reload a plan
- record_batch_hypotheses(experiments) — write hypotheses for every title in a batch
- edit_train_code_batch(edits) — save all ``train.py`` edits in one memory transaction
- run_experiment_batch(titles, concurrency=...) — parallel sandbox runs (LLM right-sized; OOM-healed)
- evaluate_batch_results(results, batch_id=...) — rank successes vs failures

Typical batch loop (aim for **≤8 code turns** before your plain-text summary):
0. If prior research exists in your directive, ``get_code_edit_history()`` then
   ``read_train_code(best_title)`` before planning new batches.
1. Turn 1: ``get_baseline_train_code()`` + ``inspect_dataset()``.
2. Turn 2: ``record_batch_plan`` then ``edit_train_code_batch(edits=[...])`` for the whole batch.
3. Turn 3: ``record_batch_hypotheses`` + ``run_experiment_batch(titles, concurrency=...)``.
4. Turn 4+: fork winners into the next batch with **train.py** edits, or reply in plain text when done.

Batch diversity (required):
- Every title in a batch must test a **distinct hypothesis** — no duplicate configs or renames.
- **Spread axes across the batch**: e.g. one edit tweaks depth/width, another changes the
  **training loop** (cosine LR, AdamW betas, weight decay), another regularization or batch size.
- Avoid LR micro-sweeps (±30% of the current best LR) after batch 1 — those rarely beat a plateau.
- Vary **one or two knobs per edit**; state the change in ``change_summary`` and
  ``record_batch_hypotheses``.
- Use ``evaluate_batch_results`` to see **which axis** helped, then explore under-tested axes.

Plateau rule (required):
- If **3 consecutive batches** fail to beat the global best val_bpb by more than **0.01**,
  stop hyperparameter micro-sweeps. Switch to **training-loop code edits** in ``train.py``
  (scheduler, optimizer, regularization, data/loss changes).

Rules:
- Use ``edit_train_code_batch`` for all code saves (including a single title: ``edits=[{{...}}]``).
- Every edit must keep ``run_training(config: ExperimentConfig) -> ExperimentResult``.
- Do NOT size compute — each run is LLM right-sized and retried automatically on OOM.
- Workshop limits: n_layer<={MAX_N_LAYER}, n_embd<={MAX_N_EMBD}, n_head<={MAX_N_HEAD},
  device_batch_size<={MAX_DEVICE_BATCH_SIZE}, seq_len=512.
- Monty sandbox: no imports, no dict mutation, no augmented assignment (`+=`).
- **Always finish with plain text (no code block)** once you have results to report.
"""

DEFAULT_MAX_TURNS = 50

def build_fanout_agent(*, max_turns: int = DEFAULT_MAX_TURNS) -> Agent:
    """Construct the fan-out agent (``code_mode=True``) with a configurable turn budget."""
    return Agent(
        name="parallelized-autoresearch",
        instructions=INSTRUCTIONS,
        model=MODEL,
        tools=[
            tools.search_arxiv,
            tools.inspect_dataset,
            tools.get_baseline_train_code,
            tools.get_code_edit_history,
            tools.edit_train_code_batch,
            tools.read_train_code,
            tools.get_promising_code,
            tools.get_leaderboard,
            tools.compare_experiments,
            tools.record_batch_plan,
            tools.get_batch_plan,
            tools.record_batch_hypotheses,
            run_experiment_batch,
            tools.evaluate_batch_results,
        ],
        max_turns=max_turns,
        call_llm=tools.call_llm,
        code_mode=True,
    )

# {{docs-fragment agent}}
@agent_env.task(report=True)
async def parallelized_autoresearch(
    n_experiments: int = 6,
    num_shards: int = DEFAULT_NUM_SHARDS,
    memory_key: str = tools.MEMORY_KEY_FANOUT,
    batch_size: int = 3,
    max_turns: int = DEFAULT_MAX_TURNS,
) -> AutoresearchOutput:
    """Drive the fan-out code-edit MLE agent with sandbox batch execution."""
    bundle = await build_bundle(num_shards=num_shards)
    profile = await profile_bundle(bundle)

    memory = await MemoryStore.get_or_create.aio(key=memory_key)
    persisted = await memory.read_json.aio("memory/leaderboard.json", default=[])
    promising = await memory.read_json.aio("memory/promising_code.json", default=[])
    history = await tools.load_research_history(memory_key)
    flyte.logger.info(
        "Fan-out agent restored %d messages, %d experiments, %d promising edits, best val_bpb=%s.",
        len(memory.messages),
        len(persisted),
        len(promising),
        history.get("best_val_bpb"),
    )

    events: list[dict[str, Any]] = []

    async def on_event(ev) -> None:
        events.append({"type": ev.type, "data": ev.data})
        if ev.type in ("tool_start", "tool_end", "tool_error", "turn_start", "agent_end"):
            tab = flyte.report.get_tab("Activity")
            tab.replace(ui.render_activity_log(events))
            await flyte.report.flush.aio()
        if ev.type == "tool_end" and ev.data.get("tool") in (
            "edit_train_code_batch",
            "<sandbox>",
        ):
            edits = await tools.load_saved_code_edits(memory_key)
            if edits:
                flyte.report.get_tab("Code edits").replace(ui.render_code_edits_panel(edits))
                await flyte.report.flush.aio()

    directive_text = ui.directive_code_edit_fanout(
        n_experiments,
        profile,
        memory_key,
        batch_size=batch_size,
        history=history,
    )

    token = agent_progress_cb.set(on_event)
    run_agent = build_fanout_agent(max_turns=max_turns)
    try:
        result = await run_agent.run.aio(directive_text, memory=memory)
    finally:
        agent_progress_cb.reset(token)

    leaderboard, best = ui.parse_leaderboard(
        memory.messages,
        promising_fallback=promising,
    )
    leaderboard_dicts = [dataclasses.asdict(e) for e in leaderboard]
    code_edits = await tools.load_saved_code_edits(memory_key)

    tab_lb = flyte.report.get_tab("Leaderboard")
    tab_lb.replace(ui.render_leaderboard(leaderboard, best))

    flyte.report.get_tab("Code edits").replace(
        ui.render_code_edits_panel(code_edits, best_title=best.title if best else None)
    )

    await memory.write_json.aio(
        "memory/leaderboard.json",
        leaderboard_dicts,
        actor="parallelized-autoresearch",
        reason=f"leaderboard after {len(leaderboard)} experiments",
    )
    await memory.save.aio()
    audit = await memory.audit_tail(20)
    hypotheses = await memory.read_json.aio("memory/hypotheses.json", default=[])
    promising = await memory.read_json.aio("memory/promising_code.json", default=[])

    tab_mem = flyte.report.get_tab("Memory")
    tab_mem.replace(
        ui.render_memory_panel(
            memory_key,
            len(memory.messages),
            leaderboard_dicts,
            audit,
            hypotheses,
            persisted_promising=promising,
            code_edits=code_edits,
        )
    )

    summary_body = result.summary or result.error or ""
    if result.error and leaderboard:
        best_line = f" Best val_bpb so far: {best.val_bpb} ({best.title})." if best and best.val_bpb else ""
        summary_body = f"{result.error}{best_line}"

    await flyte.report.replace.aio(
        ui.render_summary(
            directive_text,
            leaderboard,
            best,
            summary_body,
            code_edits=code_edits,
        )
    )
    await flyte.report.flush.aio()

    return AutoresearchOutput(
        directive=directive_text,
        dataset_profile=profile,
        best=best,
        leaderboard=leaderboard,
        summary=summary_body,
        memory_key=memory_key,
        total_experiments=len(leaderboard),
    )

# {{/docs-fragment agent}}

# {{docs-fragment main}}
if __name__ == "__main__":
    import argparse
    import asyncio
    import os

    parser = argparse.ArgumentParser(description="Parallelized autoresearch agent (CODE MODE)")
    parser.add_argument("--n-experiments", type=int, default=6)
    parser.add_argument("--batch-size", type=int, default=3)
    parser.add_argument("--max-turns", type=int, default=DEFAULT_MAX_TURNS)
    parser.add_argument("--num-shards", type=int, default=DEFAULT_NUM_SHARDS)
    parser.add_argument("--memory-key", default=tools.MEMORY_KEY_FANOUT)
    parser.add_argument(
        "--config",
        default=os.environ.get("FLYTE_CONFIG", os.path.expanduser("~/.flyte/config.yaml")),
    )
    args = parser.parse_args()

    flyte.init_from_config(args.config, image_builder="remote")

    async def main() -> None:
        run = await flyte.with_runcontext(copy_style="all").run.aio(
            parallelized_autoresearch,
            n_experiments=args.n_experiments,
            num_shards=args.num_shards,
            memory_key=args.memory_key,
            batch_size=args.batch_size,
            max_turns=args.max_turns,
        )
        print(f"View run at: {run.url}")

    asyncio.run(main())
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/parallelized_autoresearch/parallelized_autoresearch.py*

## Run the agent

### Create secrets

Register an Anthropic API key for agent LLM calls and for per-experiment resource sizing inside `call_handler`:

CODE3

### Run remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/parallelized_autoresearch):

CODE4

Use `--memory-key` to resume a prior research session (default: `parallelized-autoresearch`). Pass a unique key (for example `parallelized-autoresearch-20260622-215057`) to start with empty memory. Code mode needs more turns than JSON tool mode. Increase `--max-turns` for larger sweeps.

Or invoke the agent task directly with `flyte run` (snake_case task inputs):

CODE5

> [!NOTE]
> The first run downloads climbmix data shards and trains a BPE tokenizer. Subsequent runs reuse cached bundle tasks. Requires **Flyte SDK ≥ 2.5.5** for `call_handler` support in code mode and on `AgentTool.aio` (used by `flyte.map` fan-out).

See also the single-task [Autoresearch agent](../autoresearch/_index) tutorial for the Claude Code + pull-request workflow.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/autosec-research-agent ===

# AutoSec researcher agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autosec_research_agent).

This tutorial demonstrates an autonomous security-research agent on Flyte. The pipeline fans out across bundled C source files (each with a planted memory-corruption bug), runs static analysis, uses a `flyte.ai.agents.Agent` to hypothesize vulnerabilities, builds proof-of-concept payloads, and validates exploits inside an on-device [unionai-sandbox](../../../user-guide/agents/sandboxing/_index) user-namespace session.

Flyte provides:

- **Parallel fan-out** across every target file with `asyncio.gather`
- **Self-healing tasks**: LLM timeouts, malformed JSON, and OOM during static analysis retry with bounded resources
- **Sandbox isolation**: PoC compilation and execution never runs on the orchestration node
- **Live HTML reports** with per-target detail tabs in the Flyte UI

> [!WARNING]
> This example analyzes deliberately vulnerable C code and runs generated exploit payloads in a sandbox. Use it only in controlled environments.

## Define the task environment

The agent needs an Anthropic API key and a container image with `gcc` for sandbox compilation.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "unionai-sandbox",
#    "litellm",
# ]
# main = "run_autosec_agent"
# params = ""
# ///
"""AutoSec researcher agent — parallel vulnerability analysis with sandbox PoC validation."""

from __future__ import annotations

import asyncio
import html
import json
import os
import pathlib
import re
from typing import Any

import flyte
import flyte.errors
import flyte.report
from flyte.ai.agents import Agent

HERE = pathlib.Path(__file__).parent
TARGETS_DIR = HERE / "targets"
MODEL = os.getenv("AUTOSEC_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="autosec-research-agent", pre=True).with_apt_packages("gcc")

env = flyte.TaskEnvironment(
    name="autosec-research-agent",
    image=main_img,
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    include=[str(TARGETS_DIR)],
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

def _attempt() -> int:
    tc = flyte.ctx()
    return tc.attempt_number if tc is not None else 0

def _force(flag: str) -> bool:
    return bool(os.getenv(flag) or os.getenv("AUTOSEC_FORCE_ALL"))

def _extract_json(text: str) -> dict[str, Any]:
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        raise ValueError(f"no JSON object in model reply: {text[:200]!r}")
    blob = match.group(0)
    try:
        return json.loads(blob)
    except json.JSONDecodeError:
        fixed = re.sub(r'\\(?!["\\/bfnrtu])', r"\\\\", blob)
        return json.loads(fixed)

# --- Stage 1: static analysis (CPU, OOM-prone) ------------------------------
@env.task(retries=2, timeout=30)
async def scan_static(source: str, scope: str = "whole") -> str:
    """Cheap stand-in for whole-program analysis (Joern/CodeQL in the real system)."""
    try:
        if scope == "whole" and _force("AUTOSEC_FORCE_OOM") and _attempt() == 0:
            raise flyte.errors.OOMError("whole-program graph exceeded memory limit")
        findings = _grep_dangerous_calls(source)
        return findings or "(no dangerous-call sites found)"
    except flyte.errors.OOMError as exc:
        print(f"[scan_static] {exc}; escalating resources + narrowing scope")
        return await scan_static.override(
            short_name="scan_static_more_resources", resources=flyte.Resources(cpu=2, memory="4Gi")
        )(source, scope="file")

def _grep_dangerous_calls(source: str) -> str:
    hits = []
    for i, line in enumerate(source.splitlines(), start=1):
        for fn in ("strcpy", "strcat", "sprintf", "gets", "memcpy"):
            if fn in line:
                hits.append(f"L{i}: {fn} -> {line.strip()}")
    return "\n".join(hits)

# --- Stage 2: hypothesize the vulnerability (LLM via Agent) ------------------
ANALYSIS_INSTRUCTIONS = """\
You are a vulnerability researcher. Your job is to determine whether a given \
C source file contains an exploitable memory-corruption bug reachable from argv.

You have access to these tools during your analysis:
- scan_static: Run static analysis on the source to find dangerous function calls.
- build_poc: Build a proof-of-concept payload (do not call during analysis).
- validate_in_sandbox: Compile and run the target with a PoC input (do not call during analysis).

Focus on analyzing the source and the provided static analysis findings. Call \
scan_static only if you need additional details about dangerous function usage.

Reply with ONLY a JSON object (no prose, no markdown fences):
If vulnerable: {"vulnerable": true, "function": str, \
"buffer_size": int (bytes of the overflowable buffer), "vuln_class": str, \
"reasoning": str}.
If the code looks safe (bounded copies, length checks, snprintf/strlcpy, \
etc.): {"vulnerable": false, "reasoning": str}.
"""

# --- Stage 3: build a proof-of-concept --------------------------------------
@env.task(retries=2, timeout=90)
async def build_poc(hypothesis: dict) -> dict:
    buffer_size = int(hypothesis.get("buffer_size", 64))
    payload_len = buffer_size + 64
    return {
        "payload_len": payload_len,
        "payload_repr": f'"A" * {payload_len}',
        "target_function": hypothesis.get("function", "greet"),
    }

# --- Stage 4: validate in an on-device sandbox -------------------------------
@env.task(retries=2, timeout=300)
async def validate_in_sandbox(source: str, poc: dict) -> dict:
    """Compile + run the target with the PoC input inside an on-device sandbox.

    The exploit code runs in a user-namespace sandbox on the same machine, never
    on the Flyte orchestration node (SPEC §2.6 / §7). The session is torn down
    in __aexit__ regardless of outcome (SPEC VD-5) so a stuck or failed run
    cannot leak resources.
    """
    import tempfile

    from union import sandbox as sb

    with tempfile.TemporaryDirectory() as work:
        async with sb.on_device.session(host_work_dir=work, backend="userns") as sbx:
            await sbx.put_bytes(f"{work}/target.c", source.encode())

            compile_proc = await sbx.run(
                f"gcc -fno-stack-protector -w -o {work}/target {work}/target.c",
                stdout=True,
                stderr=True,
                timeout_s=60,
            )
            compile_out, compile_err = await compile_proc.communicate_text()
            log = compile_out + compile_err
            if "error" in log.lower():
                return {
                    "triggered": False,
                    "sandbox_exit_code": -1,
                    "log": f"COMPILE_FAILED\n{log}",
                }

            payload = "A" * int(poc["payload_len"])
            run_proc = await sbx.run(
                f"{work}/target {payload}",
                stdout=True,
                stderr=True,
                timeout_s=60,
            )
            run_out, run_err = await run_proc.communicate_text()
            log = run_out + "\n" + run_err
            triggered = "SIGSEGV" in log

            return {
                "triggered": bool(triggered),
                "sandbox_exit_code": getattr(run_proc, "returncode", 0),
                "log": log,
            }

# --- Agent + hypothesize task (depends on all tools above) ------------------
hypothesis_agent = Agent(
    name="autosec-hypothesis",
    instructions=ANALYSIS_INSTRUCTIONS,
    model=MODEL,
    tools=[scan_static, build_poc, validate_in_sandbox],
    max_turns=6,
)

@env.task(retries=3, timeout=20)
async def hypothesize(source: str, static_findings: str) -> dict:
    prompt = (
        "Analyze this C source file for memory-corruption vulnerabilities.\n\n"
        f"SOURCE:\n{source}\n\nDANGEROUS CALLS:\n{static_findings}\n"
    )

    # Beat A: hang on the first attempt -> task timeout -> retry.
    timeout_on = _force("AUTOSEC_FORCE_LLM_TIMEOUT") and _attempt() == 0
    bad_on = _force("AUTOSEC_FORCE_BAD_TOOL_CALL")

    if timeout_on:
        await asyncio.sleep(600)

    result = await hypothesis_agent.run.aio(prompt, memory=[])
    raw = result.summary or ""

    # Beat B: simulate a hallucinated/malformed tool call. When the timeout beat
    # is also active it consumes attempt 0, so defer this to attempt 1 — that way
    # both beats are actually demonstrated in a single run (e.g. AUTOSEC_FORCE_ALL).
    bad_attempt = 1 if _force("AUTOSEC_FORCE_LLM_TIMEOUT") else 0
    if bad_on and _attempt() == bad_attempt:
        raw = "Sure! The bug is somewhere around here, trust me."

    hyp = _extract_json(raw)
    if "vulnerable" not in hyp:
        hyp["vulnerable"] = "buffer_size" in hyp
    if hyp.get("vulnerable") and "buffer_size" not in hyp:
        raise ValueError(f"vulnerable hypothesis missing buffer_size: {hyp}")
    return hyp

# --- Orchestration ----------------------------------------------------------
def _load_targets() -> dict[str, str]:
    return {p.name: p.read_text() for p in sorted(TARGETS_DIR.glob("*.c"))}

@env.task
async def analyze_target(name: str, source: str) -> dict:
    findings = await scan_static(source)
    hypothesis = await hypothesize(source, findings)

    if not hypothesis.get("vulnerable"):
        poc: dict = {}
        verdict = {"triggered": False, "skipped": True}
    else:
        poc = await build_poc(hypothesis)
        verdict = await validate_in_sandbox(source, poc)

    return {
        "target": name,
        "static_findings": findings,
        "hypothesis": hypothesis,
        "poc": poc,
        "verdict": verdict,
    }

_REPORT_CSS = """
<style>
  .autosec { --bg:#ffffff; --card:#f7f8fa; --line:#e3e7ec; --muted:#5b6675;
    --text:#1b2330; --red:#c0392b; --amber:#b6791f; --green:#1e7e34; --accent:#1f6feb;
    font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;
    background:var(--bg); color:var(--text); padding:24px; border-radius:12px;
    border:1px solid var(--line); }
  .autosec h2 { margin:0 0 4px; font-size:20px; letter-spacing:.2px; }
  .autosec .sub { color:var(--muted); font-size:13px; margin:0 0 20px; }
  .autosec .cards { display:flex; gap:12px; flex-wrap:wrap; margin-bottom:22px; }
  .autosec .card { background:var(--card); border:1px solid var(--line);
    border-radius:10px; padding:14px 18px; min-width:120px; }
  .autosec .card .n { font-size:26px; font-weight:700; line-height:1; }
  .autosec .card .l { color:var(--muted); font-size:12px; margin-top:6px;
    text-transform:uppercase; letter-spacing:.6px; }
  .autosec table { width:100%; table-layout:fixed; border-collapse:collapse; font-size:13px;
    background:#fff; border:1px solid var(--line); border-radius:10px; overflow:hidden; }
  .autosec th, .autosec td { overflow-wrap:anywhere; }
  .autosec thead th { background:#eef1f5; color:var(--muted); text-align:left;
    font-weight:600; font-size:11px; text-transform:uppercase; letter-spacing:.6px;
    padding:11px 14px; border-bottom:1px solid var(--line); }
  .autosec tbody td { padding:11px 14px; border-bottom:1px solid var(--line);
    vertical-align:top; }
  .autosec tbody tr:last-child td { border-bottom:none; }
  .autosec tbody tr:nth-child(even) { background:#fafbfc; }
  .autosec tbody tr:hover { background:#eef4ff; }
  .autosec code { background:#eef1f5; border:1px solid var(--line); border-radius:5px;
    padding:1px 6px; font-family:ui-monospace,SFMono-Regular,Menlo,monospace; font-size:12px; }
  .autosec .num { text-align:right; font-variant-numeric:tabular-nums; }
  .autosec .reason { color:var(--muted); line-height:1.45; }
  .autosec .badge { display:inline-block; padding:3px 10px; border-radius:999px;
    font-size:11px; font-weight:700; letter-spacing:.4px; white-space:nowrap; }
  .autosec .b-exploited { background:rgba(192,57,43,.10); color:var(--red);
    border:1px solid rgba(192,57,43,.35); }
  .autosec .b-vuln { background:rgba(182,121,31,.12); color:var(--amber);
    border:1px solid rgba(182,121,31,.35); }
  .autosec .b-secure { background:rgba(30,126,52,.10); color:var(--green);
    border:1px solid rgba(30,126,52,.35); }
  .autosec col.c-target  { width:15%; }
  .autosec col.c-status  { width:11%; }
  .autosec col.c-class   { width:11%; }
  .autosec col.c-fn      { width:11%; }
  .autosec col.c-buf     { width:8%; }
  .autosec col.c-payload { width:8%; }
  .autosec col.c-exit    { width:6%; }
  .autosec col.c-reason  { width:30%; }
  .autosec .kv { display:flex; flex-wrap:wrap; gap:12px 28px; margin:16px 0 4px; }
  .autosec .kv .k { color:var(--muted); font-size:11px; text-transform:uppercase; letter-spacing:.6px; }
  .autosec .kv .v { font-weight:600; font-size:14px; margin-top:3px; }
  .autosec .section-label { font-size:11px; text-transform:uppercase; letter-spacing:.6px;
    color:var(--muted); margin:22px 0 7px; }
  .autosec .reason-block { background:var(--card); border:1px solid var(--line);
    border-radius:8px; padding:13px 15px; line-height:1.5; font-size:13px; }
  .autosec pre.code { background:#f6f8fa; border:1px solid var(--line); border-radius:8px;
    padding:14px 16px; overflow:auto; font-family:ui-monospace,SFMono-Regular,Menlo,monospace;
    font-size:12.5px; line-height:1.5; color:#1b2330; margin:0; }
  .autosec .subtabs > input[type=radio] { position:absolute; opacity:0; pointer-events:none; }
  .autosec .subnav { display:flex; flex-wrap:wrap; gap:4px; border-bottom:1px solid var(--line);
    margin:8px 0 18px; }
  .autosec .subnav label { display:inline-flex; align-items:center; gap:8px; padding:8px 14px;
    font-size:12.5px; cursor:pointer; border:1px solid transparent; border-bottom:none;
    border-radius:8px 8px 0 0; color:var(--muted); margin-bottom:-1px; }
  .autosec .subnav label:hover { background:var(--card); color:var(--text); }
  .autosec .subnav .dot { width:8px; height:8px; border-radius:50%; }
  .autosec .dot.b-exploited { background:var(--red); }
  .autosec .dot.b-vuln { background:var(--amber); }
  .autosec .dot.b-secure { background:var(--green); }
  .autosec .panel { display:none; }
</style>
"""

def _status(finding: dict) -> tuple[str, str]:
    hyp = finding.get("hypothesis") or {}
    verdict = finding.get("verdict") or {}
    if not hyp.get("vulnerable"):
        return "b-secure", "SECURE"
    if verdict.get("triggered"):
        return "b-exploited", "EXPLOITED"
    return "b-vuln", "VULNERABLE"

def _render_report_html(findings: list[dict]) -> str:
    exploited = sum(1 for f in findings if (f.get("verdict") or {}).get("triggered"))
    vulnerable = sum(1 for f in findings if (f.get("hypothesis") or {}).get("vulnerable"))
    secure = len(findings) - vulnerable

    rows = []
    for f in sorted(findings, key=lambda x: x["target"]):
        hyp = f.get("hypothesis") or {}
        verdict = f.get("verdict") or {}
        cls, label = _status(f)
        is_vuln = bool(hyp.get("vulnerable"))
        vuln_class = hyp.get("vuln_class", "\u2014") if is_vuln else "\u2014"
        fn = hyp.get("function", "\u2014") if is_vuln else "\u2014"
        buf = hyp.get("buffer_size", "\u2014") if is_vuln else "\u2014"
        payload = (f.get("poc") or {}).get("payload_len", "\u2014") if is_vuln else "\u2014"
        exit_code = verdict.get("sandbox_exit_code", "\u2014")
        rows.append(
            "<tr>"
            f"<td><code>{html.escape(str(f['target']))}</code></td>"
            f'<td><span class="badge {cls}">{label}</span></td>'
            f"<td>{html.escape(str(vuln_class))}</td>"
            f"<td><code>{html.escape(str(fn))}</code></td>"
            f'<td class="num">{html.escape(str(buf))}</td>'
            f'<td class="num">{html.escape(str(payload))}</td>'
            f'<td class="num">{html.escape(str(exit_code))}</td>'
            f'<td class="reason">{html.escape(str(hyp.get("reasoning", "")))}</td>'
            "</tr>"
        )

    return f"""{_REPORT_CSS}
    <div class="autosec">
      <h2>AutoSec &middot; security findings report</h2>
      <p class="sub">{len(findings)} target(s) analyzed in parallel &middot; PoCs validated in isolated sandbox.</p>
      <div class="cards">
        <div class="card"><div class="n">{len(findings)}</div><div class="l">Targets</div></div>
        <div class="card"><div class="n" style="color:#ff6b6b">{exploited}</div><div class="l">Exploited</div></div>
        <div class="card">
          <div class="n" style="color:#ffb454">{vulnerable - exploited}</div>
          <div class="l">Vuln, PoC failed</div>
        </div>
        <div class="card"><div class="n" style="color:#3fb950">{secure}</div><div class="l">Secure</div></div>
      </div>
      <table>
        <colgroup>
          <col class="c-target"><col class="c-status"><col class="c-class"><col class="c-fn">
          <col class="c-buf"><col class="c-payload"><col class="c-exit"><col class="c-reason">
        </colgroup>
        <thead>
          <tr>
            <th>Target</th><th>Status</th><th>Vuln class</th><th>Function</th>
            <th class="num">Buffer&nbsp;(B)</th><th class="num">Payload&nbsp;(B)</th>
            <th class="num">Exit</th><th>Reasoning</th>
          </tr>
        </thead>
        <tbody>{"".join(rows)}</tbody>
      </table>
    </div>
    """

def _target_detail_html(finding: dict, source: str) -> str:
    hyp = finding.get("hypothesis") or {}
    verdict = finding.get("verdict") or {}
    cls, label = _status(finding)
    is_vuln = bool(hyp.get("vulnerable"))

    def cell(k: str, v: Any) -> str:
        return f'<div><div class="k">{html.escape(k)}</div><div class="v">{html.escape(str(v))}</div></div>'

    triggered = verdict.get("triggered")
    verdict_txt = "skipped (secure)" if verdict.get("skipped") else ("triggered" if triggered else "not triggered")
    stats = "".join(
        [
            cell("Vuln class", hyp.get("vuln_class", "\u2014") if is_vuln else "\u2014"),
            cell("Function", hyp.get("function", "\u2014") if is_vuln else "\u2014"),
            cell("Buffer (B)", hyp.get("buffer_size", "\u2014") if is_vuln else "\u2014"),
            cell("Payload (B)", (finding.get("poc") or {}).get("payload_len", "\u2014") if is_vuln else "\u2014"),
            cell("Sandbox exit", verdict.get("sandbox_exit_code", "\u2014")),
            cell("PoC", verdict_txt),
        ]
    )
    reasoning = html.escape(str(hyp.get("reasoning", "")) or "\u2014")
    code = html.escape(source or "(source unavailable)")
    return f"""
      <h3 style="margin:0 0 4px"><code>{html.escape(str(finding["target"]))}</code>
          &nbsp;<span class="badge {cls}">{label}</span></h3>
      <p class="sub">Per-target detail &middot; PoCs validated in an isolated sandbox.</p>
      <div class="kv">{stats}</div>
      <div class="section-label">Reasoning</div>
      <div class="reason-block">{reasoning}</div>
      <div class="section-label">Source</div>
      <pre class="code">{code}</pre>
    """

def _render_targets_tab_html(findings: list[dict], sources: dict[str, str]) -> str:
    ordered = sorted(findings, key=lambda x: x["target"])

    radios, nav, panels, rules = [], [], [], []
    for i, f in enumerate(ordered):
        name = f["target"]
        cls, _ = _status(f)
        checked = " checked" if i == 0 else ""
        radios.append(f'<input type="radio" name="as-targets" id="as-t{i}"{checked}>')
        nav.append(f'<label for="as-t{i}"><span class="dot {cls}"></span><code>{html.escape(str(name))}</code></label>')
        panels.append(f'<div class="panel" id="as-p{i}">{_target_detail_html(f, sources.get(name, ""))}</div>')
        rules.append(
            f'.autosec #as-t{i}:checked ~ .subnav label[for="as-t{i}"]'
            "{background:#fff;color:var(--text);border-color:var(--line);font-weight:600;}"
            f".autosec #as-t{i}:checked ~ .panels #as-p{i}{{display:block;}}"
        )

    return f"""{_REPORT_CSS}
    <style>{"".join(rules)}</style>
    <div class="autosec">
      <h2>AutoSec &middot; target detail</h2>
      <p class="sub">{len(ordered)} target(s) &middot; select a file to see its status, reasoning, and source.</p>
      <div class="subtabs">
        {"".join(radios)}
        <div class="subnav">{"".join(nav)}</div>
        <div class="panels">{"".join(panels)}</div>
      </div>
    </div>
    """

@env.task(retries=1)
async def random_error() -> str:
    if _attempt() == 0:
        raise Exception("Random error")
    return "Passed!"

# {{docs-fragment pipeline}}
@env.task(report=True)
async def run_autosec_agent() -> dict:
    targets = _load_targets()
    if not targets:
        raise FileNotFoundError(f"no targets found under {TARGETS_DIR}")

    findings = list(await asyncio.gather(*(analyze_target(name, src) for name, src in targets.items())))

    await flyte.report.replace.aio(_render_report_html(findings))
    flyte.report.get_tab("targets").replace(_render_targets_tab_html(findings, targets))
    await flyte.report.flush.aio()

    await random_error()

    return {
        "targets_analyzed": len(findings),
        "triggered": sum(1 for f in findings if f["verdict"].get("triggered")),
        "findings": findings,
    }
# {{/docs-fragment pipeline}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(run_autosec_agent)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autosec_research_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "unionai-sandbox",
#    "litellm",
# ]
# ///
```

## Run the security pipeline

Each target flows through four stages: static scan, LLM hypothesis, PoC construction, and sandbox validation. The `run_autosec_agent` driver task analyzes all bundled targets in parallel and streams a findings report.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "unionai-sandbox",
#    "litellm",
# ]
# main = "run_autosec_agent"
# params = ""
# ///
"""AutoSec researcher agent — parallel vulnerability analysis with sandbox PoC validation."""

from __future__ import annotations

import asyncio
import html
import json
import os
import pathlib
import re
from typing import Any

import flyte
import flyte.errors
import flyte.report
from flyte.ai.agents import Agent

HERE = pathlib.Path(__file__).parent
TARGETS_DIR = HERE / "targets"
MODEL = os.getenv("AUTOSEC_MODEL", "claude-haiku-4-5")

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="autosec-research-agent", pre=True).with_apt_packages("gcc")

env = flyte.TaskEnvironment(
    name="autosec-research-agent",
    image=main_img,
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    include=[str(TARGETS_DIR)],
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)
# {{/docs-fragment env}}

def _attempt() -> int:
    tc = flyte.ctx()
    return tc.attempt_number if tc is not None else 0

def _force(flag: str) -> bool:
    return bool(os.getenv(flag) or os.getenv("AUTOSEC_FORCE_ALL"))

def _extract_json(text: str) -> dict[str, Any]:
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        raise ValueError(f"no JSON object in model reply: {text[:200]!r}")
    blob = match.group(0)
    try:
        return json.loads(blob)
    except json.JSONDecodeError:
        fixed = re.sub(r'\\(?!["\\/bfnrtu])', r"\\\\", blob)
        return json.loads(fixed)

# --- Stage 1: static analysis (CPU, OOM-prone) ------------------------------
@env.task(retries=2, timeout=30)
async def scan_static(source: str, scope: str = "whole") -> str:
    """Cheap stand-in for whole-program analysis (Joern/CodeQL in the real system)."""
    try:
        if scope == "whole" and _force("AUTOSEC_FORCE_OOM") and _attempt() == 0:
            raise flyte.errors.OOMError("whole-program graph exceeded memory limit")
        findings = _grep_dangerous_calls(source)
        return findings or "(no dangerous-call sites found)"
    except flyte.errors.OOMError as exc:
        print(f"[scan_static] {exc}; escalating resources + narrowing scope")
        return await scan_static.override(
            short_name="scan_static_more_resources", resources=flyte.Resources(cpu=2, memory="4Gi")
        )(source, scope="file")

def _grep_dangerous_calls(source: str) -> str:
    hits = []
    for i, line in enumerate(source.splitlines(), start=1):
        for fn in ("strcpy", "strcat", "sprintf", "gets", "memcpy"):
            if fn in line:
                hits.append(f"L{i}: {fn} -> {line.strip()}")
    return "\n".join(hits)

# --- Stage 2: hypothesize the vulnerability (LLM via Agent) ------------------
ANALYSIS_INSTRUCTIONS = """\
You are a vulnerability researcher. Your job is to determine whether a given \
C source file contains an exploitable memory-corruption bug reachable from argv.

You have access to these tools during your analysis:
- scan_static: Run static analysis on the source to find dangerous function calls.
- build_poc: Build a proof-of-concept payload (do not call during analysis).
- validate_in_sandbox: Compile and run the target with a PoC input (do not call during analysis).

Focus on analyzing the source and the provided static analysis findings. Call \
scan_static only if you need additional details about dangerous function usage.

Reply with ONLY a JSON object (no prose, no markdown fences):
If vulnerable: {"vulnerable": true, "function": str, \
"buffer_size": int (bytes of the overflowable buffer), "vuln_class": str, \
"reasoning": str}.
If the code looks safe (bounded copies, length checks, snprintf/strlcpy, \
etc.): {"vulnerable": false, "reasoning": str}.
"""

# --- Stage 3: build a proof-of-concept --------------------------------------
@env.task(retries=2, timeout=90)
async def build_poc(hypothesis: dict) -> dict:
    buffer_size = int(hypothesis.get("buffer_size", 64))
    payload_len = buffer_size + 64
    return {
        "payload_len": payload_len,
        "payload_repr": f'"A" * {payload_len}',
        "target_function": hypothesis.get("function", "greet"),
    }

# --- Stage 4: validate in an on-device sandbox -------------------------------
@env.task(retries=2, timeout=300)
async def validate_in_sandbox(source: str, poc: dict) -> dict:
    """Compile + run the target with the PoC input inside an on-device sandbox.

    The exploit code runs in a user-namespace sandbox on the same machine, never
    on the Flyte orchestration node (SPEC §2.6 / §7). The session is torn down
    in __aexit__ regardless of outcome (SPEC VD-5) so a stuck or failed run
    cannot leak resources.
    """
    import tempfile

    from union import sandbox as sb

    with tempfile.TemporaryDirectory() as work:
        async with sb.on_device.session(host_work_dir=work, backend="userns") as sbx:
            await sbx.put_bytes(f"{work}/target.c", source.encode())

            compile_proc = await sbx.run(
                f"gcc -fno-stack-protector -w -o {work}/target {work}/target.c",
                stdout=True,
                stderr=True,
                timeout_s=60,
            )
            compile_out, compile_err = await compile_proc.communicate_text()
            log = compile_out + compile_err
            if "error" in log.lower():
                return {
                    "triggered": False,
                    "sandbox_exit_code": -1,
                    "log": f"COMPILE_FAILED\n{log}",
                }

            payload = "A" * int(poc["payload_len"])
            run_proc = await sbx.run(
                f"{work}/target {payload}",
                stdout=True,
                stderr=True,
                timeout_s=60,
            )
            run_out, run_err = await run_proc.communicate_text()
            log = run_out + "\n" + run_err
            triggered = "SIGSEGV" in log

            return {
                "triggered": bool(triggered),
                "sandbox_exit_code": getattr(run_proc, "returncode", 0),
                "log": log,
            }

# --- Agent + hypothesize task (depends on all tools above) ------------------
hypothesis_agent = Agent(
    name="autosec-hypothesis",
    instructions=ANALYSIS_INSTRUCTIONS,
    model=MODEL,
    tools=[scan_static, build_poc, validate_in_sandbox],
    max_turns=6,
)

@env.task(retries=3, timeout=20)
async def hypothesize(source: str, static_findings: str) -> dict:
    prompt = (
        "Analyze this C source file for memory-corruption vulnerabilities.\n\n"
        f"SOURCE:\n{source}\n\nDANGEROUS CALLS:\n{static_findings}\n"
    )

    # Beat A: hang on the first attempt -> task timeout -> retry.
    timeout_on = _force("AUTOSEC_FORCE_LLM_TIMEOUT") and _attempt() == 0
    bad_on = _force("AUTOSEC_FORCE_BAD_TOOL_CALL")

    if timeout_on:
        await asyncio.sleep(600)

    result = await hypothesis_agent.run.aio(prompt, memory=[])
    raw = result.summary or ""

    # Beat B: simulate a hallucinated/malformed tool call. When the timeout beat
    # is also active it consumes attempt 0, so defer this to attempt 1 — that way
    # both beats are actually demonstrated in a single run (e.g. AUTOSEC_FORCE_ALL).
    bad_attempt = 1 if _force("AUTOSEC_FORCE_LLM_TIMEOUT") else 0
    if bad_on and _attempt() == bad_attempt:
        raw = "Sure! The bug is somewhere around here, trust me."

    hyp = _extract_json(raw)
    if "vulnerable" not in hyp:
        hyp["vulnerable"] = "buffer_size" in hyp
    if hyp.get("vulnerable") and "buffer_size" not in hyp:
        raise ValueError(f"vulnerable hypothesis missing buffer_size: {hyp}")
    return hyp

# --- Orchestration ----------------------------------------------------------
def _load_targets() -> dict[str, str]:
    return {p.name: p.read_text() for p in sorted(TARGETS_DIR.glob("*.c"))}

@env.task
async def analyze_target(name: str, source: str) -> dict:
    findings = await scan_static(source)
    hypothesis = await hypothesize(source, findings)

    if not hypothesis.get("vulnerable"):
        poc: dict = {}
        verdict = {"triggered": False, "skipped": True}
    else:
        poc = await build_poc(hypothesis)
        verdict = await validate_in_sandbox(source, poc)

    return {
        "target": name,
        "static_findings": findings,
        "hypothesis": hypothesis,
        "poc": poc,
        "verdict": verdict,
    }

_REPORT_CSS = """
<style>
  .autosec { --bg:#ffffff; --card:#f7f8fa; --line:#e3e7ec; --muted:#5b6675;
    --text:#1b2330; --red:#c0392b; --amber:#b6791f; --green:#1e7e34; --accent:#1f6feb;
    font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;
    background:var(--bg); color:var(--text); padding:24px; border-radius:12px;
    border:1px solid var(--line); }
  .autosec h2 { margin:0 0 4px; font-size:20px; letter-spacing:.2px; }
  .autosec .sub { color:var(--muted); font-size:13px; margin:0 0 20px; }
  .autosec .cards { display:flex; gap:12px; flex-wrap:wrap; margin-bottom:22px; }
  .autosec .card { background:var(--card); border:1px solid var(--line);
    border-radius:10px; padding:14px 18px; min-width:120px; }
  .autosec .card .n { font-size:26px; font-weight:700; line-height:1; }
  .autosec .card .l { color:var(--muted); font-size:12px; margin-top:6px;
    text-transform:uppercase; letter-spacing:.6px; }
  .autosec table { width:100%; table-layout:fixed; border-collapse:collapse; font-size:13px;
    background:#fff; border:1px solid var(--line); border-radius:10px; overflow:hidden; }
  .autosec th, .autosec td { overflow-wrap:anywhere; }
  .autosec thead th { background:#eef1f5; color:var(--muted); text-align:left;
    font-weight:600; font-size:11px; text-transform:uppercase; letter-spacing:.6px;
    padding:11px 14px; border-bottom:1px solid var(--line); }
  .autosec tbody td { padding:11px 14px; border-bottom:1px solid var(--line);
    vertical-align:top; }
  .autosec tbody tr:last-child td { border-bottom:none; }
  .autosec tbody tr:nth-child(even) { background:#fafbfc; }
  .autosec tbody tr:hover { background:#eef4ff; }
  .autosec code { background:#eef1f5; border:1px solid var(--line); border-radius:5px;
    padding:1px 6px; font-family:ui-monospace,SFMono-Regular,Menlo,monospace; font-size:12px; }
  .autosec .num { text-align:right; font-variant-numeric:tabular-nums; }
  .autosec .reason { color:var(--muted); line-height:1.45; }
  .autosec .badge { display:inline-block; padding:3px 10px; border-radius:999px;
    font-size:11px; font-weight:700; letter-spacing:.4px; white-space:nowrap; }
  .autosec .b-exploited { background:rgba(192,57,43,.10); color:var(--red);
    border:1px solid rgba(192,57,43,.35); }
  .autosec .b-vuln { background:rgba(182,121,31,.12); color:var(--amber);
    border:1px solid rgba(182,121,31,.35); }
  .autosec .b-secure { background:rgba(30,126,52,.10); color:var(--green);
    border:1px solid rgba(30,126,52,.35); }
  .autosec col.c-target  { width:15%; }
  .autosec col.c-status  { width:11%; }
  .autosec col.c-class   { width:11%; }
  .autosec col.c-fn      { width:11%; }
  .autosec col.c-buf     { width:8%; }
  .autosec col.c-payload { width:8%; }
  .autosec col.c-exit    { width:6%; }
  .autosec col.c-reason  { width:30%; }
  .autosec .kv { display:flex; flex-wrap:wrap; gap:12px 28px; margin:16px 0 4px; }
  .autosec .kv .k { color:var(--muted); font-size:11px; text-transform:uppercase; letter-spacing:.6px; }
  .autosec .kv .v { font-weight:600; font-size:14px; margin-top:3px; }
  .autosec .section-label { font-size:11px; text-transform:uppercase; letter-spacing:.6px;
    color:var(--muted); margin:22px 0 7px; }
  .autosec .reason-block { background:var(--card); border:1px solid var(--line);
    border-radius:8px; padding:13px 15px; line-height:1.5; font-size:13px; }
  .autosec pre.code { background:#f6f8fa; border:1px solid var(--line); border-radius:8px;
    padding:14px 16px; overflow:auto; font-family:ui-monospace,SFMono-Regular,Menlo,monospace;
    font-size:12.5px; line-height:1.5; color:#1b2330; margin:0; }
  .autosec .subtabs > input[type=radio] { position:absolute; opacity:0; pointer-events:none; }
  .autosec .subnav { display:flex; flex-wrap:wrap; gap:4px; border-bottom:1px solid var(--line);
    margin:8px 0 18px; }
  .autosec .subnav label { display:inline-flex; align-items:center; gap:8px; padding:8px 14px;
    font-size:12.5px; cursor:pointer; border:1px solid transparent; border-bottom:none;
    border-radius:8px 8px 0 0; color:var(--muted); margin-bottom:-1px; }
  .autosec .subnav label:hover { background:var(--card); color:var(--text); }
  .autosec .subnav .dot { width:8px; height:8px; border-radius:50%; }
  .autosec .dot.b-exploited { background:var(--red); }
  .autosec .dot.b-vuln { background:var(--amber); }
  .autosec .dot.b-secure { background:var(--green); }
  .autosec .panel { display:none; }
</style>
"""

def _status(finding: dict) -> tuple[str, str]:
    hyp = finding.get("hypothesis") or {}
    verdict = finding.get("verdict") or {}
    if not hyp.get("vulnerable"):
        return "b-secure", "SECURE"
    if verdict.get("triggered"):
        return "b-exploited", "EXPLOITED"
    return "b-vuln", "VULNERABLE"

def _render_report_html(findings: list[dict]) -> str:
    exploited = sum(1 for f in findings if (f.get("verdict") or {}).get("triggered"))
    vulnerable = sum(1 for f in findings if (f.get("hypothesis") or {}).get("vulnerable"))
    secure = len(findings) - vulnerable

    rows = []
    for f in sorted(findings, key=lambda x: x["target"]):
        hyp = f.get("hypothesis") or {}
        verdict = f.get("verdict") or {}
        cls, label = _status(f)
        is_vuln = bool(hyp.get("vulnerable"))
        vuln_class = hyp.get("vuln_class", "\u2014") if is_vuln else "\u2014"
        fn = hyp.get("function", "\u2014") if is_vuln else "\u2014"
        buf = hyp.get("buffer_size", "\u2014") if is_vuln else "\u2014"
        payload = (f.get("poc") or {}).get("payload_len", "\u2014") if is_vuln else "\u2014"
        exit_code = verdict.get("sandbox_exit_code", "\u2014")
        rows.append(
            "<tr>"
            f"<td><code>{html.escape(str(f['target']))}</code></td>"
            f'<td><span class="badge {cls}">{label}</span></td>'
            f"<td>{html.escape(str(vuln_class))}</td>"
            f"<td><code>{html.escape(str(fn))}</code></td>"
            f'<td class="num">{html.escape(str(buf))}</td>'
            f'<td class="num">{html.escape(str(payload))}</td>'
            f'<td class="num">{html.escape(str(exit_code))}</td>'
            f'<td class="reason">{html.escape(str(hyp.get("reasoning", "")))}</td>'
            "</tr>"
        )

    return f"""{_REPORT_CSS}
    <div class="autosec">
      <h2>AutoSec &middot; security findings report</h2>
      <p class="sub">{len(findings)} target(s) analyzed in parallel &middot; PoCs validated in isolated sandbox.</p>
      <div class="cards">
        <div class="card"><div class="n">{len(findings)}</div><div class="l">Targets</div></div>
        <div class="card"><div class="n" style="color:#ff6b6b">{exploited}</div><div class="l">Exploited</div></div>
        <div class="card">
          <div class="n" style="color:#ffb454">{vulnerable - exploited}</div>
          <div class="l">Vuln, PoC failed</div>
        </div>
        <div class="card"><div class="n" style="color:#3fb950">{secure}</div><div class="l">Secure</div></div>
      </div>
      <table>
        <colgroup>
          <col class="c-target"><col class="c-status"><col class="c-class"><col class="c-fn">
          <col class="c-buf"><col class="c-payload"><col class="c-exit"><col class="c-reason">
        </colgroup>
        <thead>
          <tr>
            <th>Target</th><th>Status</th><th>Vuln class</th><th>Function</th>
            <th class="num">Buffer&nbsp;(B)</th><th class="num">Payload&nbsp;(B)</th>
            <th class="num">Exit</th><th>Reasoning</th>
          </tr>
        </thead>
        <tbody>{"".join(rows)}</tbody>
      </table>
    </div>
    """

def _target_detail_html(finding: dict, source: str) -> str:
    hyp = finding.get("hypothesis") or {}
    verdict = finding.get("verdict") or {}
    cls, label = _status(finding)
    is_vuln = bool(hyp.get("vulnerable"))

    def cell(k: str, v: Any) -> str:
        return f'<div><div class="k">{html.escape(k)}</div><div class="v">{html.escape(str(v))}</div></div>'

    triggered = verdict.get("triggered")
    verdict_txt = "skipped (secure)" if verdict.get("skipped") else ("triggered" if triggered else "not triggered")
    stats = "".join(
        [
            cell("Vuln class", hyp.get("vuln_class", "\u2014") if is_vuln else "\u2014"),
            cell("Function", hyp.get("function", "\u2014") if is_vuln else "\u2014"),
            cell("Buffer (B)", hyp.get("buffer_size", "\u2014") if is_vuln else "\u2014"),
            cell("Payload (B)", (finding.get("poc") or {}).get("payload_len", "\u2014") if is_vuln else "\u2014"),
            cell("Sandbox exit", verdict.get("sandbox_exit_code", "\u2014")),
            cell("PoC", verdict_txt),
        ]
    )
    reasoning = html.escape(str(hyp.get("reasoning", "")) or "\u2014")
    code = html.escape(source or "(source unavailable)")
    return f"""
      <h3 style="margin:0 0 4px"><code>{html.escape(str(finding["target"]))}</code>
          &nbsp;<span class="badge {cls}">{label}</span></h3>
      <p class="sub">Per-target detail &middot; PoCs validated in an isolated sandbox.</p>
      <div class="kv">{stats}</div>
      <div class="section-label">Reasoning</div>
      <div class="reason-block">{reasoning}</div>
      <div class="section-label">Source</div>
      <pre class="code">{code}</pre>
    """

def _render_targets_tab_html(findings: list[dict], sources: dict[str, str]) -> str:
    ordered = sorted(findings, key=lambda x: x["target"])

    radios, nav, panels, rules = [], [], [], []
    for i, f in enumerate(ordered):
        name = f["target"]
        cls, _ = _status(f)
        checked = " checked" if i == 0 else ""
        radios.append(f'<input type="radio" name="as-targets" id="as-t{i}"{checked}>')
        nav.append(f'<label for="as-t{i}"><span class="dot {cls}"></span><code>{html.escape(str(name))}</code></label>')
        panels.append(f'<div class="panel" id="as-p{i}">{_target_detail_html(f, sources.get(name, ""))}</div>')
        rules.append(
            f'.autosec #as-t{i}:checked ~ .subnav label[for="as-t{i}"]'
            "{background:#fff;color:var(--text);border-color:var(--line);font-weight:600;}"
            f".autosec #as-t{i}:checked ~ .panels #as-p{i}{{display:block;}}"
        )

    return f"""{_REPORT_CSS}
    <style>{"".join(rules)}</style>
    <div class="autosec">
      <h2>AutoSec &middot; target detail</h2>
      <p class="sub">{len(ordered)} target(s) &middot; select a file to see its status, reasoning, and source.</p>
      <div class="subtabs">
        {"".join(radios)}
        <div class="subnav">{"".join(nav)}</div>
        <div class="panels">{"".join(panels)}</div>
      </div>
    </div>
    """

@env.task(retries=1)
async def random_error() -> str:
    if _attempt() == 0:
        raise Exception("Random error")
    return "Passed!"

# {{docs-fragment pipeline}}
@env.task(report=True)
async def run_autosec_agent() -> dict:
    targets = _load_targets()
    if not targets:
        raise FileNotFoundError(f"no targets found under {TARGETS_DIR}")

    findings = list(await asyncio.gather(*(analyze_target(name, src) for name, src in targets.items())))

    await flyte.report.replace.aio(_render_report_html(findings))
    flyte.report.get_tab("targets").replace(_render_targets_tab_html(findings, targets))
    await flyte.report.flush.aio()

    await random_error()

    return {
        "targets_analyzed": len(findings),
        "triggered": sum(1 for f in findings if f["verdict"].get("triggered")),
        "findings": findings,
    }
# {{/docs-fragment pipeline}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(run_autosec_agent)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autosec_research_agent/main.py*

## Run the agent

### Create secrets

Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/) and register it as a Flyte secret:

```
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autosec_research_agent):

```
cd v2/tutorials/autosec_research_agent
uv run --script main.py
```

Follow the printed run URL to watch each target progress through the pipeline and open the report panel for the findings table and per-target detail tabs.

Optional environment variables demonstrate self-healing behavior (`AUTOSEC_FORCE_LLM_TIMEOUT`, `AUTOSEC_FORCE_BAD_TOOL_CALL`, `AUTOSEC_FORCE_OOM`, or `AUTOSEC_FORCE_ALL=1`).


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/code-agent ===

# Coding agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/code_runner).

This example demonstrates how to run code generated by a large language model (LLM) using a `ContainerTask`.
The agent takes a user’s question, generates Flyte 2 code using the Flyte 2 documentation as context, and runs it in an isolated container.
If the execution fails, the agent reflects on the error and retries
up to a configurable limit until it succeeds.

Using `ContainerTask` ensures that all generated code runs in a secure environment.
This gives you full flexibility to execute arbitrary logic safely and reliably.

## What this example demonstrates

- How to combine LLM generation with programmatic execution.
- How to run untrusted or dynamically generated code securely.
- How to iteratively improve code using agent-like behavior.

## Setting up the agent environment

Let's start by importing the necessary libraries and setting up two environments: one for the container task and another for the agent task.
This example follows the `uv` script format to declare dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b23",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# ///
```

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

> [!NOTE]
> You can set up access to the OpenAI API using a Flyte secret.
>
> ```
> flyte create secret openai_api_key <YOUR_OPENAI_API_KEY>
> ```

We store the LLM-generated code in a structured format. This allows us to:

- Enforce consistent formatting
- Make debugging easier
- Log and analyze generations systematically

By capturing metadata alongside the raw code, we maintain transparency and make it easier to iterate or trace issues over time.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

We then define a state model to persist the agent's history across iterations. This includes previous messages,
generated code, and any errors encountered.

Maintaining this history allows the agent to reflect on past attempts, avoid repeating mistakes,
and iteratively improve the generated code.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

## Retrieve docs

We define a task to load documents from a given URL and concatenate them into a single string.
This string is then used as part of the LLM prompt.

We set `max_depth = 20` to avoid loading an excessive number of documents.
However, even with this limit, the resulting context can still be quite large.
To handle this, we use an LLM (GPT-4 in this case) that supports extended context windows.

> [!NOTE]
> Appending all documents into a single string can result in extremely large contexts, potentially exceeding the LLM’s token limit.
> If your dataset grows beyond what a single prompt can handle, there are a couple of strategies you can use.
> One option is to apply Retrieval-Augmented Generation (RAG), where you chunk the documents, embed them using a model,
> store the vectors in a vector database, and retrieve only the most relevant pieces at inference time.
>
> An alternative approach is to pass references to full files into the prompt, allowing the LLM to decide which files are most relevant based
> on natural-language search over file paths, summaries, or even contents. This method assumes that only a subset of files
> will be necessary for a given task, and the LLM is responsible for navigating the structure and identifying what to read.
> While this can be a lighter-weight solution for smaller datasets, its effectiveness depends on how well the LLM can
> reason over file references and the reliability of its internal search heuristics.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

## Code generation

Next, we define a utility function to construct the LLM chain responsible for generating Python code from user input. This chain uses
a LangChain `PromptTemplate` to structure the input and an OpenAI chat model to generate well-formed, Flyte 2-compatible Python scripts.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

We then define a `generate` task responsible for producing the code solution.
To improve clarity and testability, the output is structured in three parts:
a short summary of the generated solution, a list of necessary imports,
and the main body of executable code.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

A `ContainerTask` then executes this code in an isolated container environment.
It takes the code as input, runs it safely, and returns the program’s output and exit code.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

This task verifies that the generated code runs as expected.
It tests the import statements first, then executes the full code.
It records the output and any error messages in the agent state for further analysis.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

If an error occurs, a separate task reflects on the failure and generates a response.
This reflection is added to the agent state to guide future iterations.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

Finally, we define a `main` task that runs the code agent and orchestrates the steps above.
If the code execution fails, we reflect on the error and retry until we reach the maximum number of iterations.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "langchain-core==0.3.66",
#    "langchain-openai==0.3.24",
#    "langchain-community==0.3.26",
#    "beautifulsoup4==4.13.4",
#    "docker==7.1.0",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment code_runner_task}}
import flyte
from flyte.extras import ContainerTask
from flyte.io import File

code_runner_task = ContainerTask(
    name="run_flyte_v2",
    image=flyte.Image.from_debian_base(),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script": File},
    outputs={"result": str, "exit_code": str},
    command=[
        "/bin/bash",
        "-c",
        (
            "set -o pipefail && "
            "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "
            "echo $? > /var/outputs/exit_code"
        ),
    ],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

# {{/docs-fragment code_runner_task}}

# {{docs-fragment env}}
import tempfile
from typing import Optional

from langchain_core.runnables import Runnable
from pydantic import BaseModel, Field

container_env = flyte.TaskEnvironment.from_task(
    "code-runner-container", code_runner_task
)

env = flyte.TaskEnvironment(
    name="code_runner",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_uv_script(__file__, name="code-runner-agent"),
    resources=flyte.Resources(cpu=1),
    depends_on=[container_env],
)

# {{/docs-fragment env}}

# {{docs-fragment code_base_model}}
class Code(BaseModel):
    """Schema for code solutions to questions about Flyte v2."""

    prefix: str = Field(
        default="", description="Description of the problem and approach"
    )
    imports: str = Field(
        default="", description="Code block with just import statements"
    )
    code: str = Field(
        default="", description="Code block not including import statements"
    )

# {{/docs-fragment code_base_model}}

# {{docs-fragment agent_state}}
class AgentState(BaseModel):
    messages: list[dict[str, str]] = Field(default_factory=list)
    generation: Code = Field(default_factory=Code)
    iterations: int = 0
    error: str = "no"
    output: Optional[str] = None

# {{/docs-fragment agent_state}}

# {{docs-fragment generate_code_gen_chain}}
async def generate_code_gen_chain(debug: bool) -> Runnable:
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI

    # Grader prompt
    code_gen_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """
You are a coding assistant with expertise in Python.
You are able to execute the Flyte v2 code locally in a sandbox environment.

Use the following pattern to execute the code:

<code>
if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(...))
</code>

Your response will be shown to the user.
Here is a full set of documentation:

-------
{context}
-------

Answer the user question based on the above provided documentation.
Ensure any code you provide can be executed with all required imports and variables defined.
Structure your answer with a description of the code solution.
Then list the imports. And finally list the functioning code block.
Here is the user question:""",
            ),
            ("placeholder", "{messages}"),
        ]
    )

    expt_llm = "gpt-4o" if not debug else "gpt-4o-mini"
    llm = ChatOpenAI(temperature=0, model=expt_llm)

    code_gen_chain = code_gen_prompt | llm.with_structured_output(Code)
    return code_gen_chain

# {{/docs-fragment generate_code_gen_chain}}

# {{docs-fragment docs_retriever}}
@env.task
async def docs_retriever(url: str) -> str:
    from bs4 import BeautifulSoup
    from langchain_community.document_loaders.recursive_url_loader import (
        RecursiveUrlLoader,
    )

    loader = RecursiveUrlLoader(
        url=url, max_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text
    )
    docs = loader.load()

    # Sort the list based on the URLs and get the text
    d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
    d_reversed = list(reversed(d_sorted))

    concatenated_content = "\n\n\n --- \n\n\n".join(
        [doc.page_content for doc in d_reversed]
    )
    return concatenated_content

# {{/docs-fragment docs_retriever}}

# {{docs-fragment generate}}
@env.task
async def generate(
    question: str, state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Generate a code solution

    Args:
        question (str): The user question
        state (dict): The current graph state
        concatenated_content (str): The concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")

    messages = state.messages
    iterations = state.iterations
    error = state.error

    # We have been routed back to generation with an error
    if error == "yes":
        messages += [
            {
                "role": "user",
                "content": (
                    "Now, try again. Invoke the code tool to structure the output "
                    "with a prefix, imports, and code block:"
                ),
            }
        ]

    code_gen_chain = await generate_code_gen_chain(debug)

    # Solution
    code_solution = code_gen_chain.invoke(
        {
            "context": concatenated_content,
            "messages": (
                messages if messages else [{"role": "user", "content": question}]
            ),
        }
    )

    messages += [
        {
            "role": "assistant",
            "content": f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        }
    ]

    return AgentState(
        messages=messages,
        generation=code_solution,
        iterations=iterations + 1,
        error=error,
        output=state.output,
    )

# {{/docs-fragment generate}}

# {{docs-fragment code_check}}
@env.task
async def code_check(state: AgentState) -> AgentState:
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")

    # State
    messages = state.messages
    code_solution = state.generation
    iterations = state.iterations

    # Get solution components
    imports = code_solution.imports.strip()
    code = code_solution.code.strip()

    # Create temp file for imports
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as imports_file:
        imports_file.write(imports + "\n")
        imports_path = imports_file.name

    # Create temp file for code body
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code_file:
        code_file.write(imports + "\n" + code + "\n")
        code_path = code_file.name

    # Check imports
    import_output, import_exit_code = await code_runner_task(
        script=await File.from_local(imports_path)
    )

    if import_exit_code.strip() != "0":
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the import test: {import_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=import_output,
        )
    else:
        print("---CODE IMPORT CHECK: PASSED---")

    # Check execution
    code_output, code_exit_code = await code_runner_task(
        script=await File.from_local(code_path)
    )

    if code_exit_code.strip() != "0":
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [
            {
                "role": "user",
                "content": f"Your solution failed the code execution test: {code_output}",
            }
        ]
        messages += error_message
        return AgentState(
            generation=code_solution,
            messages=messages,
            iterations=iterations,
            error="yes",
            output=code_output,
        )
    else:
        print("---CODE BLOCK CHECK: PASSED---")

    # No errors
    print("---NO CODE TEST FAILURES---")

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error="no",
        output=code_output,
    )

# {{/docs-fragment code_check}}

# {{docs-fragment reflect}}
@env.task
async def reflect(
    state: AgentState, concatenated_content: str, debug: bool
) -> AgentState:
    """
    Reflect on errors

    Args:
        state (dict): The current graph state
        concatenated_content (str): Concatenated docs content
        debug (bool): Debug mode

    Returns:
        state (dict): New key added to state, reflection
    """

    print("---REFLECTING---")

    # State
    messages = state.messages
    iterations = state.iterations
    code_solution = state.generation

    # Prompt reflection
    code_gen_chain = await generate_code_gen_chain(debug)

    # Add reflection
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )

    messages += [
        {
            "role": "assistant",
            "content": f"Here are reflections on the error: {reflections}",
        }
    ]

    return AgentState(
        generation=code_solution,
        messages=messages,
        iterations=iterations,
        error=state.error,
        output=state.output,
    )

# {{/docs-fragment reflect}}

# {{docs-fragment main}}
@env.task
async def main(
    question: str = (
        "Define a two-task pattern where the second catches OOM from the first and retries with more memory."
    ),
    url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/",
    max_iterations: int = 3,
    debug: bool = False,
) -> str:
    concatenated_content = await docs_retriever(url=url)

    state: AgentState = AgentState()
    iterations = 0

    while True:
        with flyte.group(f"code-generation-pass-{iterations + 1}"):
            state = await generate(question, state, concatenated_content, debug)
            state = await code_check(state)

            error = state.error
            iterations = state.iterations

            if error == "no" or iterations >= max_iterations:
                print("---DECISION: FINISH---")
                code_solution = state.generation

                prefix = code_solution.prefix
                imports = code_solution.imports
                code = code_solution.code

                code_output = state.output

                return f"""{prefix}

{imports}
{code}

Result of code execution:
{code_output}
"""
            else:
                print("---DECISION: RE-TRY SOLUTION---")
                state = await reflect(state, concatenated_content, debug)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()

# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_runner/agent.py*

## Running the code agent

If things are working properly, you should see output similar to the following:

```
---GENERATING CODE SOLUTION---
---CHECKING CODE---
---CODE BLOCK CHECK: PASSED---
---NO CODE TEST FAILURES---
---DECISION: FINISH---
In this solution, we define two tasks using Flyte v2.
The first task, `oomer`, is designed to simulate an out-of-memory (OOM) error by attempting to allocate a large list.
The second task, `failure_recovery`, attempts to execute `oomer` and catches any OOM errors.
If an OOM error is caught, it retries the `oomer` task with increased memory resources.
This pattern demonstrates how to handle resource-related exceptions and dynamically adjust task configurations in Flyte workflows.

import asyncio
import flyte
import flyte.errors
env = flyte.TaskEnvironment(name="oom_example", resources=flyte.Resources(cpu=1, memory="250Mi"))

@env.task
async def oomer(x: int):
    large_list = [0] * 100000000  # Simulate OOM
    print(len(large_list))

@env.task
async def always_succeeds() -> int:
    await asyncio.sleep(1)
    return 42

...
CODE12
uv run agent.py
```


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/competitive-intelligence-agent ===

# Competitive intelligence agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/competitive_intelligence_agent).

This example demonstrates how to build a continuous competitive and market intelligence agent on Flyte. The agent fans out across a list of competitors, pulls fresh, source-cited web and news results from the [You.com Search API](https://you.com/docs/search/overview), and uses [Claude](https://docs.anthropic.com/) via [LiteLLM](https://docs.litellm.ai/) to extract structured **deltas** (pricing changes, product launches, funding events, leadership moves, and more) into a knowledge-graph-ready table.

You.com returns ranked web and news results with snippets and publication timestamps, giving the LLM attributable sources to cite. Flyte orchestrates the rest:

- **Fan-out parallelism** across competitors with `asyncio.gather`
- **`cache="auto"`** so converging parallel or repeat runs reuse prior You.com and LLM results when queries overlap
- **`@flyte.trace`** on every You.com and LLM call for full prompt → query → source lineage
- **Flyte reports** that render an HTML dashboard grouping deltas by competitor and category

![Competitive intelligence agent report](../../../_static/images/tutorials/competitive_intelligence_agent/competitive-intelligence-agent.png)

## Setting up the environment

The agent runs in a single `TaskEnvironment` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the `uv` script dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE0", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

CODE1

## Data types

The agent models search hits, deltas, and the final report as dataclasses. Each `Delta` links back to a `SearchHit` that preserves You.com metadata: domain, publication date, author, and snippet.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE2", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

## Search with the You.com Search API

The `you_search` helper calls the [You.com Search API](https://you.com/docs/search/overview) at `https://ydc-index.io/v1/search`. It requests unified web and news results with a `freshness` filter (`day`, `week`, `month`, or `year`) and returns structured hits the LLM can cite by index.

See the [Search API reference](https://you.com/docs/api-reference/search/v1-search) for all supported parameters, including `count`, `country`, and search operators.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE3", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

> [!NOTE]
> We use `@flyte.trace` to track intermediate steps within a task, like You.com API calls and LLM invocations. Each traced call appears as a span in the Flyte dashboard with its inputs and outputs captured.

## Extract deltas with Claude

A shared `llm_json` helper routes to Claude through LiteLLM and parses structured JSON from the response.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE4", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

## Watch one competitor

The `watch_competitor` task builds a category-scoped search query, calls the You.com Search API, and asks Claude to extract only changes that are supported by a specific search result. Each delta carries a confidence score and a link to its source hit.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE5", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

## Orchestration

The `competitive_intelligence` driver task fans out across all competitors with `asyncio.gather`, aggregates the results, and renders a Flyte report.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "competitive_intelligence"
# params = ""
# ///
"""Continuous competitive & market intelligence agent.

A Dragonfly-style agent that fans out across competitors, pulls fresh,
source-cited web + news results from the You.com Search API, and uses Claude to
extract structured "deltas" (pricing, features, funding, leadership, etc.) into
a knowledge-graph-ready table.
"""

# {{docs-fragment env}}
import asyncio
import json
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="competitive-intelligence",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="competitive-intelligence", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class SearchHit:
    """A You.com Search result with its full structured metadata."""

    title: str
    url: str
    domain: str
    snippet: str
    published: str  # You.com page_age timestamp
    author: str
    favicon: str  # You.com favicon_url
    thumbnail: str
    section: str  # "news" or "web" — You.com's auto classification

@dataclass
class Delta:
    competitor: str
    category: str
    summary: str
    confidence: float
    source: SearchHit | None = None

@dataclass
class CompetitorWatch:
    competitor: str
    deltas: list[Delta] = field(default_factory=list)
    sources: list[SearchHit] = field(default_factory=list)

@dataclass
class IntelReport:
    watches: list[CompetitorWatch] = field(default_factory=list)

    @property
    def deltas(self) -> list[Delta]:
        return [d for w in self.watches for d in w.deltas]
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import os
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

@flyte.trace
async def you_search(
    query: str,
    count: int = 10,
    freshness: str = "week",
    boost_domains: str = "",
) -> list[SearchHit]:
    """Call the You.com Search API and return unified web + news hits.

    ``boost_domains`` (comma-separated) gives a ranking boost to authoritative
    sources without restricting results to only those domains — useful for
    competitive intelligence where credible third-party reporting should
    surface above noise, but primary sources (e.g. a company's own blog)
    must still be reachable.
    """
    params: dict = {"query": query, "count": count, "freshness": freshness}
    if boost_domains:
        params["boost_domains"] = boost_domains
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    thumbnail=item.get("thumbnail_url", "") or "",
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict | list:
    """Call Claude via LiteLLM and parse a JSON response."""
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content
    return _parse_json(content)

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("CODE6", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min(
        (i for i in (text.find("{"), text.find("[")) if i != -1),
        default=0,
    )
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

EXTRACT_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \
search results about a competitor, extract concrete, recently-changed signals \
("deltas") in the requested categories. Only report changes that are supported \
by a specific search result. Respond with a JSON object of the form:
{"deltas": [{"category": str, "summary": str, "source_index": int (the [n] of \
the supporting search result), "confidence": float between 0 and 1}]}
If there are no clear changes, return {"deltas": []}."""

# {{docs-fragment watch_competitor}}
# Domains that consistently break AI-industry news (funding, product, leadership).
# boost_domains lifts these in ranking without excluding other sources, so a
# competitor's own blog or niche coverage still surfaces when relevant.
CI_BOOST_DOMAINS = "techcrunch.com,reuters.com,bloomberg.com,theinformation.com,venturebeat.com"

@env.task(retries=3)
async def watch_competitor(
    competitor: str,
    categories: list[str],
    freshness: str,
) -> CompetitorWatch:
    """Search for fresh signals on one competitor and extract structured deltas."""
    query = (
        f"{competitor} "
        + " OR ".join(categories)
        + " announcement OR news OR update"
    )
    hits = await you_search(
        query, count=10, freshness=freshness, boost_domains=CI_BOOST_DOMAINS
    )
    if not hits:
        return CompetitorWatch(competitor=competitor)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Competitor: {competitor}\n"
        f"Categories to watch: {', '.join(categories)}\n\n"
        f"Search results:\n{evidence}"
    )
    parsed = await llm_json(EXTRACT_SYSTEM, user)
    raw_deltas = parsed.get("deltas", []) if isinstance(parsed, dict) else []

    deltas: list[Delta] = []
    cited: list[SearchHit] = []
    for d in raw_deltas:
        idx = int(d.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        if src is not None and src not in cited:
            cited.append(src)
        deltas.append(
            Delta(
                competitor=competitor,
                category=str(d.get("category", "unknown")),
                summary=str(d.get("summary", "")),
                confidence=float(d.get("confidence", 0.0) or 0.0),
                source=src,
            )
        )
    return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited)
# {{/docs-fragment watch_competitor}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .card h2 { font-size:16px; margin:0 0 6px; color:#102a43; }
  .rpt .row { padding:11px 0; border-top:1px solid #f0f2f5; }
  .rpt .row:first-of-type { border-top:none; }
  .rpt .chip { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; white-space:nowrap;
               text-transform:uppercase; letter-spacing:.03em;
               background:#e0e8f9; color:#2b4ba0; margin-right:8px; }
  .rpt .summary { margin:6px 0 4px; font-size:14px; line-height:1.45; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .bar { display:inline-block; width:60px; height:6px; border-radius:3px;
              background:#e4e7eb; vertical-align:middle; overflow:hidden;
              margin-right:6px; }
  .rpt .bar > span { display:block; height:100%; background:#3ebd93; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:8px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .src-head { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#627d98; margin:14px 0 4px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _conf_bar(conf: float) -> str:
    pct = max(0, min(100, int(conf * 100)))
    return (
        f"<span class='bar'><span style='width:{pct}%'></span></span>"
        f"<span class='meta'>{conf:.0%} confidence</span>"
    )

def _cite(src: SearchHit) -> str:
    """Render a rich You.com citation: favicon, domain, date, author, snippet."""
    if src is None:
        return ""
    tag = (
        f"<span class='ctag news'>news</span>"
        if src.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if src.published:
        meta_bits.append(src.published[:10])
    if src.author:
        meta_bits.append(f"by {src.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{src.snippet}&rdquo;</div>" if src.snippet else ""
    return (
        f"<div class='cite'>"
        f"<img class='fav' src='{src.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{src.url}'><span class='cdom'>{src.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: IntelReport) -> str:
    watches = sorted(report.watches, key=lambda w: w.competitor)
    total_sources = sum(len(w.sources) for w in watches)

    cards = []
    for w in watches:
        deltas = sorted(w.deltas, key=lambda d: -d.confidence)
        rows = "".join(
            f"<div class='row'><span class='chip'>{d.category}</span>"
            f"<div class='summary'>{d.summary}</div>"
            f"{_conf_bar(d.confidence)}"
            f"{_cite(d.source)}"
            "</div>"
            for d in deltas
        )
        cards.append(
            f"<div class='card'><h2>{w.competitor}</h2>"
            f"<span class='meta'>{len(deltas)} signal(s) &middot; "
            f"{len(w.sources)} You.com source(s)</span>{rows or ''}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Competitive Intelligence Deltas</h1>
      <p class="sub">Fresh, source-cited market signals — every delta links back
      to a ranked, timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.deltas)}</b> signals</span>
        <span class="pill"><b>{len(watches)}</b> competitors tracked</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No signals detected in this window.</p>"}
      <p class="yoube">Sources retrieved and ranked by the You.com Search API
      (web + auto-classified news), with publication timestamps, authors, and
      snippet provenance preserved for full prompt &rarr; citation lineage.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
@env.task(report=True)
async def competitive_intelligence(
    competitors: list[str] = [
        "Anthropic",
        "OpenAI",
        "Mistral AI",
        "Google DeepMind",
        "Cohere",
        "Perplexity AI",
        "xAI",
        "Hugging Face",
        "Databricks",
        "Together AI",
    ],
    categories: list[str] = [
        "pricing",
        "product launch",
        "model release",
        "funding",
        "leadership",
        "partnership",
    ],
    freshness: str = "week",
) -> IntelReport:
    """Fan out across competitors and aggregate structured deltas."""
    with flyte.group("watch-competitors"):
        results = await asyncio.gather(
            *[watch_competitor(c, categories, freshness) for c in competitors]
        )

    report = IntelReport(watches=list(results))

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(competitive_intelligence)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive_intelligence_agent/main.py*

## Run the agent

### Create secrets

Get a You.com API key from the [You.com platform](https://you.com/platform) (see the [quickstart guide](https://you.com/docs/quickstart)). Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/).

Register both keys as Flyte secrets. The secret key names must match those declared in the `TaskEnvironment`:

```
flyte create secret youdotcom-api-key <YOUR_YOU_API_KEY>
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run locally or remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/competitive_intelligence_agent):

```
cd v2/tutorials/competitive_intelligence_agent
uv run --script main.py
```

Or pass custom competitors with the Flyte CLI:

```
flyte run main.py competitive_intelligence \
  --competitors '["Anthropic", "OpenAI"]'
```

To test locally without Flyte secrets, export the environment variables directly:

```
export YOU_API_KEY=<YOUR_YOU_API_KEY>
export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>

uv run --script main.py
```

When the run completes, open the Flyte report in the UI to review deltas grouped by competitor, each with a clickable You.com source citation.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/deep-research ===

# Deep research

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/deep_research_agent); based on work by [Together AI](https://github.com/togethercomputer/open_deep_research).

This example demonstrates how to build an agentic workflow for deep research: a multi-step reasoning system that mirrors how a human researcher explores, analyzes, and synthesizes information from the web.

Deep research refers to the iterative process of thoroughly investigating a topic: identifying relevant sources, evaluating their usefulness, refining the research direction, and ultimately producing a well-structured summary or report. It's a long-running task that requires the agent to reason over time, adapt its strategy, and chain multiple steps together, making it an ideal fit for an agentic architecture.

In this example, we use:

- [Tavily](https://www.tavily.com/) to search for and retrieve high-quality online resources.
- [LiteLLM](https://litellm.ai/) to route LLM calls that perform reasoning, evaluation, and synthesis.

The agent executes a multi-step trajectory:

- Parallel search across multiple queries.
- Evaluation of retrieved results.
- Adaptive iteration: If results are insufficient, it formulates new research queries and repeats the search-evaluate cycle.
- Synthesis: After a fixed number of iterations, it produces a comprehensive research report.

What makes this workflow compelling is its dynamic, evolving nature. The agent isn't just following a fixed plan; it's making decisions in context, using multiple prompts and reasoning steps to steer the process.

Flyte is uniquely well-suited for this kind of system. It provides:

- Structured composition of dynamic reasoning steps
- Built-in parallelism for faster search and evaluation
- Traceability and observability into each step and iteration
- Scalability for long-running or compute-intensive workloads

![Result](../../../_static/images/tutorials/deep-research/result.gif)

Throughout this guide, we'll show how to design this workflow using the Flyte SDK, and how to make the most of agentic development with tools you already know and trust.

## Setting up the environment

Let's begin by setting up the task environment. We define the following components:

- Secrets for Together and Tavily API keys
- A custom image with required Python packages and apt dependencies (`pandoc`, `texlive-xetex`)
- External YAML file with all LLM prompts baked into the container

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pydantic==2.11.5",
#    "litellm==1.72.2",
#    "tavily-python==0.7.5",
#    "together==1.5.24",
#    "markdown==3.8.2",
#    "pymdown-extensions==10.16.1",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import json
from pathlib import Path

import flyte
import yaml
from flyte.io._file import File
from libs.utils.data_types import (
    DeepResearchResult,
    DeepResearchResults,
    ResearchPlan,
    SourceList,
)
from libs.utils.generation import generate_html, generate_toc_image
from libs.utils.llms import asingle_shot_llm_call
from libs.utils.log import AgentLogger
from libs.utils.tavily_search import atavily_search_results

TIME_LIMIT_MULTIPLIER = 5
MAX_COMPLETION_TOKENS = 4096

logging = AgentLogger("together.open_deep_research")

env = flyte.TaskEnvironment(
    name="deep-researcher",
    secrets=[
        flyte.Secret(key="together_api_key", as_env_var="TOGETHER_API_KEY"),
        flyte.Secret(key="tavily_api_key", as_env_var="TAVILY_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="deep-research-agent", pre=True)
    .with_apt_packages("pandoc", "texlive-xetex")
    .with_source_file(Path("prompts.yaml"), "/root"),
    resources=flyte.Resources(cpu=1),
)
# {{/docs-fragment env}}

# {{docs-fragment generate_research_queries}}
@env.task
async def generate_research_queries(
    topic: str,
    planning_model: str,
    json_model: str,
    prompts_file: File,
) -> list[str]:
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    PLANNING_PROMPT = prompts["planning_prompt"]

    plan = ""
    logging.info(f"\n\nGenerated deep research plan for topic: {topic}\n\nPlan:")
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=PLANNING_PROMPT,
        message=f"Research Topic: {topic}",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        plan += chunk
        print(chunk, end="", flush=True)

    SEARCH_PROMPT = prompts["plan_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=SEARCH_PROMPT,
        message=f"Plan to be parsed: {plan}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    plan = json.loads(response_json)
    return plan["queries"]
# {{/docs-fragment generate_research_queries}}

async def _summarize_content_async(
    raw_content: str,
    query: str,
    prompt: str,
    summarization_model: str,
) -> str:
    """Summarize content asynchronously using the LLM"""
    logging.info("Summarizing content asynchronously using the LLM")

    result = ""
    async for chunk in asingle_shot_llm_call(
        model=summarization_model,
        system_prompt=prompt,
        message=f"<Raw Content>{raw_content}</Raw Content>\n\n<Research Topic>{query}</Research Topic>",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        result += chunk
    return result

# {{docs-fragment search_and_summarize}}
@env.task
async def search_and_summarize(
    query: str,
    prompts_file: File,
    summarization_model: str,
) -> DeepResearchResults:
    """Perform search for a single query"""

    if len(query) > 400:
        # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues
        query = query[:400]
        logging.info(f"Truncated query to 400 characters: {query}")

    response = await atavily_search_results(query)

    logging.info("Tavily Search Called.")

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    RAW_CONTENT_SUMMARIZER_PROMPT = prompts["raw_content_summarizer_prompt"]

    with flyte.group("summarize-content"):
        # Create tasks for summarization
        summarization_tasks = []
        result_info = []
        for result in response.results:
            if result.raw_content is None:
                continue

            task = _summarize_content_async(
                result.raw_content,
                query,
                RAW_CONTENT_SUMMARIZER_PROMPT,
                summarization_model,
            )
            summarization_tasks.append(task)
            result_info.append(result)

        # Use return_exceptions=True to prevent exceptions from propagating
        summarized_contents = await asyncio.gather(
            *summarization_tasks, return_exceptions=True
        )

    # Filter out exceptions
    summarized_contents = [
        result for result in summarized_contents if not isinstance(result, Exception)
    ]

    formatted_results = []
    for result, summarized_content in zip(result_info, summarized_contents):
        formatted_results.append(
            DeepResearchResult(
                title=result.title,
                link=result.link,
                content=result.content,
                raw_content=result.raw_content,
                filtered_raw_content=summarized_content,
            )
        )
    return DeepResearchResults(results=formatted_results)
# {{/docs-fragment search_and_summarize}}

@env.task
async def search_all_queries(
    queries: list[str], summarization_model: str, prompts_file: File
) -> DeepResearchResults:
    """Execute searches for all queries in parallel"""
    tasks = []
    results_list = []

    tasks = [
        search_and_summarize(query, prompts_file, summarization_model)
        for query in queries
    ]

    if tasks:
        res_list = await asyncio.gather(*tasks)

    results_list.extend(res_list)

    # Combine all results
    combined_results = DeepResearchResults(results=[])
    for results in results_list:
        combined_results = combined_results + results

    return combined_results

# {{docs-fragment evaluate_research_completeness}}
@env.task
async def evaluate_research_completeness(
    topic: str,
    results: DeepResearchResults,
    queries: list[str],
    prompts_file: File,
    planning_model: str,
    json_model: str,
) -> list[str]:
    """
    Evaluate if the current search results are sufficient or if more research is needed.
    Returns an empty list if research is complete, or a list of additional queries if more research is needed.
    """

    # Format the search results for the LLM
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)

    EVALUATION_PROMPT = prompts["evaluation_prompt"]

    logging.info("\nEvaluation: ")
    evaluation = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=EVALUATION_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Search Queries Used>{queries}</Search Queries Used>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=None,
    ):
        evaluation += chunk
        print(chunk, end="", flush=True)

    EVALUATION_PARSING_PROMPT = prompts["evaluation_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=EVALUATION_PARSING_PROMPT,
        message=f"Evaluation to be parsed: {evaluation}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    evaluation = json.loads(response_json)
    return evaluation["queries"]
# {{/docs-fragment evaluate_research_completeness}}

# {{docs-fragment filter_results}}
@env.task
async def filter_results(
    topic: str,
    results: DeepResearchResults,
    prompts_file: File,
    planning_model: str,
    json_model: str,
    max_sources: int,
) -> DeepResearchResults:
    """Filter the search results based on the research plan"""

    # Format the search results for the LLM, without the raw content
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    FILTER_PROMPT = prompts["filter_prompt"]

    logging.info("\nFilter response: ")
    filter_response = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=FILTER_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        filter_response += chunk
        print(chunk, end="", flush=True)

    logging.info(f"Filter response: {filter_response}")

    FILTER_PARSING_PROMPT = prompts["filter_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=FILTER_PARSING_PROMPT,
        message=f"Filter response to be parsed: {filter_response}",
        response_format={
            "type": "json_object",
            "schema": SourceList.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    sources = json.loads(response_json)["sources"]

    logging.info(f"Filtered sources: {sources}")

    if max_sources != -1:
        sources = sources[:max_sources]

    # Filter the results based on the source list
    filtered_results = [
        results.results[i - 1] for i in sources if i - 1 < len(results.results)
    ]

    return DeepResearchResults(results=filtered_results)
# {{/docs-fragment filter_results}}

def _remove_thinking_tags(answer: str) -> str:
    """Remove content within <think> tags"""
    while "<think>" in answer and "</think>" in answer:
        start = answer.find("<think>")
        end = answer.find("</think>") + len("</think>")
        answer = answer[:start] + answer[end:]
    return answer

# {{docs-fragment generate_research_answer}}
@env.task
async def generate_research_answer(
    topic: str,
    results: DeepResearchResults,
    remove_thinking_tags: bool,
    prompts_file: File,
    answer_model: str,
) -> str:
    """
    Generate a comprehensive answer to the research topic based on the search results.
    Returns a detailed response that synthesizes information from all search results.
    """

    formatted_results = str(results)
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    ANSWER_PROMPT = prompts["answer_prompt"]

    answer = ""
    async for chunk in asingle_shot_llm_call(
        model=answer_model,
        system_prompt=ANSWER_PROMPT,
        message=f"Research Topic: {topic}\n\nSearch Results:\n{formatted_results}",
        response_format=None,
        # NOTE: This is the max_token parameter for the LLM call on Together AI,
        # may need to be changed for other providers
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        answer += chunk

    # this is just to avoid typing complaints
    if answer is None or not isinstance(answer, str):
        logging.error("No answer generated")
        return "No answer generated"

    if remove_thinking_tags:
        # Remove content within <think> tags
        answer = _remove_thinking_tags(answer)

    # Remove markdown code block markers if they exist at the beginning
    if answer.lstrip().startswith("```"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("```"))
        if first_linebreak != -1:
            # Remove everything up to and including the first line break
            answer = answer[first_linebreak + 1 :]

        # Remove closing code block if it exists
        if answer.rstrip().endswith("```"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/agent.py*

The Python packages are declared at the top of the file using the `uv` script style:

CODE2

## Generate research queries

This task converts a user prompt into a list of focused queries. It makes two LLM calls to generate a high-level research plan and parse that plan into atomic search queries.

CODE3"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("CODE4"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/agent.py*

LLM calls use LiteLLM, and each is wrapped with `flyte.trace` for observability:

CODE5

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/libs/utils/llms.py*

> [!NOTE]
> We use `flyte.trace` to track intermediate steps within a task, like LLM calls or specific function executions. This lightweight decorator adds observability with minimal overhead and is especially useful for inspecting reasoning chains during task execution.

## Search and summarize

We submit each research query to Tavily and summarize the results using an LLM. We run all summarization tasks with `asyncio.gather`, which signals to Flyte that these tasks can be distributed across separate compute resources.

CODE6"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("CODE7"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
CODE8
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pydantic==2.11.5",
#    "litellm==1.72.2",
#    "tavily-python==0.7.5",
#    "together==1.5.24",
#    "markdown==3.8.2",
#    "pymdown-extensions==10.16.1",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import json
from pathlib import Path

import flyte
import yaml
from flyte.io._file import File
from libs.utils.data_types import (
    DeepResearchResult,
    DeepResearchResults,
    ResearchPlan,
    SourceList,
)
from libs.utils.generation import generate_html, generate_toc_image
from libs.utils.llms import asingle_shot_llm_call
from libs.utils.log import AgentLogger
from libs.utils.tavily_search import atavily_search_results

TIME_LIMIT_MULTIPLIER = 5
MAX_COMPLETION_TOKENS = 4096

logging = AgentLogger("together.open_deep_research")

env = flyte.TaskEnvironment(
    name="deep-researcher",
    secrets=[
        flyte.Secret(key="together_api_key", as_env_var="TOGETHER_API_KEY"),
        flyte.Secret(key="tavily_api_key", as_env_var="TAVILY_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="deep-research-agent", pre=True)
    .with_apt_packages("pandoc", "texlive-xetex")
    .with_source_file(Path("prompts.yaml"), "/root"),
    resources=flyte.Resources(cpu=1),
)
# {{/docs-fragment env}}

# {{docs-fragment generate_research_queries}}
@env.task
async def generate_research_queries(
    topic: str,
    planning_model: str,
    json_model: str,
    prompts_file: File,
) -> list[str]:
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    PLANNING_PROMPT = prompts["planning_prompt"]

    plan = ""
    logging.info(f"\n\nGenerated deep research plan for topic: {topic}\n\nPlan:")
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=PLANNING_PROMPT,
        message=f"Research Topic: {topic}",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        plan += chunk
        print(chunk, end="", flush=True)

    SEARCH_PROMPT = prompts["plan_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=SEARCH_PROMPT,
        message=f"Plan to be parsed: {plan}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    plan = json.loads(response_json)
    return plan["queries"]
# {{/docs-fragment generate_research_queries}}

async def _summarize_content_async(
    raw_content: str,
    query: str,
    prompt: str,
    summarization_model: str,
) -> str:
    """Summarize content asynchronously using the LLM"""
    logging.info("Summarizing content asynchronously using the LLM")

    result = ""
    async for chunk in asingle_shot_llm_call(
        model=summarization_model,
        system_prompt=prompt,
        message=f"<Raw Content>{raw_content}</Raw Content>\n\n<Research Topic>{query}</Research Topic>",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        result += chunk
    return result

# {{docs-fragment search_and_summarize}}
@env.task
async def search_and_summarize(
    query: str,
    prompts_file: File,
    summarization_model: str,
) -> DeepResearchResults:
    """Perform search for a single query"""

    if len(query) > 400:
        # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues
        query = query[:400]
        logging.info(f"Truncated query to 400 characters: {query}")

    response = await atavily_search_results(query)

    logging.info("Tavily Search Called.")

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    RAW_CONTENT_SUMMARIZER_PROMPT = prompts["raw_content_summarizer_prompt"]

    with flyte.group("summarize-content"):
        # Create tasks for summarization
        summarization_tasks = []
        result_info = []
        for result in response.results:
            if result.raw_content is None:
                continue

            task = _summarize_content_async(
                result.raw_content,
                query,
                RAW_CONTENT_SUMMARIZER_PROMPT,
                summarization_model,
            )
            summarization_tasks.append(task)
            result_info.append(result)

        # Use return_exceptions=True to prevent exceptions from propagating
        summarized_contents = await asyncio.gather(
            *summarization_tasks, return_exceptions=True
        )

    # Filter out exceptions
    summarized_contents = [
        result for result in summarized_contents if not isinstance(result, Exception)
    ]

    formatted_results = []
    for result, summarized_content in zip(result_info, summarized_contents):
        formatted_results.append(
            DeepResearchResult(
                title=result.title,
                link=result.link,
                content=result.content,
                raw_content=result.raw_content,
                filtered_raw_content=summarized_content,
            )
        )
    return DeepResearchResults(results=formatted_results)
# {{/docs-fragment search_and_summarize}}

@env.task
async def search_all_queries(
    queries: list[str], summarization_model: str, prompts_file: File
) -> DeepResearchResults:
    """Execute searches for all queries in parallel"""
    tasks = []
    results_list = []

    tasks = [
        search_and_summarize(query, prompts_file, summarization_model)
        for query in queries
    ]

    if tasks:
        res_list = await asyncio.gather(*tasks)

    results_list.extend(res_list)

    # Combine all results
    combined_results = DeepResearchResults(results=[])
    for results in results_list:
        combined_results = combined_results + results

    return combined_results

# {{docs-fragment evaluate_research_completeness}}
@env.task
async def evaluate_research_completeness(
    topic: str,
    results: DeepResearchResults,
    queries: list[str],
    prompts_file: File,
    planning_model: str,
    json_model: str,
) -> list[str]:
    """
    Evaluate if the current search results are sufficient or if more research is needed.
    Returns an empty list if research is complete, or a list of additional queries if more research is needed.
    """

    # Format the search results for the LLM
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)

    EVALUATION_PROMPT = prompts["evaluation_prompt"]

    logging.info("\nEvaluation: ")
    evaluation = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=EVALUATION_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Search Queries Used>{queries}</Search Queries Used>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=None,
    ):
        evaluation += chunk
        print(chunk, end="", flush=True)

    EVALUATION_PARSING_PROMPT = prompts["evaluation_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=EVALUATION_PARSING_PROMPT,
        message=f"Evaluation to be parsed: {evaluation}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    evaluation = json.loads(response_json)
    return evaluation["queries"]
# {{/docs-fragment evaluate_research_completeness}}

# {{docs-fragment filter_results}}
@env.task
async def filter_results(
    topic: str,
    results: DeepResearchResults,
    prompts_file: File,
    planning_model: str,
    json_model: str,
    max_sources: int,
) -> DeepResearchResults:
    """Filter the search results based on the research plan"""

    # Format the search results for the LLM, without the raw content
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    FILTER_PROMPT = prompts["filter_prompt"]

    logging.info("\nFilter response: ")
    filter_response = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=FILTER_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        filter_response += chunk
        print(chunk, end="", flush=True)

    logging.info(f"Filter response: {filter_response}")

    FILTER_PARSING_PROMPT = prompts["filter_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=FILTER_PARSING_PROMPT,
        message=f"Filter response to be parsed: {filter_response}",
        response_format={
            "type": "json_object",
            "schema": SourceList.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    sources = json.loads(response_json)["sources"]

    logging.info(f"Filtered sources: {sources}")

    if max_sources != -1:
        sources = sources[:max_sources]

    # Filter the results based on the source list
    filtered_results = [
        results.results[i - 1] for i in sources if i - 1 < len(results.results)
    ]

    return DeepResearchResults(results=filtered_results)
# {{/docs-fragment filter_results}}

def _remove_thinking_tags(answer: str) -> str:
    """Remove content within <think> tags"""
    while "<think>" in answer and "</think>" in answer:
        start = answer.find("<think>")
        end = answer.find("</think>") + len("</think>")
        answer = answer[:start] + answer[end:]
    return answer

# {{docs-fragment generate_research_answer}}
@env.task
async def generate_research_answer(
    topic: str,
    results: DeepResearchResults,
    remove_thinking_tags: bool,
    prompts_file: File,
    answer_model: str,
) -> str:
    """
    Generate a comprehensive answer to the research topic based on the search results.
    Returns a detailed response that synthesizes information from all search results.
    """

    formatted_results = str(results)
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    ANSWER_PROMPT = prompts["answer_prompt"]

    answer = ""
    async for chunk in asingle_shot_llm_call(
        model=answer_model,
        system_prompt=ANSWER_PROMPT,
        message=f"Research Topic: {topic}\n\nSearch Results:\n{formatted_results}",
        response_format=None,
        # NOTE: This is the max_token parameter for the LLM call on Together AI,
        # may need to be changed for other providers
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        answer += chunk

    # this is just to avoid typing complaints
    if answer is None or not isinstance(answer, str):
        logging.error("No answer generated")
        return "No answer generated"

    if remove_thinking_tags:
        # Remove content within <think> tags
        answer = _remove_thinking_tags(answer)

    # Remove markdown code block markers if they exist at the beginning
    if answer.lstrip().startswith("CODE9"))
        if first_linebreak != -1:
            # Remove everything up to and including the first line break
            answer = answer[first_linebreak + 1 :]

        # Remove closing code block if it exists
        if answer.rstrip().endswith("CODE10

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/agent.py*

## Filter results

In this step, we evaluate the relevance of search results and rank them. This task returns the most useful sources for the final synthesis.

CODE11"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("CODE12"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
CODE13
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pydantic==2.11.5",
#    "litellm==1.72.2",
#    "tavily-python==0.7.5",
#    "together==1.5.24",
#    "markdown==3.8.2",
#    "pymdown-extensions==10.16.1",
# ]
# main = "main"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import json
from pathlib import Path

import flyte
import yaml
from flyte.io._file import File
from libs.utils.data_types import (
    DeepResearchResult,
    DeepResearchResults,
    ResearchPlan,
    SourceList,
)
from libs.utils.generation import generate_html, generate_toc_image
from libs.utils.llms import asingle_shot_llm_call
from libs.utils.log import AgentLogger
from libs.utils.tavily_search import atavily_search_results

TIME_LIMIT_MULTIPLIER = 5
MAX_COMPLETION_TOKENS = 4096

logging = AgentLogger("together.open_deep_research")

env = flyte.TaskEnvironment(
    name="deep-researcher",
    secrets=[
        flyte.Secret(key="together_api_key", as_env_var="TOGETHER_API_KEY"),
        flyte.Secret(key="tavily_api_key", as_env_var="TAVILY_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="deep-research-agent", pre=True)
    .with_apt_packages("pandoc", "texlive-xetex")
    .with_source_file(Path("prompts.yaml"), "/root"),
    resources=flyte.Resources(cpu=1),
)
# {{/docs-fragment env}}

# {{docs-fragment generate_research_queries}}
@env.task
async def generate_research_queries(
    topic: str,
    planning_model: str,
    json_model: str,
    prompts_file: File,
) -> list[str]:
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    PLANNING_PROMPT = prompts["planning_prompt"]

    plan = ""
    logging.info(f"\n\nGenerated deep research plan for topic: {topic}\n\nPlan:")
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=PLANNING_PROMPT,
        message=f"Research Topic: {topic}",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        plan += chunk
        print(chunk, end="", flush=True)

    SEARCH_PROMPT = prompts["plan_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=SEARCH_PROMPT,
        message=f"Plan to be parsed: {plan}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    plan = json.loads(response_json)
    return plan["queries"]
# {{/docs-fragment generate_research_queries}}

async def _summarize_content_async(
    raw_content: str,
    query: str,
    prompt: str,
    summarization_model: str,
) -> str:
    """Summarize content asynchronously using the LLM"""
    logging.info("Summarizing content asynchronously using the LLM")

    result = ""
    async for chunk in asingle_shot_llm_call(
        model=summarization_model,
        system_prompt=prompt,
        message=f"<Raw Content>{raw_content}</Raw Content>\n\n<Research Topic>{query}</Research Topic>",
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        result += chunk
    return result

# {{docs-fragment search_and_summarize}}
@env.task
async def search_and_summarize(
    query: str,
    prompts_file: File,
    summarization_model: str,
) -> DeepResearchResults:
    """Perform search for a single query"""

    if len(query) > 400:
        # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues
        query = query[:400]
        logging.info(f"Truncated query to 400 characters: {query}")

    response = await atavily_search_results(query)

    logging.info("Tavily Search Called.")

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    RAW_CONTENT_SUMMARIZER_PROMPT = prompts["raw_content_summarizer_prompt"]

    with flyte.group("summarize-content"):
        # Create tasks for summarization
        summarization_tasks = []
        result_info = []
        for result in response.results:
            if result.raw_content is None:
                continue

            task = _summarize_content_async(
                result.raw_content,
                query,
                RAW_CONTENT_SUMMARIZER_PROMPT,
                summarization_model,
            )
            summarization_tasks.append(task)
            result_info.append(result)

        # Use return_exceptions=True to prevent exceptions from propagating
        summarized_contents = await asyncio.gather(
            *summarization_tasks, return_exceptions=True
        )

    # Filter out exceptions
    summarized_contents = [
        result for result in summarized_contents if not isinstance(result, Exception)
    ]

    formatted_results = []
    for result, summarized_content in zip(result_info, summarized_contents):
        formatted_results.append(
            DeepResearchResult(
                title=result.title,
                link=result.link,
                content=result.content,
                raw_content=result.raw_content,
                filtered_raw_content=summarized_content,
            )
        )
    return DeepResearchResults(results=formatted_results)
# {{/docs-fragment search_and_summarize}}

@env.task
async def search_all_queries(
    queries: list[str], summarization_model: str, prompts_file: File
) -> DeepResearchResults:
    """Execute searches for all queries in parallel"""
    tasks = []
    results_list = []

    tasks = [
        search_and_summarize(query, prompts_file, summarization_model)
        for query in queries
    ]

    if tasks:
        res_list = await asyncio.gather(*tasks)

    results_list.extend(res_list)

    # Combine all results
    combined_results = DeepResearchResults(results=[])
    for results in results_list:
        combined_results = combined_results + results

    return combined_results

# {{docs-fragment evaluate_research_completeness}}
@env.task
async def evaluate_research_completeness(
    topic: str,
    results: DeepResearchResults,
    queries: list[str],
    prompts_file: File,
    planning_model: str,
    json_model: str,
) -> list[str]:
    """
    Evaluate if the current search results are sufficient or if more research is needed.
    Returns an empty list if research is complete, or a list of additional queries if more research is needed.
    """

    # Format the search results for the LLM
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)

    EVALUATION_PROMPT = prompts["evaluation_prompt"]

    logging.info("\nEvaluation: ")
    evaluation = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=EVALUATION_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Search Queries Used>{queries}</Search Queries Used>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=None,
    ):
        evaluation += chunk
        print(chunk, end="", flush=True)

    EVALUATION_PARSING_PROMPT = prompts["evaluation_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=EVALUATION_PARSING_PROMPT,
        message=f"Evaluation to be parsed: {evaluation}",
        response_format={
            "type": "json_object",
            "schema": ResearchPlan.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    evaluation = json.loads(response_json)
    return evaluation["queries"]
# {{/docs-fragment evaluate_research_completeness}}

# {{docs-fragment filter_results}}
@env.task
async def filter_results(
    topic: str,
    results: DeepResearchResults,
    prompts_file: File,
    planning_model: str,
    json_model: str,
    max_sources: int,
) -> DeepResearchResults:
    """Filter the search results based on the research plan"""

    # Format the search results for the LLM, without the raw content
    formatted_results = str(results)

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    FILTER_PROMPT = prompts["filter_prompt"]

    logging.info("\nFilter response: ")
    filter_response = ""
    async for chunk in asingle_shot_llm_call(
        model=planning_model,
        system_prompt=FILTER_PROMPT,
        message=(
            f"<Research Topic>{topic}</Research Topic>\n\n"
            f"<Current Search Results>{formatted_results}</Current Search Results>"
        ),
        response_format=None,
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        filter_response += chunk
        print(chunk, end="", flush=True)

    logging.info(f"Filter response: {filter_response}")

    FILTER_PARSING_PROMPT = prompts["filter_parsing_prompt"]

    response_json = ""
    async for chunk in asingle_shot_llm_call(
        model=json_model,
        system_prompt=FILTER_PARSING_PROMPT,
        message=f"Filter response to be parsed: {filter_response}",
        response_format={
            "type": "json_object",
            "schema": SourceList.model_json_schema(),
        },
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        response_json += chunk

    sources = json.loads(response_json)["sources"]

    logging.info(f"Filtered sources: {sources}")

    if max_sources != -1:
        sources = sources[:max_sources]

    # Filter the results based on the source list
    filtered_results = [
        results.results[i - 1] for i in sources if i - 1 < len(results.results)
    ]

    return DeepResearchResults(results=filtered_results)
# {{/docs-fragment filter_results}}

def _remove_thinking_tags(answer: str) -> str:
    """Remove content within <think> tags"""
    while "<think>" in answer and "</think>" in answer:
        start = answer.find("<think>")
        end = answer.find("</think>") + len("</think>")
        answer = answer[:start] + answer[end:]
    return answer

# {{docs-fragment generate_research_answer}}
@env.task
async def generate_research_answer(
    topic: str,
    results: DeepResearchResults,
    remove_thinking_tags: bool,
    prompts_file: File,
    answer_model: str,
) -> str:
    """
    Generate a comprehensive answer to the research topic based on the search results.
    Returns a detailed response that synthesizes information from all search results.
    """

    formatted_results = str(results)
    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    prompts = yaml.safe_load(yaml_contents)
    ANSWER_PROMPT = prompts["answer_prompt"]

    answer = ""
    async for chunk in asingle_shot_llm_call(
        model=answer_model,
        system_prompt=ANSWER_PROMPT,
        message=f"Research Topic: {topic}\n\nSearch Results:\n{formatted_results}",
        response_format=None,
        # NOTE: This is the max_token parameter for the LLM call on Together AI,
        # may need to be changed for other providers
        max_completion_tokens=MAX_COMPLETION_TOKENS,
    ):
        answer += chunk

    # this is just to avoid typing complaints
    if answer is None or not isinstance(answer, str):
        logging.error("No answer generated")
        return "No answer generated"

    if remove_thinking_tags:
        # Remove content within <think> tags
        answer = _remove_thinking_tags(answer)

    # Remove markdown code block markers if they exist at the beginning
    if answer.lstrip().startswith("CODE14"))
        if first_linebreak != -1:
            # Remove everything up to and including the first line break
            answer = answer[first_linebreak + 1 :]

        # Remove closing code block if it exists
        if answer.rstrip().endswith("CODE15

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/agent.py*

## Orchestration

Next, we define a `research_topic` task to orchestrate the entire deep research workflow. It runs the core stages in sequence: generating research queries, performing search and summarization, evaluating the completeness of results, and producing the final report.

CODE16"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("CODE17"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep_research_agent/agent.py*

The `main` task wraps this entire pipeline and adds report generation in HTML format as the final step.
It also serves as the main entry point to the workflow, allowing us to pass in all configuration parameters, including which LLMs to use at each stage.
This flexibility lets us mix and match models for planning, summarization, and final synthesis, helping us optimize for both cost and quality.

CODE18"):
        # Find the first line break after the opening backticks
        first_linebreak = answer.find("\n", answer.find("CODE19"):
            answer = answer.rstrip()[:-3].rstrip()

    return answer.strip()
# {{/docs-fragment generate_research_answer}}

# {{docs-fragment research_topic}}
@env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff_factor=2))
async def research_topic(
    topic: str,
    budget: int = 3,
    remove_thinking_tags: bool = True,
    max_queries: int = 5,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 40,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
    prompts_file: File | str = "prompts.yaml",
) -> str:
    """Main method to conduct research on a topic. Will be used for weave evals."""
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    # Step 1: Generate initial queries
    queries = await generate_research_queries(
        topic=topic,
        planning_model=planning_model,
        json_model=json_model,
        prompts_file=prompts_file,
    )
    queries = [topic, *queries[: max_queries - 1]]
    all_queries = queries.copy()
    logging.info(f"Initial queries: {queries}")

    if len(queries) == 0:
        logging.error("No initial queries generated")
        return "No initial queries generated"

    # Step 2: Perform initial search
    results = await search_all_queries(queries, summarization_model, prompts_file)
    logging.info(f"Initial search complete, found {len(results.results)} results")

    # Step 3: Conduct iterative research within budget
    for iteration in range(budget):
        with flyte.group(f"eval_iteration_{iteration}"):
            # Evaluate if more research is needed
            additional_queries = await evaluate_research_completeness(
                topic=topic,
                results=results,
                queries=all_queries,
                prompts_file=prompts_file,
                planning_model=planning_model,
                json_model=json_model,
            )

            # Filter out empty strings and check if any queries remain
            additional_queries = [q for q in additional_queries if q]
            if not additional_queries:
                logging.info("No need for additional research")
                break

            # for debugging purposes we limit the number of queries
            additional_queries = additional_queries[:max_queries]
            logging.info(f"Additional queries: {additional_queries}")

            # Expand research with new queries
            new_results = await search_all_queries(
                additional_queries, summarization_model, prompts_file
            )
            logging.info(
                f"Follow-up search complete, found {len(new_results.results)} results"
            )

            results = results + new_results
            all_queries.extend(additional_queries)

    # Step 4: Generate final answer
    logging.info(f"Generating final answer for topic: {topic}")
    results = results.dedup()
    logging.info(f"Deduplication complete, kept {len(results.results)} results")
    filtered_results = await filter_results(
        topic=topic,
        results=results,
        prompts_file=prompts_file,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
    )
    logging.info(
        f"LLM Filtering complete, kept {len(filtered_results.results)} results"
    )

    # Generate final answer
    answer = await generate_research_answer(
        topic=topic,
        results=filtered_results,
        remove_thinking_tags=remove_thinking_tags,
        prompts_file=prompts_file,
        answer_model=answer_model,
    )

    return answer
# {{/docs-fragment research_topic}}

# {{docs-fragment main}}
@env.task(report=True)
async def main(
    topic: str = (
        "List the essential requirements for a developer-focused agent orchestration system."
    ),
    prompts_file: File | str = "/root/prompts.yaml",
    budget: int = 2,
    remove_thinking_tags: bool = True,
    max_queries: int = 3,
    answer_model: str = "together_ai/deepseek-ai/DeepSeek-V3",
    planning_model: str = "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
    json_model: str = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    max_sources: int = 10,
    summarization_model: str = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
) -> str:
    if isinstance(prompts_file, str):
        prompts_file = await File.from_local(prompts_file)

    answer = await research_topic(
        topic=topic,
        budget=budget,
        remove_thinking_tags=remove_thinking_tags,
        max_queries=max_queries,
        answer_model=answer_model,
        planning_model=planning_model,
        json_model=json_model,
        max_sources=max_sources,
        summarization_model=summarization_model,
        prompts_file=prompts_file,
    )

    async with prompts_file.open() as fh:
        data = await fh.read()
        yaml_contents = str(data, "utf-8")

    toc_image_url = await generate_toc_image(
        yaml.safe_load(yaml_contents)["data_visualization_prompt"],
        planning_model,
        topic,
    )

    html_content = await generate_html(answer, toc_image_url)
    await flyte.report.replace.aio(html_content, do_flush=True)
    await flyte.report.flush.aio()

    return html_content
# {{/docs-fragment main}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
CODE20
flyte create secret TOGETHER_API_KEY <>
flyte create secret TAVILY_API_KEY <>
CODE21
uv run agent.py
CODE22
brew install pandoc
brew install basictex # restart your terminal after install

export TOGETHER_API_KEY=<>
export TAVILY_API_KEY=<>

uv run agent.py
CODE23
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "weave==0.51.51",
#    "datasets==3.6.0",
#    "huggingface-hub==0.32.6",
#    "litellm==1.72.2",
#    "tavily-python==0.7.5",
# ]
# ///

import os

import weave
from agent import research_topic
from datasets import load_dataset
from huggingface_hub import login
from libs.utils.log import AgentLogger
from litellm import completion

import flyte

logging = AgentLogger()

weave.init(project_name="deep-researcher")

env = flyte.TaskEnvironment(name="deep-researcher-eval")

@weave.op
def llm_as_a_judge_scoring(answer: str, output: str, question: str) -> bool:
    prompt = f"""
    Given the following question and answer, evaluate the answer against the correct answer:

    <question>
    {question}
    </question>

    <agent_answer>
    {output}
    </agent_answer>

    <correct_answer>
    {answer}
    </correct_answer>

    Note that the agent answer might be a long text containing a lot of information or it might be a short answer.

    You should read the entire text and think if the agent answers the question somewhere
    in the text. You should try to be flexible with the answer but careful.

    For example, answering with names instead of name and surname is fine.

    The important thing is that the answer of the agent either contains the correct answer or is equal to
    the correct answer.

    <reasoning>
    The agent answer is correct because I can read that ....
    </reasoning>

    <answer>
    1
    </answer>

    Otherwise, return

    <reasoning>
    The agent answer is incorrect because there is ...
    </reasoning>

    <answer>
    0
    </answer>

    """

    messages = [
        {
            "role": "system",
            "content": "You are an helpful assistant that returns a number between 0 and 1.",
        },
        {"role": "user", "content": prompt},
    ]
    answer = (
        completion(
            model="together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
            messages=messages,
            max_tokens=1000,
            temperature=0.0,
        )
        .choices[0]  # type: ignore
        .message["content"]  # type: ignore
    )

    return bool(int(answer.split("<answer>")[1].split("</answer>")[0].strip()))

def authenticate_huggingface():
    """Authenticate with Hugging Face Hub using token from environment variable."""
    token = os.getenv("HUGGINGFACE_TOKEN")
    if not token:
        raise ValueError(
            "HUGGINGFACE_TOKEN environment variable not set. "
            "Please set it with your token from https://huggingface.co/settings/tokens"
        )

    try:
        login(token=token)
        print("Successfully authenticated with Hugging Face Hub")
    except Exception as e:
        raise RuntimeError(f"Failed to authenticate with Hugging Face Hub: {e!s}")

@env.task
async def load_questions(
    dataset_names: list[str] | None = None,
) -> list[dict[str, str]]:
    """
    Load questions from the specified Hugging Face dataset configurations.

    Args:
        dataset_names: List of dataset configurations to load
                      Options:
                          "smolagents:simpleqa",
                          "hotpotqa",
                          "simpleqa",
                          "together-search-bench"
                      If None, all available configurations except hotpotqa will be loaded

    Returns:
        List of question-answer pairs
    """
    if dataset_names is None:
        dataset_names = ["smolagents:simpleqa"]

    all_questions = []

    # Authenticate with Hugging Face Hub (once and for all)
    authenticate_huggingface()

    for dataset_name in dataset_names:
        print(f"Loading dataset: {dataset_name}")

        try:
            if dataset_name == "together-search-bench":
                # Load Together-Search-Bench dataset
                dataset_path = "togethercomputer/together-search-bench"
                ds = load_dataset(dataset_path)
                if "test" in ds:
                    split_data = ds["test"]
                else:
                    print(f"No 'test' split found in dataset at {dataset_path}")
                    continue

                for i in range(len(split_data)):
                    item = split_data[i]
                    question_data = {
                        "question": item["question"],
                        "answer": item["answer"],
                        "dataset": item.get("dataset", "together-search-bench"),
                    }
                    all_questions.append(question_data)

                print(f"Loaded {len(split_data)} questions from together-search-bench dataset")
                continue

            elif dataset_name == "hotpotqa":
                # Load HotpotQA dataset (using distractor version for validation)
                ds = load_dataset("hotpotqa/hotpot_qa", "distractor", trust_remote_code=True)
                split_name = "validation"
            elif dataset_name == "simpleqa":
                ds = load_dataset("basicv8vc/SimpleQA")
                split_name = "test"
            else:
                # Strip "smolagents:" prefix when loading the dataset
                actual_dataset = dataset_name.split(":")[-1]
                ds = load_dataset("smolagents/benchmark-v1", actual_dataset)
                split_name = "test"

        except Exception as e:
            print(f"Failed to load dataset {dataset_name}: {e!s}")
            continue  # Skip this dataset if it fails to load

        print(f"Dataset structure for {dataset_name}: {ds}")
        print(f"Available splits: {list(ds)}")

        split_data = ds[split_name]  # type: ignore

        for i in range(len(split_data)):
            item = split_data[i]

            if dataset_name == "hotpotqa":
                # we remove questions that are easy or medium (if any) just to reduce the number of questions
                if item["level"] != "hard":
                    continue

                question_data = {
                    "question": item["question"],
                    "answer": item["answer"],
                    "dataset": dataset_name,
                }
            elif dataset_name == "simpleqa":
                # Handle SimpleQA dataset format
                question_data = {
                    "question": item["problem"],
                    "answer": item["answer"],
                    "dataset": dataset_name,
                }
            else:
                question_data = {
                    "question": item["question"],
                    "answer": item["true_answer"],
                    "dataset": dataset_name,
                }

            all_questions.append(question_data)

    print(f"Loaded {len(all_questions)} questions in total")
    return all_questions

@weave.op
async def predict(question: str):
    return await research_topic(topic=str(question))

@env.task
async def main(datasets: list[str] = ["together-search-bench"], limit: int | None = 1):
    questions = await load_questions(datasets)

    if limit is not None:
        questions = questions[:limit]
        print(f"Limited to {len(questions)} question(s)")

    evaluation = weave.Evaluation(dataset=questions, scorers=[llm_as_a_judge_scoring])
    await evaluation.evaluate(predict)

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.with_runcontext(raw_data_path="data").run(main)
CODE24
export HUGGINGFACE_TOKEN=<> # https://huggingface.co/settings/tokens
export WANDB_API_KEY=<> # https://wandb.ai/settings

uv run weave_evals.py
```

The script will run all tasks in the pipeline and log the evaluation results to Weights & Biases.
While you can also evaluate individual tasks, this script focuses on end-to-end evaluation of the end-to-end deep research workflow.

![Weave evaluations](../../../_static/images/tutorials/deep-research/weave_evals.png)


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/langgraph-agent-research ===

# LangGraph research agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/langgraph_agent_research).

This tutorial combines [LangGraph](https://langchain-ai.github.io/langgraph/) for agentic control flow with Flyte for durable compute. A research pipeline plans sub-topics, fans out ReAct agents that search the web with [Tavily](https://tavily.com/), synthesizes findings, and loops on quality gaps until the report is good enough. Each LangGraph step dispatches to a separate Flyte task so planning, research, synthesis, and quality checks appear independently in the run UI.

Flyte provides:

- **Per-step tasks** visible in the Flyte UI while LangGraph orchestrates the graph.
- **Secrets** for OpenAI and Tavily API keys.
- **Live HTML reports** with Mermaid graph visualizations and the final synthesized report.

## Define the task environment

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "langgraph>=1.0.7",
#    "langchain-anthropic",
#    "tavily-python",
#    "markdown",
#    "pydantic",
# ]
# main = "research_pipeline"
# params = ""
# ///
import json
import os
import base64
import logging
import markdown

import flyte
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="langgraph-agent-research", pre=True)

env = flyte.TaskEnvironment(
    name="langgraph-agent-research",
    image=main_img,
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
        flyte.Secret(key="tavily_api_key", as_env_var="TAVILY_API_KEY"),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
)
# {{/docs-fragment env}}

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage

from models import TopicReport, QualityResult, PipelineResult
from graph import build_pipeline_graph, build_research_subgraph

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
logging.getLogger("graph").setLevel(logging.INFO)
logging.getLogger("tools.search").setLevel(logging.INFO)

MODEL = "claude-3-5-haiku-latest"

def md_to_html(text: str) -> str:
    """Convert markdown to HTML for Flyte reports."""
    return markdown.markdown(text, extensions=["tables", "fenced_code"])

# ------------------------------------------------------------------
# Flyte tasks — each step is visible in the UI while running
# ------------------------------------------------------------------

@env.task(report=True)
async def plan_topics(query: str, num_topics: int = 3) -> list[str]:
    """Break a research query into focused sub-topics."""
    log.info(f"Planning {num_topics} sub-topics for: {query}")

    await flyte.report.replace.aio(
        f"<h2>Planning</h2><p>Breaking query into {num_topics} sub-topics...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    response = llm.invoke(
        f"Break this research question into exactly {num_topics} focused sub-topics. "
        f"Return ONLY a JSON array of strings, nothing else.\n\nQuestion: {query}"
    )
    try:
        topics = json.loads(response.content)
    except json.JSONDecodeError:
        topics = [query]

    topics = topics[:num_topics]
    log.info(f"Sub-topics: {topics}")

    topic_html = "".join(f"<li>{t}</li>" for t in topics)
    await flyte.report.replace.aio(
        f"<h2>Planning</h2><p>Sub-topics:</p><ul>{topic_html}</ul>"
    )
    await flyte.report.flush.aio()

    return topics

@env.task(report=True)
async def research_topic(topic: str, max_searches: int = 2) -> TopicReport:
    """Run the ReAct research agent on a single sub-topic."""
    log.info(f"[Research Task] Starting: {topic}")

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    tavily_api_key = os.getenv("TAVILY_API_KEY")

    await flyte.report.replace.aio(f"<h2>Researching: {topic}</h2><p>Running searches...</p>")
    await flyte.report.flush.aio()

    graph = build_research_subgraph(
        anthropic_api_key=anthropic_api_key,
        tavily_api_key=tavily_api_key,
        max_searches=max_searches,
        model=MODEL,
    )
    result = await graph.ainvoke({"messages": [HumanMessage(content=f"Research this topic: {topic}")]})
    report = result["messages"][-1].content
    log.info(f"[Research Task] Done: {topic}")

    await flyte.report.replace.aio(f"<h2>{topic}</h2>{md_to_html(report)}")
    await flyte.report.flush.aio()

    return TopicReport(topic=topic, report=report)

@env.task(report=True)
async def synthesize(query: str, results: list[TopicReport]) -> str:
    """Combine sub-topic research reports into a unified synthesis."""
    log.info(f"Synthesizing {len(results)} report(s)...")

    await flyte.report.replace.aio(
        f"<h2>Synthesis</h2><p>Combining {len(results)} reports...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    sections = "\n\n---\n\n".join(
        f"## {r.topic}\n\n{r.report}" for r in results
    )

    response = llm.invoke(
        f"You have research reports on sub-topics of this question:\n\n{query}\n\n"
        f"Sub-topic reports:\n\n{sections}\n\n"
        f"Write a comprehensive report that synthesizes all findings. "
        f"Organize by theme, highlight connections between sub-topics, "
        f"and end with key takeaways."
    )
    synthesis = response.content
    log.info(f"Synthesis complete: {len(synthesis)} chars")

    await flyte.report.replace.aio(f"<h2>Synthesis</h2>{md_to_html(synthesis)}")
    await flyte.report.flush.aio()

    return synthesis

@env.task(report=True)
async def quality_check(query: str, synthesis: str) -> QualityResult:
    """Evaluate report quality and identify gaps."""
    log.info("Evaluating quality...")

    await flyte.report.replace.aio(
        "<h2>Quality Check</h2><p>Evaluating report quality...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    response = llm.invoke(
        f'Evaluate this research report for the question: {query}\n\n'
        f'Report:\n{synthesis}\n\n'
        f'Rate the report quality from 1-10 and identify any gaps or missing perspectives. '
        f'Return JSON: {{"score": <int>, "gaps": [<string>, ...]}}\n'
        f'If the report is comprehensive (score >= 8) or there are no significant gaps, '
        f'return an empty gaps list.'
    )

    try:
        evaluation = json.loads(response.content)
        score = evaluation.get("score", 8)
        gaps = evaluation.get("gaps", [])
    except json.JSONDecodeError:
        score = 8
        gaps = []

    result = QualityResult(score=score, gaps=gaps)
    log.info(f"Score: {result.score}/10, Gaps: {len(result.gaps)}")

    gap_html = "".join(f"<li>{g}</li>" for g in result.gaps) if result.gaps else "<li>None</li>"
    await flyte.report.replace.aio(
        f"<h2>Quality Check</h2>"
        f"<p><b>Score:</b> {result.score}/10</p>"
        f"<p><b>Gaps:</b></p><ul>{gap_html}</ul>"
    )
    await flyte.report.flush.aio()

    return result

# ------------------------------------------------------------------
# Orchestrator: runs the LangGraph pipeline, backed by Flyte tasks
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def research_pipeline(
    query: str,
    num_topics: int = 3,
    max_searches: int = 2,
    max_iterations: int = 2,
) -> PipelineResult:
    """
    Research pipeline workflow:
    1. LangGraph plans sub-topics via plan_topics Flyte task
    2. LangGraph fans out research via Send → each dispatches to research_topic Flyte task
    3. LangGraph synthesizes results via synthesize Flyte task
    4. LangGraph evaluates quality via quality_check Flyte task
    5. If gaps found, loops back to step 2
    """
    log.info(f"Starting research pipeline: {query}")

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    tavily_api_key = os.getenv("TAVILY_API_KEY")

    # Build the pipeline graph, passing all Flyte tasks as compute backends
    pipeline = build_pipeline_graph(
        anthropic_api_key=anthropic_api_key,
        tavily_api_key=tavily_api_key,
        plan_task=plan_topics,
        research_task=research_topic,
        synthesize_task=synthesize,
        quality_check_task=quality_check,
        model=MODEL,
    )

    # Visualize the graphs in report tabs
    graph_tab = flyte.report.get_tab("Agent Graphs")

    png_bytes = pipeline.get_graph().draw_mermaid_png()
    img_b64 = base64.b64encode(png_bytes).decode()
    graph_tab.log(f"""\
<h2>Research Pipeline</h2>\
<img src="data:image/png;base64,{img_b64}" alt="Research pipeline" />""")

    subgraph = build_research_subgraph(anthropic_api_key, tavily_api_key, max_searches, model=MODEL)
    sub_png = subgraph.get_graph().draw_mermaid_png()
    sub_b64 = base64.b64encode(sub_png).decode()
    graph_tab.log(f"""\
<h2>Research Agent (ReAct)</h2>\
<img src="data:image/png;base64,{sub_b64}" alt="ReAct research agent" />""")
    await flyte.report.flush.aio()

    # Run the pipeline — LangGraph controls the flow, Flyte tasks run the compute
    result = await pipeline.ainvoke({
        "query": query,
        "num_topics": num_topics,
        "max_searches": max_searches,
        "max_iterations": max_iterations,
        "iteration": 0,
        "topics": [],
        "research_results": [],
        "synthesis": "",
        "score": 0,
        "gaps": [],
        "final_report": "",
    })

    # Build the final report
    final_report = result["final_report"]
    sub_reports = [TopicReport(**r) for r in result["research_results"]]
    score = result.get("score", 0)
    iteration = result.get("iteration", 1) - 1

    await flyte.report.replace.aio(f"""\
<h2>Research Report</h2>\
<p><b>Query:</b> {query}</p>\
<p><b>Quality:</b> {score}/10 after {iteration} iteration(s)</p>\
<hr/>{md_to_html(final_report)}""")
    await flyte.report.flush.aio()

    log.info(f"Research pipeline complete. Score: {score}/10, Iterations: {iteration}")
    return PipelineResult(
        query=query,
        report=final_report,
        sub_reports=sub_reports,
        score=score,
        iterations=iteration,
    )

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(research_pipeline(query="Compare quantum computing approaches"))
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/langgraph_agent_research/langgraph_agent_research.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "langgraph>=1.0.7",
#    "langchain-openai",
#    "tavily-python",
#    ...
# ]
# ///
```

## Orchestrate the pipeline

The `research_pipeline` task builds the LangGraph workflow, renders graph diagrams in a report tab, and runs the full plan → research → synthesize → quality-check loop.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "langgraph>=1.0.7",
#    "langchain-anthropic",
#    "tavily-python",
#    "markdown",
#    "pydantic",
# ]
# main = "research_pipeline"
# params = ""
# ///
import json
import os
import base64
import logging
import markdown

import flyte
import flyte.report

# {{docs-fragment env}}
main_img = flyte.Image.from_uv_script(__file__, name="langgraph-agent-research", pre=True)

env = flyte.TaskEnvironment(
    name="langgraph-agent-research",
    image=main_img,
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
        flyte.Secret(key="tavily_api_key", as_env_var="TAVILY_API_KEY"),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
)
# {{/docs-fragment env}}

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage

from models import TopicReport, QualityResult, PipelineResult
from graph import build_pipeline_graph, build_research_subgraph

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
logging.getLogger("graph").setLevel(logging.INFO)
logging.getLogger("tools.search").setLevel(logging.INFO)

MODEL = "claude-3-5-haiku-latest"

def md_to_html(text: str) -> str:
    """Convert markdown to HTML for Flyte reports."""
    return markdown.markdown(text, extensions=["tables", "fenced_code"])

# ------------------------------------------------------------------
# Flyte tasks — each step is visible in the UI while running
# ------------------------------------------------------------------

@env.task(report=True)
async def plan_topics(query: str, num_topics: int = 3) -> list[str]:
    """Break a research query into focused sub-topics."""
    log.info(f"Planning {num_topics} sub-topics for: {query}")

    await flyte.report.replace.aio(
        f"<h2>Planning</h2><p>Breaking query into {num_topics} sub-topics...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    response = llm.invoke(
        f"Break this research question into exactly {num_topics} focused sub-topics. "
        f"Return ONLY a JSON array of strings, nothing else.\n\nQuestion: {query}"
    )
    try:
        topics = json.loads(response.content)
    except json.JSONDecodeError:
        topics = [query]

    topics = topics[:num_topics]
    log.info(f"Sub-topics: {topics}")

    topic_html = "".join(f"<li>{t}</li>" for t in topics)
    await flyte.report.replace.aio(
        f"<h2>Planning</h2><p>Sub-topics:</p><ul>{topic_html}</ul>"
    )
    await flyte.report.flush.aio()

    return topics

@env.task(report=True)
async def research_topic(topic: str, max_searches: int = 2) -> TopicReport:
    """Run the ReAct research agent on a single sub-topic."""
    log.info(f"[Research Task] Starting: {topic}")

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    tavily_api_key = os.getenv("TAVILY_API_KEY")

    await flyte.report.replace.aio(f"<h2>Researching: {topic}</h2><p>Running searches...</p>")
    await flyte.report.flush.aio()

    graph = build_research_subgraph(
        anthropic_api_key=anthropic_api_key,
        tavily_api_key=tavily_api_key,
        max_searches=max_searches,
        model=MODEL,
    )
    result = await graph.ainvoke({"messages": [HumanMessage(content=f"Research this topic: {topic}")]})
    report = result["messages"][-1].content
    log.info(f"[Research Task] Done: {topic}")

    await flyte.report.replace.aio(f"<h2>{topic}</h2>{md_to_html(report)}")
    await flyte.report.flush.aio()

    return TopicReport(topic=topic, report=report)

@env.task(report=True)
async def synthesize(query: str, results: list[TopicReport]) -> str:
    """Combine sub-topic research reports into a unified synthesis."""
    log.info(f"Synthesizing {len(results)} report(s)...")

    await flyte.report.replace.aio(
        f"<h2>Synthesis</h2><p>Combining {len(results)} reports...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    sections = "\n\n---\n\n".join(
        f"## {r.topic}\n\n{r.report}" for r in results
    )

    response = llm.invoke(
        f"You have research reports on sub-topics of this question:\n\n{query}\n\n"
        f"Sub-topic reports:\n\n{sections}\n\n"
        f"Write a comprehensive report that synthesizes all findings. "
        f"Organize by theme, highlight connections between sub-topics, "
        f"and end with key takeaways."
    )
    synthesis = response.content
    log.info(f"Synthesis complete: {len(synthesis)} chars")

    await flyte.report.replace.aio(f"<h2>Synthesis</h2>{md_to_html(synthesis)}")
    await flyte.report.flush.aio()

    return synthesis

@env.task(report=True)
async def quality_check(query: str, synthesis: str) -> QualityResult:
    """Evaluate report quality and identify gaps."""
    log.info("Evaluating quality...")

    await flyte.report.replace.aio(
        "<h2>Quality Check</h2><p>Evaluating report quality...</p>"
    )
    await flyte.report.flush.aio()

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = ChatAnthropic(model=MODEL, api_key=anthropic_api_key)

    response = llm.invoke(
        f'Evaluate this research report for the question: {query}\n\n'
        f'Report:\n{synthesis}\n\n'
        f'Rate the report quality from 1-10 and identify any gaps or missing perspectives. '
        f'Return JSON: {{"score": <int>, "gaps": [<string>, ...]}}\n'
        f'If the report is comprehensive (score >= 8) or there are no significant gaps, '
        f'return an empty gaps list.'
    )

    try:
        evaluation = json.loads(response.content)
        score = evaluation.get("score", 8)
        gaps = evaluation.get("gaps", [])
    except json.JSONDecodeError:
        score = 8
        gaps = []

    result = QualityResult(score=score, gaps=gaps)
    log.info(f"Score: {result.score}/10, Gaps: {len(result.gaps)}")

    gap_html = "".join(f"<li>{g}</li>" for g in result.gaps) if result.gaps else "<li>None</li>"
    await flyte.report.replace.aio(
        f"<h2>Quality Check</h2>"
        f"<p><b>Score:</b> {result.score}/10</p>"
        f"<p><b>Gaps:</b></p><ul>{gap_html}</ul>"
    )
    await flyte.report.flush.aio()

    return result

# ------------------------------------------------------------------
# Orchestrator: runs the LangGraph pipeline, backed by Flyte tasks
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@env.task(report=True)
async def research_pipeline(
    query: str,
    num_topics: int = 3,
    max_searches: int = 2,
    max_iterations: int = 2,
) -> PipelineResult:
    """
    Research pipeline workflow:
    1. LangGraph plans sub-topics via plan_topics Flyte task
    2. LangGraph fans out research via Send → each dispatches to research_topic Flyte task
    3. LangGraph synthesizes results via synthesize Flyte task
    4. LangGraph evaluates quality via quality_check Flyte task
    5. If gaps found, loops back to step 2
    """
    log.info(f"Starting research pipeline: {query}")

    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    tavily_api_key = os.getenv("TAVILY_API_KEY")

    # Build the pipeline graph, passing all Flyte tasks as compute backends
    pipeline = build_pipeline_graph(
        anthropic_api_key=anthropic_api_key,
        tavily_api_key=tavily_api_key,
        plan_task=plan_topics,
        research_task=research_topic,
        synthesize_task=synthesize,
        quality_check_task=quality_check,
        model=MODEL,
    )

    # Visualize the graphs in report tabs
    graph_tab = flyte.report.get_tab("Agent Graphs")

    png_bytes = pipeline.get_graph().draw_mermaid_png()
    img_b64 = base64.b64encode(png_bytes).decode()
    graph_tab.log(f"""\
<h2>Research Pipeline</h2>\
<img src="data:image/png;base64,{img_b64}" alt="Research pipeline" />""")

    subgraph = build_research_subgraph(anthropic_api_key, tavily_api_key, max_searches, model=MODEL)
    sub_png = subgraph.get_graph().draw_mermaid_png()
    sub_b64 = base64.b64encode(sub_png).decode()
    graph_tab.log(f"""\
<h2>Research Agent (ReAct)</h2>\
<img src="data:image/png;base64,{sub_b64}" alt="ReAct research agent" />""")
    await flyte.report.flush.aio()

    # Run the pipeline — LangGraph controls the flow, Flyte tasks run the compute
    result = await pipeline.ainvoke({
        "query": query,
        "num_topics": num_topics,
        "max_searches": max_searches,
        "max_iterations": max_iterations,
        "iteration": 0,
        "topics": [],
        "research_results": [],
        "synthesis": "",
        "score": 0,
        "gaps": [],
        "final_report": "",
    })

    # Build the final report
    final_report = result["final_report"]
    sub_reports = [TopicReport(**r) for r in result["research_results"]]
    score = result.get("score", 0)
    iteration = result.get("iteration", 1) - 1

    await flyte.report.replace.aio(f"""\
<h2>Research Report</h2>\
<p><b>Query:</b> {query}</p>\
<p><b>Quality:</b> {score}/10 after {iteration} iteration(s)</p>\
<hr/>{md_to_html(final_report)}""")
    await flyte.report.flush.aio()

    log.info(f"Research pipeline complete. Score: {score}/10, Iterations: {iteration}")
    return PipelineResult(
        query=query,
        report=final_report,
        sub_reports=sub_reports,
        score=score,
        iterations=iteration,
    )

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(research_pipeline(query="Compare quantum computing approaches"))
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/langgraph_agent_research/langgraph_agent_research.py*

Inside each research task, a ReAct subgraph (`graph.py`) uses `@flyte.trace` on Tavily search calls for observability.

## Run the agent

Create secrets for Anthropic and Tavily:

```
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
flyte create secret tavily_api_key <YOUR_TAVILY_API_KEY>
```

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/langgraph_agent_research):

```
cd v2/tutorials/langgraph_agent_research
uv run --script langgraph_agent_research.py
```

Or pass a custom query:

```
flyte run langgraph_agent_research.py research_pipeline --query "Compare quantum computing approaches"
```

Check the **Agent Graphs** report tab for the LangGraph diagram and the main report for the synthesized research output.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/mle-bot ===

# MLE bot: an autonomous ML engineer

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/mle_bot).

You have a dataset and a business question. Today, going from a raw CSV to a trained, evaluated model with a written report takes an ML engineer hours of experimentation: profiling the data, picking algorithms, engineering features, tuning hyperparameters, analyzing results, and iterating. What if you could describe the problem in plain English and let an agent handle the rest?

This tutorial walks you through building exactly that. You'll construct an autonomous ML engineer that takes a problem description and a dataset, designs experiments, runs them on cloud infrastructure, analyzes results, iterates, and produces a report summarizing the best model it found.

## TL;DR

- You'll build an agent that takes a natural language problem description and a CSV file, then produces a trained model and a detailed report comparing the results.
- The LLM reasons over dataset statistics, never raw data. Trusted tools compute statistics in the cloud, and only those statistics reach the LLM.
- LLM-generated orchestration code runs inside Flyte's sandbox: no imports, no network access, no filesystem. It can only call pre-approved tool functions.
- Each tool function runs as a durable Flyte task in the cloud, with retries, observability, and full traceability.

## The problem with LLMs and ML pipelines

If you ask an LLM to "train a model on this dataset," you run into a few issues fast. The LLM might hallucinate sklearn APIs that don't exist. It has no access to real compute, so it can't actually train anything. It runs everything in a single context with no way to handle large datasets. And if something goes wrong, there's no structured way to iterate.

The core tension is that LLMs are genuinely good at reasoning about *what* to try. Given a dataset profile showing class imbalance and temporal structure, a capable model will suggest rolling window features and appropriate resampling strategies. But LLMs are unreliable at *executing* those plans. They generate buggy code, lose track of variable names, and have no way to dispatch real compute.

The solution is to separate the two concerns. Let the LLM handle the planning: which algorithms to try, what feature engineering to apply, which hyperparameters to tune. Then hand the execution to trusted tool functions that run on real infrastructure. The LLM controls *what* happens. The tools control *how*.

Think of it like giving a junior engineer access to a curated set of approved tools and reviewing their work. They can compose those tools in creative ways, but they can't go off-script and install random packages or hit arbitrary endpoints.

## How it works

The agent runs in five phases:

1. **Profile** the dataset using a trusted tool. The tool returns statistics (shape, class balance, feature correlations, missing values). The LLM never touches the raw data.
2. **Design** a batch of experiments. The LLM reads the profile and proposes 2 to 3 experiments, each with an algorithm, hyperparameters, and a feature engineering strategy.
3. **Execute** each experiment in parallel. For each one, the LLM generates Python orchestration code that chains together pre-approved tool functions. That code runs inside a restricted sandbox, and each tool call dispatches as a durable Flyte task on cloud compute.
4. **Analyze** the results. The LLM reviews metrics across experiments, optionally requests targeted data explorations (e.g., "are failures concentrated on specific machines?"), and decides whether to iterate with new experiments.
5. **Produce a report** summarizing the winning model, the experiment journey, and deployment recommendations.

Two things make this work. First, the LLM never sees raw data. The profiling tool runs in the cloud on managed compute and returns only aggregated statistics. This keeps prompt sizes manageable and avoids leaking sensitive data into LLM API calls. Second, the LLM-generated code runs inside Flyte's sandbox where the only thing it can do is call your pre-approved tool functions. More on that shortly.

### What to expect

Here's what an actual run looks like on a synthetic predictive maintenance dataset (175k rows of sensor data from 20 industrial pumps, ~3% failure rate).

In the first iteration, the agent designed three experiments: a logistic regression baseline, an XGBoost model with rolling window features, and a random forest with lag features. After reviewing the results, it decided to continue. It requested two targeted explorations ("do failure cases show meaningfully higher vibration?" and "how do feature-target correlations vary by pump?"), then used those findings to design a second round of experiments with tuned feature engineering and class weighting.

After two iterations and five total experiments, the final rankings looked like this:

| Rank | Experiment | ROC-AUC | F1 | Recall | Precision |
|------|-----------|---------|------|--------|-----------|
| 1 | **Random Forest with Balanced Class Weights** | 0.7983 | 0.4284 | 0.4561 | 0.4038 |
| 2 | XGBoost with Feature Engineering | 0.7847 | 0.4568 | 0.4722 | 0.4425 |
| 3 | Enhanced XGBoost with Focused Feature Engineering | 0.7821 | 0.3565 | 0.4973 | 0.2778 |
| 4 | Random Forest with Lag Features | 0.7651 | 0.5206 | 0.4104 | 0.7116 |
| 5 | Baseline Logistic Regression | 0.7528 | 0.118 | 0.6496 | 0.0649 |

The agent autonomously explored different algorithms, feature strategies, and class imbalance techniques, then ranked everything by ROC-AUC. The full report includes the LLM's reasoning and generated code for every experiment, so you can trace exactly why it chose each approach and what code it wrote to implement it. Since the LLM makes different decisions each run, your results will vary, but the overall pattern (profile, design, execute, analyze, iterate) stays the same.

## Declaring task environments

Before writing any tasks, you need to declare *where* and *how* they run. In Flyte v2, a `TaskEnvironment` bundles together the container image, resource requirements, secrets, and dependencies for a group of tasks.

The MLE Bot uses two environments. One for the ML tools (pandas, sklearn, xgboost) and one for the agent itself (the OpenAI client and the sandbox runtime):

```
"""Flyte TaskEnvironment definitions for mle-bot.

Two environments:
- tool_env: Runs the ML tools (data loading, feature engineering, training, evaluation).
            Has sklearn, xgboost, pandas, numpy, joblib.
- agent_env: Runs the orchestrating agent (OpenAI calls, sandbox orchestration).
             Has openai, pydantic-monty. Depends on tool_env.
"""

# {{docs-fragment environments}}
import flyte

tool_env = flyte.TaskEnvironment(
    "mle-tools",
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    image=(
        flyte.Image.from_debian_base(name="mle-tools-image").with_pip_packages(
            "pandas>=2.0.0",
            "scikit-learn>=1.3.0",
            "xgboost>=2.0.0",
            "numpy>=1.24.0",
            "joblib>=1.3.0",
        )
    ),
)

agent_env = flyte.TaskEnvironment(
    "mle-agent",
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    secrets=[flyte.Secret(key="OPENAI_API_KEY", as_env_var="OPENAI_API_KEY")],
    env_vars={"PYTHONUNBUFFERED": "1"},
    image=(
        flyte.Image.from_debian_base(name="mle-agent-image")
        .with_apt_packages("git")
        .with_pip_packages(
            "openai>=1.0.0",
            "flyte[sandbox]",
        )
    ),
    depends_on=[tool_env],
)
# {{/docs-fragment environments}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/environments.py*

A few things to note. `flyte.Resources` sets the CPU and memory for every task in that environment. `flyte.Image.from_debian_base()` builds a container image on the fly with the packages you declare, so you never need to manage Dockerfiles. `flyte.Secret` injects a secret from your cluster's secret store as an environment variable. And `depends_on=[tool_env]` tells Flyte that the agent environment needs to be able to dispatch tasks in the tool environment. This is what enables the sandbox to call tool functions that run on separate, appropriately-resourced compute.

## Building durable tool functions

Each tool is a regular Python async function decorated with `@env.task`. That decorator turns it into a durable Flyte task: it runs in its own container with the resources declared on the environment, it's automatically retried on transient failures, and every invocation is tracked in the Flyte UI.

Data flows between tasks as `flyte.io.File` objects. A `File` is a reference to data in cloud storage. When a task needs the actual bytes, it calls `await data.download()` to pull them into the container's local filesystem. When it produces output, it creates a `File` from a local path and returns it. Flyte handles the upload to cloud storage when the task completes. The data itself never passes through the agent or the LLM.

Here's what the training tool looks like:

```
"""Model training tools.

A single unified interface for training classifiers with different algorithms.
The tool handles serialization, class imbalance, and basic hyperparameter passing.
"""

from flyte.io import File

from mle_bot.environments import tool_env
from mle_bot.schemas import (
    GradientBoostingParams,
    LogisticRegressionParams,
    RandomForestParams,
    XGBoostParams,
)

# {{docs-fragment train_model}}
@tool_env.task
async def train_model(
    data: File,
    target_column: str,
    algorithm: str,
    hyperparams: dict,
) -> File:
    """Train a classification model and return the serialized model and training metrics.

    Supports multiple algorithms through a single interface so the agent can
    dispatch different approaches without knowing implementation details.

    Args:
        data: CSV file with training data (features + target column).
        target_column: Name of the column to predict.
        algorithm: One of:
            "xgboost"            — Gradient boosted trees. Good default for tabular data.
                                   Handles missing values and class imbalance natively.
            "random_forest"      — Ensemble of decision trees. More robust to outliers.
            "logistic_regression"— Linear model. Fast baseline, good for linearly separable problems.
            "gradient_boosting"  — Sklearn GradientBoostingClassifier. Slower than xgboost
                                   but sometimes better on small datasets.
        hyperparams: Dict of algorithm-specific hyperparameters. Common keys:
            For xgboost / gradient_boosting:
                n_estimators (int, default 100): Number of trees.
                max_depth (int, default 6): Maximum tree depth.
                learning_rate (float, default 0.1): Step size shrinkage.
                scale_pos_weight (float): Ratio negative/positive — use for imbalanced data.
                                          Set to (n_negative / n_positive) to upweight minority class.
                subsample (float, default 1.0): Fraction of samples used per tree.
                colsample_bytree (float, default 1.0): Fraction of features per tree.
            For random_forest:
                n_estimators (int, default 100): Number of trees.
                max_depth (int or null, default null): Maximum tree depth (null = unlimited).
                min_samples_leaf (int, default 1): Minimum samples at a leaf node.
                class_weight (str, default "balanced"): "balanced" reweights by class frequency.
            For logistic_regression:
                C (float, default 1.0): Inverse regularization strength (higher = less regularization).
                max_iter (int, default 1000): Maximum iterations for solver.
                class_weight (str, default "balanced"): "balanced" reweights by class frequency.

    Returns:
        File — serialized model (joblib format, contains model + feature columns + target column).
    """
    # {{/docs-fragment train_model}}
    import tempfile

    import joblib
    import numpy as np
    import pandas as pd
    from flyte.io import File as FlyteFile
    from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score

    path = await data.download()
    df = pd.read_csv(path)

    # Only use numeric columns — drop strings like machine_id automatically
    feature_cols = [c for c in df.select_dtypes(include=[np.number]).columns if c != target_column]
    X = df[feature_cols].values
    y = df[target_column].values

    class_dist = {str(k): int(v) for k, v in zip(*np.unique(y, return_counts=True))}
    n_positive = int((y == 1).sum())
    n_negative = int((y == 0).sum())
    default_scale = max(1.0, n_negative / n_positive) if n_positive > 0 else 1.0

    if algorithm == "xgboost":
        from xgboost import XGBClassifier
        p = XGBoostParams.model_validate({**hyperparams, "scale_pos_weight": hyperparams.get("scale_pos_weight", default_scale)})
        params = {**p.model_dump(), "eval_metric": "logloss", "random_state": 42}
        model = XGBClassifier(**params)

    elif algorithm == "random_forest":
        p = RandomForestParams.model_validate(hyperparams)
        params = {**p.model_dump(), "random_state": 42, "n_jobs": -1}
        model = RandomForestClassifier(**params)

    elif algorithm == "gradient_boosting":
        p = GradientBoostingParams.model_validate(hyperparams)
        params = {**p.model_dump(), "random_state": 42}
        model = GradientBoostingClassifier(**params)

    elif algorithm == "logistic_regression":
        p = LogisticRegressionParams.model_validate(hyperparams)
        params = {**p.model_dump(), "random_state": 42}
        model = LogisticRegression(**params)

    else:
        raise ValueError(f"Unknown algorithm: {algorithm!r}. Choose from: xgboost, random_forest, gradient_boosting, logistic_regression")

    model.fit(X, y)
    y_pred = model.predict(X)
    y_prob = model.predict_proba(X)[:, 1] if hasattr(model, "predict_proba") else y_pred

    train_metrics = {
        "accuracy": round(float(accuracy_score(y, y_pred)), 4),
        "f1": round(float(f1_score(y, y_pred, average="binary", zero_division=0)), 4),
        "precision": round(float(precision_score(y, y_pred, average="binary", zero_division=0)), 4),
        "recall": round(float(recall_score(y, y_pred, average="binary", zero_division=0)), 4),
        "roc_auc": round(float(roc_auc_score(y, y_prob)), 4),
    }

    # Feature importance (top 20)
    if hasattr(model, "feature_importances_"):
        importances = model.feature_importances_
        importance_dict = {feature_cols[i]: round(float(importances[i]), 4) for i in range(len(feature_cols))}
        importance_dict = dict(sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:20])
    elif hasattr(model, "coef_"):
        importances = abs(model.coef_[0])
        importance_dict = {feature_cols[i]: round(float(importances[i]), 4) for i in range(len(feature_cols))}
        importance_dict = dict(sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:20])
    else:
        importance_dict = {}

    model_file = tempfile.NamedTemporaryFile(suffix=".joblib", delete=False)
    joblib.dump({"model": model, "feature_columns": feature_cols, "target_column": target_column}, model_file.name)
    model_file.close()

    return await FlyteFile.from_local(model_file.name)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/tools/training.py*

And here's the profiling tool, which is the first thing the agent calls. It computes dataset statistics that the LLM will use to design experiments:

```
"""Data loading, profiling, and splitting tools.

These tools are safe, general-purpose, and side-effect free.
They run as durable Flyte tasks so they execute in the cloud on managed compute.
"""

from flyte.io import File

from mle_bot.environments import tool_env

# {{docs-fragment profile_dataset}}
@tool_env.task
async def profile_dataset(data: File, target_column: str) -> dict:
    """Profile a dataset and return statistics that inform ML problem design.

    Call this first before designing any experiments. The returned profile tells
    you the shape, column types, class balance, missing values, and numeric
    statistics — everything needed to choose algorithms and feature strategies.

    Args:
        data: CSV file to profile.
        target_column: Name of the column to predict.

    Returns a dict with keys:
        - shape: [n_rows, n_cols]
        - columns: list of all column names
        - dtypes: {col: dtype_string, ...}
        - numeric_columns: list of numeric column names (excluding target)
        - categorical_columns: list of non-numeric column names (excluding target)
        - target_distribution: {class_value: count, ...}
        - class_balance: {class_value: pct, ...}  (proportions, sum to 100)
        - missing_pct: {col: pct_missing, ...}
        - numeric_stats: {col: {mean, std, min, max, median}, ...}
        - n_classes: int — number of unique target values
        - is_imbalanced: bool — True if minority class < 20% of data
        - sample: list of 3 example rows as dicts
    """
    import numpy as np
    import pandas as pd

    path = await data.download()
    df = pd.read_csv(path)

    target_counts = df[target_column].value_counts()
    class_balance = (df[target_column].value_counts(normalize=True) * 100).round(2).to_dict()
    minority_pct = float(min(class_balance.values()))

    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    categorical_cols = df.select_dtypes(exclude=[np.number]).columns.tolist()

    numeric_stats = {}
    for col in numeric_cols:
        if col == target_column:
            continue
        numeric_stats[col] = {
            "mean": round(float(df[col].mean()), 4),
            "std": round(float(df[col].std()), 4),
            "min": round(float(df[col].min()), 4),
            "max": round(float(df[col].max()), 4),
            "median": round(float(df[col].median()), 4),
        }

    # Point-biserial correlation between each numeric feature and the target
    feature_target_corr = {}
    for col in numeric_cols:
        if col == target_column:
            continue
        corr = float(df[col].corr(df[target_column]))
        if not np.isnan(corr):
            feature_target_corr[col] = round(corr, 4)
    # Sort by absolute correlation descending
    feature_target_corr = dict(
        sorted(feature_target_corr.items(), key=lambda x: abs(x[1]), reverse=True)
    )

    return {
        "shape": list(df.shape),
        "columns": list(df.columns),
        "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
        "numeric_columns": [c for c in numeric_cols if c != target_column],
        "categorical_columns": [c for c in categorical_cols if c != target_column],
        "target_distribution": {str(k): int(v) for k, v in target_counts.items()},
        "class_balance": {str(k): float(v) for k, v in class_balance.items()},
        "missing_pct": {col: round(float(pct * 100), 2) for col, pct in df.isnull().mean().items()},
        "numeric_stats": numeric_stats,
        "feature_target_corr": feature_target_corr,
        "n_classes": int(df[target_column].nunique()),
        "is_imbalanced": minority_pct < 20.0,
        "sample": df.head(3).fillna("").to_dict(orient="records"),
    }
# {{/docs-fragment profile_dataset}}

@tool_env.task
async def split_dataset(
    data: File,
    target_column: str,
    test_size: float,
    time_column: str,
    split_type: str,
) -> File:
    """Split a dataset and return either the train or test half.

    Call this twice — once with split_type="train" and once with split_type="test" —
    to get both halves. Always split before feature engineering to prevent data leakage.

    Args:
        data: CSV file to split.
        target_column: Name of the column to predict.
        test_size: Fraction of data to use for testing (e.g. 0.2 for 20%).
        time_column: If non-empty, sort by this column and take the last
                     `test_size` fraction as test (time-based split, no shuffling).
                     If empty string "", use stratified random split.
        split_type: Which half to return — "train" or "test".

    Returns:
        File — CSV file containing the requested split (train or test rows).
    """
    import tempfile

    import pandas as pd
    from flyte.io import File as FlyteFile
    from sklearn.model_selection import train_test_split

    path = await data.download()
    df = pd.read_csv(path)

    if time_column:
        df = df.sort_values(time_column).reset_index(drop=True)
        split_idx = int(len(df) * (1 - test_size))
        train_df = df.iloc[:split_idx]
        test_df = df.iloc[split_idx:]
    else:
        train_df, test_df = train_test_split(
            df,
            test_size=test_size,
            stratify=df[target_column],
            random_state=42,
        )

    selected_df = train_df if split_type == "train" else test_df

    out = tempfile.NamedTemporaryFile(suffix=".csv", delete=False)
    selected_df.to_csv(out.name, index=False)
    out.close()
    return await FlyteFile.from_local(out.name)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/tools/data.py*

The full tool inventory includes ten functions: `profile_dataset`, `split_dataset`, `explore_dataset`, `engineer_features`, `select_features`, `resample_dataset`, `train_model`, `get_predictions`, `evaluate_model`, and `rank_experiments`. Each one does exactly one thing. The LLM composes them into pipelines, but each tool enforces its own correctness guarantees internally. For example, `resample_dataset` only applies resampling to training data, never test data, regardless of what the LLM asks for.

## Guiding the LLM with domain knowledge

The quality of the agent's experiments depends heavily on what you tell it. The MLE Bot bakes ML best practices directly into its system prompts, so the LLM starts from a solid foundation rather than relying on whatever it picked up during pretraining.

The orchestration prompt, for example, includes guidance on feature engineering strategies, class imbalance handling, and algorithm selection. It's dynamically built from the dataset profile, so the LLM sees concrete context alongside the general advice:

CODE2

This means the LLM doesn't just get a blank canvas. It gets a structured briefing that combines the actual dataset characteristics with best practices for handling them. When the profile shows class imbalance, the prompt tells it which hyperparameters to adjust and which resampling strategies to consider. When there's a timestamp column, the prompt suggests rolling window features with guidance on window sizing.

The user's problem description also has a significant impact on the agent's behavior. A query like "Predict pump failures 24 hours before they happen based on sensor readings" tells the LLM that this is a time-series classification problem with a specific prediction horizon. That shapes everything: the LLM will favor temporal feature engineering (rolling windows sized relative to that 24-hour horizon), pick algorithms that handle imbalanced binary classification well, and focus on recall as a key metric because missing a failure is worse than a false alarm. Change the query to something like "Classify machine health status from the latest sensor snapshot" and the same dataset would produce a completely different set of experiments, with less emphasis on temporal features and more on cross-sectional patterns.

## The agent loop: profile, design, execute, iterate

The agent's main function orchestrates five phases. Let's walk through each one.

**Phase 1: Profile.** The agent calls `profile_dataset` directly as a trusted tool. This isn't sandboxed because there's nothing to protect against here: the function is your code, running on your compute. The `flyte.group` call organizes this step in the Flyte UI so you can inspect it later.

CODE3

**Phase 2: Design.** The profile dict goes to the LLM along with the problem description. The LLM returns a structured response matching the `InitialDesign` schema:

```
"""Pydantic schemas for tool inputs and agent data structures.

These models define the expected shape of configs and results throughout the agent.

Important: Tool functions that are called from the Monty sandbox must accept plain
`dict` at the boundary (Monty can't import or instantiate classes). Each tool parses
its incoming dict into the appropriate model internally for validation. In agent.py,
use `.model_dump()` to convert models back to dicts before passing to the sandbox.
"""

from typing import Literal

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Feature engineering
# ---------------------------------------------------------------------------

class FeatureConfig(BaseModel):
    """Configuration for the engineer_features tool."""

    group_column: str = Field(
        default="",
        description="Column to group by for rolling/lag features (e.g. 'machine_id'). "
                    "Required when rolling_columns or lag_columns is specified.",
    )
    time_column: str = Field(
        default="",
        description="Timestamp column to sort by before computing rolling/lag features.",
    )
    rolling_columns: list[str] = Field(
        default_factory=list,
        description="Numeric columns to compute rolling statistics for (mean, std, min, max).",
    )
    windows: list[int] = Field(
        default_factory=list,
        description="Rolling window sizes in rows (e.g. [6, 12, 24]).",
    )
    lag_columns: list[str] = Field(
        default_factory=list,
        description="Numeric columns to create lag features for.",
    )
    lags: list[int] = Field(
        default_factory=list,
        description="Lag steps in rows (e.g. [1, 3, 6]).",
    )
    normalize: bool = Field(
        default=False,
        description="If true, z-score normalize all numeric columns except target_column.",
    )
    target_column: str = Field(
        default="",
        description="Column to exclude from normalization. Required when normalize=True.",
    )
    drop_columns: list[str] = Field(
        default_factory=list,
        description="Columns to remove from output (e.g. raw timestamp after rolling).",
    )
    fillna_method: Literal["forward", "zero", "drop"] = Field(
        default="forward",
        description="How to fill NaN values introduced by rolling/lag. "
                    "'forward' forward-fills then fills remaining with 0. "
                    "'zero' fills all NaN with 0. 'drop' drops rows with NaN.",
    )

# ---------------------------------------------------------------------------
# Training hyperparameters (per algorithm)
# ---------------------------------------------------------------------------

class XGBoostParams(BaseModel):
    n_estimators: int = Field(default=100, ge=1)
    max_depth: int = Field(default=6, ge=1, le=20)
    learning_rate: float = Field(default=0.1, gt=0, le=1)
    scale_pos_weight: float = Field(
        default=1.0, ge=0,
        description="Set to n_negative/n_positive for imbalanced datasets.",
    )
    subsample: float = Field(default=1.0, gt=0, le=1)
    colsample_bytree: float = Field(default=1.0, gt=0, le=1)

class RandomForestParams(BaseModel):
    n_estimators: int = Field(default=100, ge=1)
    max_depth: int | None = Field(
        default=None,
        description="Maximum tree depth. None means unlimited.",
    )
    min_samples_leaf: int = Field(default=1, ge=1)
    class_weight: Literal["balanced", "balanced_subsample"] | None = Field(default="balanced")

class GradientBoostingParams(BaseModel):
    n_estimators: int = Field(default=100, ge=1)
    max_depth: int = Field(default=3, ge=1, le=10)
    learning_rate: float = Field(default=0.1, gt=0, le=1)
    subsample: float = Field(default=1.0, gt=0, le=1)

class LogisticRegressionParams(BaseModel):
    C: float = Field(default=1.0, gt=0, description="Inverse regularization strength.")
    max_iter: int = Field(default=1000, ge=100)
    class_weight: Literal["balanced"] | None = Field(default="balanced")

# ---------------------------------------------------------------------------
# Experiment design (used by agent.py, validated when parsing LLM JSON)
# ---------------------------------------------------------------------------

Algorithm = Literal["xgboost", "random_forest", "gradient_boosting", "logistic_regression"]

# {{docs-fragment schemas}}
class ExperimentConfig(BaseModel):
    """One experiment to run — produced by the LLM and executed by the agent."""

    name: str = Field(description="Short descriptive name for this experiment.")
    algorithm: Algorithm
    hyperparams: dict = Field(
        default_factory=dict,
        description="Algorithm-specific hyperparameters. Will be validated inside train_model.",
    )
    feature_config: FeatureConfig = Field(default_factory=FeatureConfig)
    rationale: str = Field(default="", description="Why this experiment is worth running.")

class InitialDesign(BaseModel):
    """LLM response for initial experiment design."""

    problem_type: str = Field(default="binary_classification")
    primary_metric: Literal["roc_auc", "f1", "recall"] = Field(default="roc_auc")
    reasoning: str
    experiments: list[ExperimentConfig]

class IterationDecision(BaseModel):
    """LLM response after analyzing experiment results."""

    should_continue: bool
    reasoning: str
    exploration_requests: list[dict] = Field(
        default_factory=list,
        description="Optional list of explore_dataset config dicts to run before designing "
                    "the next batch. Each dict is passed directly to explore_dataset.",
    )
    next_experiments: list[ExperimentConfig] = Field(default_factory=list)
# {{/docs-fragment schemas}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/schemas.py*

The LLM typically proposes 2 to 3 experiments: a baseline with minimal feature engineering, an experiment with rolling window features for temporal data, and perhaps one testing a different algorithm or resampling strategy.

**Phase 3: Execute in parallel.** All experiments in a batch run simultaneously using `asyncio.gather()`. Each experiment dispatches its own set of durable Flyte tasks:

CODE4python
[your orchestration code]
CODE5python" in text:
        start = text.index("CODE6python")
        end = text.index("CODE7" in text:
        start = text.index("CODE8", start)
        return text[start:end].strip()
    return text.strip()

def _extract_reasoning(text: str) -> str:
    """Extract the ## Reasoning section from LLM response."""
    if "## Reasoning" in text:
        start = text.index("## Reasoning") + len("## Reasoning")
        if "## Code" in text:
            end = text.index("## Code")
            return text[start:end].strip()
        return text[start:].strip()
    return ""

def _parse_json(text: str) -> dict:
    """Extract and parse JSON from LLM response."""
    text = text.strip()
    if "CODE9json") + 7
        end = text.index("CODE10" in text:
        start = text.index("CODE11", start)
        text = text[start:end].strip()
    return json.loads(text)

# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------

def _recommend_threshold(threshold_analysis: list, min_precision: float = 0.70) -> dict | None:
    """Find the threshold that maximises recall subject to precision >= min_precision."""
    candidates = [t for t in threshold_analysis if t["precision"] >= min_precision]
    if not candidates:
        return None
    return max(candidates, key=lambda t: t["recall"])

def _print_experiment_table(results: list["ExperimentResult"], best_name: str) -> None:
    """Print a ranked comparison table of all experiments."""
    sorted_results = sorted(results, key=lambda r: r.metrics.get("roc_auc", 0), reverse=True)
    print("\n" + "─" * 78)
    print(f"  {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}")
    print("─" * 78)
    for rank, r in enumerate(sorted_results, 1):
        note = "◀ winner" if r.name == best_name else ""
        roc = r.metrics.get("roc_auc", 0)
        f1 = r.metrics.get("f1", 0)
        recall = r.metrics.get("recall", 0)
        print(f"  {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}")
    print("─" * 78)

def _print_threshold_recommendation(threshold_analysis: list, default_metrics: dict) -> None:
    """Print the operational threshold recommendation."""
    rec = _recommend_threshold(threshold_analysis)
    if not rec:
        return
    default_recall = default_metrics.get("recall", 0)
    default_precision = default_metrics.get("precision", 0)
    missed_pct = round((1 - rec["recall"]) * 100, 1)
    false_alarm_pct = round((1 - rec["precision"]) * 100, 1)

    print(f"\n  Recommended decision threshold: {rec['threshold']}")
    print(f"  ├─ Precision : {rec['precision']:.0%}   ({false_alarm_pct}% of alerts are false alarms)")
    print(f"  ├─ Recall    : {rec['recall']:.0%}   (catches {rec['recall']*100:.0f}% of actual failures)")
    print(f"  └─ F1        : {rec['f1']:.4f}")
    print(f"  Default threshold (0.5): Precision={default_precision:.0%}, Recall={default_recall:.0%}")
    if rec["recall"] > default_recall:
        extra = round((rec["recall"] - default_recall) * 100, 1)
        print(f"  → Lowering threshold catches {extra}% more failures at cost of more alerts")

# ---------------------------------------------------------------------------
# Orchestration code generation (durable Flyte task with Flyte report)
# ---------------------------------------------------------------------------

@agent_env.task
async def plan_experiment(
    experiment_json: str,
    profile_json: str,
    target_column: str,
    time_column: str,
    previous_error: str = "",
    previous_code: str = "",
    llm_model: str = "gpt-4o",
) -> str:
    """LLM plans a single experiment: reasons about the pipeline and generates Monty code.

    Runs as a durable Flyte task so each experiment's planning step is traceable.
    Returns a JSON string: {"code": "...", "reasoning": "..."}.

    Args:
        experiment_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...).
        profile_json: JSON string of the dataset profile from profile_dataset.
        target_column: Name of the target column.
        time_column: Time column for temporal splitting, or empty string.
        previous_error: Error message from the previous attempt (empty on first try).
        previous_code: Code that failed on the previous attempt (empty on first try).
        llm_model: OpenAI model identifier.

    Returns:
        str — JSON string with keys "code" and "reasoning".
    """
    experiment = json.loads(experiment_json)
    profile = json.loads(profile_json)
    exp_name = experiment.get("name", "experiment")

    # Strip rationale — it was written by the design LLM to explain *why* this
    # experiment was chosen. Passing it here causes plan_experiment to parrot it
    # back as "reasoning" instead of independently thinking about *how* to build
    # the best pipeline. Keep only the technical spec.
    pipeline_spec = {
        k: v for k, v in experiment.items()
        if k not in ("rationale",)
    }

    system = _build_orchestration_system_prompt(profile)

    user_content = textwrap.dedent(f"""
        Design and implement the best pipeline for this experiment:

        Name: {exp_name}
        Algorithm: {pipeline_spec.get("algorithm")}
        Hyperparams: {json.dumps(pipeline_spec.get("hyperparams", {}), indent=2)}
        Feature config hint: {json.dumps(pipeline_spec.get("feature_config", {}), indent=2)}

        Available sandbox inputs:
        - data: File  — the full dataset CSV
        - target_column: str = "{target_column}"
        - time_column: str = "{time_column}"  (empty string means no time ordering)
        - experiment_name: str = "{exp_name}"

        The feature config hint is a suggestion from the experiment designer — you can
        follow it, improve on it, or override it if the dataset context and your ML
        judgment suggest a better approach. In your ## Reasoning, explain your actual
        pipeline decisions: what you chose to do (or not do) and why, based on the
        dataset profile above. Do not restate the experiment name or why it was chosen.
    """).strip()

    messages = [{"role": "user", "content": user_content}]
    if previous_code and previous_error:
        messages = [
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": f"CODE12"},
            {"role": "user", "content": f"That code failed with this error:\n\n{previous_error}\n\nPlease fix it."},
        ]

    response = await _call_llm(system, messages, llm_model)
    reasoning = _extract_reasoning(response)
    code = _extract_code(response)
    return json.dumps({"code": code, "reasoning": reasoning})

@flyte.trace
async def design_experiments(
    problem_description: str,
    profile_json: str,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs the initial batch of experiments given problem + dataset profile.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching InitialDesign schema).
    """
    design_prompt = textwrap.dedent(f"""
        Problem description: {problem_description}

        Dataset profile:
        {profile_json}

        Design the first batch of experiments.
    """).strip()
    return await _call_llm(
        _build_initial_design_system_prompt(),
        [{"role": "user", "content": design_prompt}],
        llm_model,
    )

@flyte.trace
async def analyze_iteration(
    analysis_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM analyzes experiment results and decides whether/how to continue.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching IterationDecision schema).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [{"role": "user", "content": analysis_prompt}],
        llm_model,
    )

@flyte.trace
async def plan_followup(
    analysis_prompt: str,
    analysis_response: str,
    followup_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs next experiments after targeted data explorations.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string with {"next_experiments": [...]}).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [
            {"role": "user", "content": analysis_prompt},
            {"role": "assistant", "content": analysis_response},
            {"role": "user", "content": followup_prompt},
        ],
        llm_model,
    )

def _corrupt_experiment_for_demo(exp_dict: dict) -> dict:
    """Introduce a deliberate error into the first experiment for demo purposes.

    Corrupts the algorithm name so the LLM must recover from a known-bad value.
    The retry loop will catch this, regenerate with the error message, and fix it.
    """
    corrupted = dict(exp_dict)
    corrupted["algorithm"] = corrupted["algorithm"] + "_INVALID"
    return corrupted

# ---------------------------------------------------------------------------
# Main agent loop
# ---------------------------------------------------------------------------

@dataclass
class ExperimentResult:
    name: str
    algorithm: str
    metrics: dict
    confusion_matrix: dict
    threshold_analysis: list
    n_samples: int
    code: str
    attempts: int
    reasoning: str = ""
    error: str = ""

@dataclass
class AgentResult:
    model_card: str
    best_experiment: str
    best_metrics: dict
    all_results: list[ExperimentResult]
    iterations: int
    total_experiments: int

async def _run_experiment(
    exp: "ExperimentConfig",
    exp_dict: dict,
    inject_failure: bool,
    data: File,
    target_column: str,
    time_column: str,
    profile: dict,
    llm_model: str,
    max_retries: int,
) -> "ExperimentResult | None":
    """Run a single experiment with retries. Returns None on total failure."""
    exp_name = exp.name
    profile_json = json.dumps(profile)
    print(f"\n   ┌─ {exp_name}  [{exp.algorithm}]")
    if exp.rationale:
        for line in textwrap.wrap(exp.rationale, width=58):
            print(f"   │  {line}")
    if inject_failure:
        print(f"   │  [injecting failure for demo: algorithm='{exp_dict['algorithm']}']")

    code = ""
    error = ""
    result = None
    attempt = 0

    reasoning = ""
    # {{docs-fragment retry_loop}}
    for attempt in range(max_retries):
        try:
            with flyte.group(exp_name):
                plan_json = await plan_experiment.aio(
                    experiment_json=json.dumps(exp_dict),
                    profile_json=profile_json,
                    target_column=target_column,
                    time_column=time_column,
                    previous_error=error,
                    previous_code=code,
                    llm_model=llm_model,
                )
                plan = json.loads(plan_json)
                code = plan["code"]
                reasoning = plan.get("reasoning", "")
                result = await flyte.sandbox.orchestrate_local(
                    code,
                    inputs={"data": data, "target_column": target_column,
                            "time_column": time_column, "experiment_name": exp_name},
                    tasks=TOOLS,
                )
            error = ""
            break
        except Exception as exc:
            error = str(exc)
            short_error = error[:100] + "..." if len(error) > 100 else error
            print(f"   │  attempt {attempt + 1} failed: {short_error}")
            print(f"   │  → asking LLM to fix and retry...")
            if inject_failure and attempt == 0:
                exp_dict = exp.model_dump()
    # {{/docs-fragment retry_loop}}

    if result and not error:
        exp_result = ExperimentResult(
            name=exp_name,
            algorithm=exp.algorithm,
            metrics=result.get("metrics", {}),
            confusion_matrix=result.get("confusion_matrix", {}),
            threshold_analysis=result.get("threshold_analysis", []),
            n_samples=result.get("n_samples", 0),
            code=code,
            reasoning=reasoning,
            attempts=attempt + 1,
        )
        m = exp_result.metrics
        attempts_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else ""
        print(f"   └─ ROC-AUC={m.get('roc_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts_note}")
        return exp_result

    print(f"   └─ FAILED after {max_retries} attempts — skipping.")
    return None

async def run_agent(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
    max_retries_per_experiment: int = 3,
    llm_model: str = "gpt-4o",
    inject_failure: bool = False,
) -> AgentResult:
    """Run the MLE agent end-to-end.

    Args:
        data: CSV file containing the dataset.
        problem_description: Natural language description of the ML problem.
        target_column: Name of the target column to predict.
        time_column: Optional column to use for time-based train/test split.
        max_iterations: Maximum number of experiment iterations to run.
        max_retries_per_experiment: Max times to retry a failed sandbox execution.
        llm_model: OpenAI model to use (default: gpt-4o).
        inject_failure: If True, corrupts the first experiment to demonstrate self-healing.
    """
    print(f"\n{'='*60}")
    print(f"MLE Agent starting")
    print(f"Problem: {problem_description}")
    print(f"Target: {target_column}")
    if inject_failure:
        print(f"[demo mode: failure injection enabled]")
    print(f"{'='*60}\n")

    # {{docs-fragment phase1_profile}}
    # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) ---
    print(">> Phase 1: Profiling dataset...")
    with flyte.group("profile"):
        profile = await profile_dataset(data, target_column)
    # {{/docs-fragment phase1_profile}}
    print(f"   Shape: {profile['shape']}, Classes: {profile['target_distribution']}")
    print(f"   Imbalanced: {profile['is_imbalanced']}, Columns: {len(profile['columns'])}")
    corr = profile.get("feature_target_corr", {})
    top_corr = list(corr.items())[:5]
    print(f"   Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top_corr)}")

    # Stream report: dataset summary
    await flyte.report.log.aio(
        f"<h1>MLE Agent Run</h1>"
        f"<p><b>Problem:</b> {problem_description}</p>"
        f"<p><b>Dataset:</b> {profile['shape'][0]:,} rows × {profile['shape'][1]} cols &nbsp;|&nbsp; "
        f"Class balance: {profile['class_balance']} &nbsp;|&nbsp; Imbalanced: {profile['is_imbalanced']}</p>"
        f"<p><b>Top feature-target correlations (raw):</b> "
        + ", ".join(f"{k}: {v:+.3f}" for k, v in top_corr) +
        f"</p><hr>",
        do_flush=True,
    )

    # --- Phase 2: LLM designs initial experiments ---
    print("\n>> Phase 2: Designing initial experiments...")
    design_response = await design_experiments(
        problem_description=problem_description,
        profile_json=json.dumps(profile),
        llm_model=llm_model,
    )
    design = InitialDesign.model_validate(_parse_json(design_response))
    print(f"   Primary metric: {design.primary_metric}")
    print(f"   Strategy: {design.reasoning}")
    print(f"   Experiments planned: {len(design.experiments)}")

    all_results: list[ExperimentResult] = []
    iteration_log: list[dict] = []  # tracks per-iteration decisions + explorations for summary
    current_experiments: list[ExperimentConfig] = design.experiments
    first_experiment = True

    # --- Phase 3: Iterative experiment loop ---
    for iteration in range(max_iterations):
        experiments = current_experiments

        if not experiments:
            print(f"\n>> No experiments to run in iteration {iteration + 1}. Stopping.")
            break

        print(f"\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...")

        # Assign names and prepare dicts before launching in parallel
        exp_batch = []
        for i, exp in enumerate(experiments):
            if not exp.name:
                exp.name = f"experiment_{len(all_results) + i + 1}"
            exp_dict = exp.model_dump()
            inject_this = inject_failure and first_experiment and i == 0
            if inject_this:
                exp_dict = _corrupt_experiment_for_demo(exp_dict)
            first_experiment = False
            exp_batch.append((exp, exp_dict, inject_this))

        # {{docs-fragment parallel_execute}}
        batch_results = await asyncio.gather(*[
            _run_experiment(
                exp=exp,
                exp_dict=exp_dict,
                inject_failure=inject_this,
                data=data,
                target_column=target_column,
                time_column=time_column,
                profile=profile,
                llm_model=llm_model,
                max_retries=max_retries_per_experiment,
            )
            for exp, exp_dict, inject_this in exp_batch
        ])
        # {{/docs-fragment parallel_execute}}

        for exp_result in batch_results:
            if exp_result is not None:
                all_results.append(exp_result)
                # Stream report: each experiment as it completes
                m = exp_result.metrics
                html = (
                    f"<h3>Iteration {iteration + 1} — {exp_result.name}</h3>"
                    f"<p><b>Algorithm:</b> {exp_result.algorithm} &nbsp;|&nbsp; "
                    f"<b>ROC-AUC:</b> {m.get('roc_auc')} &nbsp;|&nbsp; "
                    f"<b>F1:</b> {m.get('f1')} &nbsp;|&nbsp; "
                    f"<b>Recall:</b> {m.get('recall')} &nbsp;|&nbsp; "
                    f"<b>Attempts:</b> {exp_result.attempts}</p>"
                )
                if exp_result.reasoning:
                    html += f"<details><summary>Reasoning</summary><pre>{exp_result.reasoning}</pre></details>"
                html += f"<details><summary>Generated Code</summary><pre>{exp_result.code}</pre></details>"
                await flyte.report.log.aio(html, do_flush=True)

        # --- Phase 4: Analyze results, decide whether to iterate ---
        if all_results and iteration < max_iterations - 1:
            print(f"\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...")
            results_summary = [
                {
                    "experiment_name": r.name,
                    "algorithm": r.algorithm,
                    "metrics": r.metrics,
                    "confusion_matrix": r.confusion_matrix,
                    "used_feature_engineering": "engineer_features" in r.code,
                    "used_rolling_features": "rolling_columns" in r.code,
                    "used_lag_features": "lag_columns" in r.code,
                }
                for r in all_results
            ]
            analysis_prompt = textwrap.dedent(f"""
                Problem: {problem_description}
                Dataset profile: shape={profile['shape']}, imbalanced={profile['is_imbalanced']}
                Feature-target correlations (raw): {json.dumps(profile.get('feature_target_corr', {}), indent=2)}

                Experiment results so far (iteration {iteration + 1}):
                {json.dumps(results_summary, indent=2)}

                Should we run more experiments? If yes, request any data explorations
                you need, then specify what experiments to run next.
            """).strip()

            analysis_response = await analyze_iteration(
                analysis_prompt=analysis_prompt,
                max_iterations=max_iterations,
                current_iteration=iteration,
                llm_model=llm_model,
            )
            decision = IterationDecision.model_validate(_parse_json(analysis_response))
            verdict = "continuing" if decision.should_continue else "stopping"
            print(f"   Decision: {verdict}")
            print(f"   Reasoning: {decision.reasoning}")

            # Stream report: analysis decision
            await flyte.report.log.aio(
                f"<h3>Analysis — Iteration {iteration + 1}</h3>"
                f"<p><b>Decision:</b> {verdict}</p>"
                f"<p><b>Reasoning:</b> {decision.reasoning}</p>",
                do_flush=True,
            )

            # Track this iteration for the experiment journey summary
            iter_entry = {
                "iteration": iteration + 1,
                "experiments": [r.name for r in batch_results if r is not None],
                "best_roc_auc": max(
                    (r.metrics.get("roc_auc", 0) for r in all_results), default=0
                ),
                "reasoning": decision.reasoning,
                "explorations": [],
            }

            # --- Targeted exploration before next iteration ---
            if decision.should_continue and decision.exploration_requests:
                print(f"   Running {len(decision.exploration_requests)} exploration request(s)...")
                exploration_questions = []
                exploration_results = []

                for i, req in enumerate(decision.exploration_requests):
                    question = req.get("question", f"Exploration {i + 1}")
                    # Strip agent-level metadata — tool only needs the analysis config
                    tool_config = {k: v for k, v in req.items() if k not in ("question", "analysis_type")}

                    print(f"   Q: {question}")
                    with flyte.group(f"explore_{iteration + 1}_{i + 1}"):
                        result = await explore_dataset(data, tool_config)
                    exploration_questions.append(question)
                    exploration_results.append(result)
                    iter_entry["explorations"].append({"question": question})

                    await flyte.report.log.aio(
                        f"<h4>Exploration {i + 1}</h4>"
                        f"<p><b>Question:</b> {question}</p>"
                        f"<details><summary>Results</summary><pre>{json.dumps(result, indent=2)}</pre></details>",
                        do_flush=True,
                    )

                # Build follow-up that explicitly connects each question to its answer
                qa_pairs = "\n\n".join(
                    f'Question {i + 1}: "{q}"\nResult:\n{json.dumps(r, indent=2)}'
                    for i, (q, r) in enumerate(zip(exploration_questions, exploration_results))
                )
                followup_prompt = textwrap.dedent(f"""
                    You requested {len(exploration_results)} targeted exploration(s).
                    Here is what you asked and what you learned:

                    {qa_pairs}

                    Given what you learned and your earlier reasoning:
                    "{decision.reasoning}"

                    Now specify the next experiments. For each experiment, briefly state
                    which exploration insight informed your choice.
                    Respond with valid JSON: {{"next_experiments": [...same schema as before...]}}
                """).strip()
                followup_response = await plan_followup(
                    analysis_prompt=analysis_prompt,
                    analysis_response=analysis_response,
                    followup_prompt=followup_prompt,
                    max_iterations=max_iterations,
                    current_iteration=iteration,
                    llm_model=llm_model,
                )
                followup = _parse_json(followup_response)
                current_experiments = IterationDecision.model_validate({
                    "should_continue": True,
                    "reasoning": decision.reasoning,
                    "next_experiments": followup.get("next_experiments", []),
                }).next_experiments
                print(f"   Post-exploration: {len(current_experiments)} experiment(s) planned")
            else:
                current_experiments = decision.next_experiments

            iteration_log.append(iter_entry)

            if not decision.should_continue:
                break

    # --- Phase 5: Rank all results and generate model card ---
    print(f"\n>> Phase 5: Ranking {len(all_results)} experiment(s) and generating model card...")

    if not all_results:
        return AgentResult(
            model_card="No experiments completed successfully.",
            best_experiment="",
            best_metrics={},
            all_results=[],
            iterations=iteration + 1,
            total_experiments=0,
        )

    ranking_input = [
        {
            "experiment_name": r.name,
            "metrics": r.metrics,
            "confusion_matrix": r.confusion_matrix,
        }
        for r in all_results
    ]
    with flyte.group("rank"):
        ranking = await rank_experiments(json.dumps(ranking_input))
    best_name = ranking["best_experiment"]
    best_result = next(r for r in all_results if r.name == best_name)

    _print_experiment_table(all_results, best_name)
    _print_threshold_recommendation(best_result.threshold_analysis, best_result.metrics)

    # Stream report: final rankings table
    rows = "".join(
        f"<tr><td>{row['rank']}</td>"
        f"<td>{'<b>' if row['experiment_name'] == best_name else ''}"
        f"{row['experiment_name']}"
        f"{'</b>' if row['experiment_name'] == best_name else ''}</td>"
        f"<td>{row['roc_auc']}</td><td>{row['f1']}</td>"
        f"<td>{row['recall']}</td><td>{row['precision']}</td></tr>"
        for row in ranking.get("ranking", [])
    )
    await flyte.report.log.aio(
        f"<hr><h2>Final Rankings</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0'>"
        f"<tr><th>Rank</th><th>Experiment</th><th>ROC-AUC</th><th>F1</th><th>Recall</th><th>Precision</th></tr>"
        f"{rows}</table>"
        f"<p>{ranking.get('summary', '')}</p>",
        do_flush=True,
    )

    # Stream report: experiment journey summary
    journey_rows = ""
    for entry in iteration_log:
        exps = ", ".join(entry["experiments"]) if entry["experiments"] else "—"
        explorations = "; ".join(e["question"] for e in entry["explorations"]) if entry["explorations"] else "—"
        short_reasoning = (entry["reasoning"][:120] + "…") if len(entry["reasoning"]) > 120 else entry["reasoning"]
        journey_rows += (
            f"<tr>"
            f"<td style='text-align:center'>{entry['iteration']}</td>"
            f"<td>{exps}</td>"
            f"<td style='text-align:center'>{entry['best_roc_auc']:.4f}</td>"
            f"<td>{short_reasoning}</td>"
            f"<td>{explorations}</td>"
            f"</tr>"
        )
    await flyte.report.log.aio(
        f"<hr><h2>Experiment Journey</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0' style='width:100%;border-collapse:collapse'>"
        f"<tr><th>Iter</th><th>Experiments</th><th>Best ROC-AUC</th><th>Key insight</th><th>Explorations</th></tr>"
        f"{journey_rows}"
        f"</table>",
        do_flush=True,
    )

    model_card = await _generate_model_card(
        problem_description=problem_description,
        profile=profile,
        all_results=all_results,
        best_result=best_result,
        ranking=ranking,
        iteration_log=iteration_log,
        llm_model=llm_model,
    )

    print(f"\n{'='*60}")
    print(f"DONE — Best model: {best_name}")
    print(f"       ROC-AUC={best_result.metrics.get('roc_auc')}, F1={best_result.metrics.get('f1')}")
    print(f"{'='*60}\n")

    return AgentResult(
        model_card=model_card,
        best_experiment=best_name,
        best_metrics=best_result.metrics,
        all_results=all_results,
        iterations=iteration + 1,
        total_experiments=len(all_results),
    )

async def _generate_model_card(
    problem_description: str,
    profile: dict,
    all_results: list[ExperimentResult],
    best_result: ExperimentResult,
    ranking: dict,
    iteration_log: list[dict],
    llm_model: str,
) -> str:
    """Generate a markdown model card summarizing the winning model."""
    system = textwrap.dedent("""
        You are an ML engineer writing a model card for a trained model.
        Write in markdown. Be concise but informative. Include:
        - Problem statement
        - Dataset summary
        - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed)
        - Experiment summary (table of all experiments with metrics)
        - Winning model details (algorithm, key hyperparams, metrics, threshold analysis)
        - Recommendations for deployment (decision threshold, monitoring)
    """).strip()

    results_text = "\n".join(
        f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc_auc')}, "
        f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}"
        for r in all_results
    )

    journey_text = ""
    if iteration_log:
        journey_text = "\n\nIteration log:\n" + "\n".join(
            f"  Iteration {e['iteration']}: ran [{', '.join(e['experiments'])}], "
            f"best ROC-AUC so far={e['best_roc_auc']:.4f}. "
            f"Key insight: {e['reasoning'][:200]}. "
            + (f"Explorations: {'; '.join(x['question'] for x in e['explorations'])}" if e['explorations'] else "")
            for e in iteration_log
        )

    user_content = textwrap.dedent(f"""
        Problem: {problem_description}

        Dataset: {profile['shape'][0]} rows × {profile['shape'][1]} cols.
        Class balance: {profile['class_balance']}
        Imbalanced: {profile['is_imbalanced']}
        {journey_text}

        All experiments:
        {results_text}

        Best model: {best_result.name} ({best_result.algorithm})
        Metrics: {json.dumps(best_result.metrics, indent=2)}
        Confusion matrix: {json.dumps(best_result.confusion_matrix, indent=2)}
        Threshold analysis: {json.dumps(best_result.threshold_analysis, indent=2)}

        Ranking summary: {ranking['summary']}
    """).strip()

    response = await _call_llm(system, [{"role": "user", "content": user_content}], llm_model)
    return response

# ---------------------------------------------------------------------------
# Durable entrypoint (runs the agent as a Flyte task in the cloud)
# ---------------------------------------------------------------------------

# {{docs-fragment entrypoint}}
@agent_env.task(retries=1, report=True)
async def mle_agent_task(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
) -> str:
    """Durable Flyte task entrypoint for the MLE agent."""
    result = await run_agent(
        data=data,
        problem_description=problem_description,
        target_column=target_column,
        time_column=time_column,
        max_iterations=max_iterations,
    )
    return result.model_card
# {{/docs-fragment entrypoint}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/agent.py*

**Phase 4: Analyze and iterate.** After each batch completes, the LLM reviews the results and decides whether to continue. It can optionally request targeted data explorations before designing the next round. If the LLM requests explorations (e.g., "do failure cases show higher vibration readings?"), the agent runs `explore_dataset` with those configurations, feeds the results back to the LLM, and lets it refine the next batch of experiments based on what it learned. The loop continues until the LLM decides to stop, the target metric threshold is reached, or the maximum number of iterations is exhausted.

## Running LLM-generated code in Flyte's sandbox

This is where it gets interesting. The LLM doesn't just pick parameters from a dropdown. For each experiment, it writes actual Python code that decides how to compose the tool functions into a pipeline. Maybe it splits the data, engineers rolling window features, applies SMOTE resampling on the training split, trains an XGBoost model, and evaluates it. Or maybe it skips feature engineering entirely for a baseline. The LLM decides the structure.

That code runs inside Flyte's sandbox, a restricted execution environment that enforces strict constraints:

- No `import` statements. The only callable functions are the ones you explicitly provide.
- No network access and no filesystem access.
- No `try`/`except`, no `class` definitions, no augmented assignment (`+=`).
- No `with` statements, no generators, no `global`/`nonlocal`.

The sandbox sees your pre-approved tool functions as plain function calls. When the code calls `train_model(...)`, the sandbox pauses execution, dispatches the call to Flyte (which runs it as a durable task on cloud compute with the resources declared on `tool_env`), waits for the result, and resumes. The LLM-generated code looks like synchronous Python, but under the hood each tool call is a full Flyte task execution.

Here's how the sandbox is invoked:

CODE13

The `code` parameter is a string of Python generated by the LLM. `inputs` provides the variables that the code can reference. `tasks` is the allowlist: a list of Flyte task functions that the code is permitted to call. Nothing else is available.

Here's an example of what the LLM might generate for a single experiment:

CODE14

Each function call in that snippet dispatches a separate Flyte task. The `split_dataset` calls run on the tool environment's compute (2 CPU, 4Gi memory). The `train_model` call trains an actual XGBoost model. The last expression (a dict literal) is returned as the sandbox result.

Sometimes the LLM generates code with bugs, like a wrong variable name or a missing argument. The agent handles this with a retry loop. If the sandbox raises an exception, the error message and the failing code are fed back to the LLM, which gets a chance to fix the issue:

CODE15python
[your orchestration code]
CODE16python" in text:
        start = text.index("CODE17python")
        end = text.index("CODE18" in text:
        start = text.index("CODE19", start)
        return text[start:end].strip()
    return text.strip()

def _extract_reasoning(text: str) -> str:
    """Extract the ## Reasoning section from LLM response."""
    if "## Reasoning" in text:
        start = text.index("## Reasoning") + len("## Reasoning")
        if "## Code" in text:
            end = text.index("## Code")
            return text[start:end].strip()
        return text[start:].strip()
    return ""

def _parse_json(text: str) -> dict:
    """Extract and parse JSON from LLM response."""
    text = text.strip()
    if "CODE20json") + 7
        end = text.index("CODE21" in text:
        start = text.index("CODE22", start)
        text = text[start:end].strip()
    return json.loads(text)

# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------

def _recommend_threshold(threshold_analysis: list, min_precision: float = 0.70) -> dict | None:
    """Find the threshold that maximises recall subject to precision >= min_precision."""
    candidates = [t for t in threshold_analysis if t["precision"] >= min_precision]
    if not candidates:
        return None
    return max(candidates, key=lambda t: t["recall"])

def _print_experiment_table(results: list["ExperimentResult"], best_name: str) -> None:
    """Print a ranked comparison table of all experiments."""
    sorted_results = sorted(results, key=lambda r: r.metrics.get("roc_auc", 0), reverse=True)
    print("\n" + "─" * 78)
    print(f"  {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}")
    print("─" * 78)
    for rank, r in enumerate(sorted_results, 1):
        note = "◀ winner" if r.name == best_name else ""
        roc = r.metrics.get("roc_auc", 0)
        f1 = r.metrics.get("f1", 0)
        recall = r.metrics.get("recall", 0)
        print(f"  {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}")
    print("─" * 78)

def _print_threshold_recommendation(threshold_analysis: list, default_metrics: dict) -> None:
    """Print the operational threshold recommendation."""
    rec = _recommend_threshold(threshold_analysis)
    if not rec:
        return
    default_recall = default_metrics.get("recall", 0)
    default_precision = default_metrics.get("precision", 0)
    missed_pct = round((1 - rec["recall"]) * 100, 1)
    false_alarm_pct = round((1 - rec["precision"]) * 100, 1)

    print(f"\n  Recommended decision threshold: {rec['threshold']}")
    print(f"  ├─ Precision : {rec['precision']:.0%}   ({false_alarm_pct}% of alerts are false alarms)")
    print(f"  ├─ Recall    : {rec['recall']:.0%}   (catches {rec['recall']*100:.0f}% of actual failures)")
    print(f"  └─ F1        : {rec['f1']:.4f}")
    print(f"  Default threshold (0.5): Precision={default_precision:.0%}, Recall={default_recall:.0%}")
    if rec["recall"] > default_recall:
        extra = round((rec["recall"] - default_recall) * 100, 1)
        print(f"  → Lowering threshold catches {extra}% more failures at cost of more alerts")

# ---------------------------------------------------------------------------
# Orchestration code generation (durable Flyte task with Flyte report)
# ---------------------------------------------------------------------------

@agent_env.task
async def plan_experiment(
    experiment_json: str,
    profile_json: str,
    target_column: str,
    time_column: str,
    previous_error: str = "",
    previous_code: str = "",
    llm_model: str = "gpt-4o",
) -> str:
    """LLM plans a single experiment: reasons about the pipeline and generates Monty code.

    Runs as a durable Flyte task so each experiment's planning step is traceable.
    Returns a JSON string: {"code": "...", "reasoning": "..."}.

    Args:
        experiment_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...).
        profile_json: JSON string of the dataset profile from profile_dataset.
        target_column: Name of the target column.
        time_column: Time column for temporal splitting, or empty string.
        previous_error: Error message from the previous attempt (empty on first try).
        previous_code: Code that failed on the previous attempt (empty on first try).
        llm_model: OpenAI model identifier.

    Returns:
        str — JSON string with keys "code" and "reasoning".
    """
    experiment = json.loads(experiment_json)
    profile = json.loads(profile_json)
    exp_name = experiment.get("name", "experiment")

    # Strip rationale — it was written by the design LLM to explain *why* this
    # experiment was chosen. Passing it here causes plan_experiment to parrot it
    # back as "reasoning" instead of independently thinking about *how* to build
    # the best pipeline. Keep only the technical spec.
    pipeline_spec = {
        k: v for k, v in experiment.items()
        if k not in ("rationale",)
    }

    system = _build_orchestration_system_prompt(profile)

    user_content = textwrap.dedent(f"""
        Design and implement the best pipeline for this experiment:

        Name: {exp_name}
        Algorithm: {pipeline_spec.get("algorithm")}
        Hyperparams: {json.dumps(pipeline_spec.get("hyperparams", {}), indent=2)}
        Feature config hint: {json.dumps(pipeline_spec.get("feature_config", {}), indent=2)}

        Available sandbox inputs:
        - data: File  — the full dataset CSV
        - target_column: str = "{target_column}"
        - time_column: str = "{time_column}"  (empty string means no time ordering)
        - experiment_name: str = "{exp_name}"

        The feature config hint is a suggestion from the experiment designer — you can
        follow it, improve on it, or override it if the dataset context and your ML
        judgment suggest a better approach. In your ## Reasoning, explain your actual
        pipeline decisions: what you chose to do (or not do) and why, based on the
        dataset profile above. Do not restate the experiment name or why it was chosen.
    """).strip()

    messages = [{"role": "user", "content": user_content}]
    if previous_code and previous_error:
        messages = [
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": f"CODE23"},
            {"role": "user", "content": f"That code failed with this error:\n\n{previous_error}\n\nPlease fix it."},
        ]

    response = await _call_llm(system, messages, llm_model)
    reasoning = _extract_reasoning(response)
    code = _extract_code(response)
    return json.dumps({"code": code, "reasoning": reasoning})

@flyte.trace
async def design_experiments(
    problem_description: str,
    profile_json: str,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs the initial batch of experiments given problem + dataset profile.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching InitialDesign schema).
    """
    design_prompt = textwrap.dedent(f"""
        Problem description: {problem_description}

        Dataset profile:
        {profile_json}

        Design the first batch of experiments.
    """).strip()
    return await _call_llm(
        _build_initial_design_system_prompt(),
        [{"role": "user", "content": design_prompt}],
        llm_model,
    )

@flyte.trace
async def analyze_iteration(
    analysis_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM analyzes experiment results and decides whether/how to continue.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching IterationDecision schema).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [{"role": "user", "content": analysis_prompt}],
        llm_model,
    )

@flyte.trace
async def plan_followup(
    analysis_prompt: str,
    analysis_response: str,
    followup_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs next experiments after targeted data explorations.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string with {"next_experiments": [...]}).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [
            {"role": "user", "content": analysis_prompt},
            {"role": "assistant", "content": analysis_response},
            {"role": "user", "content": followup_prompt},
        ],
        llm_model,
    )

def _corrupt_experiment_for_demo(exp_dict: dict) -> dict:
    """Introduce a deliberate error into the first experiment for demo purposes.

    Corrupts the algorithm name so the LLM must recover from a known-bad value.
    The retry loop will catch this, regenerate with the error message, and fix it.
    """
    corrupted = dict(exp_dict)
    corrupted["algorithm"] = corrupted["algorithm"] + "_INVALID"
    return corrupted

# ---------------------------------------------------------------------------
# Main agent loop
# ---------------------------------------------------------------------------

@dataclass
class ExperimentResult:
    name: str
    algorithm: str
    metrics: dict
    confusion_matrix: dict
    threshold_analysis: list
    n_samples: int
    code: str
    attempts: int
    reasoning: str = ""
    error: str = ""

@dataclass
class AgentResult:
    model_card: str
    best_experiment: str
    best_metrics: dict
    all_results: list[ExperimentResult]
    iterations: int
    total_experiments: int

async def _run_experiment(
    exp: "ExperimentConfig",
    exp_dict: dict,
    inject_failure: bool,
    data: File,
    target_column: str,
    time_column: str,
    profile: dict,
    llm_model: str,
    max_retries: int,
) -> "ExperimentResult | None":
    """Run a single experiment with retries. Returns None on total failure."""
    exp_name = exp.name
    profile_json = json.dumps(profile)
    print(f"\n   ┌─ {exp_name}  [{exp.algorithm}]")
    if exp.rationale:
        for line in textwrap.wrap(exp.rationale, width=58):
            print(f"   │  {line}")
    if inject_failure:
        print(f"   │  [injecting failure for demo: algorithm='{exp_dict['algorithm']}']")

    code = ""
    error = ""
    result = None
    attempt = 0

    reasoning = ""
    # {{docs-fragment retry_loop}}
    for attempt in range(max_retries):
        try:
            with flyte.group(exp_name):
                plan_json = await plan_experiment.aio(
                    experiment_json=json.dumps(exp_dict),
                    profile_json=profile_json,
                    target_column=target_column,
                    time_column=time_column,
                    previous_error=error,
                    previous_code=code,
                    llm_model=llm_model,
                )
                plan = json.loads(plan_json)
                code = plan["code"]
                reasoning = plan.get("reasoning", "")
                result = await flyte.sandbox.orchestrate_local(
                    code,
                    inputs={"data": data, "target_column": target_column,
                            "time_column": time_column, "experiment_name": exp_name},
                    tasks=TOOLS,
                )
            error = ""
            break
        except Exception as exc:
            error = str(exc)
            short_error = error[:100] + "..." if len(error) > 100 else error
            print(f"   │  attempt {attempt + 1} failed: {short_error}")
            print(f"   │  → asking LLM to fix and retry...")
            if inject_failure and attempt == 0:
                exp_dict = exp.model_dump()
    # {{/docs-fragment retry_loop}}

    if result and not error:
        exp_result = ExperimentResult(
            name=exp_name,
            algorithm=exp.algorithm,
            metrics=result.get("metrics", {}),
            confusion_matrix=result.get("confusion_matrix", {}),
            threshold_analysis=result.get("threshold_analysis", []),
            n_samples=result.get("n_samples", 0),
            code=code,
            reasoning=reasoning,
            attempts=attempt + 1,
        )
        m = exp_result.metrics
        attempts_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else ""
        print(f"   └─ ROC-AUC={m.get('roc_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts_note}")
        return exp_result

    print(f"   └─ FAILED after {max_retries} attempts — skipping.")
    return None

async def run_agent(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
    max_retries_per_experiment: int = 3,
    llm_model: str = "gpt-4o",
    inject_failure: bool = False,
) -> AgentResult:
    """Run the MLE agent end-to-end.

    Args:
        data: CSV file containing the dataset.
        problem_description: Natural language description of the ML problem.
        target_column: Name of the target column to predict.
        time_column: Optional column to use for time-based train/test split.
        max_iterations: Maximum number of experiment iterations to run.
        max_retries_per_experiment: Max times to retry a failed sandbox execution.
        llm_model: OpenAI model to use (default: gpt-4o).
        inject_failure: If True, corrupts the first experiment to demonstrate self-healing.
    """
    print(f"\n{'='*60}")
    print(f"MLE Agent starting")
    print(f"Problem: {problem_description}")
    print(f"Target: {target_column}")
    if inject_failure:
        print(f"[demo mode: failure injection enabled]")
    print(f"{'='*60}\n")

    # {{docs-fragment phase1_profile}}
    # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) ---
    print(">> Phase 1: Profiling dataset...")
    with flyte.group("profile"):
        profile = await profile_dataset(data, target_column)
    # {{/docs-fragment phase1_profile}}
    print(f"   Shape: {profile['shape']}, Classes: {profile['target_distribution']}")
    print(f"   Imbalanced: {profile['is_imbalanced']}, Columns: {len(profile['columns'])}")
    corr = profile.get("feature_target_corr", {})
    top_corr = list(corr.items())[:5]
    print(f"   Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top_corr)}")

    # Stream report: dataset summary
    await flyte.report.log.aio(
        f"<h1>MLE Agent Run</h1>"
        f"<p><b>Problem:</b> {problem_description}</p>"
        f"<p><b>Dataset:</b> {profile['shape'][0]:,} rows × {profile['shape'][1]} cols &nbsp;|&nbsp; "
        f"Class balance: {profile['class_balance']} &nbsp;|&nbsp; Imbalanced: {profile['is_imbalanced']}</p>"
        f"<p><b>Top feature-target correlations (raw):</b> "
        + ", ".join(f"{k}: {v:+.3f}" for k, v in top_corr) +
        f"</p><hr>",
        do_flush=True,
    )

    # --- Phase 2: LLM designs initial experiments ---
    print("\n>> Phase 2: Designing initial experiments...")
    design_response = await design_experiments(
        problem_description=problem_description,
        profile_json=json.dumps(profile),
        llm_model=llm_model,
    )
    design = InitialDesign.model_validate(_parse_json(design_response))
    print(f"   Primary metric: {design.primary_metric}")
    print(f"   Strategy: {design.reasoning}")
    print(f"   Experiments planned: {len(design.experiments)}")

    all_results: list[ExperimentResult] = []
    iteration_log: list[dict] = []  # tracks per-iteration decisions + explorations for summary
    current_experiments: list[ExperimentConfig] = design.experiments
    first_experiment = True

    # --- Phase 3: Iterative experiment loop ---
    for iteration in range(max_iterations):
        experiments = current_experiments

        if not experiments:
            print(f"\n>> No experiments to run in iteration {iteration + 1}. Stopping.")
            break

        print(f"\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...")

        # Assign names and prepare dicts before launching in parallel
        exp_batch = []
        for i, exp in enumerate(experiments):
            if not exp.name:
                exp.name = f"experiment_{len(all_results) + i + 1}"
            exp_dict = exp.model_dump()
            inject_this = inject_failure and first_experiment and i == 0
            if inject_this:
                exp_dict = _corrupt_experiment_for_demo(exp_dict)
            first_experiment = False
            exp_batch.append((exp, exp_dict, inject_this))

        # {{docs-fragment parallel_execute}}
        batch_results = await asyncio.gather(*[
            _run_experiment(
                exp=exp,
                exp_dict=exp_dict,
                inject_failure=inject_this,
                data=data,
                target_column=target_column,
                time_column=time_column,
                profile=profile,
                llm_model=llm_model,
                max_retries=max_retries_per_experiment,
            )
            for exp, exp_dict, inject_this in exp_batch
        ])
        # {{/docs-fragment parallel_execute}}

        for exp_result in batch_results:
            if exp_result is not None:
                all_results.append(exp_result)
                # Stream report: each experiment as it completes
                m = exp_result.metrics
                html = (
                    f"<h3>Iteration {iteration + 1} — {exp_result.name}</h3>"
                    f"<p><b>Algorithm:</b> {exp_result.algorithm} &nbsp;|&nbsp; "
                    f"<b>ROC-AUC:</b> {m.get('roc_auc')} &nbsp;|&nbsp; "
                    f"<b>F1:</b> {m.get('f1')} &nbsp;|&nbsp; "
                    f"<b>Recall:</b> {m.get('recall')} &nbsp;|&nbsp; "
                    f"<b>Attempts:</b> {exp_result.attempts}</p>"
                )
                if exp_result.reasoning:
                    html += f"<details><summary>Reasoning</summary><pre>{exp_result.reasoning}</pre></details>"
                html += f"<details><summary>Generated Code</summary><pre>{exp_result.code}</pre></details>"
                await flyte.report.log.aio(html, do_flush=True)

        # --- Phase 4: Analyze results, decide whether to iterate ---
        if all_results and iteration < max_iterations - 1:
            print(f"\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...")
            results_summary = [
                {
                    "experiment_name": r.name,
                    "algorithm": r.algorithm,
                    "metrics": r.metrics,
                    "confusion_matrix": r.confusion_matrix,
                    "used_feature_engineering": "engineer_features" in r.code,
                    "used_rolling_features": "rolling_columns" in r.code,
                    "used_lag_features": "lag_columns" in r.code,
                }
                for r in all_results
            ]
            analysis_prompt = textwrap.dedent(f"""
                Problem: {problem_description}
                Dataset profile: shape={profile['shape']}, imbalanced={profile['is_imbalanced']}
                Feature-target correlations (raw): {json.dumps(profile.get('feature_target_corr', {}), indent=2)}

                Experiment results so far (iteration {iteration + 1}):
                {json.dumps(results_summary, indent=2)}

                Should we run more experiments? If yes, request any data explorations
                you need, then specify what experiments to run next.
            """).strip()

            analysis_response = await analyze_iteration(
                analysis_prompt=analysis_prompt,
                max_iterations=max_iterations,
                current_iteration=iteration,
                llm_model=llm_model,
            )
            decision = IterationDecision.model_validate(_parse_json(analysis_response))
            verdict = "continuing" if decision.should_continue else "stopping"
            print(f"   Decision: {verdict}")
            print(f"   Reasoning: {decision.reasoning}")

            # Stream report: analysis decision
            await flyte.report.log.aio(
                f"<h3>Analysis — Iteration {iteration + 1}</h3>"
                f"<p><b>Decision:</b> {verdict}</p>"
                f"<p><b>Reasoning:</b> {decision.reasoning}</p>",
                do_flush=True,
            )

            # Track this iteration for the experiment journey summary
            iter_entry = {
                "iteration": iteration + 1,
                "experiments": [r.name for r in batch_results if r is not None],
                "best_roc_auc": max(
                    (r.metrics.get("roc_auc", 0) for r in all_results), default=0
                ),
                "reasoning": decision.reasoning,
                "explorations": [],
            }

            # --- Targeted exploration before next iteration ---
            if decision.should_continue and decision.exploration_requests:
                print(f"   Running {len(decision.exploration_requests)} exploration request(s)...")
                exploration_questions = []
                exploration_results = []

                for i, req in enumerate(decision.exploration_requests):
                    question = req.get("question", f"Exploration {i + 1}")
                    # Strip agent-level metadata — tool only needs the analysis config
                    tool_config = {k: v for k, v in req.items() if k not in ("question", "analysis_type")}

                    print(f"   Q: {question}")
                    with flyte.group(f"explore_{iteration + 1}_{i + 1}"):
                        result = await explore_dataset(data, tool_config)
                    exploration_questions.append(question)
                    exploration_results.append(result)
                    iter_entry["explorations"].append({"question": question})

                    await flyte.report.log.aio(
                        f"<h4>Exploration {i + 1}</h4>"
                        f"<p><b>Question:</b> {question}</p>"
                        f"<details><summary>Results</summary><pre>{json.dumps(result, indent=2)}</pre></details>",
                        do_flush=True,
                    )

                # Build follow-up that explicitly connects each question to its answer
                qa_pairs = "\n\n".join(
                    f'Question {i + 1}: "{q}"\nResult:\n{json.dumps(r, indent=2)}'
                    for i, (q, r) in enumerate(zip(exploration_questions, exploration_results))
                )
                followup_prompt = textwrap.dedent(f"""
                    You requested {len(exploration_results)} targeted exploration(s).
                    Here is what you asked and what you learned:

                    {qa_pairs}

                    Given what you learned and your earlier reasoning:
                    "{decision.reasoning}"

                    Now specify the next experiments. For each experiment, briefly state
                    which exploration insight informed your choice.
                    Respond with valid JSON: {{"next_experiments": [...same schema as before...]}}
                """).strip()
                followup_response = await plan_followup(
                    analysis_prompt=analysis_prompt,
                    analysis_response=analysis_response,
                    followup_prompt=followup_prompt,
                    max_iterations=max_iterations,
                    current_iteration=iteration,
                    llm_model=llm_model,
                )
                followup = _parse_json(followup_response)
                current_experiments = IterationDecision.model_validate({
                    "should_continue": True,
                    "reasoning": decision.reasoning,
                    "next_experiments": followup.get("next_experiments", []),
                }).next_experiments
                print(f"   Post-exploration: {len(current_experiments)} experiment(s) planned")
            else:
                current_experiments = decision.next_experiments

            iteration_log.append(iter_entry)

            if not decision.should_continue:
                break

    # --- Phase 5: Rank all results and generate model card ---
    print(f"\n>> Phase 5: Ranking {len(all_results)} experiment(s) and generating model card...")

    if not all_results:
        return AgentResult(
            model_card="No experiments completed successfully.",
            best_experiment="",
            best_metrics={},
            all_results=[],
            iterations=iteration + 1,
            total_experiments=0,
        )

    ranking_input = [
        {
            "experiment_name": r.name,
            "metrics": r.metrics,
            "confusion_matrix": r.confusion_matrix,
        }
        for r in all_results
    ]
    with flyte.group("rank"):
        ranking = await rank_experiments(json.dumps(ranking_input))
    best_name = ranking["best_experiment"]
    best_result = next(r for r in all_results if r.name == best_name)

    _print_experiment_table(all_results, best_name)
    _print_threshold_recommendation(best_result.threshold_analysis, best_result.metrics)

    # Stream report: final rankings table
    rows = "".join(
        f"<tr><td>{row['rank']}</td>"
        f"<td>{'<b>' if row['experiment_name'] == best_name else ''}"
        f"{row['experiment_name']}"
        f"{'</b>' if row['experiment_name'] == best_name else ''}</td>"
        f"<td>{row['roc_auc']}</td><td>{row['f1']}</td>"
        f"<td>{row['recall']}</td><td>{row['precision']}</td></tr>"
        for row in ranking.get("ranking", [])
    )
    await flyte.report.log.aio(
        f"<hr><h2>Final Rankings</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0'>"
        f"<tr><th>Rank</th><th>Experiment</th><th>ROC-AUC</th><th>F1</th><th>Recall</th><th>Precision</th></tr>"
        f"{rows}</table>"
        f"<p>{ranking.get('summary', '')}</p>",
        do_flush=True,
    )

    # Stream report: experiment journey summary
    journey_rows = ""
    for entry in iteration_log:
        exps = ", ".join(entry["experiments"]) if entry["experiments"] else "—"
        explorations = "; ".join(e["question"] for e in entry["explorations"]) if entry["explorations"] else "—"
        short_reasoning = (entry["reasoning"][:120] + "…") if len(entry["reasoning"]) > 120 else entry["reasoning"]
        journey_rows += (
            f"<tr>"
            f"<td style='text-align:center'>{entry['iteration']}</td>"
            f"<td>{exps}</td>"
            f"<td style='text-align:center'>{entry['best_roc_auc']:.4f}</td>"
            f"<td>{short_reasoning}</td>"
            f"<td>{explorations}</td>"
            f"</tr>"
        )
    await flyte.report.log.aio(
        f"<hr><h2>Experiment Journey</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0' style='width:100%;border-collapse:collapse'>"
        f"<tr><th>Iter</th><th>Experiments</th><th>Best ROC-AUC</th><th>Key insight</th><th>Explorations</th></tr>"
        f"{journey_rows}"
        f"</table>",
        do_flush=True,
    )

    model_card = await _generate_model_card(
        problem_description=problem_description,
        profile=profile,
        all_results=all_results,
        best_result=best_result,
        ranking=ranking,
        iteration_log=iteration_log,
        llm_model=llm_model,
    )

    print(f"\n{'='*60}")
    print(f"DONE — Best model: {best_name}")
    print(f"       ROC-AUC={best_result.metrics.get('roc_auc')}, F1={best_result.metrics.get('f1')}")
    print(f"{'='*60}\n")

    return AgentResult(
        model_card=model_card,
        best_experiment=best_name,
        best_metrics=best_result.metrics,
        all_results=all_results,
        iterations=iteration + 1,
        total_experiments=len(all_results),
    )

async def _generate_model_card(
    problem_description: str,
    profile: dict,
    all_results: list[ExperimentResult],
    best_result: ExperimentResult,
    ranking: dict,
    iteration_log: list[dict],
    llm_model: str,
) -> str:
    """Generate a markdown model card summarizing the winning model."""
    system = textwrap.dedent("""
        You are an ML engineer writing a model card for a trained model.
        Write in markdown. Be concise but informative. Include:
        - Problem statement
        - Dataset summary
        - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed)
        - Experiment summary (table of all experiments with metrics)
        - Winning model details (algorithm, key hyperparams, metrics, threshold analysis)
        - Recommendations for deployment (decision threshold, monitoring)
    """).strip()

    results_text = "\n".join(
        f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc_auc')}, "
        f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}"
        for r in all_results
    )

    journey_text = ""
    if iteration_log:
        journey_text = "\n\nIteration log:\n" + "\n".join(
            f"  Iteration {e['iteration']}: ran [{', '.join(e['experiments'])}], "
            f"best ROC-AUC so far={e['best_roc_auc']:.4f}. "
            f"Key insight: {e['reasoning'][:200]}. "
            + (f"Explorations: {'; '.join(x['question'] for x in e['explorations'])}" if e['explorations'] else "")
            for e in iteration_log
        )

    user_content = textwrap.dedent(f"""
        Problem: {problem_description}

        Dataset: {profile['shape'][0]} rows × {profile['shape'][1]} cols.
        Class balance: {profile['class_balance']}
        Imbalanced: {profile['is_imbalanced']}
        {journey_text}

        All experiments:
        {results_text}

        Best model: {best_result.name} ({best_result.algorithm})
        Metrics: {json.dumps(best_result.metrics, indent=2)}
        Confusion matrix: {json.dumps(best_result.confusion_matrix, indent=2)}
        Threshold analysis: {json.dumps(best_result.threshold_analysis, indent=2)}

        Ranking summary: {ranking['summary']}
    """).strip()

    response = await _call_llm(system, [{"role": "user", "content": user_content}], llm_model)
    return response

# ---------------------------------------------------------------------------
# Durable entrypoint (runs the agent as a Flyte task in the cloud)
# ---------------------------------------------------------------------------

# {{docs-fragment entrypoint}}
@agent_env.task(retries=1, report=True)
async def mle_agent_task(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
) -> str:
    """Durable Flyte task entrypoint for the MLE agent."""
    result = await run_agent(
        data=data,
        problem_description=problem_description,
        target_column=target_column,
        time_column=time_column,
        max_iterations=max_iterations,
    )
    return result.model_card
# {{/docs-fragment entrypoint}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle_bot/mle_bot/agent.py*

On the first attempt, `previous_error` and `previous_code` are empty. On subsequent attempts, the LLM sees exactly what went wrong and can fix it. In practice, most experiments succeed on the first try, with occasional recoveries on the second.

## Streaming results to a live report

While the agent runs, it streams results to the Flyte UI in real time using `flyte.report.log.aio()`. You don't have to wait for the full run to finish to see how experiments are performing.

The entrypoint task enables this with `report=True`:

CODE24python
[your orchestration code]
CODE25python" in text:
        start = text.index("CODE26python")
        end = text.index("CODE27" in text:
        start = text.index("CODE28", start)
        return text[start:end].strip()
    return text.strip()

def _extract_reasoning(text: str) -> str:
    """Extract the ## Reasoning section from LLM response."""
    if "## Reasoning" in text:
        start = text.index("## Reasoning") + len("## Reasoning")
        if "## Code" in text:
            end = text.index("## Code")
            return text[start:end].strip()
        return text[start:].strip()
    return ""

def _parse_json(text: str) -> dict:
    """Extract and parse JSON from LLM response."""
    text = text.strip()
    if "CODE29json") + 7
        end = text.index("CODE30" in text:
        start = text.index("CODE31", start)
        text = text[start:end].strip()
    return json.loads(text)

# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------

def _recommend_threshold(threshold_analysis: list, min_precision: float = 0.70) -> dict | None:
    """Find the threshold that maximises recall subject to precision >= min_precision."""
    candidates = [t for t in threshold_analysis if t["precision"] >= min_precision]
    if not candidates:
        return None
    return max(candidates, key=lambda t: t["recall"])

def _print_experiment_table(results: list["ExperimentResult"], best_name: str) -> None:
    """Print a ranked comparison table of all experiments."""
    sorted_results = sorted(results, key=lambda r: r.metrics.get("roc_auc", 0), reverse=True)
    print("\n" + "─" * 78)
    print(f"  {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}")
    print("─" * 78)
    for rank, r in enumerate(sorted_results, 1):
        note = "◀ winner" if r.name == best_name else ""
        roc = r.metrics.get("roc_auc", 0)
        f1 = r.metrics.get("f1", 0)
        recall = r.metrics.get("recall", 0)
        print(f"  {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}")
    print("─" * 78)

def _print_threshold_recommendation(threshold_analysis: list, default_metrics: dict) -> None:
    """Print the operational threshold recommendation."""
    rec = _recommend_threshold(threshold_analysis)
    if not rec:
        return
    default_recall = default_metrics.get("recall", 0)
    default_precision = default_metrics.get("precision", 0)
    missed_pct = round((1 - rec["recall"]) * 100, 1)
    false_alarm_pct = round((1 - rec["precision"]) * 100, 1)

    print(f"\n  Recommended decision threshold: {rec['threshold']}")
    print(f"  ├─ Precision : {rec['precision']:.0%}   ({false_alarm_pct}% of alerts are false alarms)")
    print(f"  ├─ Recall    : {rec['recall']:.0%}   (catches {rec['recall']*100:.0f}% of actual failures)")
    print(f"  └─ F1        : {rec['f1']:.4f}")
    print(f"  Default threshold (0.5): Precision={default_precision:.0%}, Recall={default_recall:.0%}")
    if rec["recall"] > default_recall:
        extra = round((rec["recall"] - default_recall) * 100, 1)
        print(f"  → Lowering threshold catches {extra}% more failures at cost of more alerts")

# ---------------------------------------------------------------------------
# Orchestration code generation (durable Flyte task with Flyte report)
# ---------------------------------------------------------------------------

@agent_env.task
async def plan_experiment(
    experiment_json: str,
    profile_json: str,
    target_column: str,
    time_column: str,
    previous_error: str = "",
    previous_code: str = "",
    llm_model: str = "gpt-4o",
) -> str:
    """LLM plans a single experiment: reasons about the pipeline and generates Monty code.

    Runs as a durable Flyte task so each experiment's planning step is traceable.
    Returns a JSON string: {"code": "...", "reasoning": "..."}.

    Args:
        experiment_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...).
        profile_json: JSON string of the dataset profile from profile_dataset.
        target_column: Name of the target column.
        time_column: Time column for temporal splitting, or empty string.
        previous_error: Error message from the previous attempt (empty on first try).
        previous_code: Code that failed on the previous attempt (empty on first try).
        llm_model: OpenAI model identifier.

    Returns:
        str — JSON string with keys "code" and "reasoning".
    """
    experiment = json.loads(experiment_json)
    profile = json.loads(profile_json)
    exp_name = experiment.get("name", "experiment")

    # Strip rationale — it was written by the design LLM to explain *why* this
    # experiment was chosen. Passing it here causes plan_experiment to parrot it
    # back as "reasoning" instead of independently thinking about *how* to build
    # the best pipeline. Keep only the technical spec.
    pipeline_spec = {
        k: v for k, v in experiment.items()
        if k not in ("rationale",)
    }

    system = _build_orchestration_system_prompt(profile)

    user_content = textwrap.dedent(f"""
        Design and implement the best pipeline for this experiment:

        Name: {exp_name}
        Algorithm: {pipeline_spec.get("algorithm")}
        Hyperparams: {json.dumps(pipeline_spec.get("hyperparams", {}), indent=2)}
        Feature config hint: {json.dumps(pipeline_spec.get("feature_config", {}), indent=2)}

        Available sandbox inputs:
        - data: File  — the full dataset CSV
        - target_column: str = "{target_column}"
        - time_column: str = "{time_column}"  (empty string means no time ordering)
        - experiment_name: str = "{exp_name}"

        The feature config hint is a suggestion from the experiment designer — you can
        follow it, improve on it, or override it if the dataset context and your ML
        judgment suggest a better approach. In your ## Reasoning, explain your actual
        pipeline decisions: what you chose to do (or not do) and why, based on the
        dataset profile above. Do not restate the experiment name or why it was chosen.
    """).strip()

    messages = [{"role": "user", "content": user_content}]
    if previous_code and previous_error:
        messages = [
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": f"CODE32"},
            {"role": "user", "content": f"That code failed with this error:\n\n{previous_error}\n\nPlease fix it."},
        ]

    response = await _call_llm(system, messages, llm_model)
    reasoning = _extract_reasoning(response)
    code = _extract_code(response)
    return json.dumps({"code": code, "reasoning": reasoning})

@flyte.trace
async def design_experiments(
    problem_description: str,
    profile_json: str,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs the initial batch of experiments given problem + dataset profile.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching InitialDesign schema).
    """
    design_prompt = textwrap.dedent(f"""
        Problem description: {problem_description}

        Dataset profile:
        {profile_json}

        Design the first batch of experiments.
    """).strip()
    return await _call_llm(
        _build_initial_design_system_prompt(),
        [{"role": "user", "content": design_prompt}],
        llm_model,
    )

@flyte.trace
async def analyze_iteration(
    analysis_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM analyzes experiment results and decides whether/how to continue.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string matching IterationDecision schema).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [{"role": "user", "content": analysis_prompt}],
        llm_model,
    )

@flyte.trace
async def plan_followup(
    analysis_prompt: str,
    analysis_response: str,
    followup_prompt: str,
    max_iterations: int,
    current_iteration: int,
    llm_model: str = "gpt-4o",
) -> str:
    """LLM designs next experiments after targeted data explorations.

    Traced so the prompt/response is visible in the Flyte UI and results are
    cached for deterministic replay on crash/retry.
    Returns raw LLM response (JSON string with {"next_experiments": [...]}).
    """
    return await _call_llm(
        _build_analysis_system_prompt(max_iterations, current_iteration),
        [
            {"role": "user", "content": analysis_prompt},
            {"role": "assistant", "content": analysis_response},
            {"role": "user", "content": followup_prompt},
        ],
        llm_model,
    )

def _corrupt_experiment_for_demo(exp_dict: dict) -> dict:
    """Introduce a deliberate error into the first experiment for demo purposes.

    Corrupts the algorithm name so the LLM must recover from a known-bad value.
    The retry loop will catch this, regenerate with the error message, and fix it.
    """
    corrupted = dict(exp_dict)
    corrupted["algorithm"] = corrupted["algorithm"] + "_INVALID"
    return corrupted

# ---------------------------------------------------------------------------
# Main agent loop
# ---------------------------------------------------------------------------

@dataclass
class ExperimentResult:
    name: str
    algorithm: str
    metrics: dict
    confusion_matrix: dict
    threshold_analysis: list
    n_samples: int
    code: str
    attempts: int
    reasoning: str = ""
    error: str = ""

@dataclass
class AgentResult:
    model_card: str
    best_experiment: str
    best_metrics: dict
    all_results: list[ExperimentResult]
    iterations: int
    total_experiments: int

async def _run_experiment(
    exp: "ExperimentConfig",
    exp_dict: dict,
    inject_failure: bool,
    data: File,
    target_column: str,
    time_column: str,
    profile: dict,
    llm_model: str,
    max_retries: int,
) -> "ExperimentResult | None":
    """Run a single experiment with retries. Returns None on total failure."""
    exp_name = exp.name
    profile_json = json.dumps(profile)
    print(f"\n   ┌─ {exp_name}  [{exp.algorithm}]")
    if exp.rationale:
        for line in textwrap.wrap(exp.rationale, width=58):
            print(f"   │  {line}")
    if inject_failure:
        print(f"   │  [injecting failure for demo: algorithm='{exp_dict['algorithm']}']")

    code = ""
    error = ""
    result = None
    attempt = 0

    reasoning = ""
    # {{docs-fragment retry_loop}}
    for attempt in range(max_retries):
        try:
            with flyte.group(exp_name):
                plan_json = await plan_experiment.aio(
                    experiment_json=json.dumps(exp_dict),
                    profile_json=profile_json,
                    target_column=target_column,
                    time_column=time_column,
                    previous_error=error,
                    previous_code=code,
                    llm_model=llm_model,
                )
                plan = json.loads(plan_json)
                code = plan["code"]
                reasoning = plan.get("reasoning", "")
                result = await flyte.sandbox.orchestrate_local(
                    code,
                    inputs={"data": data, "target_column": target_column,
                            "time_column": time_column, "experiment_name": exp_name},
                    tasks=TOOLS,
                )
            error = ""
            break
        except Exception as exc:
            error = str(exc)
            short_error = error[:100] + "..." if len(error) > 100 else error
            print(f"   │  attempt {attempt + 1} failed: {short_error}")
            print(f"   │  → asking LLM to fix and retry...")
            if inject_failure and attempt == 0:
                exp_dict = exp.model_dump()
    # {{/docs-fragment retry_loop}}

    if result and not error:
        exp_result = ExperimentResult(
            name=exp_name,
            algorithm=exp.algorithm,
            metrics=result.get("metrics", {}),
            confusion_matrix=result.get("confusion_matrix", {}),
            threshold_analysis=result.get("threshold_analysis", []),
            n_samples=result.get("n_samples", 0),
            code=code,
            reasoning=reasoning,
            attempts=attempt + 1,
        )
        m = exp_result.metrics
        attempts_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else ""
        print(f"   └─ ROC-AUC={m.get('roc_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts_note}")
        return exp_result

    print(f"   └─ FAILED after {max_retries} attempts — skipping.")
    return None

async def run_agent(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
    max_retries_per_experiment: int = 3,
    llm_model: str = "gpt-4o",
    inject_failure: bool = False,
) -> AgentResult:
    """Run the MLE agent end-to-end.

    Args:
        data: CSV file containing the dataset.
        problem_description: Natural language description of the ML problem.
        target_column: Name of the target column to predict.
        time_column: Optional column to use for time-based train/test split.
        max_iterations: Maximum number of experiment iterations to run.
        max_retries_per_experiment: Max times to retry a failed sandbox execution.
        llm_model: OpenAI model to use (default: gpt-4o).
        inject_failure: If True, corrupts the first experiment to demonstrate self-healing.
    """
    print(f"\n{'='*60}")
    print(f"MLE Agent starting")
    print(f"Problem: {problem_description}")
    print(f"Target: {target_column}")
    if inject_failure:
        print(f"[demo mode: failure injection enabled]")
    print(f"{'='*60}\n")

    # {{docs-fragment phase1_profile}}
    # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) ---
    print(">> Phase 1: Profiling dataset...")
    with flyte.group("profile"):
        profile = await profile_dataset(data, target_column)
    # {{/docs-fragment phase1_profile}}
    print(f"   Shape: {profile['shape']}, Classes: {profile['target_distribution']}")
    print(f"   Imbalanced: {profile['is_imbalanced']}, Columns: {len(profile['columns'])}")
    corr = profile.get("feature_target_corr", {})
    top_corr = list(corr.items())[:5]
    print(f"   Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top_corr)}")

    # Stream report: dataset summary
    await flyte.report.log.aio(
        f"<h1>MLE Agent Run</h1>"
        f"<p><b>Problem:</b> {problem_description}</p>"
        f"<p><b>Dataset:</b> {profile['shape'][0]:,} rows × {profile['shape'][1]} cols &nbsp;|&nbsp; "
        f"Class balance: {profile['class_balance']} &nbsp;|&nbsp; Imbalanced: {profile['is_imbalanced']}</p>"
        f"<p><b>Top feature-target correlations (raw):</b> "
        + ", ".join(f"{k}: {v:+.3f}" for k, v in top_corr) +
        f"</p><hr>",
        do_flush=True,
    )

    # --- Phase 2: LLM designs initial experiments ---
    print("\n>> Phase 2: Designing initial experiments...")
    design_response = await design_experiments(
        problem_description=problem_description,
        profile_json=json.dumps(profile),
        llm_model=llm_model,
    )
    design = InitialDesign.model_validate(_parse_json(design_response))
    print(f"   Primary metric: {design.primary_metric}")
    print(f"   Strategy: {design.reasoning}")
    print(f"   Experiments planned: {len(design.experiments)}")

    all_results: list[ExperimentResult] = []
    iteration_log: list[dict] = []  # tracks per-iteration decisions + explorations for summary
    current_experiments: list[ExperimentConfig] = design.experiments
    first_experiment = True

    # --- Phase 3: Iterative experiment loop ---
    for iteration in range(max_iterations):
        experiments = current_experiments

        if not experiments:
            print(f"\n>> No experiments to run in iteration {iteration + 1}. Stopping.")
            break

        print(f"\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...")

        # Assign names and prepare dicts before launching in parallel
        exp_batch = []
        for i, exp in enumerate(experiments):
            if not exp.name:
                exp.name = f"experiment_{len(all_results) + i + 1}"
            exp_dict = exp.model_dump()
            inject_this = inject_failure and first_experiment and i == 0
            if inject_this:
                exp_dict = _corrupt_experiment_for_demo(exp_dict)
            first_experiment = False
            exp_batch.append((exp, exp_dict, inject_this))

        # {{docs-fragment parallel_execute}}
        batch_results = await asyncio.gather(*[
            _run_experiment(
                exp=exp,
                exp_dict=exp_dict,
                inject_failure=inject_this,
                data=data,
                target_column=target_column,
                time_column=time_column,
                profile=profile,
                llm_model=llm_model,
                max_retries=max_retries_per_experiment,
            )
            for exp, exp_dict, inject_this in exp_batch
        ])
        # {{/docs-fragment parallel_execute}}

        for exp_result in batch_results:
            if exp_result is not None:
                all_results.append(exp_result)
                # Stream report: each experiment as it completes
                m = exp_result.metrics
                html = (
                    f"<h3>Iteration {iteration + 1} — {exp_result.name}</h3>"
                    f"<p><b>Algorithm:</b> {exp_result.algorithm} &nbsp;|&nbsp; "
                    f"<b>ROC-AUC:</b> {m.get('roc_auc')} &nbsp;|&nbsp; "
                    f"<b>F1:</b> {m.get('f1')} &nbsp;|&nbsp; "
                    f"<b>Recall:</b> {m.get('recall')} &nbsp;|&nbsp; "
                    f"<b>Attempts:</b> {exp_result.attempts}</p>"
                )
                if exp_result.reasoning:
                    html += f"<details><summary>Reasoning</summary><pre>{exp_result.reasoning}</pre></details>"
                html += f"<details><summary>Generated Code</summary><pre>{exp_result.code}</pre></details>"
                await flyte.report.log.aio(html, do_flush=True)

        # --- Phase 4: Analyze results, decide whether to iterate ---
        if all_results and iteration < max_iterations - 1:
            print(f"\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...")
            results_summary = [
                {
                    "experiment_name": r.name,
                    "algorithm": r.algorithm,
                    "metrics": r.metrics,
                    "confusion_matrix": r.confusion_matrix,
                    "used_feature_engineering": "engineer_features" in r.code,
                    "used_rolling_features": "rolling_columns" in r.code,
                    "used_lag_features": "lag_columns" in r.code,
                }
                for r in all_results
            ]
            analysis_prompt = textwrap.dedent(f"""
                Problem: {problem_description}
                Dataset profile: shape={profile['shape']}, imbalanced={profile['is_imbalanced']}
                Feature-target correlations (raw): {json.dumps(profile.get('feature_target_corr', {}), indent=2)}

                Experiment results so far (iteration {iteration + 1}):
                {json.dumps(results_summary, indent=2)}

                Should we run more experiments? If yes, request any data explorations
                you need, then specify what experiments to run next.
            """).strip()

            analysis_response = await analyze_iteration(
                analysis_prompt=analysis_prompt,
                max_iterations=max_iterations,
                current_iteration=iteration,
                llm_model=llm_model,
            )
            decision = IterationDecision.model_validate(_parse_json(analysis_response))
            verdict = "continuing" if decision.should_continue else "stopping"
            print(f"   Decision: {verdict}")
            print(f"   Reasoning: {decision.reasoning}")

            # Stream report: analysis decision
            await flyte.report.log.aio(
                f"<h3>Analysis — Iteration {iteration + 1}</h3>"
                f"<p><b>Decision:</b> {verdict}</p>"
                f"<p><b>Reasoning:</b> {decision.reasoning}</p>",
                do_flush=True,
            )

            # Track this iteration for the experiment journey summary
            iter_entry = {
                "iteration": iteration + 1,
                "experiments": [r.name for r in batch_results if r is not None],
                "best_roc_auc": max(
                    (r.metrics.get("roc_auc", 0) for r in all_results), default=0
                ),
                "reasoning": decision.reasoning,
                "explorations": [],
            }

            # --- Targeted exploration before next iteration ---
            if decision.should_continue and decision.exploration_requests:
                print(f"   Running {len(decision.exploration_requests)} exploration request(s)...")
                exploration_questions = []
                exploration_results = []

                for i, req in enumerate(decision.exploration_requests):
                    question = req.get("question", f"Exploration {i + 1}")
                    # Strip agent-level metadata — tool only needs the analysis config
                    tool_config = {k: v for k, v in req.items() if k not in ("question", "analysis_type")}

                    print(f"   Q: {question}")
                    with flyte.group(f"explore_{iteration + 1}_{i + 1}"):
                        result = await explore_dataset(data, tool_config)
                    exploration_questions.append(question)
                    exploration_results.append(result)
                    iter_entry["explorations"].append({"question": question})

                    await flyte.report.log.aio(
                        f"<h4>Exploration {i + 1}</h4>"
                        f"<p><b>Question:</b> {question}</p>"
                        f"<details><summary>Results</summary><pre>{json.dumps(result, indent=2)}</pre></details>",
                        do_flush=True,
                    )

                # Build follow-up that explicitly connects each question to its answer
                qa_pairs = "\n\n".join(
                    f'Question {i + 1}: "{q}"\nResult:\n{json.dumps(r, indent=2)}'
                    for i, (q, r) in enumerate(zip(exploration_questions, exploration_results))
                )
                followup_prompt = textwrap.dedent(f"""
                    You requested {len(exploration_results)} targeted exploration(s).
                    Here is what you asked and what you learned:

                    {qa_pairs}

                    Given what you learned and your earlier reasoning:
                    "{decision.reasoning}"

                    Now specify the next experiments. For each experiment, briefly state
                    which exploration insight informed your choice.
                    Respond with valid JSON: {{"next_experiments": [...same schema as before...]}}
                """).strip()
                followup_response = await plan_followup(
                    analysis_prompt=analysis_prompt,
                    analysis_response=analysis_response,
                    followup_prompt=followup_prompt,
                    max_iterations=max_iterations,
                    current_iteration=iteration,
                    llm_model=llm_model,
                )
                followup = _parse_json(followup_response)
                current_experiments = IterationDecision.model_validate({
                    "should_continue": True,
                    "reasoning": decision.reasoning,
                    "next_experiments": followup.get("next_experiments", []),
                }).next_experiments
                print(f"   Post-exploration: {len(current_experiments)} experiment(s) planned")
            else:
                current_experiments = decision.next_experiments

            iteration_log.append(iter_entry)

            if not decision.should_continue:
                break

    # --- Phase 5: Rank all results and generate model card ---
    print(f"\n>> Phase 5: Ranking {len(all_results)} experiment(s) and generating model card...")

    if not all_results:
        return AgentResult(
            model_card="No experiments completed successfully.",
            best_experiment="",
            best_metrics={},
            all_results=[],
            iterations=iteration + 1,
            total_experiments=0,
        )

    ranking_input = [
        {
            "experiment_name": r.name,
            "metrics": r.metrics,
            "confusion_matrix": r.confusion_matrix,
        }
        for r in all_results
    ]
    with flyte.group("rank"):
        ranking = await rank_experiments(json.dumps(ranking_input))
    best_name = ranking["best_experiment"]
    best_result = next(r for r in all_results if r.name == best_name)

    _print_experiment_table(all_results, best_name)
    _print_threshold_recommendation(best_result.threshold_analysis, best_result.metrics)

    # Stream report: final rankings table
    rows = "".join(
        f"<tr><td>{row['rank']}</td>"
        f"<td>{'<b>' if row['experiment_name'] == best_name else ''}"
        f"{row['experiment_name']}"
        f"{'</b>' if row['experiment_name'] == best_name else ''}</td>"
        f"<td>{row['roc_auc']}</td><td>{row['f1']}</td>"
        f"<td>{row['recall']}</td><td>{row['precision']}</td></tr>"
        for row in ranking.get("ranking", [])
    )
    await flyte.report.log.aio(
        f"<hr><h2>Final Rankings</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0'>"
        f"<tr><th>Rank</th><th>Experiment</th><th>ROC-AUC</th><th>F1</th><th>Recall</th><th>Precision</th></tr>"
        f"{rows}</table>"
        f"<p>{ranking.get('summary', '')}</p>",
        do_flush=True,
    )

    # Stream report: experiment journey summary
    journey_rows = ""
    for entry in iteration_log:
        exps = ", ".join(entry["experiments"]) if entry["experiments"] else "—"
        explorations = "; ".join(e["question"] for e in entry["explorations"]) if entry["explorations"] else "—"
        short_reasoning = (entry["reasoning"][:120] + "…") if len(entry["reasoning"]) > 120 else entry["reasoning"]
        journey_rows += (
            f"<tr>"
            f"<td style='text-align:center'>{entry['iteration']}</td>"
            f"<td>{exps}</td>"
            f"<td style='text-align:center'>{entry['best_roc_auc']:.4f}</td>"
            f"<td>{short_reasoning}</td>"
            f"<td>{explorations}</td>"
            f"</tr>"
        )
    await flyte.report.log.aio(
        f"<hr><h2>Experiment Journey</h2>"
        f"<table border='1' cellpadding='6' cellspacing='0' style='width:100%;border-collapse:collapse'>"
        f"<tr><th>Iter</th><th>Experiments</th><th>Best ROC-AUC</th><th>Key insight</th><th>Explorations</th></tr>"
        f"{journey_rows}"
        f"</table>",
        do_flush=True,
    )

    model_card = await _generate_model_card(
        problem_description=problem_description,
        profile=profile,
        all_results=all_results,
        best_result=best_result,
        ranking=ranking,
        iteration_log=iteration_log,
        llm_model=llm_model,
    )

    print(f"\n{'='*60}")
    print(f"DONE — Best model: {best_name}")
    print(f"       ROC-AUC={best_result.metrics.get('roc_auc')}, F1={best_result.metrics.get('f1')}")
    print(f"{'='*60}\n")

    return AgentResult(
        model_card=model_card,
        best_experiment=best_name,
        best_metrics=best_result.metrics,
        all_results=all_results,
        iterations=iteration + 1,
        total_experiments=len(all_results),
    )

async def _generate_model_card(
    problem_description: str,
    profile: dict,
    all_results: list[ExperimentResult],
    best_result: ExperimentResult,
    ranking: dict,
    iteration_log: list[dict],
    llm_model: str,
) -> str:
    """Generate a markdown model card summarizing the winning model."""
    system = textwrap.dedent("""
        You are an ML engineer writing a model card for a trained model.
        Write in markdown. Be concise but informative. Include:
        - Problem statement
        - Dataset summary
        - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed)
        - Experiment summary (table of all experiments with metrics)
        - Winning model details (algorithm, key hyperparams, metrics, threshold analysis)
        - Recommendations for deployment (decision threshold, monitoring)
    """).strip()

    results_text = "\n".join(
        f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc_auc')}, "
        f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}"
        for r in all_results
    )

    journey_text = ""
    if iteration_log:
        journey_text = "\n\nIteration log:\n" + "\n".join(
            f"  Iteration {e['iteration']}: ran [{', '.join(e['experiments'])}], "
            f"best ROC-AUC so far={e['best_roc_auc']:.4f}. "
            f"Key insight: {e['reasoning'][:200]}. "
            + (f"Explorations: {'; '.join(x['question'] for x in e['explorations'])}" if e['explorations'] else "")
            for e in iteration_log
        )

    user_content = textwrap.dedent(f"""
        Problem: {problem_description}

        Dataset: {profile['shape'][0]} rows × {profile['shape'][1]} cols.
        Class balance: {profile['class_balance']}
        Imbalanced: {profile['is_imbalanced']}
        {journey_text}

        All experiments:
        {results_text}

        Best model: {best_result.name} ({best_result.algorithm})
        Metrics: {json.dumps(best_result.metrics, indent=2)}
        Confusion matrix: {json.dumps(best_result.confusion_matrix, indent=2)}
        Threshold analysis: {json.dumps(best_result.threshold_analysis, indent=2)}

        Ranking summary: {ranking['summary']}
    """).strip()

    response = await _call_llm(system, [{"role": "user", "content": user_content}], llm_model)
    return response

# ---------------------------------------------------------------------------
# Durable entrypoint (runs the agent as a Flyte task in the cloud)
# ---------------------------------------------------------------------------

# {{docs-fragment entrypoint}}
@agent_env.task(retries=1, report=True)
async def mle_agent_task(
    data: File,
    problem_description: str,
    target_column: str,
    time_column: str = "",
    max_iterations: int = 3,
) -> str:
    """Durable Flyte task entrypoint for the MLE agent."""
    result = await run_agent(
        data=data,
        problem_description=problem_description,
        target_column=target_column,
        time_column=time_column,
        max_iterations=max_iterations,
    )
    return result.model_card
# {{/docs-fragment entrypoint}}
CODE33python
await flyte.report.log.aio(
    f"<h3>Iteration {iteration + 1}: {exp_result.name}</h3>"
    f"<p><b>Algorithm:</b> {exp_result.algorithm} &nbsp;|&nbsp; "
    f"<b>ROC-AUC:</b> {m.get('roc_auc')} &nbsp;|&nbsp; "
    f"<b>F1:</b> {m.get('f1')}</p>",
    do_flush=True,
)
CODE34bash
uv run main.py generate-data
CODE35bash
uv run main.py run \
    --data data/predictive_maintenance.csv \
    --problem "Predict pump failures 24 hours before they happen" \
    --target failure_24h \
    --time-column timestamp \
    --max-iterations 3 \
    --output results/report.md
```

The agent connects to your cluster via `~/.flyte/config.yaml`, uploads the CSV, and submits the agent task. You'll see a URL to track the execution in the Flyte UI, and logs will stream to your terminal.

> [!NOTE]
> You'll need to register your OpenAI API key as a cluster secret before running:
> `flyte create secret openai-api-key <YOUR_KEY>`

If you want to see the self-healing retry loop in action, add the `--inject-failure` flag. This deliberately corrupts the first experiment so the agent has to detect the error and recover, which makes for a nice demo of the durability guarantees.

## Why Flyte?

You could build something similar with plain Python and `exec()`. But there are a few things you'd lose.

**Safety.** Flyte's sandbox restricts LLM-generated code to calling your pre-approved functions and nothing else. No imports, no network, no filesystem. If you wouldn't give an intern root access to your production cluster, you probably shouldn't give an LLM unrestricted code execution either.

**Durability.** Every tool call is a Flyte task. If the agent process crashes halfway through iteration 3, the experiments that already completed are cached. You restart and pick up where you left off instead of retraining models from scratch. For long-running ML experiments, this matters.

**Observability.** You can see every LLM prompt, every generated code snippet, every tool invocation, and every result in the Flyte UI. When the agent makes a questionable decision (like skipping feature engineering on temporal data), you can trace exactly why: the prompt it received, the profile it read, the reasoning it generated.

**Compute isolation.** The ML tools run on cloud instances with the CPU and memory they need. The agent itself runs on a small 1-CPU instance since all it does is call the LLM and dispatch tool tasks. You're not bottlenecked by your laptop, and you're not paying for GPU-class compute to run an orchestration loop.

**Parallelism.** Multiple experiments run simultaneously via `asyncio.gather()`, each dispatching its own durable tasks. Flyte handles the scheduling. If you have three experiments in a batch and each involves training + evaluation, that's six tasks running concurrently on cloud compute.

The MLE Bot is a specific example of a more general pattern: giving an LLM the ability to reason about *what* work should be done, while Flyte handles *how* that work gets executed safely, durably, and at scale. The sandbox is the boundary between the two. Everything above the boundary is LLM-generated and untrusted. Everything below it is your code, running on your infrastructure, with all the guarantees you'd expect from a production orchestrator.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/compliance-monitoring-agent ===

# Compliance monitoring agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/compliance_monitoring_agent).

This example demonstrates how to build a regulatory and compliance monitoring agent on Flyte. The agent watches trusted regulatory sources (FDA guidance, SEC filings, sanctions lists, state-level privacy laws) and routes structured, **citation-precise** findings to the right downstream team (compliance, legal, or clinical ops).

Compliance monitoring requires **citation precision and recency** so every finding can be verified. The [You.com Research API](https://you.com/docs/research/overview) returns a grounded, synthesized answer plus structured sources (URL, title, snippet). Use `source_control` to restrict research to trusted government and regulator domains within a recency window, and `output_schema` when you need machine-readable findings. [Claude](https://docs.anthropic.com/) via [LiteLLM](https://docs.litellm.ai/) triages each finding for severity and routing. Combined with Flyte's audit lineage, you get end-to-end traceability from query to citation.

Flyte provides:

- **Fan-out parallelism** across watch items
- **`@flyte.trace`** on every You.com Research and LLM call
- **Retries** on monitoring tasks for robustness
- **Flyte reports** grouped by team and severity

![Compliance monitoring agent report](../../../_static/images/tutorials/compliance_monitoring_agent/compliance-monitoring-agent.png)

## Setting up the environment

The agent runs in a `TaskEnvironment` with secrets for the You.com and Anthropic API keys and a container image built from the `uv` script dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# ///
```

## Data types

Each `WatchItem` specifies a regulatory topic, a list of trusted domains for `source_control`, and a routing destination team. Findings carry citation metadata (source URL, published date, and snippet) so every claim can be verified.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

## Research with the You.com Research API

The `you_research` helper calls the [You.com Research API](https://you.com/docs/research/overview) at `https://api.you.com/v1/research`. It passes `source_control` with an `include_domains` allowlist and a `freshness` filter, and requests structured output via `output_schema`.

See the [Research API reference](https://you.com/docs/api-reference/research/v1-research) for `research_effort` levels (`lite`, `standard`, `deep`, `exhaustive`), `source_control`, and `output_schema` parameters.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

## Triage findings with Claude

After the Research API returns structured findings, Claude assigns a severity (`info`, `watch`, or `action`) and a routing rationale for each one.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

## Monitor one watch item

The `monitor_watch_item` task researches a single regulatory topic, enriches findings with source metadata from the Research API response, and triages each finding for severity and routing.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

## Orchestration

The `compliance_monitoring` driver task fans out across all watch items, aggregates findings, and renders a Flyte report sorted by severity and team.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "compliance_monitoring"
# params = ""
# ///
"""Regulatory & compliance monitoring agent.

Watches trusted regulatory sources via the You.com Research API (with
domain/freshness source controls and a structured output schema), then uses
Claude to assign severity and route citation-precise findings to the right team.
Every external call is traced so Flyte's audit lineage extends to the web layer.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="compliance-monitoring",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="compliance-monitoring", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class WatchItem:
    topic: str
    trusted_domains: list[str]
    team: str

@dataclass
class Finding:
    topic: str
    team: str
    title: str
    summary: str
    source_url: str
    published_date: str
    snippet: str
    domain: str = ""
    favicon: str = ""
    severity: str = "info"
    rationale: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class ComplianceReport:
    findings: list[Finding] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

FINDINGS_SCHEMA = {
    "type": "object",
    "properties": {
        "findings": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "source_url": {"type": "string"},
                    "published_date": {"type": "string"},
                    "snippet": {"type": "string"},
                },
                "required": [
                    "title",
                    "summary",
                    "source_url",
                    "published_date",
                    "snippet",
                ],
                "additionalProperties": False,
            },
        }
    },
    "required": ["findings"],
    "additionalProperties": False,
}

async def _you_post(url: str, body: dict, timeout: float = 300.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(
    question: str,
    include_domains: list[str],
    freshness: str,
    research_effort: str = "standard",
) -> dict:
    """Call the You.com Research API with domain + freshness source controls."""
    body = {
        "input": question,
        "research_effort": research_effort,
        "source_control": {
            "include_domains": include_domains,
            "freshness": freshness,
        },
        "output_schema": FINDINGS_SCHEMA,
    }
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment llm}}
@flyte.trace
async def triage(topic: str, findings: list[dict]) -> list[dict]:
    """Use Claude to assign a severity + rationale to each finding."""
    from litellm import acompletion

    if not findings:
        return []

    system = (
        "You are a regulatory-compliance triage analyst. For each finding, "
        "assign a severity of 'info' (FYI), 'watch' (monitor closely), or "
        "'action' (requires a concrete compliance/legal response), and a one-"
        "sentence rationale. Respond ONLY with JSON: "
        '{"triage": [{"severity": str, "rationale": str}]} with one entry per '
        "finding, in order."
    )
    listing = "\n".join(
        f"[{i + 1}] {f.get('title', '')}: {f.get('summary', '')}"
        for i, f in enumerate(findings)
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{listing}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed.get("triage", []) if isinstance(parsed, dict) else []

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

# {{docs-fragment monitor_watch_item}}
@env.task(retries=3)
async def monitor_watch_item(item: WatchItem, freshness: str) -> list[Finding]:
    """Research one regulatory topic and produce triaged, cited findings."""
    question = (
        f"What are the most recent changes, updates, or new guidance regarding "
        f"'{item.topic}'? Report concrete, dated changes with their sources."
    )
    result = await you_research(question, item.trusted_domains, freshness)
    output = result.get("output", {})

    # Build a lookup from the Research API's full source list (url -> metadata).
    src_by_url: dict[str, dict] = {}
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        if url:
            src_by_url[url] = s

    content = output.get("content", {})
    if isinstance(content, str):
        content = _parse_json(content) if content.strip() else {}
    raw_findings = content.get("findings", []) if isinstance(content, dict) else []

    triage_results = await triage(item.topic, raw_findings)

    findings: list[Finding] = []
    for i, f in enumerate(raw_findings):
        t = triage_results[i] if i < len(triage_results) else {}
        url = str(f.get("source_url", ""))
        meta = src_by_url.get(url, {})
        snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or [""])[0])
        findings.append(
            Finding(
                topic=item.topic,
                team=item.team,
                title=str(f.get("title", "") or meta.get("title", "")),
                summary=str(f.get("summary", "")),
                source_url=url,
                published_date=str(f.get("published_date", "")),
                snippet=snippet,
                domain=_domain(url),
                favicon=_favicon_for(url),
                severity=str(t.get("severity", "info")),
                rationale=str(t.get("rationale", "")),
            )
        )
    return findings
# {{/docs-fragment monitor_watch_item}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"action": 0, "watch": 1, "info": 2}
_SEVERITY_STYLE = {
    "action": ("#fdecea", "#c0392b"),
    "watch": ("#fdf3e1", "#b7791f"),
    "info": ("#e3f1fb", "#2b6cb0"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.action { border-left-color:#c0392b; }
  .rpt .card.watch { border-left-color:#b7791f; }
  .rpt .card.info { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .team { display:inline-block; font-size:11px; font-weight:600;
               padding:3px 9px; border-radius:6px; background:#edf0f3;
               color:#52606d; text-transform:uppercase; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .rationale { font-size:13px; color:#486581; font-style:italic; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:8px 10px;
               margin-top:10px; }
  .rpt .cite img.fav { width:16px; height:16px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.45; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:3px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#edf0f3", "#52606d"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(f: Finding) -> str:
    """Render a rich You.com Research citation with domain, date, and snippet."""
    if not f.source_url:
        return ""
    meta = f.published_date[:10] if f.published_date else ""
    snip = f"<div class='csnip'>&ldquo;{f.snippet}&rdquo;</div>" if f.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{f.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{f.source_url}'><span class='cdom'>{f.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{meta} &middot; {f.title}</div>{snip}</div></div>"
    )

def _render_report(report: ComplianceReport) -> str:
    findings = sorted(
        report.findings,
        key=lambda f: (_SEVERITY_ORDER.get(f.severity, 3), f.team),
    )
    counts = {s: sum(1 for f in findings if f.severity == s) for s in _SEVERITY_ORDER}
    cited = sum(1 for f in findings if f.source_url)

    cards = []
    for f in findings:
        cards.append(
            f"<div class='card {f.severity}'>"
            f"<div>{_sev_badge(f.severity)}<span class='team'>{f.team}</span></div>"
            f"<h2>{f.title or f.topic}</h2>"
            f"<div class='summary'>{f.summary}</div>"
            f"<div class='rationale'>{f.rationale}</div>"
            f"<div class='meta' style='margin-top:6px'>{f.topic}</div>"
            f"{_cite(f)}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Compliance Monitoring Findings</h1>
      <p class="sub">Citation-precise regulatory changes from trusted domains —
      every finding links to a You.com Research source with snippet provenance.</p>
      <div class="stats">
        <span class="pill"><b>{len(findings)}</b> findings</span>
        <span class="pill"><b>{cited}</b> cited You.com sources</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{counts['action']}</b> action</span>
        <span class="pill" style="background:#fdf3e1;color:#b7791f">
          <b>{counts['watch']}</b> watch</span>
        <span class="pill" style="background:#e3f1fb;color:#2b6cb0">
          <b>{counts['info']}</b> info</span>
      </div>
      {''.join(cards) or "<p class='empty'>No findings in this window.</p>"}
      <p class="yoube">Findings retrieved via the You.com Research API with
      <code>source_control</code> domain allowlists and freshness filters.
      Flyte logs which agent called which query and got which document — full
      prompt &rarr; citation lineage for audit.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_watch_items() -> list[WatchItem]:
    return [
        WatchItem(
            topic="FDA guidance on AI/ML-enabled medical device software",
            trusted_domains=["fda.gov", "federalregister.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="SEC climate-related disclosure rules for public companies",
            trusted_domains=["sec.gov", "federalregister.gov"],
            team="legal",
        ),
        WatchItem(
            topic="OFAC sanctions list additions and updates",
            trusted_domains=["treasury.gov", "ofac.treasury.gov"],
            team="compliance",
        ),
        WatchItem(
            topic="State-level consumer data privacy laws and amendments",
            trusted_domains=["iapp.org", "oag.ca.gov"],
            team="legal",
        ),
        WatchItem(
            topic="FDA drug recalls and safety communications",
            trusted_domains=["fda.gov"],
            team="clinical",
        ),
        WatchItem(
            topic="HIPAA enforcement actions and guidance updates",
            trusted_domains=["hhs.gov"],
            team="compliance",
        ),
    ]

@env.task(report=True)
async def compliance_monitoring(
    watch_items: list[WatchItem] | None = None,
    freshness: str = "month",
) -> ComplianceReport:
    """Fan out across regulatory watch items and aggregate triaged findings."""
    if watch_items is None:
        watch_items = _default_watch_items()

    with flyte.group("monitor-watch-items"):
        results = await asyncio.gather(
            *[monitor_watch_item(item, freshness) for item in watch_items]
        )

    report = ComplianceReport(findings=[f for fs in results for f in fs])

    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compliance_monitoring)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance_monitoring_agent/main.py*

## Run the agent

### Create secrets

Get a You.com API key from the [You.com platform](https://you.com/platform) (see the [quickstart guide](https://you.com/docs/quickstart)). Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/).

Register both keys as Flyte secrets. The secret key names must match those declared in the `TaskEnvironment`:

```
flyte create secret youdotcom-api-key <YOUR_YOU_API_KEY>
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run locally or remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/compliance_monitoring_agent):

```
cd v2/tutorials/compliance_monitoring_agent
uv run --script main.py
```

To test locally without Flyte secrets:

```
export YOU_API_KEY=<YOUR_YOU_API_KEY>
export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>

uv run --script main.py
```

When the run completes, open the Flyte report to review findings grouped by severity, each with a verifiable You.com Research citation.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/field-data-enrichment-agent ===

# Field data enrichment agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/field_data_enrichment_agent).

This example demonstrates how to build an autonomous systems and field-data enrichment agent on Flyte. The agent enriches geo-tagged operational events (from autonomous vehicles, aircraft, satellites, or field sensors) with **real-world public context**: road closures, weather events, airspace changes, or local incidents tied to a geofence.

Operational data stays in your environment while public-web grounding queries go to the [You.com Search API](https://you.com/docs/search/overview). The API provides unified web and news results with `freshness` and `country` targeting, and [Claude](https://docs.anthropic.com/) via [LiteLLM](https://docs.litellm.ai/) summarizes the relevant context for each geo-tagged event.

Flyte provides:

- **Fan-out parallelism** across geo-tagged events
- **`cache="auto"`** so repeated geofence checks within the cache window reuse prior results
- **`@flyte.trace`** on every external call for lineage
- **Flyte reports** with operational severity and per-incident citations

![Field data enrichment agent report](../../../_static/images/tutorials/field_data_enrichment_agent/field-data-enrichment-data.png)

## Setting up the environment

The agent runs in a `TaskEnvironment` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the `uv` script dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "field_data_enrichment"
# params = ""
# ///
"""Autonomous systems & field-data enrichment agent.

Enriches geo-tagged operational events with real-world public context (road
closures, weather, incidents) using the You.com Search API with country +
freshness targeting, then uses Claude to summarize the relevant context. Only
public-web grounding queries leave the customer's cloud, never operational data.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="field-data-enrichment",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="field-data-enrichment", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class GeoEvent:
    event_id: str
    location: str
    country: str
    event_type: str

@dataclass
class Incident:
    description: str
    source_url: str
    published: str
    domain: str = ""
    author: str = ""
    favicon: str = ""
    snippet: str = ""
    section: str = "web"

@dataclass
class EnrichedEvent:
    event_id: str
    location: str
    context_summary: str
    severity: str
    incidents: list[Incident] = field(default_factory=list)

@dataclass
class EnrichmentReport:
    events: list[EnrichedEvent] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

@dataclass
class SearchHit:
    title: str
    url: str
    domain: str
    snippet: str
    published: str
    author: str
    favicon: str
    section: str

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_search(
    query: str, country: str, freshness: str = "day", count: int = 8
) -> list[SearchHit]:
    """Search the public web + news for context near a geofenced location."""
    params = {
        "query": query,
        "count": count,
        "freshness": freshness,
        "country": country,
    }
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict:
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

ENRICH_SYSTEM = """You are an operational-context analyst for autonomous and \
field systems. Given fresh local search results near a geofenced location, \
summarize the real-world context relevant to operations, extract discrete \
incidents (road closures, weather events, regulatory/airspace changes, local \
incidents), and assign an operational severity of 'none', 'low', 'medium', or \
'high'. Each incident must reference the supporting search result by its index. \
Respond ONLY with JSON:
{"context_summary": str, "severity": str, "incidents": [{"description": str, \
"source_index": int (the [n] of the supporting search result)}]}"""

# {{docs-fragment enrich_event}}
@env.task(retries=3)
async def enrich_event(event: GeoEvent, freshness: str) -> EnrichedEvent:
    """Ground one geo-tagged event in fresh public context."""
    query = f"{event.location} {event.event_type.replace('_', ' ')} road closure weather incident"
    hits = await you_search(query, country=event.country, freshness=freshness)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Location: {event.location}\n"
        f"Event type: {event.event_type}\n\n"
        f"Search results:\n{evidence or 'No results.'}"
    )
    parsed = await llm_json(ENRICH_SYSTEM, user)

    def _incident(it: dict) -> Incident:
        idx = int(it.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        return Incident(
            description=str(it.get("description", "")),
            source_url=src.url if src else "",
            published=src.published if src else "",
            domain=src.domain if src else "",
            author=src.author if src else "",
            favicon=src.favicon if src else "",
            snippet=src.snippet if src else "",
            section=src.section if src else "web",
        )

    incidents = [_incident(it) for it in (parsed.get("incidents", []) or [])]
    return EnrichedEvent(
        event_id=event.event_id,
        location=event.location,
        context_summary=str(parsed.get("context_summary", "")),
        severity=str(parsed.get("severity", "none")),
        incidents=incidents,
    )
# {{/docs-fragment enrich_event}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3}
_SEVERITY_STYLE = {
    "high": ("#fdecea", "#c0392b"),
    "medium": ("#fdf3e1", "#b7791f"),
    "low": ("#e3f1fb", "#2b6cb0"),
    "none": ("#eef1f4", "#627d98"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.high { border-left-color:#c0392b; }
  .rpt .card.medium { border-left-color:#b7791f; }
  .rpt .card.low { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .loc { font-size:13px; color:#52606d; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .inc { font-size:13px; color:#334e68; padding:6px 0; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:5px 0 2px 14px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#eef1f4", "#627d98"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(it: Incident) -> str:
    """Render a rich You.com citation for an incident's supporting source."""
    if not it.source_url:
        return ""
    tag = (
        "<span class='ctag news'>news</span>"
        if it.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if it.published:
        meta_bits.append(it.published[:10])
    if it.author:
        meta_bits.append(f"by {it.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{it.snippet}&rdquo;</div>" if it.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{it.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{it.source_url}'><span class='cdom'>{it.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: EnrichmentReport) -> str:
    events = sorted(report.events, key=lambda e: _SEVERITY_ORDER.get(e.severity, 4))
    flagged = sum(1 for e in events if e.severity in ("high", "medium"))
    total_sources = sum(len(e.incidents) for e in events)

    cards = []
    for e in events:
        incidents = "".join(
            f"<div class='inc'>&bull; {it.description}{_cite(it)}</div>"
            for it in e.incidents
        )
        cards.append(
            f"<div class='card {e.severity}'>"
            f"<div>{_sev_badge(e.severity)}"
            f"<span class='loc'><b>{e.event_id}</b> &middot; {e.location}</span></div>"
            f"<div class='summary'>{e.context_summary or 'No relevant public context found.'}</div>"
            f"{incidents}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Field-Data Enrichment</h1>
      <p class="sub">Geo-tagged events grounded in fresh public context — each
      incident cites a timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(events)}</b> events</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{flagged}</b> flagged (high/medium)</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No events processed.</p>"}
      <p class="yoube">Public context retrieved via the You.com Search API with
      country + freshness targeting. Operational data never leaves the BYOC
      boundary — only public-web queries go out.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
DEFAULT_EVENTS = [
    GeoEvent("evt-1", "Mountain View, CA", "US", "road_closure_check"),
    GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),
    GeoEvent("evt-3", "Austin, TX", "US", "road_closure_check"),
    GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),
    GeoEvent("evt-5", "London, UK", "GB", "incident"),
    GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),
    GeoEvent("evt-7", "Seattle, WA", "US", "weather"),
    GeoEvent("evt-8", "Miami, FL", "US", "weather"),
    GeoEvent("evt-9", "Denver, CO", "US", "road_closure_check"),
    GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),
]

@env.task(report=True)
async def field_data_enrichment(
    events: list[GeoEvent] = DEFAULT_EVENTS,
    freshness: str = "day",
) -> EnrichmentReport:
    """Fan out across geo-tagged events and enrich each with public context."""
    with flyte.group("enrich-events"):
        enriched = await asyncio.gather(
            *[enrich_event(e, freshness) for e in events]
        )

    report = EnrichmentReport(events=list(enriched))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(field_data_enrichment)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field_data_enrichment_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# ///
```

## Data types

Each `GeoEvent` carries an event ID, location, ISO country code for geo-targeting, and an event type. Enriched events include a context summary, operational severity, and discrete incidents with source citations.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "field_data_enrichment"
# params = ""
# ///
"""Autonomous systems & field-data enrichment agent.

Enriches geo-tagged operational events with real-world public context (road
closures, weather, incidents) using the You.com Search API with country +
freshness targeting, then uses Claude to summarize the relevant context. Only
public-web grounding queries leave the customer's cloud, never operational data.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="field-data-enrichment",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="field-data-enrichment", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class GeoEvent:
    event_id: str
    location: str
    country: str
    event_type: str

@dataclass
class Incident:
    description: str
    source_url: str
    published: str
    domain: str = ""
    author: str = ""
    favicon: str = ""
    snippet: str = ""
    section: str = "web"

@dataclass
class EnrichedEvent:
    event_id: str
    location: str
    context_summary: str
    severity: str
    incidents: list[Incident] = field(default_factory=list)

@dataclass
class EnrichmentReport:
    events: list[EnrichedEvent] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

@dataclass
class SearchHit:
    title: str
    url: str
    domain: str
    snippet: str
    published: str
    author: str
    favicon: str
    section: str

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_search(
    query: str, country: str, freshness: str = "day", count: int = 8
) -> list[SearchHit]:
    """Search the public web + news for context near a geofenced location."""
    params = {
        "query": query,
        "count": count,
        "freshness": freshness,
        "country": country,
    }
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict:
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

ENRICH_SYSTEM = """You are an operational-context analyst for autonomous and \
field systems. Given fresh local search results near a geofenced location, \
summarize the real-world context relevant to operations, extract discrete \
incidents (road closures, weather events, regulatory/airspace changes, local \
incidents), and assign an operational severity of 'none', 'low', 'medium', or \
'high'. Each incident must reference the supporting search result by its index. \
Respond ONLY with JSON:
{"context_summary": str, "severity": str, "incidents": [{"description": str, \
"source_index": int (the [n] of the supporting search result)}]}"""

# {{docs-fragment enrich_event}}
@env.task(retries=3)
async def enrich_event(event: GeoEvent, freshness: str) -> EnrichedEvent:
    """Ground one geo-tagged event in fresh public context."""
    query = f"{event.location} {event.event_type.replace('_', ' ')} road closure weather incident"
    hits = await you_search(query, country=event.country, freshness=freshness)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Location: {event.location}\n"
        f"Event type: {event.event_type}\n\n"
        f"Search results:\n{evidence or 'No results.'}"
    )
    parsed = await llm_json(ENRICH_SYSTEM, user)

    def _incident(it: dict) -> Incident:
        idx = int(it.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        return Incident(
            description=str(it.get("description", "")),
            source_url=src.url if src else "",
            published=src.published if src else "",
            domain=src.domain if src else "",
            author=src.author if src else "",
            favicon=src.favicon if src else "",
            snippet=src.snippet if src else "",
            section=src.section if src else "web",
        )

    incidents = [_incident(it) for it in (parsed.get("incidents", []) or [])]
    return EnrichedEvent(
        event_id=event.event_id,
        location=event.location,
        context_summary=str(parsed.get("context_summary", "")),
        severity=str(parsed.get("severity", "none")),
        incidents=incidents,
    )
# {{/docs-fragment enrich_event}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3}
_SEVERITY_STYLE = {
    "high": ("#fdecea", "#c0392b"),
    "medium": ("#fdf3e1", "#b7791f"),
    "low": ("#e3f1fb", "#2b6cb0"),
    "none": ("#eef1f4", "#627d98"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.high { border-left-color:#c0392b; }
  .rpt .card.medium { border-left-color:#b7791f; }
  .rpt .card.low { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .loc { font-size:13px; color:#52606d; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .inc { font-size:13px; color:#334e68; padding:6px 0; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:5px 0 2px 14px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#eef1f4", "#627d98"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(it: Incident) -> str:
    """Render a rich You.com citation for an incident's supporting source."""
    if not it.source_url:
        return ""
    tag = (
        "<span class='ctag news'>news</span>"
        if it.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if it.published:
        meta_bits.append(it.published[:10])
    if it.author:
        meta_bits.append(f"by {it.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{it.snippet}&rdquo;</div>" if it.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{it.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{it.source_url}'><span class='cdom'>{it.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: EnrichmentReport) -> str:
    events = sorted(report.events, key=lambda e: _SEVERITY_ORDER.get(e.severity, 4))
    flagged = sum(1 for e in events if e.severity in ("high", "medium"))
    total_sources = sum(len(e.incidents) for e in events)

    cards = []
    for e in events:
        incidents = "".join(
            f"<div class='inc'>&bull; {it.description}{_cite(it)}</div>"
            for it in e.incidents
        )
        cards.append(
            f"<div class='card {e.severity}'>"
            f"<div>{_sev_badge(e.severity)}"
            f"<span class='loc'><b>{e.event_id}</b> &middot; {e.location}</span></div>"
            f"<div class='summary'>{e.context_summary or 'No relevant public context found.'}</div>"
            f"{incidents}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Field-Data Enrichment</h1>
      <p class="sub">Geo-tagged events grounded in fresh public context — each
      incident cites a timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(events)}</b> events</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{flagged}</b> flagged (high/medium)</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No events processed.</p>"}
      <p class="yoube">Public context retrieved via the You.com Search API with
      country + freshness targeting. Operational data never leaves the BYOC
      boundary — only public-web queries go out.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
DEFAULT_EVENTS = [
    GeoEvent("evt-1", "Mountain View, CA", "US", "road_closure_check"),
    GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),
    GeoEvent("evt-3", "Austin, TX", "US", "road_closure_check"),
    GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),
    GeoEvent("evt-5", "London, UK", "GB", "incident"),
    GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),
    GeoEvent("evt-7", "Seattle, WA", "US", "weather"),
    GeoEvent("evt-8", "Miami, FL", "US", "weather"),
    GeoEvent("evt-9", "Denver, CO", "US", "road_closure_check"),
    GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),
]

@env.task(report=True)
async def field_data_enrichment(
    events: list[GeoEvent] = DEFAULT_EVENTS,
    freshness: str = "day",
) -> EnrichmentReport:
    """Fan out across geo-tagged events and enrich each with public context."""
    with flyte.group("enrich-events"):
        enriched = await asyncio.gather(
            *[enrich_event(e, freshness) for e in events]
        )

    report = EnrichmentReport(events=list(enriched))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(field_data_enrichment)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field_data_enrichment_agent/main.py*

## Search with the You.com Search API

The `you_search` helper calls the [You.com Search API](https://you.com/docs/search/overview) with `freshness` and `country` parameters to retrieve location-relevant web and news results. See the [Search API reference](https://you.com/docs/api-reference/search/v1-search) for supported country codes and freshness values.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "field_data_enrichment"
# params = ""
# ///
"""Autonomous systems & field-data enrichment agent.

Enriches geo-tagged operational events with real-world public context (road
closures, weather, incidents) using the You.com Search API with country +
freshness targeting, then uses Claude to summarize the relevant context. Only
public-web grounding queries leave the customer's cloud, never operational data.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="field-data-enrichment",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="field-data-enrichment", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class GeoEvent:
    event_id: str
    location: str
    country: str
    event_type: str

@dataclass
class Incident:
    description: str
    source_url: str
    published: str
    domain: str = ""
    author: str = ""
    favicon: str = ""
    snippet: str = ""
    section: str = "web"

@dataclass
class EnrichedEvent:
    event_id: str
    location: str
    context_summary: str
    severity: str
    incidents: list[Incident] = field(default_factory=list)

@dataclass
class EnrichmentReport:
    events: list[EnrichedEvent] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

@dataclass
class SearchHit:
    title: str
    url: str
    domain: str
    snippet: str
    published: str
    author: str
    favicon: str
    section: str

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_search(
    query: str, country: str, freshness: str = "day", count: int = 8
) -> list[SearchHit]:
    """Search the public web + news for context near a geofenced location."""
    params = {
        "query": query,
        "count": count,
        "freshness": freshness,
        "country": country,
    }
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict:
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

ENRICH_SYSTEM = """You are an operational-context analyst for autonomous and \
field systems. Given fresh local search results near a geofenced location, \
summarize the real-world context relevant to operations, extract discrete \
incidents (road closures, weather events, regulatory/airspace changes, local \
incidents), and assign an operational severity of 'none', 'low', 'medium', or \
'high'. Each incident must reference the supporting search result by its index. \
Respond ONLY with JSON:
{"context_summary": str, "severity": str, "incidents": [{"description": str, \
"source_index": int (the [n] of the supporting search result)}]}"""

# {{docs-fragment enrich_event}}
@env.task(retries=3)
async def enrich_event(event: GeoEvent, freshness: str) -> EnrichedEvent:
    """Ground one geo-tagged event in fresh public context."""
    query = f"{event.location} {event.event_type.replace('_', ' ')} road closure weather incident"
    hits = await you_search(query, country=event.country, freshness=freshness)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Location: {event.location}\n"
        f"Event type: {event.event_type}\n\n"
        f"Search results:\n{evidence or 'No results.'}"
    )
    parsed = await llm_json(ENRICH_SYSTEM, user)

    def _incident(it: dict) -> Incident:
        idx = int(it.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        return Incident(
            description=str(it.get("description", "")),
            source_url=src.url if src else "",
            published=src.published if src else "",
            domain=src.domain if src else "",
            author=src.author if src else "",
            favicon=src.favicon if src else "",
            snippet=src.snippet if src else "",
            section=src.section if src else "web",
        )

    incidents = [_incident(it) for it in (parsed.get("incidents", []) or [])]
    return EnrichedEvent(
        event_id=event.event_id,
        location=event.location,
        context_summary=str(parsed.get("context_summary", "")),
        severity=str(parsed.get("severity", "none")),
        incidents=incidents,
    )
# {{/docs-fragment enrich_event}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3}
_SEVERITY_STYLE = {
    "high": ("#fdecea", "#c0392b"),
    "medium": ("#fdf3e1", "#b7791f"),
    "low": ("#e3f1fb", "#2b6cb0"),
    "none": ("#eef1f4", "#627d98"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.high { border-left-color:#c0392b; }
  .rpt .card.medium { border-left-color:#b7791f; }
  .rpt .card.low { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .loc { font-size:13px; color:#52606d; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .inc { font-size:13px; color:#334e68; padding:6px 0; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:5px 0 2px 14px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#eef1f4", "#627d98"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(it: Incident) -> str:
    """Render a rich You.com citation for an incident's supporting source."""
    if not it.source_url:
        return ""
    tag = (
        "<span class='ctag news'>news</span>"
        if it.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if it.published:
        meta_bits.append(it.published[:10])
    if it.author:
        meta_bits.append(f"by {it.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{it.snippet}&rdquo;</div>" if it.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{it.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{it.source_url}'><span class='cdom'>{it.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: EnrichmentReport) -> str:
    events = sorted(report.events, key=lambda e: _SEVERITY_ORDER.get(e.severity, 4))
    flagged = sum(1 for e in events if e.severity in ("high", "medium"))
    total_sources = sum(len(e.incidents) for e in events)

    cards = []
    for e in events:
        incidents = "".join(
            f"<div class='inc'>&bull; {it.description}{_cite(it)}</div>"
            for it in e.incidents
        )
        cards.append(
            f"<div class='card {e.severity}'>"
            f"<div>{_sev_badge(e.severity)}"
            f"<span class='loc'><b>{e.event_id}</b> &middot; {e.location}</span></div>"
            f"<div class='summary'>{e.context_summary or 'No relevant public context found.'}</div>"
            f"{incidents}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Field-Data Enrichment</h1>
      <p class="sub">Geo-tagged events grounded in fresh public context — each
      incident cites a timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(events)}</b> events</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{flagged}</b> flagged (high/medium)</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No events processed.</p>"}
      <p class="yoube">Public context retrieved via the You.com Search API with
      country + freshness targeting. Operational data never leaves the BYOC
      boundary — only public-web queries go out.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
DEFAULT_EVENTS = [
    GeoEvent("evt-1", "Mountain View, CA", "US", "road_closure_check"),
    GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),
    GeoEvent("evt-3", "Austin, TX", "US", "road_closure_check"),
    GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),
    GeoEvent("evt-5", "London, UK", "GB", "incident"),
    GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),
    GeoEvent("evt-7", "Seattle, WA", "US", "weather"),
    GeoEvent("evt-8", "Miami, FL", "US", "weather"),
    GeoEvent("evt-9", "Denver, CO", "US", "road_closure_check"),
    GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),
]

@env.task(report=True)
async def field_data_enrichment(
    events: list[GeoEvent] = DEFAULT_EVENTS,
    freshness: str = "day",
) -> EnrichmentReport:
    """Fan out across geo-tagged events and enrich each with public context."""
    with flyte.group("enrich-events"):
        enriched = await asyncio.gather(
            *[enrich_event(e, freshness) for e in events]
        )

    report = EnrichmentReport(events=list(enriched))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(field_data_enrichment)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field_data_enrichment_agent/main.py*

## Enrich one event

The `enrich_event` task builds a location- and type-scoped query, calls the You.com Search API, and asks Claude to summarize relevant real-world context, extract discrete incidents, and assign an operational severity, all grounded in the returned sources.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "field_data_enrichment"
# params = ""
# ///
"""Autonomous systems & field-data enrichment agent.

Enriches geo-tagged operational events with real-world public context (road
closures, weather, incidents) using the You.com Search API with country +
freshness targeting, then uses Claude to summarize the relevant context. Only
public-web grounding queries leave the customer's cloud, never operational data.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="field-data-enrichment",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="field-data-enrichment", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class GeoEvent:
    event_id: str
    location: str
    country: str
    event_type: str

@dataclass
class Incident:
    description: str
    source_url: str
    published: str
    domain: str = ""
    author: str = ""
    favicon: str = ""
    snippet: str = ""
    section: str = "web"

@dataclass
class EnrichedEvent:
    event_id: str
    location: str
    context_summary: str
    severity: str
    incidents: list[Incident] = field(default_factory=list)

@dataclass
class EnrichmentReport:
    events: list[EnrichedEvent] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

@dataclass
class SearchHit:
    title: str
    url: str
    domain: str
    snippet: str
    published: str
    author: str
    favicon: str
    section: str

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_search(
    query: str, country: str, freshness: str = "day", count: int = 8
) -> list[SearchHit]:
    """Search the public web + news for context near a geofenced location."""
    params = {
        "query": query,
        "count": count,
        "freshness": freshness,
        "country": country,
    }
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict:
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

ENRICH_SYSTEM = """You are an operational-context analyst for autonomous and \
field systems. Given fresh local search results near a geofenced location, \
summarize the real-world context relevant to operations, extract discrete \
incidents (road closures, weather events, regulatory/airspace changes, local \
incidents), and assign an operational severity of 'none', 'low', 'medium', or \
'high'. Each incident must reference the supporting search result by its index. \
Respond ONLY with JSON:
{"context_summary": str, "severity": str, "incidents": [{"description": str, \
"source_index": int (the [n] of the supporting search result)}]}"""

# {{docs-fragment enrich_event}}
@env.task(retries=3)
async def enrich_event(event: GeoEvent, freshness: str) -> EnrichedEvent:
    """Ground one geo-tagged event in fresh public context."""
    query = f"{event.location} {event.event_type.replace('_', ' ')} road closure weather incident"
    hits = await you_search(query, country=event.country, freshness=freshness)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Location: {event.location}\n"
        f"Event type: {event.event_type}\n\n"
        f"Search results:\n{evidence or 'No results.'}"
    )
    parsed = await llm_json(ENRICH_SYSTEM, user)

    def _incident(it: dict) -> Incident:
        idx = int(it.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        return Incident(
            description=str(it.get("description", "")),
            source_url=src.url if src else "",
            published=src.published if src else "",
            domain=src.domain if src else "",
            author=src.author if src else "",
            favicon=src.favicon if src else "",
            snippet=src.snippet if src else "",
            section=src.section if src else "web",
        )

    incidents = [_incident(it) for it in (parsed.get("incidents", []) or [])]
    return EnrichedEvent(
        event_id=event.event_id,
        location=event.location,
        context_summary=str(parsed.get("context_summary", "")),
        severity=str(parsed.get("severity", "none")),
        incidents=incidents,
    )
# {{/docs-fragment enrich_event}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3}
_SEVERITY_STYLE = {
    "high": ("#fdecea", "#c0392b"),
    "medium": ("#fdf3e1", "#b7791f"),
    "low": ("#e3f1fb", "#2b6cb0"),
    "none": ("#eef1f4", "#627d98"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.high { border-left-color:#c0392b; }
  .rpt .card.medium { border-left-color:#b7791f; }
  .rpt .card.low { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .loc { font-size:13px; color:#52606d; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .inc { font-size:13px; color:#334e68; padding:6px 0; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:5px 0 2px 14px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#eef1f4", "#627d98"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(it: Incident) -> str:
    """Render a rich You.com citation for an incident's supporting source."""
    if not it.source_url:
        return ""
    tag = (
        "<span class='ctag news'>news</span>"
        if it.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if it.published:
        meta_bits.append(it.published[:10])
    if it.author:
        meta_bits.append(f"by {it.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{it.snippet}&rdquo;</div>" if it.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{it.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{it.source_url}'><span class='cdom'>{it.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: EnrichmentReport) -> str:
    events = sorted(report.events, key=lambda e: _SEVERITY_ORDER.get(e.severity, 4))
    flagged = sum(1 for e in events if e.severity in ("high", "medium"))
    total_sources = sum(len(e.incidents) for e in events)

    cards = []
    for e in events:
        incidents = "".join(
            f"<div class='inc'>&bull; {it.description}{_cite(it)}</div>"
            for it in e.incidents
        )
        cards.append(
            f"<div class='card {e.severity}'>"
            f"<div>{_sev_badge(e.severity)}"
            f"<span class='loc'><b>{e.event_id}</b> &middot; {e.location}</span></div>"
            f"<div class='summary'>{e.context_summary or 'No relevant public context found.'}</div>"
            f"{incidents}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Field-Data Enrichment</h1>
      <p class="sub">Geo-tagged events grounded in fresh public context — each
      incident cites a timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(events)}</b> events</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{flagged}</b> flagged (high/medium)</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No events processed.</p>"}
      <p class="yoube">Public context retrieved via the You.com Search API with
      country + freshness targeting. Operational data never leaves the BYOC
      boundary — only public-web queries go out.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
DEFAULT_EVENTS = [
    GeoEvent("evt-1", "Mountain View, CA", "US", "road_closure_check"),
    GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),
    GeoEvent("evt-3", "Austin, TX", "US", "road_closure_check"),
    GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),
    GeoEvent("evt-5", "London, UK", "GB", "incident"),
    GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),
    GeoEvent("evt-7", "Seattle, WA", "US", "weather"),
    GeoEvent("evt-8", "Miami, FL", "US", "weather"),
    GeoEvent("evt-9", "Denver, CO", "US", "road_closure_check"),
    GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),
]

@env.task(report=True)
async def field_data_enrichment(
    events: list[GeoEvent] = DEFAULT_EVENTS,
    freshness: str = "day",
) -> EnrichmentReport:
    """Fan out across geo-tagged events and enrich each with public context."""
    with flyte.group("enrich-events"):
        enriched = await asyncio.gather(
            *[enrich_event(e, freshness) for e in events]
        )

    report = EnrichmentReport(events=list(enriched))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(field_data_enrichment)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field_data_enrichment_agent/main.py*

## Orchestration

The `field_data_enrichment` driver task fans out across all events and renders a Flyte report sorted by severity.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "field_data_enrichment"
# params = ""
# ///
"""Autonomous systems & field-data enrichment agent.

Enriches geo-tagged operational events with real-world public context (road
closures, weather, incidents) using the You.com Search API with country +
freshness targeting, then uses Claude to summarize the relevant context. Only
public-web grounding queries leave the customer's cloud, never operational data.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="field-data-enrichment",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="field-data-enrichment", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    cache="auto",
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class GeoEvent:
    event_id: str
    location: str
    country: str
    event_type: str

@dataclass
class Incident:
    description: str
    source_url: str
    published: str
    domain: str = ""
    author: str = ""
    favicon: str = ""
    snippet: str = ""
    section: str = "web"

@dataclass
class EnrichedEvent:
    event_id: str
    location: str
    context_summary: str
    severity: str
    incidents: list[Incident] = field(default_factory=list)

@dataclass
class EnrichmentReport:
    events: list[EnrichedEvent] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_search}}
YOU_SEARCH_URL = "https://ydc-index.io/v1/search"

@dataclass
class SearchHit:
    title: str
    url: str
    domain: str
    snippet: str
    published: str
    author: str
    favicon: str
    section: str

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon(item: dict, url: str) -> str:
    return item.get("favicon_url") or (
        f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"
    )

async def _you_get(url: str, params: dict, timeout: float = 60.0) -> dict:
    """GET with exponential backoff + jitter on 429 rate limits."""
    import asyncio
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {"X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"]}
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.get(url, headers=headers, params=params)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_search(
    query: str, country: str, freshness: str = "day", count: int = 8
) -> list[SearchHit]:
    """Search the public web + news for context near a geofenced location."""
    params = {
        "query": query,
        "count": count,
        "freshness": freshness,
        "country": country,
    }
    data = await _you_get(YOU_SEARCH_URL, params)

    results = data.get("results", {})
    hits: list[SearchHit] = []
    for section in ("news", "web"):
        for item in results.get(section, []) or []:
            snippets = item.get("snippets") or []
            url = item.get("url", "")
            hits.append(
                SearchHit(
                    title=item.get("title", ""),
                    url=url,
                    domain=_domain(url),
                    snippet=(snippets[0] if snippets else item.get("description", "")),
                    published=item.get("page_age", "") or "",
                    author=", ".join(item.get("authors") or []),
                    favicon=_favicon(item, url),
                    section=section,
                )
            )
    return hits
# {{/docs-fragment you_search}}

# {{docs-fragment llm}}
@flyte.trace
async def llm_json(system: str, user: str) -> dict:
    from litellm import acompletion

    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.0,
        max_tokens=1536,
    )
    parsed = _parse_json(resp.choices[0].message.content)
    return parsed if isinstance(parsed, dict) else {}

def _parse_json(text: str) -> dict | list:
    text = text.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start = min((i for i in (text.find("{"), text.find("[")) if i != -1), default=0)
    end = max(text.rfind("}"), text.rfind("]")) + 1
    return json.loads(text[start:end])
# {{/docs-fragment llm}}

ENRICH_SYSTEM = """You are an operational-context analyst for autonomous and \
field systems. Given fresh local search results near a geofenced location, \
summarize the real-world context relevant to operations, extract discrete \
incidents (road closures, weather events, regulatory/airspace changes, local \
incidents), and assign an operational severity of 'none', 'low', 'medium', or \
'high'. Each incident must reference the supporting search result by its index. \
Respond ONLY with JSON:
{"context_summary": str, "severity": str, "incidents": [{"description": str, \
"source_index": int (the [n] of the supporting search result)}]}"""

# {{docs-fragment enrich_event}}
@env.task(retries=3)
async def enrich_event(event: GeoEvent, freshness: str) -> EnrichedEvent:
    """Ground one geo-tagged event in fresh public context."""
    query = f"{event.location} {event.event_type.replace('_', ' ')} road closure weather incident"
    hits = await you_search(query, country=event.country, freshness=freshness)

    evidence = "\n\n".join(
        f"[{i + 1}] {h.title} ({h.published}) — {h.domain}\n{h.url}\n{h.snippet}"
        for i, h in enumerate(hits)
    )
    user = (
        f"Location: {event.location}\n"
        f"Event type: {event.event_type}\n\n"
        f"Search results:\n{evidence or 'No results.'}"
    )
    parsed = await llm_json(ENRICH_SYSTEM, user)

    def _incident(it: dict) -> Incident:
        idx = int(it.get("source_index", 0) or 0)
        src = hits[idx - 1] if 1 <= idx <= len(hits) else None
        return Incident(
            description=str(it.get("description", "")),
            source_url=src.url if src else "",
            published=src.published if src else "",
            domain=src.domain if src else "",
            author=src.author if src else "",
            favicon=src.favicon if src else "",
            snippet=src.snippet if src else "",
            section=src.section if src else "web",
        )

    incidents = [_incident(it) for it in (parsed.get("incidents", []) or [])]
    return EnrichedEvent(
        event_id=event.event_id,
        location=event.location,
        context_summary=str(parsed.get("context_summary", "")),
        severity=str(parsed.get("severity", "none")),
        incidents=incidents,
    )
# {{/docs-fragment enrich_event}}

# {{docs-fragment report}}
_SEVERITY_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3}
_SEVERITY_STYLE = {
    "high": ("#fdecea", "#c0392b"),
    "medium": ("#fdf3e1", "#b7791f"),
    "low": ("#e3f1fb", "#2b6cb0"),
    "none": ("#eef1f4", "#627d98"),
}

REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:16px 18px;
               margin:0 0 14px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; border-left:4px solid #cbd2d9; }
  .rpt .card.high { border-left-color:#c0392b; }
  .rpt .card.medium { border-left-color:#b7791f; }
  .rpt .card.low { border-left-color:#2b6cb0; }
  .rpt .card h2 { font-size:15px; margin:0 0 6px; color:#102a43; }
  .rpt .sev { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; text-transform:uppercase;
              letter-spacing:.03em; margin-right:8px; }
  .rpt .loc { font-size:13px; color:#52606d; }
  .rpt .summary { margin:8px 0; font-size:14px; line-height:1.45; }
  .rpt .inc { font-size:13px; color:#334e68; padding:6px 0; }
  .rpt .meta { color:#829ab1; font-size:12px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:5px 0 2px 14px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#bcccdc; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .ctag.news { background:#e8833a; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _sev_badge(sev: str) -> str:
    bg, fg = _SEVERITY_STYLE.get(sev, ("#eef1f4", "#627d98"))
    return f"<span class='sev' style='background:{bg};color:{fg}'>{sev}</span>"

def _cite(it: Incident) -> str:
    """Render a rich You.com citation for an incident's supporting source."""
    if not it.source_url:
        return ""
    tag = (
        "<span class='ctag news'>news</span>"
        if it.section == "news"
        else "<span class='ctag'>web</span>"
    )
    meta_bits = []
    if it.published:
        meta_bits.append(it.published[:10])
    if it.author:
        meta_bits.append(f"by {it.author}")
    meta = " &middot; ".join(meta_bits)
    snip = f"<div class='csnip'>&ldquo;{it.snippet}&rdquo;</div>" if it.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{it.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{it.source_url}'><span class='cdom'>{it.domain or 'source'}</span></a>{tag}"
        f"<div class='cmeta'>{meta}</div>{snip}</div></div>"
    )

def _render_report(report: EnrichmentReport) -> str:
    events = sorted(report.events, key=lambda e: _SEVERITY_ORDER.get(e.severity, 4))
    flagged = sum(1 for e in events if e.severity in ("high", "medium"))
    total_sources = sum(len(e.incidents) for e in events)

    cards = []
    for e in events:
        incidents = "".join(
            f"<div class='inc'>&bull; {it.description}{_cite(it)}</div>"
            for it in e.incidents
        )
        cards.append(
            f"<div class='card {e.severity}'>"
            f"<div>{_sev_badge(e.severity)}"
            f"<span class='loc'><b>{e.event_id}</b> &middot; {e.location}</span></div>"
            f"<div class='summary'>{e.context_summary or 'No relevant public context found.'}</div>"
            f"{incidents}</div>"
        )

    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Field-Data Enrichment</h1>
      <p class="sub">Geo-tagged events grounded in fresh public context — each
      incident cites a timestamped You.com Search result.</p>
      <div class="stats">
        <span class="pill"><b>{len(events)}</b> events</span>
        <span class="pill" style="background:#fdecea;color:#c0392b">
          <b>{flagged}</b> flagged (high/medium)</span>
        <span class="pill"><b>{total_sources}</b> cited You.com sources</span>
      </div>
      {''.join(cards) or "<p class='empty'>No events processed.</p>"}
      <p class="yoube">Public context retrieved via the You.com Search API with
      country + freshness targeting. Operational data never leaves the BYOC
      boundary — only public-web queries go out.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
DEFAULT_EVENTS = [
    GeoEvent("evt-1", "Mountain View, CA", "US", "road_closure_check"),
    GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),
    GeoEvent("evt-3", "Austin, TX", "US", "road_closure_check"),
    GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),
    GeoEvent("evt-5", "London, UK", "GB", "incident"),
    GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),
    GeoEvent("evt-7", "Seattle, WA", "US", "weather"),
    GeoEvent("evt-8", "Miami, FL", "US", "weather"),
    GeoEvent("evt-9", "Denver, CO", "US", "road_closure_check"),
    GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),
]

@env.task(report=True)
async def field_data_enrichment(
    events: list[GeoEvent] = DEFAULT_EVENTS,
    freshness: str = "day",
) -> EnrichmentReport:
    """Fan out across geo-tagged events and enrich each with public context."""
    with flyte.group("enrich-events"):
        enriched = await asyncio.gather(
            *[enrich_event(e, freshness) for e in events]
        )

    report = EnrichmentReport(events=list(enriched))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(field_data_enrichment)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field_data_enrichment_agent/main.py*

## Run the agent

### Create secrets

Get a You.com API key from the [You.com platform](https://you.com/platform) (see the [quickstart guide](https://you.com/docs/quickstart)). Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/).

Register both keys as Flyte secrets. The secret key names must match those declared in the `TaskEnvironment`:

```
flyte create secret youdotcom-api-key <YOUR_YOU_API_KEY>
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run locally or remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/field_data_enrichment_agent):

```
cd v2/tutorials/field_data_enrichment_agent
uv run --script main.py
```

To test locally without Flyte secrets:

```
export YOU_API_KEY=<YOUR_YOU_API_KEY>
export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>

uv run --script main.py
```

When the run completes, open the Flyte report to review enriched events with operational severity and timestamped You.com source citations for each incident.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/support-resolution-agent ===

# Support resolution agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/support_resolution_agent).

This example demonstrates how to build a customer-support and field-service resolution agent on Flyte. The agent resolves tickets that need current public information (return policies, weather advisories, product recalls, manufacturer specs) and drafts a customer-ready reply with sources a human agent can verify before sending.

The [You.com Research API](https://you.com/docs/research/overview) grounds each ticket in fresh, citable sources. [Claude](https://docs.anthropic.com/) via [LiteLLM](https://docs.litellm.ai/) turns that research into a reply draft. With `research_effort="lite"`, the research step stays fast enough for human-in-the-loop support flows.

Flyte provides:

- **Fan-out parallelism** across support tickets
- **`@flyte.trace`** on every external call for lineage
- A **two-step pipeline** per ticket: ground the answer, then draft the reply
- **Flyte reports** with draft replies and verifiable source citations

![Support resolution agent report](../../../_static/images/tutorials/support_resolution_agent/support-resolutions-agent.png)

## Setting up the environment

The agent runs in a `TaskEnvironment` with secrets for the You.com and Anthropic API keys and a container image built from the `uv` script dependencies.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

The Python packages are declared at the top of the file using the `uv` script style:

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# ///
```

## Data types

Each `Ticket` carries a ticket ID, a customer question, and optional product or vendor context. The final `Resolution` includes the grounded answer, a draft reply, and the list of You.com sources.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Ground answers with the You.com Research API

The `you_research` helper calls the [You.com Research API](https://you.com/docs/research/overview) with a configurable `research_effort`. For support use cases, `lite` provides a fast, citation-backed answer suitable for real-time, human-in-the-loop flows. See the [Research API reference](https://you.com/docs/api-reference/research/v1-research) for effort levels and parameters.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Ground one ticket

The `ground_answer` task combines the ticket question and context into a research query and collects the grounded answer plus structured sources from the Research API response.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Draft a customer-ready reply

The `draft_reply` task turns the grounded answer into a concise, friendly reply that cites source URLs inline so a human agent can verify before sending.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Resolve one ticket

Each ticket runs `ground_answer` followed by `draft_reply` in sequence.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Orchestration

The `support_resolution` driver task fans out across all tickets and renders a Flyte report with every draft reply and its sources.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#     "flyte>=2.4.0",
#     "httpx>=0.27.0",
#     "litellm>=1.72.0",
# ]
# main = "support_resolution"
# params = ""
# ///
"""Customer-support & field-service resolution agent.

Grounds a support ticket in fresh, public, citable sources via the You.com
Research API (low effort for low latency, human-in-the-loop use), then uses
Claude to draft a customer-ready reply that cites its sources inline so a human
agent can verify before sending.
"""

# {{docs-fragment env}}
import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YDC_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)
# {{/docs-fragment env}}

# {{docs-fragment data_types}}
@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""

def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""

def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"

@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""

@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)

@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)
# {{/docs-fragment data_types}}

# {{docs-fragment you_research}}
YOU_RESEARCH_URL = "https://api.you.com/v1/research"

async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    # YDC_API_KEY is canonical; YOU_API_KEY accepted as a backwards-compatible fallback.
    headers = {
        "X-API-Key": os.environ.get("YDC_API_KEY") or os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()

@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)
# {{/docs-fragment you_research}}

# {{docs-fragment ground_answer}}
@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)
# {{/docs-fragment ground_answer}}

# {{docs-fragment draft_reply}}
@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content

@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )
# {{/docs-fragment draft_reply}}

# {{docs-fragment resolve_ticket}}
async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)
# {{/docs-fragment resolve_ticket}}

# {{docs-fragment report}}
REPORT_CSS = """
<style>
  .rpt { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
         Helvetica, Arial, sans-serif; color:#1f2933; max-width:1040px;
         margin:0 auto; }
  .rpt h1 { font-size:22px; margin:0 0 4px; color:#102a43; }
  .rpt .sub { color:#647488; font-size:13px; margin:0 0 18px; }
  .rpt .stats { display:flex; gap:10px; flex-wrap:wrap; margin:0 0 22px; }
  .rpt .pill { background:#f0f4f8; border-radius:999px; padding:6px 14px;
               font-size:13px; color:#334e68; }
  .rpt .pill b { color:#102a43; }
  .rpt .card { border:1px solid #e4e7eb; border-radius:12px; padding:18px 20px;
               margin:0 0 16px; box-shadow:0 1px 3px rgba(16,42,67,0.06);
               background:#fff; }
  .rpt .tid { display:inline-block; font-size:11px; font-weight:700;
              padding:3px 9px; border-radius:6px; background:#e0e8f9;
              color:#2b4ba0; margin-right:8px; }
  .rpt .q { font-size:15px; font-weight:600; color:#102a43; margin:8px 0 12px; }
  .rpt .reply { background:#f7faf7; border:1px solid #e1ece1; border-radius:8px;
                padding:12px 14px; font-size:14px; line-height:1.55; }
  .rpt .reply h3 { font-size:11px; text-transform:uppercase; letter-spacing:.04em;
                   color:#3c8a5e; margin:0 0 8px; }
  .rpt .sources { margin-top:12px; }
  .rpt .sources h3 { font-size:11px; text-transform:uppercase; color:#627d98;
                     margin:0 0 8px; }
  .rpt a { color:#2b6cb0; text-decoration:none; }
  .rpt a:hover { text-decoration:underline; }
  .rpt .empty { color:#829ab1; font-style:italic; padding:8px 0; }
  .rpt .cite { display:flex; gap:9px; align-items:flex-start; background:#f7f9fb;
               border:1px solid #eef1f4; border-radius:8px; padding:7px 10px;
               margin:0 0 6px; }
  .rpt .cite img.fav { width:15px; height:15px; border-radius:3px; margin-top:2px;
                       flex:0 0 auto; background:#e4e7eb; }
  .rpt .cite .cb { font-size:12px; line-height:1.4; }
  .rpt .cite .cdom { font-weight:600; color:#334e68; }
  .rpt .cite .ctag { font-size:10px; font-weight:700; text-transform:uppercase;
                     color:#fff; background:#5b8def; border-radius:4px;
                     padding:1px 5px; margin-left:6px; }
  .rpt .cite .cmeta { color:#829ab1; }
  .rpt .cite .csnip { color:#52606d; font-style:italic; margin-top:2px; }
  .rpt .yoube { font-size:11px; color:#9aa5b1; margin-top:4px; }
</style>
"""

def _cite(s: Source) -> str:
    """Render a rich You.com Research citation for a support source."""
    if not s.url:
        return ""
    snip = f"<div class='csnip'>&ldquo;{s.snippet}&rdquo;</div>" if s.snippet else ""
    return (
        f"<div class='cite'><img class='fav' src='{s.favicon}' alt=''/>"
        f"<div class='cb'>"
        f"<a href='{s.url}'><span class='cdom'>{s.domain or 'source'}</span></a>"
        f"<span class='ctag'>research</span>"
        f"<div class='cmeta'>{s.title}</div>{snip}</div></div>"
    )

def _render_report(report: ResolutionReport) -> str:
    cards = []
    for res in report.resolutions:
        src = "".join(_cite(s) for s in res.sources[:8])
        reply_html = res.draft_reply.replace("\n", "<br/>")
        cards.append(
            f"<div class='card'>"
            f"<div><span class='tid'>{res.ticket_id}</span></div>"
            f"<div class='q'>{res.ticket}</div>"
            f"<div class='reply'><h3>Draft reply (for human review)</h3>{reply_html}</div>"
            + (f"<div class='sources'><h3>You.com sources ({len(res.sources)})</h3>{src}</div>" if src else "")
            + "</div>"
        )

    total_sources = sum(len(r.sources) for r in report.resolutions)
    return f"""
    {REPORT_CSS}
    <div class="rpt">
      <h1>Support Resolutions</h1>
      <p class="sub">Tickets grounded in fresh public sources via the You.com
      Research API — draft replies cite sources a human agent can verify.</p>
      <div class="stats">
        <span class="pill"><b>{len(report.resolutions)}</b> tickets</span>
        <span class="pill"><b>{total_sources}</b> You.com sources cited</span>
      </div>
      {''.join(cards) or "<p class='empty'>No tickets processed.</p>"}
      <p class="yoube">Each ticket grounded by the You.com Research API
      (<code>lite</code> effort for low-latency, human-in-the-loop use). Sources
      include domain, title, and snippet provenance — ready to paste into a
      customer reply with verification links.</p>
    </div>
    """
# {{/docs-fragment report}}

# {{docs-fragment driver}}
def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]

@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report
# {{/docs-fragment driver}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(support_resolution)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support_resolution_agent/main.py*

## Run the agent

### Create secrets

Get a You.com API key from the [You.com platform](https://you.com/platform) (see the [quickstart guide](https://you.com/docs/quickstart)). Get an Anthropic API key from the [Anthropic console](https://console.anthropic.com/).

Register both keys as Flyte secrets. The secret key names must match those declared in the `TaskEnvironment`:

```
flyte create secret youdotcom-api-key <YOUR_YOU_API_KEY>
flyte create secret internal-anthropic-api-key <YOUR_ANTHROPIC_API_KEY>
```

See **Tasks > Configure tasks > Secrets** for scoping and file-based secrets.

### Run locally or remotely

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/support_resolution_agent):

```
cd v2/tutorials/support_resolution_agent
uv run --script main.py
```

To test locally without Flyte secrets:

```
export YOU_API_KEY=<YOUR_YOU_API_KEY>
export ANTHROPIC_API_KEY=<YOUR_ANTHROPIC_API_KEY>

uv run --script main.py
```

When the run completes, open the Flyte report to review draft replies for each ticket, with You.com source citations ready for a human agent to verify and paste into a customer response.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/code-mode-agent ===

# Code mode analytics agent

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/code_mode_agent).

This tutorial builds a "chat with live market data" app on Flyte's native AI stack. You ask a question in the browser, and the app launches a Flyte run to answer it. Inside that run, `flyte.ai.agents.Agent` running in code mode has Claude write one small Python program, the program executes in Flyte's **Agents > Sandboxing > Code sandboxing**, and the only things it can touch are the tools you registered. It fetches daily stock prices from a Yahoo Finance server plugged in over MCP, and hands them to a DuckDB `query` (a durable Flyte task), so every query the model writes shows up as a tracked, retryable child task you can open in the UI. The cheap tools that render metrics, charts, and tables run in-process. And the web layer is not hand-built either: `flyte.ai.chat.AgentChatAppEnvironment` provides the chat UI, the streaming, and the run-per-question wiring in one declaration.

![Code mode analytics agent](../../../_static/images/tutorials/code_mode_agent/demo.gif)

## Why code mode

Most tool-using agents call tools one at a time. The model asks for a tool, the result comes back, it reasons, it asks for the next one. For anything multi-step that turns into a lot of round-trips, and the orchestration logic lives in a loop you have to babysit.

In **Agents > Sandboxing > Programmatic tool calling for agents**, the model writes a single program that orchestrates the tools instead, with real control flow and composition. A question like "compare three tickers, indexed to 100 at the start, then rank them by volatility" becomes one script that does a few fetches and runs one query. It doesn't glue together a dozen tools with model turns.

The code runs inside Monty, a restricted interpreter with no imports, no filesystem access, no network access, and near-instant startup. It can only use the tools you explicitly make available to the sandbox. That means the model isn't running arbitrary Python with unrestricted access. It can only work within the boundaries you've defined.

## What runs where

The example splits work by how expensive and how worth-tracking each piece is.

| Piece | Where it runs | Why |
|---|---|---|
| the chat app | a CPU app pod | The native chat UI. Streams progress and launches one analysis run per question. |
| `analyze` | a Flyte task (the run) | Starts a report, runs the agent loop, returns the report blocks and a summary. |
| `yf_get_historical_stock_prices` | the MCP server subprocess | Live price fetch: the agent's only path to the network. Loaded over MCP; the model calls it from its code like any other function. |
| `query(sql, series)` | a durable child task | Parses the fetched price JSON into a `prices` table and runs read-only DuckDB SQL. Real work worth tracking, retrying, and caching. Dispatched from inside the sandbox. |
| `create_metric`, `create_chart`, `create_table`, `calculate_statistics` | in-process in `analyze` | Microseconds of pure Python. Making them tasks would add a round-trip for nothing. |
| the model's code | the Monty sandbox | Untrusted LLM code, confined to calling the tools above. |

## The heavy tool: a durable query

The agent fetches prices at runtime from a Yahoo Finance MCP server (covered below). Each fetch returns one ticker's closing prices as a JSON string; `query` parses those into a `prices(ticker, date, close)` table and runs read-only DuckDB SQL against it, coercing dates to ISO strings on the way out. The reshape lives in the task, not the sandbox, because the sandbox has no `json` or `pandas`:

```
"""Tools and data access for the Code Mode stock-analysis agent.

The agent (``flyte.ai.agents.Agent`` in ``code_mode``) writes Python
orchestration code that calls these tools; that code runs in the Monty
sandbox, which allows no imports, no IO, and no network, so the only things the
generated code can touch are the tools registered in ``analysis.py``.

Two kinds of tools, on purpose:

* The **fetch** is a Yahoo Finance MCP tool (``yf_get_historical_stock_prices``),
  registered on the agent via ``mcp_servers`` in ``analysis.py``. It is the only
  path to the network — the sandbox has none — so it is the agent's live data
  source. It returns a raw JSON *string* of closing prices; the sandbox does not
  parse it (it has no ``json``), it just hands it to ``query``.
* ``query`` runs read-only DuckDB SQL over the fetched series. In ``analysis.py``
  it is a durable ``@env.task``, so the heavy analytics (moving averages,
  volatility, drawdowns, cross-ticker joins) run as a tracked, cached Flyte task.
  It parses the raw MCP strings into a ``prices`` table before running the SQL —
  the messy reshape lives here, where pandas is available, not in the sandbox.
* ``create_metric``, ``create_chart``, ``create_table``, and
  ``calculate_statistics`` are cheap, pure-Python helpers that run in-process.
  The ``create_*`` ones render HTML blocks into a per-run report collector.

To add a tool: write a function with type annotations and a docstring, then add
it to the agent's ``tools`` list in ``analysis.py``. The agent generates its
system prompt from the signatures and docstrings, so there is nothing else to
wire up.
"""

from __future__ import annotations

import contextvars
import datetime as _dt
import json as _json
import math
import uuid

CHART_COLORS = [
    "rgba(14, 165, 233, 0.8)",  # #0ea5e9 — sky
    "rgba(37, 99, 235, 0.8)",  # #2563eb — blue
    "rgba(6, 182, 212, 0.8)",  # #06b6d4 — cyan
    "rgba(99, 102, 241, 0.8)",  # #6366f1 — indigo
    "rgba(8, 145, 178, 0.8)",  # #0891b2 — deep cyan
]

CHART_BORDERS = ["#0ea5e9", "#2563eb", "#06b6d4", "#6366f1", "#0891b2"]

# Dataset — live stock closing prices, fetched via the Yahoo Finance MCP server
#
# There is no local data to fetch: the agent pulls prices at runtime from the
# `mcp-yahoo-finance` server (registered in `analysis.py`). This description is
# injected into the system prompt so the model knows how the two heavy tools fit
# together without a round-trip.

DATA_DESCRIPTION = (
    "You analyze daily stock closing prices. There are two heavy tools.\n"
    "\n"
    "Fetching (one ticker per call, via the Yahoo Finance MCP server):\n"
    "  yf_get_historical_stock_prices(symbol=..., period='1y', interval='1d')\n"
    "  returns a JSON *string* of closing prices keyed by timestamp. Do NOT parse\n"
    "  it in your code — the sandbox has no json or datetime module. Pass the\n"
    "  string straight to query(). Call it once per ticker (await each call) and\n"
    "  collect the returned strings into a dict for query(). Valid period: 1mo,\n"
    "  3mo, 6mo, 1y, 2y, 5y, ytd, max. Valid interval: 1d, 1wk, 1mo.\n"
    "\n"
    "Analyzing (durable DuckDB task):\n"
    "  query(sql, series) where `series` maps each ticker symbol to the JSON\n"
    "  string returned by yf_get_historical_stock_prices for it. The task parses\n"
    "  those into one table:\n"
    "     prices(ticker TEXT, date DATE, close DOUBLE)\n"
    "  Write a single read-only SELECT against `prices`. Do the math in SQL:\n"
    "  window functions (AVG(...) OVER (PARTITION BY ticker ORDER BY date ...))\n"
    "  for moving averages, LAG(...) for daily returns, STDDEV for volatility,\n"
    "  and GROUP BY / self-joins for cross-ticker comparisons."
)

def _jsonable(value: object) -> object:
    """Coerce DuckDB scalars to JSON-friendly Python types."""
    if isinstance(value, (_dt.date, _dt.datetime)):
        return value.isoformat()
    return value

# {{docs-fragment collector}}
# The native code-mode loop ends in a plain-text answer, but the UI renders
# structured HTML blocks. A per-run collector bridges the two: each render tool
# appends its HTML here as a side effect, and the `analyze` task reads the blocks
# back after the agent finishes. A ContextVar keeps concurrent runs isolated.
_REPORT: contextvars.ContextVar[list | None] = contextvars.ContextVar(
    "report", default=None
)

def start_report() -> None:
    """Begin a fresh report for this run (called by `analyze` before the agent)."""
    _REPORT.set([])

def collect_report() -> list[str]:
    """Return the HTML blocks rendered so far, in the order they were created."""
    return list(_REPORT.get() or [])

def _add_block(html: str) -> None:
    blocks = _REPORT.get()
    if blocks is not None:
        blocks.append(html)

# {{/docs-fragment collector}}

# {{docs-fragment sql_guard}}
# The tool is a safety boundary. The model can only call the tools you register, so
# narrowing what a tool accepts shrinks the blast radius. `query` allows a single
# read-only SELECT and nothing else. DuckDB's own parser classifies the statement, so
# there is no brittle keyword matching to trip over identifiers or string literals.
def _ensure_read_only(con, sql: str) -> None:
    import duckdb

    statements = con.extract_statements(sql)
    if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
        raise ValueError("Only a single read-only SELECT query is allowed.")

# {{/docs-fragment sql_guard}}

# {{docs-fragment query_tool}}
async def run_sql(sql: str, series: dict[str, str]) -> list:
    """Parse raw Yahoo Finance price JSON per ticker, then run a read-only query.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it (closing prices keyed by
                epoch-millisecond timestamp).

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """
    import duckdb
    import pandas as pd

    # Parse each ticker's raw MCP payload into rows and stack them into one table.
    # This reshape needs json + pandas, which the Monty sandbox lacks — so it runs
    # here, in the durable task, not in the model's generated code.
    frames = []
    for ticker, raw in series.items():
        data = _json.loads(raw) if raw else {}
        if not data:
            continue
        frame = pd.DataFrame({"ts": list(data.keys()), "close": list(data.values())})
        # The MCP keys its close prices by timestamp, but the format varies by
        # pandas version inside the server: ISO date strings ("2025-07-03") or
        # epoch-millisecond integers. Detect which and parse accordingly.
        ts = frame["ts"].astype(str)
        if ts.str.fullmatch(r"\d+").all():
            frame["date"] = pd.to_datetime(ts.astype("int64"), unit="ms").dt.date
        else:
            frame["date"] = pd.to_datetime(ts).dt.date
        frame["ticker"] = ticker
        frames.append(frame[["ticker", "date", "close"]])

    prices = (
        pd.concat(frames, ignore_index=True)
        if frames
        else pd.DataFrame(columns=["ticker", "date", "close"])
    )

    # Lock the engine down: no reading or writing files, no extensions, no network.
    con = duckdb.connect(config={"enable_external_access": "false"})
    _ensure_read_only(con, sql)

    con.register("prices", prices)
    rel = con.execute(sql)
    columns = [d[0] for d in rel.description]
    return [{c: _jsonable(v) for c, v in zip(columns, row)} for row in rel.fetchall()]

# {{/docs-fragment query_tool}}

async def calculate_statistics(rows: list, column: str) -> dict:
    """Calculate descriptive statistics for a numeric column of query rows.

    Args:
        rows: A list of row dicts, e.g. the output of query().
        column: Name of the numeric column to analyze.

    Returns:
        Dict with keys: count, mean, median, min, max, std_dev.
    """
    vals = [row[column] for row in rows if column in row and row[column] is not None]
    if not vals:
        return {"count": 0, "mean": 0, "median": 0, "min": 0, "max": 0, "std_dev": 0}
    n = len(vals)
    mean = sum(vals) / n
    ordered = sorted(vals)
    median = (
        ordered[n // 2] if n % 2 == 1 else (ordered[n // 2 - 1] + ordered[n // 2]) / 2
    )
    variance = sum((v - mean) ** 2 for v in vals) / n
    return {
        "count": n,
        "mean": round(mean, 2),
        "median": round(median, 2),
        "min": min(vals),
        "max": max(vals),
        "std_dev": round(math.sqrt(variance), 2),
    }

async def create_chart(chart_type: str, title: str, labels: list, values: list) -> str:
    """Add a chart to the report (rendered with Chart.js in the UI).

    Blocks appear in the report in the order the create_* tools are called.

    Args:
        chart_type: One of "bar", "line", "pie", "doughnut".
        title: Chart title displayed above the canvas.
        labels: X-axis labels (or slice labels for pie/doughnut).
        values: Either a flat list of numbers, or a list of
                {"label": str, "data": list[number]} dicts for multi-series.

    Returns:
        A short confirmation string.
    """
    if not values:
        return f"chart {title!r} skipped: no data to plot"

    if isinstance(values[0], dict):
        datasets = []
        for i, series in enumerate(values):
            idx = i % len(CHART_COLORS)
            datasets.append(
                {
                    "label": series["label"],
                    "data": series["data"],
                    "backgroundColor": CHART_COLORS[idx],
                    "borderColor": CHART_BORDERS[idx],
                    "borderWidth": 2,
                    "tension": 0.3,
                    "fill": False,
                }
            )
    else:
        bg = [CHART_COLORS[i % len(CHART_COLORS)] for i in range(len(values))]
        border = [CHART_BORDERS[i % len(CHART_BORDERS)] for i in range(len(values))]
        datasets = [
            {
                "label": title,
                "data": values,
                "backgroundColor": (
                    bg if chart_type in ("pie", "doughnut") else CHART_COLORS[0]
                ),
                "borderColor": (
                    border if chart_type in ("pie", "doughnut") else CHART_BORDERS[0]
                ),
                "borderWidth": 2,
                "tension": 0.3,
                "fill": chart_type == "line",
            }
        ]

    # Light text and faint grid lines so the chart reads on the chat UI's dark theme
    # (Chart.js defaults to dark grey text, which disappears on a near-black page).
    options: dict = {
        "responsive": True,
        "maintainAspectRatio": False,
        "plugins": {
            "title": {
                "display": True,
                "text": title,
                "font": {"size": 16},
                "color": "#e5e7eb",
            },
            "legend": {"labels": {"color": "#cbd5e1"}},
        },
    }
    if chart_type in ("bar", "line"):
        options["scales"] = {
            axis: {
                "ticks": {"color": "#94a3b8"},
                "grid": {"color": "rgba(148,163,184,0.15)"},
            }
            for axis in ("x", "y")
        }
    config = {
        "type": chart_type,
        "data": {"labels": labels, "datasets": datasets},
        "options": options,
    }

    # A self-contained canvas plus the script that instantiates it. The chat UI injects
    # each block's HTML and re-runs its <script>, so the chart draws itself. A unique id
    # keeps two charts in one report (or a repeated title) from colliding, and escaping
    # </ stops any string in the config from closing the <script> early.
    canvas_id = "cm-chart-" + uuid.uuid4().hex[:8]
    config_json = _json.dumps(config).replace("</", "<\\/")
    _add_block(
        '<div class="block chart-block" style="position:relative;height:340px;margin:18px 0;">'
        f'<canvas id="{canvas_id}"></canvas></div>'
        f"<script>try{{new Chart(document.getElementById('{canvas_id}'),"
        f"{config_json});}}catch(e){{console.error('chart {canvas_id}',e);}}</script>"
    )
    return f"chart {title!r} added to the report"

async def create_metric(label: str, value: str, delta: str = "") -> str:
    """Add a single KPI card (a big number with a label) to the report.

    Use for headline figures, e.g. latest price or period return. Consecutive
    metric cards lay out in a row. Blocks appear in the order the tools are called.

    Args:
        label: Short caption, e.g. "AAPL return".
        value: The formatted value to display, e.g. "$185.64" or "+12%".
        delta: Optional change note, e.g. "+8% vs last month".

    Returns:
        A short confirmation string.
    """
    # Always render the delta line (blank when there is no delta) so every card is the
    # same height whether or not a delta was passed, and a row of cards stays aligned.
    # Colors are tuned for the chat UI's dark theme.
    delta_html = f'<div style="font-size:12px;color:#94a3b8;margin-top:4px;">{delta or "&nbsp;"}</div>'
    _add_block(
        '<div class="block metric-card" style="display:inline-block;min-width:150px;margin:8px 10px 8px 0;'
        "padding:16px 20px;background:rgba(14,165,233,0.12);border:1px solid rgba(14,165,233,0.35);"
        'border-radius:14px;vertical-align:top;">'
        f'<div style="font-size:11px;color:#94a3b8;text-transform:uppercase;letter-spacing:0.06em;">{label}</div>'
        f'<div style="font-size:26px;font-weight:700;color:#7dd3fc;margin-top:4px;">{value}</div>'
        f"{delta_html}</div>"
    )
    return f"metric {label!r} added to the report"

async def create_table(title: str, headers: list, rows: list) -> str:
    """Add a data table to the report.

    Use for tabular breakdowns (e.g. per-ticker detail) where a chart would lose
    the exact numbers. Blocks appear in the order the tools are called.

    Args:
        title: Table caption shown above it.
        headers: Column names.
        rows: List of rows, each a list of cell values (same length as headers).

    Returns:
        A short confirmation string.
    """
    # Colors tuned for the chat UI's dark theme.
    head = "".join(
        f'<th style="text-align:left;padding:8px 12px;color:#7dd3fc;border-bottom:1px solid '
        f'rgba(14,165,233,0.4);">{h}</th>'
        for h in headers
    )
    body = "".join(
        "<tr>"
        + "".join(
            f'<td style="padding:8px 12px;border-bottom:1px solid rgba(148,163,184,0.15);">{c}</td>'
            for c in row
        )
        + "</tr>"
        for row in rows
    )
    _add_block(
        '<div class="block table-block" style="margin:16px 0;overflow-x:auto;">'
        f'<div style="font-size:13px;color:#94a3b8;margin-bottom:8px;">{title}</div>'
        '<table style="width:100%;border-collapse:collapse;font-size:14px;font-variant-numeric:tabular-nums;">'
        f"<thead><tr>{head}</tr></thead><tbody>{body}</tbody></table></div>"
    )
    return f"table {title!r} added to the report ({len(rows)} rows)"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/tools.py*

The model does the analytics in SQL. It uses window functions for moving averages, `LAG` for daily returns, `STDDEV` for volatility, rather than looping in Python.

## Tools are a safety boundary

Restricting the model to a fixed set of tools is one layer of safety. The tool itself is a second. Because the model can only ever call the tools you register, narrowing what each tool accepts shrinks the whole system's surface area.

`query` is a good example. The sandbox confines the orchestration code, but the SQL string still runs against real DuckDB, which can read and write local files, install extensions, and reach the network. So the tool adds two of DuckDB's own controls: it opens the connection with external access disabled (no files, no extensions, no network), and it uses DuckDB's parser to classify the statement and reject anything that is not a single read-only SELECT. A `DELETE` or `DROP` is rejected as a non-SELECT; a `read_csv('/etc/passwd')` parses as a SELECT but is stopped by the disabled external access.

```
"""Tools and data access for the Code Mode stock-analysis agent.

The agent (``flyte.ai.agents.Agent`` in ``code_mode``) writes Python
orchestration code that calls these tools; that code runs in the Monty
sandbox, which allows no imports, no IO, and no network, so the only things the
generated code can touch are the tools registered in ``analysis.py``.

Two kinds of tools, on purpose:

* The **fetch** is a Yahoo Finance MCP tool (``yf_get_historical_stock_prices``),
  registered on the agent via ``mcp_servers`` in ``analysis.py``. It is the only
  path to the network — the sandbox has none — so it is the agent's live data
  source. It returns a raw JSON *string* of closing prices; the sandbox does not
  parse it (it has no ``json``), it just hands it to ``query``.
* ``query`` runs read-only DuckDB SQL over the fetched series. In ``analysis.py``
  it is a durable ``@env.task``, so the heavy analytics (moving averages,
  volatility, drawdowns, cross-ticker joins) run as a tracked, cached Flyte task.
  It parses the raw MCP strings into a ``prices`` table before running the SQL —
  the messy reshape lives here, where pandas is available, not in the sandbox.
* ``create_metric``, ``create_chart``, ``create_table``, and
  ``calculate_statistics`` are cheap, pure-Python helpers that run in-process.
  The ``create_*`` ones render HTML blocks into a per-run report collector.

To add a tool: write a function with type annotations and a docstring, then add
it to the agent's ``tools`` list in ``analysis.py``. The agent generates its
system prompt from the signatures and docstrings, so there is nothing else to
wire up.
"""

from __future__ import annotations

import contextvars
import datetime as _dt
import json as _json
import math
import uuid

CHART_COLORS = [
    "rgba(14, 165, 233, 0.8)",  # #0ea5e9 — sky
    "rgba(37, 99, 235, 0.8)",  # #2563eb — blue
    "rgba(6, 182, 212, 0.8)",  # #06b6d4 — cyan
    "rgba(99, 102, 241, 0.8)",  # #6366f1 — indigo
    "rgba(8, 145, 178, 0.8)",  # #0891b2 — deep cyan
]

CHART_BORDERS = ["#0ea5e9", "#2563eb", "#06b6d4", "#6366f1", "#0891b2"]

# Dataset — live stock closing prices, fetched via the Yahoo Finance MCP server
#
# There is no local data to fetch: the agent pulls prices at runtime from the
# `mcp-yahoo-finance` server (registered in `analysis.py`). This description is
# injected into the system prompt so the model knows how the two heavy tools fit
# together without a round-trip.

DATA_DESCRIPTION = (
    "You analyze daily stock closing prices. There are two heavy tools.\n"
    "\n"
    "Fetching (one ticker per call, via the Yahoo Finance MCP server):\n"
    "  yf_get_historical_stock_prices(symbol=..., period='1y', interval='1d')\n"
    "  returns a JSON *string* of closing prices keyed by timestamp. Do NOT parse\n"
    "  it in your code — the sandbox has no json or datetime module. Pass the\n"
    "  string straight to query(). Call it once per ticker (await each call) and\n"
    "  collect the returned strings into a dict for query(). Valid period: 1mo,\n"
    "  3mo, 6mo, 1y, 2y, 5y, ytd, max. Valid interval: 1d, 1wk, 1mo.\n"
    "\n"
    "Analyzing (durable DuckDB task):\n"
    "  query(sql, series) where `series` maps each ticker symbol to the JSON\n"
    "  string returned by yf_get_historical_stock_prices for it. The task parses\n"
    "  those into one table:\n"
    "     prices(ticker TEXT, date DATE, close DOUBLE)\n"
    "  Write a single read-only SELECT against `prices`. Do the math in SQL:\n"
    "  window functions (AVG(...) OVER (PARTITION BY ticker ORDER BY date ...))\n"
    "  for moving averages, LAG(...) for daily returns, STDDEV for volatility,\n"
    "  and GROUP BY / self-joins for cross-ticker comparisons."
)

def _jsonable(value: object) -> object:
    """Coerce DuckDB scalars to JSON-friendly Python types."""
    if isinstance(value, (_dt.date, _dt.datetime)):
        return value.isoformat()
    return value

# {{docs-fragment collector}}
# The native code-mode loop ends in a plain-text answer, but the UI renders
# structured HTML blocks. A per-run collector bridges the two: each render tool
# appends its HTML here as a side effect, and the `analyze` task reads the blocks
# back after the agent finishes. A ContextVar keeps concurrent runs isolated.
_REPORT: contextvars.ContextVar[list | None] = contextvars.ContextVar(
    "report", default=None
)

def start_report() -> None:
    """Begin a fresh report for this run (called by `analyze` before the agent)."""
    _REPORT.set([])

def collect_report() -> list[str]:
    """Return the HTML blocks rendered so far, in the order they were created."""
    return list(_REPORT.get() or [])

def _add_block(html: str) -> None:
    blocks = _REPORT.get()
    if blocks is not None:
        blocks.append(html)

# {{/docs-fragment collector}}

# {{docs-fragment sql_guard}}
# The tool is a safety boundary. The model can only call the tools you register, so
# narrowing what a tool accepts shrinks the blast radius. `query` allows a single
# read-only SELECT and nothing else. DuckDB's own parser classifies the statement, so
# there is no brittle keyword matching to trip over identifiers or string literals.
def _ensure_read_only(con, sql: str) -> None:
    import duckdb

    statements = con.extract_statements(sql)
    if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
        raise ValueError("Only a single read-only SELECT query is allowed.")

# {{/docs-fragment sql_guard}}

# {{docs-fragment query_tool}}
async def run_sql(sql: str, series: dict[str, str]) -> list:
    """Parse raw Yahoo Finance price JSON per ticker, then run a read-only query.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it (closing prices keyed by
                epoch-millisecond timestamp).

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """
    import duckdb
    import pandas as pd

    # Parse each ticker's raw MCP payload into rows and stack them into one table.
    # This reshape needs json + pandas, which the Monty sandbox lacks — so it runs
    # here, in the durable task, not in the model's generated code.
    frames = []
    for ticker, raw in series.items():
        data = _json.loads(raw) if raw else {}
        if not data:
            continue
        frame = pd.DataFrame({"ts": list(data.keys()), "close": list(data.values())})
        # The MCP keys its close prices by timestamp, but the format varies by
        # pandas version inside the server: ISO date strings ("2025-07-03") or
        # epoch-millisecond integers. Detect which and parse accordingly.
        ts = frame["ts"].astype(str)
        if ts.str.fullmatch(r"\d+").all():
            frame["date"] = pd.to_datetime(ts.astype("int64"), unit="ms").dt.date
        else:
            frame["date"] = pd.to_datetime(ts).dt.date
        frame["ticker"] = ticker
        frames.append(frame[["ticker", "date", "close"]])

    prices = (
        pd.concat(frames, ignore_index=True)
        if frames
        else pd.DataFrame(columns=["ticker", "date", "close"])
    )

    # Lock the engine down: no reading or writing files, no extensions, no network.
    con = duckdb.connect(config={"enable_external_access": "false"})
    _ensure_read_only(con, sql)

    con.register("prices", prices)
    rel = con.execute(sql)
    columns = [d[0] for d in rel.description]
    return [{c: _jsonable(v) for c, v in zip(columns, row)} for row in rel.fetchall()]

# {{/docs-fragment query_tool}}

async def calculate_statistics(rows: list, column: str) -> dict:
    """Calculate descriptive statistics for a numeric column of query rows.

    Args:
        rows: A list of row dicts, e.g. the output of query().
        column: Name of the numeric column to analyze.

    Returns:
        Dict with keys: count, mean, median, min, max, std_dev.
    """
    vals = [row[column] for row in rows if column in row and row[column] is not None]
    if not vals:
        return {"count": 0, "mean": 0, "median": 0, "min": 0, "max": 0, "std_dev": 0}
    n = len(vals)
    mean = sum(vals) / n
    ordered = sorted(vals)
    median = (
        ordered[n // 2] if n % 2 == 1 else (ordered[n // 2 - 1] + ordered[n // 2]) / 2
    )
    variance = sum((v - mean) ** 2 for v in vals) / n
    return {
        "count": n,
        "mean": round(mean, 2),
        "median": round(median, 2),
        "min": min(vals),
        "max": max(vals),
        "std_dev": round(math.sqrt(variance), 2),
    }

async def create_chart(chart_type: str, title: str, labels: list, values: list) -> str:
    """Add a chart to the report (rendered with Chart.js in the UI).

    Blocks appear in the report in the order the create_* tools are called.

    Args:
        chart_type: One of "bar", "line", "pie", "doughnut".
        title: Chart title displayed above the canvas.
        labels: X-axis labels (or slice labels for pie/doughnut).
        values: Either a flat list of numbers, or a list of
                {"label": str, "data": list[number]} dicts for multi-series.

    Returns:
        A short confirmation string.
    """
    if not values:
        return f"chart {title!r} skipped: no data to plot"

    if isinstance(values[0], dict):
        datasets = []
        for i, series in enumerate(values):
            idx = i % len(CHART_COLORS)
            datasets.append(
                {
                    "label": series["label"],
                    "data": series["data"],
                    "backgroundColor": CHART_COLORS[idx],
                    "borderColor": CHART_BORDERS[idx],
                    "borderWidth": 2,
                    "tension": 0.3,
                    "fill": False,
                }
            )
    else:
        bg = [CHART_COLORS[i % len(CHART_COLORS)] for i in range(len(values))]
        border = [CHART_BORDERS[i % len(CHART_BORDERS)] for i in range(len(values))]
        datasets = [
            {
                "label": title,
                "data": values,
                "backgroundColor": (
                    bg if chart_type in ("pie", "doughnut") else CHART_COLORS[0]
                ),
                "borderColor": (
                    border if chart_type in ("pie", "doughnut") else CHART_BORDERS[0]
                ),
                "borderWidth": 2,
                "tension": 0.3,
                "fill": chart_type == "line",
            }
        ]

    # Light text and faint grid lines so the chart reads on the chat UI's dark theme
    # (Chart.js defaults to dark grey text, which disappears on a near-black page).
    options: dict = {
        "responsive": True,
        "maintainAspectRatio": False,
        "plugins": {
            "title": {
                "display": True,
                "text": title,
                "font": {"size": 16},
                "color": "#e5e7eb",
            },
            "legend": {"labels": {"color": "#cbd5e1"}},
        },
    }
    if chart_type in ("bar", "line"):
        options["scales"] = {
            axis: {
                "ticks": {"color": "#94a3b8"},
                "grid": {"color": "rgba(148,163,184,0.15)"},
            }
            for axis in ("x", "y")
        }
    config = {
        "type": chart_type,
        "data": {"labels": labels, "datasets": datasets},
        "options": options,
    }

    # A self-contained canvas plus the script that instantiates it. The chat UI injects
    # each block's HTML and re-runs its <script>, so the chart draws itself. A unique id
    # keeps two charts in one report (or a repeated title) from colliding, and escaping
    # </ stops any string in the config from closing the <script> early.
    canvas_id = "cm-chart-" + uuid.uuid4().hex[:8]
    config_json = _json.dumps(config).replace("</", "<\\/")
    _add_block(
        '<div class="block chart-block" style="position:relative;height:340px;margin:18px 0;">'
        f'<canvas id="{canvas_id}"></canvas></div>'
        f"<script>try{{new Chart(document.getElementById('{canvas_id}'),"
        f"{config_json});}}catch(e){{console.error('chart {canvas_id}',e);}}</script>"
    )
    return f"chart {title!r} added to the report"

async def create_metric(label: str, value: str, delta: str = "") -> str:
    """Add a single KPI card (a big number with a label) to the report.

    Use for headline figures, e.g. latest price or period return. Consecutive
    metric cards lay out in a row. Blocks appear in the order the tools are called.

    Args:
        label: Short caption, e.g. "AAPL return".
        value: The formatted value to display, e.g. "$185.64" or "+12%".
        delta: Optional change note, e.g. "+8% vs last month".

    Returns:
        A short confirmation string.
    """
    # Always render the delta line (blank when there is no delta) so every card is the
    # same height whether or not a delta was passed, and a row of cards stays aligned.
    # Colors are tuned for the chat UI's dark theme.
    delta_html = f'<div style="font-size:12px;color:#94a3b8;margin-top:4px;">{delta or "&nbsp;"}</div>'
    _add_block(
        '<div class="block metric-card" style="display:inline-block;min-width:150px;margin:8px 10px 8px 0;'
        "padding:16px 20px;background:rgba(14,165,233,0.12);border:1px solid rgba(14,165,233,0.35);"
        'border-radius:14px;vertical-align:top;">'
        f'<div style="font-size:11px;color:#94a3b8;text-transform:uppercase;letter-spacing:0.06em;">{label}</div>'
        f'<div style="font-size:26px;font-weight:700;color:#7dd3fc;margin-top:4px;">{value}</div>'
        f"{delta_html}</div>"
    )
    return f"metric {label!r} added to the report"

async def create_table(title: str, headers: list, rows: list) -> str:
    """Add a data table to the report.

    Use for tabular breakdowns (e.g. per-ticker detail) where a chart would lose
    the exact numbers. Blocks appear in the order the tools are called.

    Args:
        title: Table caption shown above it.
        headers: Column names.
        rows: List of rows, each a list of cell values (same length as headers).

    Returns:
        A short confirmation string.
    """
    # Colors tuned for the chat UI's dark theme.
    head = "".join(
        f'<th style="text-align:left;padding:8px 12px;color:#7dd3fc;border-bottom:1px solid '
        f'rgba(14,165,233,0.4);">{h}</th>'
        for h in headers
    )
    body = "".join(
        "<tr>"
        + "".join(
            f'<td style="padding:8px 12px;border-bottom:1px solid rgba(148,163,184,0.15);">{c}</td>'
            for c in row
        )
        + "</tr>"
        for row in rows
    )
    _add_block(
        '<div class="block table-block" style="margin:16px 0;overflow-x:auto;">'
        f'<div style="font-size:13px;color:#94a3b8;margin-bottom:8px;">{title}</div>'
        '<table style="width:100%;border-collapse:collapse;font-size:14px;font-variant-numeric:tabular-nums;">'
        f"<thead><tr>{head}</tr></thead><tbody>{body}</tbody></table></div>"
    )
    return f"table {title!r} added to the report ({len(rows)} rows)"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/tools.py*

Letting DuckDB parse the SQL beats hand-matching keywords, which trips over identifiers and string literals. A rejected query comes back as an error that the agent loop feeds back to the model, so it rewrites the query as a proper SELECT. The same idea applies to any tool you add: accept the narrowest input that does the job.

## The durable environment

`query` becomes durable by living in a `flyte.TaskEnvironment`. The environment carries the image (DuckDB, pandas, the Anthropic SDK, the Monty sandbox package, the MCP client, and the Yahoo Finance MCP server) and the Anthropic API key as a `flyte.Secret`. The MCP server reads public market data, so it needs no credentials and the key is the only secret:

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE0\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

The `query` task itself is a thin wrapper over the tool function, and it carries `cache="auto"`:

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE1\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

Caching is scoped to `query`, not the whole environment. Given the same SQL over the same fetched series, the result is deterministic, so identical queries return instantly from cache and dedupe across conversations. The fetch itself is deliberately left off the cache: prices are live, so an MCP call the agent makes (not this task) is the right place for it. `analyze` is uncached too: it is a non-deterministic model turn keyed on the whole conversation, so it rarely repeats, and caching it could freeze a transient failure.

## The native agent

The agent is Flyte's built-in `Agent` from the [agent framework](../../../user-guide/agents/build-agent/_index) with `code_mode=True`. On each turn the model writes a Python program, the program runs in the sandbox via `orchestrate_local`, and the render tools populate the report as a side effect. When the report is done the model writes a one-line plain-text summary; a small shim stops there so a second turn can't re-run the same queries. Sandbox errors are fed back to the model automatically, so it fixes its own code within the turn budget.

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE2\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

Two things about that `tools` list are worth pausing on.

First, it mixes bindings. `query` is an `@env.task`, so the code-mode runtime passes it through as a task and every call the model writes dispatches as a durable child task. The render helpers are plain callables and run in-process. And the price fetch is a third kind: a tool the MCP server contributes, yet the model calls all three the same way. The split is about more than observability: if `query` ran in-process, a burst of questions, or one analysis firing several queries, would pile onto the single process handling the request. As a durable task, each query fans out to its own worker on the cluster, with retries for free, while the microsecond render helpers stay in-process where a round-trip would only add latency.

Second, the tools and the instructions are the whole definition of what this agent does, so what you get is a stock analyst, not a general assistant. Ask it to write a sorting function and it will not hand you one: it produces results by running code against the tools it has. That narrowness is a feature for a served app, since the behavior stays predictable and the surface stays small. Code mode itself does not impose the scope; the tools and the prompt do, so to widen or narrow the agent, you change those, not the machinery.

The agent generates its system prompt from the registry, introspecting each function's signature and docstring, so adding a tool is a matter of writing a function. The `instructions` add the data description and the report guidance on top.

The LLM callback uses the official Anthropic SDK. The agent's default callback goes through litellm, which works fine; supplying our own keeps the image lean and the API surface explicit:

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE3\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

> [!NOTE]
> This does not have to be Claude. The `call_llm` callback is the only model-specific code; everything around it is model-agnostic. Point it at any chat-completion endpoint, including an open model you host yourself, for example [an LLM served with vLLM](../../../user-guide/apps/native-app-integrations/vllm-app/page.md) as its own app right alongside this one. That keeps the data and the model on your own infrastructure and drops the per-call API cost in exchange for running the inference yourself.

## The report collector

The code-mode loop ends in a plain-text answer, but the app renders structured HTML blocks: metric cards, charts, tables. A per-run collector bridges the two. Each render tool appends its HTML as a side effect and returns a short confirmation (which keeps the sandbox observations small), and `analyze` reads the blocks back after the agent finishes. A `ContextVar` keeps concurrent runs isolated:

```
"""Tools and data access for the Code Mode stock-analysis agent.

The agent (``flyte.ai.agents.Agent`` in ``code_mode``) writes Python
orchestration code that calls these tools; that code runs in the Monty
sandbox, which allows no imports, no IO, and no network, so the only things the
generated code can touch are the tools registered in ``analysis.py``.

Two kinds of tools, on purpose:

* The **fetch** is a Yahoo Finance MCP tool (``yf_get_historical_stock_prices``),
  registered on the agent via ``mcp_servers`` in ``analysis.py``. It is the only
  path to the network — the sandbox has none — so it is the agent's live data
  source. It returns a raw JSON *string* of closing prices; the sandbox does not
  parse it (it has no ``json``), it just hands it to ``query``.
* ``query`` runs read-only DuckDB SQL over the fetched series. In ``analysis.py``
  it is a durable ``@env.task``, so the heavy analytics (moving averages,
  volatility, drawdowns, cross-ticker joins) run as a tracked, cached Flyte task.
  It parses the raw MCP strings into a ``prices`` table before running the SQL —
  the messy reshape lives here, where pandas is available, not in the sandbox.
* ``create_metric``, ``create_chart``, ``create_table``, and
  ``calculate_statistics`` are cheap, pure-Python helpers that run in-process.
  The ``create_*`` ones render HTML blocks into a per-run report collector.

To add a tool: write a function with type annotations and a docstring, then add
it to the agent's ``tools`` list in ``analysis.py``. The agent generates its
system prompt from the signatures and docstrings, so there is nothing else to
wire up.
"""

from __future__ import annotations

import contextvars
import datetime as _dt
import json as _json
import math
import uuid

CHART_COLORS = [
    "rgba(14, 165, 233, 0.8)",  # #0ea5e9 — sky
    "rgba(37, 99, 235, 0.8)",  # #2563eb — blue
    "rgba(6, 182, 212, 0.8)",  # #06b6d4 — cyan
    "rgba(99, 102, 241, 0.8)",  # #6366f1 — indigo
    "rgba(8, 145, 178, 0.8)",  # #0891b2 — deep cyan
]

CHART_BORDERS = ["#0ea5e9", "#2563eb", "#06b6d4", "#6366f1", "#0891b2"]

# Dataset — live stock closing prices, fetched via the Yahoo Finance MCP server
#
# There is no local data to fetch: the agent pulls prices at runtime from the
# `mcp-yahoo-finance` server (registered in `analysis.py`). This description is
# injected into the system prompt so the model knows how the two heavy tools fit
# together without a round-trip.

DATA_DESCRIPTION = (
    "You analyze daily stock closing prices. There are two heavy tools.\n"
    "\n"
    "Fetching (one ticker per call, via the Yahoo Finance MCP server):\n"
    "  yf_get_historical_stock_prices(symbol=..., period='1y', interval='1d')\n"
    "  returns a JSON *string* of closing prices keyed by timestamp. Do NOT parse\n"
    "  it in your code — the sandbox has no json or datetime module. Pass the\n"
    "  string straight to query(). Call it once per ticker (await each call) and\n"
    "  collect the returned strings into a dict for query(). Valid period: 1mo,\n"
    "  3mo, 6mo, 1y, 2y, 5y, ytd, max. Valid interval: 1d, 1wk, 1mo.\n"
    "\n"
    "Analyzing (durable DuckDB task):\n"
    "  query(sql, series) where `series` maps each ticker symbol to the JSON\n"
    "  string returned by yf_get_historical_stock_prices for it. The task parses\n"
    "  those into one table:\n"
    "     prices(ticker TEXT, date DATE, close DOUBLE)\n"
    "  Write a single read-only SELECT against `prices`. Do the math in SQL:\n"
    "  window functions (AVG(...) OVER (PARTITION BY ticker ORDER BY date ...))\n"
    "  for moving averages, LAG(...) for daily returns, STDDEV for volatility,\n"
    "  and GROUP BY / self-joins for cross-ticker comparisons."
)

def _jsonable(value: object) -> object:
    """Coerce DuckDB scalars to JSON-friendly Python types."""
    if isinstance(value, (_dt.date, _dt.datetime)):
        return value.isoformat()
    return value

# {{docs-fragment collector}}
# The native code-mode loop ends in a plain-text answer, but the UI renders
# structured HTML blocks. A per-run collector bridges the two: each render tool
# appends its HTML here as a side effect, and the `analyze` task reads the blocks
# back after the agent finishes. A ContextVar keeps concurrent runs isolated.
_REPORT: contextvars.ContextVar[list | None] = contextvars.ContextVar(
    "report", default=None
)

def start_report() -> None:
    """Begin a fresh report for this run (called by `analyze` before the agent)."""
    _REPORT.set([])

def collect_report() -> list[str]:
    """Return the HTML blocks rendered so far, in the order they were created."""
    return list(_REPORT.get() or [])

def _add_block(html: str) -> None:
    blocks = _REPORT.get()
    if blocks is not None:
        blocks.append(html)

# {{/docs-fragment collector}}

# {{docs-fragment sql_guard}}
# The tool is a safety boundary. The model can only call the tools you register, so
# narrowing what a tool accepts shrinks the blast radius. `query` allows a single
# read-only SELECT and nothing else. DuckDB's own parser classifies the statement, so
# there is no brittle keyword matching to trip over identifiers or string literals.
def _ensure_read_only(con, sql: str) -> None:
    import duckdb

    statements = con.extract_statements(sql)
    if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
        raise ValueError("Only a single read-only SELECT query is allowed.")

# {{/docs-fragment sql_guard}}

# {{docs-fragment query_tool}}
async def run_sql(sql: str, series: dict[str, str]) -> list:
    """Parse raw Yahoo Finance price JSON per ticker, then run a read-only query.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it (closing prices keyed by
                epoch-millisecond timestamp).

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """
    import duckdb
    import pandas as pd

    # Parse each ticker's raw MCP payload into rows and stack them into one table.
    # This reshape needs json + pandas, which the Monty sandbox lacks — so it runs
    # here, in the durable task, not in the model's generated code.
    frames = []
    for ticker, raw in series.items():
        data = _json.loads(raw) if raw else {}
        if not data:
            continue
        frame = pd.DataFrame({"ts": list(data.keys()), "close": list(data.values())})
        # The MCP keys its close prices by timestamp, but the format varies by
        # pandas version inside the server: ISO date strings ("2025-07-03") or
        # epoch-millisecond integers. Detect which and parse accordingly.
        ts = frame["ts"].astype(str)
        if ts.str.fullmatch(r"\d+").all():
            frame["date"] = pd.to_datetime(ts.astype("int64"), unit="ms").dt.date
        else:
            frame["date"] = pd.to_datetime(ts).dt.date
        frame["ticker"] = ticker
        frames.append(frame[["ticker", "date", "close"]])

    prices = (
        pd.concat(frames, ignore_index=True)
        if frames
        else pd.DataFrame(columns=["ticker", "date", "close"])
    )

    # Lock the engine down: no reading or writing files, no extensions, no network.
    con = duckdb.connect(config={"enable_external_access": "false"})
    _ensure_read_only(con, sql)

    con.register("prices", prices)
    rel = con.execute(sql)
    columns = [d[0] for d in rel.description]
    return [{c: _jsonable(v) for c, v in zip(columns, row)} for row in rel.fetchall()]

# {{/docs-fragment query_tool}}

async def calculate_statistics(rows: list, column: str) -> dict:
    """Calculate descriptive statistics for a numeric column of query rows.

    Args:
        rows: A list of row dicts, e.g. the output of query().
        column: Name of the numeric column to analyze.

    Returns:
        Dict with keys: count, mean, median, min, max, std_dev.
    """
    vals = [row[column] for row in rows if column in row and row[column] is not None]
    if not vals:
        return {"count": 0, "mean": 0, "median": 0, "min": 0, "max": 0, "std_dev": 0}
    n = len(vals)
    mean = sum(vals) / n
    ordered = sorted(vals)
    median = (
        ordered[n // 2] if n % 2 == 1 else (ordered[n // 2 - 1] + ordered[n // 2]) / 2
    )
    variance = sum((v - mean) ** 2 for v in vals) / n
    return {
        "count": n,
        "mean": round(mean, 2),
        "median": round(median, 2),
        "min": min(vals),
        "max": max(vals),
        "std_dev": round(math.sqrt(variance), 2),
    }

async def create_chart(chart_type: str, title: str, labels: list, values: list) -> str:
    """Add a chart to the report (rendered with Chart.js in the UI).

    Blocks appear in the report in the order the create_* tools are called.

    Args:
        chart_type: One of "bar", "line", "pie", "doughnut".
        title: Chart title displayed above the canvas.
        labels: X-axis labels (or slice labels for pie/doughnut).
        values: Either a flat list of numbers, or a list of
                {"label": str, "data": list[number]} dicts for multi-series.

    Returns:
        A short confirmation string.
    """
    if not values:
        return f"chart {title!r} skipped: no data to plot"

    if isinstance(values[0], dict):
        datasets = []
        for i, series in enumerate(values):
            idx = i % len(CHART_COLORS)
            datasets.append(
                {
                    "label": series["label"],
                    "data": series["data"],
                    "backgroundColor": CHART_COLORS[idx],
                    "borderColor": CHART_BORDERS[idx],
                    "borderWidth": 2,
                    "tension": 0.3,
                    "fill": False,
                }
            )
    else:
        bg = [CHART_COLORS[i % len(CHART_COLORS)] for i in range(len(values))]
        border = [CHART_BORDERS[i % len(CHART_BORDERS)] for i in range(len(values))]
        datasets = [
            {
                "label": title,
                "data": values,
                "backgroundColor": (
                    bg if chart_type in ("pie", "doughnut") else CHART_COLORS[0]
                ),
                "borderColor": (
                    border if chart_type in ("pie", "doughnut") else CHART_BORDERS[0]
                ),
                "borderWidth": 2,
                "tension": 0.3,
                "fill": chart_type == "line",
            }
        ]

    # Light text and faint grid lines so the chart reads on the chat UI's dark theme
    # (Chart.js defaults to dark grey text, which disappears on a near-black page).
    options: dict = {
        "responsive": True,
        "maintainAspectRatio": False,
        "plugins": {
            "title": {
                "display": True,
                "text": title,
                "font": {"size": 16},
                "color": "#e5e7eb",
            },
            "legend": {"labels": {"color": "#cbd5e1"}},
        },
    }
    if chart_type in ("bar", "line"):
        options["scales"] = {
            axis: {
                "ticks": {"color": "#94a3b8"},
                "grid": {"color": "rgba(148,163,184,0.15)"},
            }
            for axis in ("x", "y")
        }
    config = {
        "type": chart_type,
        "data": {"labels": labels, "datasets": datasets},
        "options": options,
    }

    # A self-contained canvas plus the script that instantiates it. The chat UI injects
    # each block's HTML and re-runs its <script>, so the chart draws itself. A unique id
    # keeps two charts in one report (or a repeated title) from colliding, and escaping
    # </ stops any string in the config from closing the <script> early.
    canvas_id = "cm-chart-" + uuid.uuid4().hex[:8]
    config_json = _json.dumps(config).replace("</", "<\\/")
    _add_block(
        '<div class="block chart-block" style="position:relative;height:340px;margin:18px 0;">'
        f'<canvas id="{canvas_id}"></canvas></div>'
        f"<script>try{{new Chart(document.getElementById('{canvas_id}'),"
        f"{config_json});}}catch(e){{console.error('chart {canvas_id}',e);}}</script>"
    )
    return f"chart {title!r} added to the report"

async def create_metric(label: str, value: str, delta: str = "") -> str:
    """Add a single KPI card (a big number with a label) to the report.

    Use for headline figures, e.g. latest price or period return. Consecutive
    metric cards lay out in a row. Blocks appear in the order the tools are called.

    Args:
        label: Short caption, e.g. "AAPL return".
        value: The formatted value to display, e.g. "$185.64" or "+12%".
        delta: Optional change note, e.g. "+8% vs last month".

    Returns:
        A short confirmation string.
    """
    # Always render the delta line (blank when there is no delta) so every card is the
    # same height whether or not a delta was passed, and a row of cards stays aligned.
    # Colors are tuned for the chat UI's dark theme.
    delta_html = f'<div style="font-size:12px;color:#94a3b8;margin-top:4px;">{delta or "&nbsp;"}</div>'
    _add_block(
        '<div class="block metric-card" style="display:inline-block;min-width:150px;margin:8px 10px 8px 0;'
        "padding:16px 20px;background:rgba(14,165,233,0.12);border:1px solid rgba(14,165,233,0.35);"
        'border-radius:14px;vertical-align:top;">'
        f'<div style="font-size:11px;color:#94a3b8;text-transform:uppercase;letter-spacing:0.06em;">{label}</div>'
        f'<div style="font-size:26px;font-weight:700;color:#7dd3fc;margin-top:4px;">{value}</div>'
        f"{delta_html}</div>"
    )
    return f"metric {label!r} added to the report"

async def create_table(title: str, headers: list, rows: list) -> str:
    """Add a data table to the report.

    Use for tabular breakdowns (e.g. per-ticker detail) where a chart would lose
    the exact numbers. Blocks appear in the order the tools are called.

    Args:
        title: Table caption shown above it.
        headers: Column names.
        rows: List of rows, each a list of cell values (same length as headers).

    Returns:
        A short confirmation string.
    """
    # Colors tuned for the chat UI's dark theme.
    head = "".join(
        f'<th style="text-align:left;padding:8px 12px;color:#7dd3fc;border-bottom:1px solid '
        f'rgba(14,165,233,0.4);">{h}</th>'
        for h in headers
    )
    body = "".join(
        "<tr>"
        + "".join(
            f'<td style="padding:8px 12px;border-bottom:1px solid rgba(148,163,184,0.15);">{c}</td>'
            for c in row
        )
        + "</tr>"
        for row in rows
    )
    _add_block(
        '<div class="block table-block" style="margin:16px 0;overflow-x:auto;">'
        f'<div style="font-size:13px;color:#94a3b8;margin-bottom:8px;">{title}</div>'
        '<table style="width:100%;border-collapse:collapse;font-size:14px;font-variant-numeric:tabular-nums;">'
        f"<thead><tr>{head}</tr></thead><tbody>{body}</tbody></table></div>"
    )
    return f"table {title!r} added to the report ({len(rows)} rows)"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/tools.py*

## Live prices over MCP

External tools plug in over MCP with `MCPServerSpec`. The agent connects on first use, lists the server's tools, and registers each one alongside the local tools; the model calls them from its generated code like any other function. The example wires the `mcp-yahoo-finance` server, launched as a subprocess in the task pod, with no credentials since it reads public market data:

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE4\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

`tool_prefix` namespaces the server's tools (`yf_get_historical_stock_prices`) to avoid collisions, and `tool_filter` narrows the server's twelve tools down to the one the analytics needs, the same surface-shrinking move as the SQL guard: the model never even sees the tools it should not use, and the prompt stays small. Because the sandbox has no `json`, the model never parses the tool's raw output. It passes the string straight to `query`, which does the reshape. Fetching a ticker becomes one more function call in the model's program, sitting alongside the durable query.

## The analysis task

`analyze` ties it together: start a fresh report, run the agent with the chat history as its memory (so a follow-up like "now compare it with two peers" refers back to earlier turns), collect the blocks, and return them with the summary. It also appends a link to its own run, so every answer carries a click-through to the task graph that produced it:

```
"""The durable half: the code-mode analysis task, built on ``flyte.ai.agents``.

The agent is Flyte's native :class:`~flyte.ai.agents.Agent` with ``code_mode=True``:
on each turn the model writes a small Python program, the program runs in the Monty
sandbox, and the tools are exposed to it as plain functions. The Yahoo Finance MCP
server supplies the live price fetch; the durable ``query`` task runs the DuckDB
analytics on the cluster; the render helpers run in-process and stream their HTML
into the report collector in ``tools.py``.

Kept separate from ``app.py`` on purpose. This module runs in the task image (which
has anthropic / duckdb / monty / the MCP client but not the web layer), and
``app.py`` serves it.
"""

from __future__ import annotations

from typing import Any

import flyte
import flyte.remote
from flyte.ai.agents import Agent, LLMMessage, MCPServerSpec
from flyte.ai.agents._code import build_sandbox_tools, extract_python_code
from flyte.ai.agents._tools import _abbreviate
from flyte.ai.agents.agent import AgentEvent, _TurnResult, _emit
from flyte.ai.agents.protocol import AgentResult

import tools

# {{docs-fragment env}}
ANTHROPIC_SECRET = flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")

# The analysis environment: the agent runs here, and `query` is a durable task in
# the same environment, so the sandboxed code's query calls dispatch as child tasks.
# The Yahoo Finance MCP server needs no credentials (public data), so the only secret
# is the Anthropic key.
env = flyte.TaskEnvironment(
    name="code-mode",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyte[mcp]",
        "anthropic",
        "pydantic-monty",
        "duckdb>=1.1.0",
        "pandas",
        "mcp-yahoo-finance",
    ),
    secrets=[ANTHROPIC_SECRET],
)
# {{/docs-fragment env}}

# {{docs-fragment query_task}}
# Cache the analytics: given the same SQL over the same fetched series, the result is
# deterministic, so identical queries dedupe across conversations. The fetch itself is
# live and is not cached — it is an MCP tool call the agent makes, not this task.
@env.task(cache="auto")
async def _query_task(sql: str, series: dict[str, str]) -> list[dict]:
    return await tools.run_sql(sql, series)

async def query(sql: str, series: dict[str, str]) -> list[dict]:
    """Run a read-only SQL query over fetched stock prices and return rows.

    Args:
        sql: A DuckDB SELECT statement against the table `prices`
             (columns: ticker, date, close). Aggregate in SQL where you can.
        series: Maps ticker symbol -> the JSON string returned by
                yf_get_historical_stock_prices for it. Pass the raw strings; the
                durable task parses them into the `prices` table.

    Returns:
        A list of row dicts (one per result row), with dates as ISO strings.
    """

    return await _query_task(sql, series)

# {{/docs-fragment query_task}}

# {{docs-fragment llm}}
async def call_llm(
    model: str, system: str, messages: list[dict], tools_schema: list[dict] | None
) -> LLMMessage:
    """LLM callback for the agent, using the official Anthropic SDK.

    The agent's default callback goes through litellm; supplying our own keeps the
    image lean and the API surface explicit. In code mode `tools_schema` is None
    (tools are called from generated code, not via JSON tool-calling).
    """

    from anthropic import AsyncAnthropic

    client = AsyncAnthropic()  # reads ANTHROPIC_API_KEY, injected as a Flyte secret
    resp = await client.messages.create(
        model=model,
        max_tokens=4096,
        thinking={"type": "adaptive"},
        system=system,
        messages=messages,
    )
    text = "".join(block.text for block in resp.content if block.type == "text")
    return LLMMessage(content=text)

# {{/docs-fragment llm}}

# {{docs-fragment mcp}}
def _mcp_servers() -> list[MCPServerSpec]:
    """The Yahoo Finance MCP server — the agent's live price source.

    The agent connects on first use, lists the server's tools, and registers each
    one alongside the local tools; the model calls them from its generated code
    like any other function. `tool_prefix` namespaces them, and `tool_filter`
    narrows the server's 12 tools down to the one the analytics needs, the same
    surface-shrinking move as the SQL guard. No auth: the server reads public
    Yahoo Finance data, so there is no secret to inject.
    """
    return [
        MCPServerSpec(
            name="yahoo-finance",
            command=["mcp-yahoo-finance"],
            tool_prefix="yf_",
            tool_filter=["get_historical_stock_prices"],
        )
    ]

# {{/docs-fragment mcp}}

INSTRUCTIONS = f"""\
You are a stock-market data analyst in a chat. Answer questions by writing one
complete Python program that fetches the price history you need and assembles a report.

{tools.DATA_DESCRIPTION}

How to build the report:
- IMPORTANT: the user only sees what you RENDER. The value your code returns and the
  rows from query(...) are NOT shown to them. You must turn your findings into report
  blocks with create_metric / create_chart / create_table, or the user sees nothing. A
  reply that describes a chart without calling create_chart shows an empty answer.
- Fetch each ticker the question needs with yf_get_historical_stock_prices(...). When
  you need more than one, call it once per ticker (await each call). Pass the raw JSON
  strings straight through — do not parse them (the sandbox has no json/datetime).
- Do all the analytics in SQL via query(sql, series): build series as
  {{"AAPL": aapl_json, "MSFT": msft_json}} and write one SELECT against the `prices`
  table (columns ticker, date, close). Use window functions, LAG, and STDDEV — do not
  compute these by hand in Python.
- Build a report, not just one chart. Lead with one or two headline numbers via
  create_metric(...), then a create_chart(...) for the trend, and a create_table(...)
  when the exact figures matter. Use the tools that fit the question.
- The create_* tools add blocks to the report in the order you call them; they return
  short confirmations, not HTML.
- Fetch each ticker once and run each query once; reuse the returned rows for every
  metric, chart, and table.
- After one successful code block has created the report, stop writing code and give
  the final plain-text summary. Do not re-run the analysis in a second code block
  unless the prior code failed.
- Format numbers with f-strings, e.g. f"${{x:.2f}}" or f"{{r:.1%}}". The format() builtin
  and the {{:,}} thousands separator are not available in the sandbox.
- Prefer ONE code block that does everything: fetch, query, render. After it runs,
  reply with a one-or-two-sentence plain-text summary of what the data shows.
- Your final reply is rendered as Markdown. Write "about 12%", never "~12%": a pair of
  ~ characters renders as strikethrough.
- For a greeting or a question that needs no data, just reply in plain text.
"""

class CodeModeAgent(Agent):
    """Agent shim for flyte 2.5.7 code-mode edge cases used by this tutorial."""

    async def _run_code_mode(
        self,
        message: str,
        memory: Any = None,
    ) -> AgentResult:
        import flyte.sandbox

        # flyte 2.5.7 loads MCP inside _run_loop, after code mode has already
        # snapshotted sandbox_tools. Load first so the yf_* MCP tools enter Monty's
        # namespace and the generated code can call them.
        await self._ensure_mcp_loaded()
        sandbox_tools = build_sandbox_tools(
            self._registry, call_llm=self.call_llm, model=self.model
        )
        last_code = ""
        sandbox_runs = 0
        report_created = False
        render_nudged = False

        async def step(
            llm_msg: LLMMessage, messages: list[dict[str, Any]], attempts: int
        ) -> _TurnResult:
            nonlocal last_code, sandbox_runs, report_created, render_nudged
            text = llm_msg.content or ""
            messages.append({"role": "assistant", "content": text})
            await _emit(AgentEvent("message", {"role": "assistant", "content": text}))

            code = extract_python_code(text)

            # Once the report exists, the model's next message is its plain-text
            # summary. Ignore any further code — the report is done and we do not
            # want a second run re-executing the same queries.
            if report_created:
                summary = text if not code else "Done. The report is above."
                await _emit(AgentEvent("turn_end", {"turn": attempts, "summary": True}))
                return _TurnResult(done=True, final_text=summary)

            if not code:
                # No report and no code: a greeting or a question needing no data.
                await _emit(
                    AgentEvent(
                        "turn_end",
                        {"turn": attempts, "had_code": False, "text_len": len(text)},
                    )
                )
                return _TurnResult(done=True, final_text=text)

            last_code = code
            sandbox_runs += 1
            await _emit(AgentEvent("tool_start", {"tool": "<sandbox>", "code": code}))
            try:
                with flyte.group(f"{self.name}-sandbox-{sandbox_runs}"):
                    result = await flyte.sandbox.orchestrate_local(
                        code,
                        inputs={"_unused": 0},
                        tasks=sandbox_tools,
                    )
            except Exception as exc:
                await _emit(
                    AgentEvent("tool_error", {"tool": "<sandbox>", "error": str(exc)})
                )
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            f"Your code raised an error:\n\nCODE5\n\n"
                            "Fix the code and try again, respecting the Monty sandbox restrictions."
                        ),
                    }
                )
                await _emit(
                    AgentEvent(
                        "turn_end", {"turn": attempts, "had_code": True, "error": True}
                    )
                )
                return _TurnResult(done=False)

            await _emit(
                AgentEvent(
                    "tool_end", {"tool": "<sandbox>", "result": _abbreviate(result)}
                )
            )
            await _emit(
                AgentEvent(
                    "turn_end",
                    {"turn": attempts, "had_code": True, "final_after_code": True},
                )
            )
            # Only treat the turn as done when the render tools actually produced
            # report blocks. If the model computed a result but rendered nothing,
            # the user would see an empty answer (the query rows are not shown), and
            # asking for a summary here would make the model narrate a report that
            # does not exist. So check the collector and nudge it to render first.
            blocks = tools.collect_report()
            if blocks:
                report_created = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "The report has been created and is shown to the user. "
                            "Reply with a one or two sentence plain-text summary of "
                            "what the data shows. Do not write any more code."
                        ),
                    }
                )
                return _TurnResult(done=False)
            if not render_nudged:
                render_nudged = True
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "Your code ran but added nothing to the report, so the "
                            "user sees no result — the query rows are not displayed "
                            "automatically. Call create_metric / create_chart / "
                            "create_table to render the findings, then stop."
                        ),
                    }
                )
                return _TurnResult(done=False)
            # Rendered nothing even after a nudge: end honestly rather than claim a
            # report that was never built.
            return _TurnResult(
                done=True,
                final_text="I ran the analysis but did not produce a visual report.",
            )

        outcome = await self._run_loop(
            message, memory, tools_schema=None, step=step, mode="code"
        )
        return AgentResult(
            code=last_code,
            summary=outcome.last_text,
            error=outcome.error_msg,
            attempts=outcome.attempts,
            memory=outcome.memory,
        )

# {{docs-fragment agent}}
agent = CodeModeAgent(
    name="code-mode-analyst",
    instructions=INSTRUCTIONS,
    model="claude-opus-4-8",
    # One list, two kinds of local tools: `query` awaits an @env.task, so the sandbox
    # dispatches the DuckDB analytics as a durable child task; the render helpers are
    # plain callables and run in-process. The live price fetch is a *third* kind — an
    # MCP tool contributed by `mcp_servers` below — but the model calls all of them the
    # same way. The agent introspects signatures and docstrings to build its prompt.
    tools=[
        query,
        tools.create_metric,
        tools.create_chart,
        tools.create_table,
        tools.calculate_statistics,
    ],
    mcp_servers=_mcp_servers(),
    code_mode=True,
    # Turn 1 writes the program; the next turn is the plain-text summary. The
    # spare turns let the agent fix its code if the sandbox rejects it.
    max_turns=5,
    call_llm=call_llm,
)
# {{/docs-fragment agent}}

# Shows up in the task logs, so a deployment is easy to spot as MCP-enabled.
print(f"Yahoo Finance MCP: {len(agent.mcp_servers)} server(s) configured")

async def _run_link_block() -> str:
    """A small HTML block linking to this run in the UI (best effort)."""
    tctx = flyte.ctx()
    if tctx is None or not tctx.action.run_name:
        return ""
    try:
        run = await flyte.remote.Run.get.aio(tctx.action.run_name)
        url = run.url
    except Exception:
        return ""
    # Inline light-sky color so the link stays readable on the chat UI's dark theme
    # (an unstyled anchor inherits a dark blue that disappears on the near-black page).
    return (
        '<div class="block" style="margin:10px 0;font-size:13px;">'
        f'<a href="{url}" target="_blank" rel="noopener" '
        'style="color:#7dd3fc;text-decoration:underline;">'
        "View this analysis run in the Union UI &#8599;</a></div>"
    )

# {{docs-fragment analyze}}
@env.task
async def analyze(message: str, history: list[dict[str, str]]) -> dict:
    """Run one analysis: start a report, run the agent, return blocks + summary.

    `history` is the prior conversation, which `Agent.run` takes as its memory, so
    follow-ups can refer back to earlier turns. This task is the chat app's
    `task_entrypoint`: each question becomes a run, and inside it the sandbox's
    `query` calls dispatch as durable child tasks.
    """

    tools.start_report()
    result = await agent.run.aio(message, memory=list(history))
    blocks = tools.collect_report()
    if link := await _run_link_block():
        blocks.append(link)
    # The UI renders the summary as Markdown, where a pair of ~ characters becomes
    # strikethrough. Models like ~ as shorthand for "approximately", so escape it.
    summary = result.summary.replace("~", "\\~")
    return {
        "summary": summary,
        "charts": blocks,
        "code": result.code,
        "error": result.error,
        "attempts": result.attempts,
    }

# {{/docs-fragment analyze}}

if __name__ == "__main__":
    # Run one analysis as a durable flyte.run (no app) — handy for testing the
    # analysis half on its own. Remote image builder so no local Docker is needed.
    flyte.init_from_config(image_builder="remote")
    run = flyte.run(
        analyze, message="Compare AAPL and MSFT over the last 6 months", history=[]
    )
    print(f"View at: {run.url}")
    run.wait()
    print(run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/analysis.py*

Because `analyze` runs inside a task context, the `query` calls made by the sandboxed code dispatch as durable child tasks. You can run this half on its own with `python analysis.py`, which submits one analysis as a `flyte.run` and prints the run URL, no app required.

## Serving it: the native chat app

The web layer is one declaration. `AgentChatAppEnvironment` from the [agent chat UI](../../../user-guide/agents/build-agent/agent-chat-ui/page.md) provides the chat interface, the tools sidebar, progress streaming, and the chat endpoint:

```
"""Serve the Code Mode stock analyst with Flyte's native chat app.

``AgentChatAppEnvironment`` provides the whole web layer: the chat UI, the
``/api/chat`` endpoint, progress streaming, and the tools sidebar. Pointing its
``task_entrypoint`` at the ``analyze`` task makes every question a durable Flyte
run, and ``passthrough_auth=True`` forwards the caller's credentials so those
runs launch as the signed-in user (no service identity or org plumbing needed).

The agent pulls live prices from the Yahoo Finance MCP server (no credentials
needed) and runs the DuckDB analytics as a durable task.

Run::

    flyte create secret anthropic_api_key <your-anthropic-key>
    python app.py
"""

import flyte
import flyte.app
from flyte.ai.chat import AgentChatAppEnvironment, CustomTheme

from analysis import agent, analyze, env as agent_env

_prompt_nudges = [
    {
        "label": "Compare two stocks",
        "prompt": "Compare AAPL and MSFT over the last year — normalized price trend and volatility.",
    },
    {
        "label": "Trend + moving average",
        "prompt": "Show NVDA's closing price with a 50-day moving average for the last year.",
    },
    {
        "label": "Best performer",
        "prompt": "Which of AAPL, MSFT, GOOGL and AMZN had the best 6-month return?",
    },
    {
        "label": "Volatility ranking",
        "prompt": "Rank AAPL, TSLA and NVDA by 3-month volatility.",
    },
]

# {{docs-fragment chat_app}}
env = AgentChatAppEnvironment(
    name="code-mode-analytics",
    agent=agent,  # powers the tools sidebar
    # Each question is launched as a durable run of `analyze` (with the chat
    # history), so the sandbox's query calls dispatch as tracked child tasks.
    task_entrypoint=analyze,
    # Run those tasks with the caller's forwarded credentials.
    passthrough_auth=True,
    title="Code Mode Stock Analytics",
    subtitle=(
        "Chat with live stock prices. The model writes one Python program per "
        "question; it fetches prices from the Yahoo Finance MCP server and runs "
        "the heavy queries as durable Flyte tasks."
    ),
    prompt_nudges=_prompt_nudges,
    theme=CustomTheme(
        accent_color="#0ea5e9",
        accent_hover_color="#0284c7",
        button_text_color="#ffffff",
    ),
    image=flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn"),
    scaling=flyte.app.Scaling(replicas=1),
    depends_on=[agent_env],
    # Every request launches a run (compute + a paid LLM call), so gate the app
    # behind platform auth.
    requires_auth=True,
)
# {{/docs-fragment chat_app}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    # Remote image builder so no local Docker is needed to build the app + task images.
    flyte.init_from_config(image_builder="remote")

    handle = flyte.serve(env)
    print(f"Deployed Code Mode Stock Analytics: {handle.url}")
# {{/docs-fragment deploy}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code_mode_agent/app.py*

Two parameters do the architectural work.

`task_entrypoint=analyze` makes each question a durable run. An app's request handler has no task context, so calling a task directly from it would run the task locally in the app pod and you would lose durability and the child-task graph. With a task entrypoint, the chat endpoint launches `analyze` with `flyte.run` (passing the message and the history), streams the run's phase changes to the UI as progress, and renders the returned blocks. For more on apps and tasks calling each other, see [hybrid app-task graphs](../../../user-guide/apps/build-apps/hybrid-graphs/page.md).

`passthrough_auth=True` forwards each caller's credentials to those runs, so the analysis executes as the signed-in user rather than as a shared service identity, and the app needs no credential plumbing of its own. Together with `requires_auth=True`, which gates the app at the platform gateway, every request is authenticated end to end, which matters because each one launches real compute and a paid LLM call.

The rest is presentation: a theme for the accent colors, prompt nudges shown before the first message (ready-made comparisons like "Compare AAPL and MSFT over the last year"), and a title and subtitle.

## Deploy and run

Deploying is one command. The entry point uses the remote image builder so no local Docker is needed, and serves the app and its task environment together:

```
"""Serve the Code Mode stock analyst with Flyte's native chat app.

``AgentChatAppEnvironment`` provides the whole web layer: the chat UI, the
``/api/chat`` endpoint, progress streaming, and the tools sidebar. Pointing its
``task_entrypoint`` at the ``analyze`` task makes every question a durable Flyte
run, and ``passthrough_auth=True`` forwards the caller's credentials so those
runs launch as the signed-in user (no service identity or org plumbing needed).

The agent pulls live prices from the Yahoo Finance MCP server (no credentials
needed) and runs the DuckDB analytics as a durable task.

Run::

    flyte create secret anthropic_api_key <your-anthropic-key>
    python app.py
"""

import flyte
import flyte.app
from flyte.ai.chat import AgentChatAppEnvironment, CustomTheme

from analysis import agent, analyze, env as agent_env

_prompt_nudges = [
    {
        "label": "Compare two stocks",
        "prompt": "Compare AAPL and MSFT over the last year — normalized price trend and volatility.",
    },
    {
        "label": "Trend + moving average",
        "prompt": "Show NVDA's closing price with a 50-day moving average for the last year.",
    },
    {
        "label": "Best performer",
        "prompt": "Which of AAPL, MSFT, GOOGL and AMZN had the best 6-month return?",
    },
    {
        "label": "Volatility ranking",
        "prompt": "Rank AAPL, TSLA and NVDA by 3-month volatility.",
    },
]

# {{docs-fragment chat_app}}
env = AgentChatAppEnvironment(
    name="code-mode-analytics",
    agent=agent,  # powers the tools sidebar
    # Each question is launched as a durable run of `analyze` (with the chat
    # history), so the sandbox's query calls dispatch as tracked child tasks.
    task_entrypoint=analyze,
    # Run those tasks with the caller's forwarded credentials.
    passthrough_auth=True,
    title="Code Mode Stock Analytics",
    subtitle=(
        "Chat with live stock prices. The model writes one Python program per "
        "question; it fetches prices from the Yahoo Finance MCP server and runs "
        "the heavy queries as durable Flyte tasks."
    ),
    prompt_nudges=_prompt_nudges,
    theme=CustomTheme(
        accent_color="#0ea5e9",
        accent_hover_color="#0284c7",
        button_text_color="#ffffff",
    ),
    image=flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn"),
    scaling=flyte.app.Scaling(replicas=1),
    depends_on=[agent_env],
    # Every request launches a run (compute + a paid LLM call), so gate the app
    # behind platform auth.
    requires_auth=True,
)
# {{/docs-fragment chat_app}}

# {{docs-fragment deploy}}
if __name__ == "__main__":
    # Remote image builder so no local Docker is needed to build the app + task images.
    flyte.init_from_config(image_builder="remote")

    handle = flyte.serve(env)
    print(f"Deployed Code Mode Stock Analytics: {handle.url}")
# {{/docs-fragment deploy}}
CODE6shell
flyte create secret anthropic_api_key <your-anthropic-key>
python app.py
```

Open the printed URL and ask something like "Compare AAPL and MSFT over the last 6 months" or "Rank the FAANG stocks by 6-month return." The first question is slower as the task image builds and the MCP server cold-starts, then each answer streams progress while the run executes, and comes back as a short report of headline numbers, a chart, and sometimes a table, with the generated code and a link to the run so you can see the query tasks it dispatched.

## Going further

- **More history, more tickers.** The queries are ordinary DuckDB SQL over whatever the model fetches. Because `query` is a durable task, large or slow queries get retries and caching for free.
- **More MCP servers.** Add another `MCPServerSpec` for web search, Slack, or a ticketing system. Use `tool_filter` to expose only the tools the agent should have.
- **Add a model-based tool.** A tool that calls another model, such as an LLM judge or an embedder, registers like any other. Cheap tools stay in-process, and expensive ones become tasks.
- **More tools.** Write a function with a docstring and add it to the agent's `tools` list. The prompt regenerates from the signatures, so there is nothing else to wire up.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering ===

# Context engineering

Tutorials for prompt engineering, prompt optimization, and context construction.

### **Context engineering > Automatic prompt engineering**

Easily run prompt optimization with real-time observability, traceability, and automatic recovery.

### **Context engineering > Text-to-SQL prompt optimization**

Learn how to turn natural language questions into SQL queries with Flyte and LlamaIndex, and explore prompt optimization in practice.

## Subpages
- **Context engineering > Text-to-SQL prompt optimization**
- **Context engineering > Automatic prompt engineering**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/text_to_sql ===

# Text-to-SQL prompt optimization

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/text_to_sql); based on work by [LlamaIndex](https://docs.llamaindex.ai/en/stable/examples/workflow/advanced_text_to_sql/).

Data analytics drives modern decision-making, but SQL often creates a bottleneck. Writing queries requires technical expertise, so non-technical stakeholders must rely on data teams. That translation layer slows everyone down.

Text-to-SQL narrows this gap by turning natural language into executable SQL queries. It lowers the barrier to structured data and makes databases accessible to more people.

In this tutorial, we build a Text-to-SQL workflow using LlamaIndex and evaluate it on the [WikiTableQuestions dataset](https://ppasupat.github.io/WikiTableQuestions/) (a benchmark of natural language questions over semi-structured tables). We then explore prompt optimization to see whether it improves accuracy and show how to track prompts and results over time. Along the way, we'll see what worked, what didn't, and what we learned about building durable evaluation pipelines. The pattern here can be adapted to your own datasets and workflows.

![Evaluation](../../../_static/images/tutorials/text-to-sql/evaluation.png)

## Ingesting data

We start by ingesting the WikiTableQuestions dataset, which comes as CSV files, into a SQLite database. This database serves as the source of truth for our Text-to-SQL pipeline.

```
import asyncio
import fnmatch
import os
import re
import zipfile

import flyte
import pandas as pd
import requests
from flyte.io import Dir, File
from llama_index.core.llms import ChatMessage
from llama_index.core.prompts import ChatPromptTemplate
from llama_index.llms.openai import OpenAI
from pydantic import BaseModel, Field
from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine
from utils import env

# {{docs-fragment table_info}}
class TableInfo(BaseModel):
    """Information regarding a structured table."""

    table_name: str = Field(..., description="table name (underscores only, no spaces)")
    table_summary: str = Field(
        ..., description="short, concise summary/caption of the table"
    )

# {{/docs-fragment table_info}}

@env.task
async def download_and_extract(zip_path: str, search_glob: str) -> Dir:
    """Download and extract the dataset zip file if not already available."""
    output_zip = "data.zip"
    extract_dir = "wiki_table_questions"

    if not os.path.exists(zip_path):
        response = requests.get(zip_path, stream=True)
        with open(output_zip, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
    else:
        output_zip = zip_path
        print(f"Using existing file {output_zip}")

    os.makedirs(extract_dir, exist_ok=True)
    with zipfile.ZipFile(output_zip, "r") as zip_ref:
        for member in zip_ref.namelist():
            if fnmatch.fnmatch(member, search_glob):
                zip_ref.extract(member, extract_dir)

    remote_dir = await Dir.from_local(extract_dir)
    return remote_dir

async def read_csv_file(
    csv_file: File, nrows: int | None = None
) -> pd.DataFrame | None:
    """Safely download and parse a CSV file into a DataFrame."""
    try:
        local_csv_file = await csv_file.download()
        return pd.read_csv(local_csv_file, nrows=nrows)
    except Exception as e:
        print(f"Error parsing {csv_file.path}: {e}")
        return None

def sanitize_column_name(col_name: str) -> str:
    """Sanitize column names by replacing spaces/special chars with underscores."""
    return re.sub(r"\W+", "_", col_name)

async def create_table_from_dataframe(
    df: pd.DataFrame, table_name: str, engine, metadata_obj
):
    """Create a SQL table from a Pandas DataFrame."""
    # Sanitize column names
    sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}
    df = df.rename(columns=sanitized_columns)

    # Define table columns based on DataFrame dtypes
    columns = [
        Column(col, String if dtype == "object" else Integer)
        for col, dtype in zip(df.columns, df.dtypes)
    ]

    table = Table(table_name, metadata_obj, *columns)

    # Create table in database
    metadata_obj.create_all(engine)

    # Insert data into table
    with engine.begin() as conn:
        for _, row in df.iterrows():
            conn.execute(table.insert().values(**row.to_dict()))

@flyte.trace
async def create_table(
    csv_file: File, table_info: TableInfo, database_path: str
) -> str:
    """Safely create a table from CSV if parsing succeeds."""
    df = await read_csv_file(csv_file)
    if df is None:
        return "false"

    print(f"Creating table: {table_info.table_name}")

    engine = create_engine(f"sqlite:///{database_path}")
    metadata_obj = MetaData()

    await create_table_from_dataframe(df, table_info.table_name, engine, metadata_obj)
    return "true"

@flyte.trace
async def llm_structured_predict(
    df_str: str,
    table_names: list[str],
    prompt_tmpl: ChatPromptTemplate,
    feedback: str,
    llm: OpenAI,
) -> TableInfo:
    return llm.structured_predict(
        TableInfo,
        prompt_tmpl,
        feedback=feedback,
        table_str=df_str,
        exclude_table_name_list=str(list(table_names)),
    )

async def generate_unique_table_info(
    df_str: str,
    table_names: list[str],
    prompt_tmpl: ChatPromptTemplate,
    llm: OpenAI,
    tablename_lock: asyncio.Lock,
    retries: int = 3,
) -> TableInfo | None:
    """Process a single CSV file to generate a unique TableInfo."""
    last_table_name = None
    for attempt in range(retries):
        feedback = ""
        if attempt > 0:
            feedback = f"Note: '{last_table_name}' already exists. Please pick a new name not in {table_names}."

        table_info = await llm_structured_predict(
            df_str, table_names, prompt_tmpl, feedback, llm
        )
        last_table_name = table_info.table_name

        async with tablename_lock:
            if table_info.table_name not in table_names:
                table_names.append(table_info.table_name)
                return table_info

        print(f"Table name {table_info.table_name} already exists, retrying...")

    return None

async def process_csv_file(
    csv_file: File,
    table_names: list[str],
    semaphore: asyncio.Semaphore,
    tablename_lock: asyncio.Lock,
    llm: OpenAI,
    prompt_tmpl: ChatPromptTemplate,
) -> TableInfo | None:
    """Process a single CSV file to generate a unique TableInfo."""
    async with semaphore:
        df = await read_csv_file(csv_file, nrows=10)
        if df is None:
            return None
        return await generate_unique_table_info(
            df.to_csv(), table_names, prompt_tmpl, llm, tablename_lock
        )

@env.task
async def extract_table_info(
    data_dir: Dir, model: str, concurrency: int
) -> list[TableInfo | None]:
    """Extract structured table information from CSV files."""
    table_names: list[str] = []
    semaphore = asyncio.Semaphore(concurrency)
    tablename_lock = asyncio.Lock()
    llm = OpenAI(model=model)

    prompt_str = """\
    Provide a JSON object with the following fields:

    - `table_name`: must be unique and descriptive (underscores only, no generic names).
    - `table_summary`: short and concise summary of the table.

    Do NOT use any of these table names: {exclude_table_name_list}

    Table:
    {table_str}

    {feedback}
    """
    prompt_tmpl = ChatPromptTemplate(
        message_templates=[ChatMessage.from_str(prompt_str, role="user")]
    )

    tasks = [
        process_csv_file(
            csv_file, table_names, semaphore, tablename_lock, llm, prompt_tmpl
        )
        async for csv_file in data_dir.walk()
    ]

    return await asyncio.gather(*tasks)

# {{docs-fragment data_ingestion}}
@env.task
async def data_ingestion(
    csv_zip_path: str = "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip",
    search_glob: str = "WikiTableQuestions/csv/200-csv/*.csv",
    concurrency: int = 5,
    model: str = "gpt-4o-mini",
) -> tuple[File, list[TableInfo | None]]:
    """Main data ingestion pipeline: download → extract → analyze → create DB."""
    data_dir = await download_and_extract(csv_zip_path, search_glob)
    table_infos = await extract_table_info(data_dir, model, concurrency)

    database_path = "wiki_table_questions.db"

    i = 0
    async for csv_file in data_dir.walk():
        table_info = table_infos[i]
        if table_info:
            ok = await create_table(csv_file, table_info, database_path)
            if ok == "false":
                table_infos[i] = None
        else:
            print(f"Skipping table creation for {csv_file} due to missing TableInfo.")
        i += 1

    db_file = await File.from_local(database_path)
    return db_file, table_infos

# {{/docs-fragment data_ingestion}}
CODE0
import asyncio
import fnmatch
import os
import re
import zipfile

import flyte
import pandas as pd
import requests
from flyte.io import Dir, File
from llama_index.core.llms import ChatMessage
from llama_index.core.prompts import ChatPromptTemplate
from llama_index.llms.openai import OpenAI
from pydantic import BaseModel, Field
from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine
from utils import env

# {{docs-fragment table_info}}
class TableInfo(BaseModel):
    """Information regarding a structured table."""

    table_name: str = Field(..., description="table name (underscores only, no spaces)")
    table_summary: str = Field(
        ..., description="short, concise summary/caption of the table"
    )

# {{/docs-fragment table_info}}

@env.task
async def download_and_extract(zip_path: str, search_glob: str) -> Dir:
    """Download and extract the dataset zip file if not already available."""
    output_zip = "data.zip"
    extract_dir = "wiki_table_questions"

    if not os.path.exists(zip_path):
        response = requests.get(zip_path, stream=True)
        with open(output_zip, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
    else:
        output_zip = zip_path
        print(f"Using existing file {output_zip}")

    os.makedirs(extract_dir, exist_ok=True)
    with zipfile.ZipFile(output_zip, "r") as zip_ref:
        for member in zip_ref.namelist():
            if fnmatch.fnmatch(member, search_glob):
                zip_ref.extract(member, extract_dir)

    remote_dir = await Dir.from_local(extract_dir)
    return remote_dir

async def read_csv_file(
    csv_file: File, nrows: int | None = None
) -> pd.DataFrame | None:
    """Safely download and parse a CSV file into a DataFrame."""
    try:
        local_csv_file = await csv_file.download()
        return pd.read_csv(local_csv_file, nrows=nrows)
    except Exception as e:
        print(f"Error parsing {csv_file.path}: {e}")
        return None

def sanitize_column_name(col_name: str) -> str:
    """Sanitize column names by replacing spaces/special chars with underscores."""
    return re.sub(r"\W+", "_", col_name)

async def create_table_from_dataframe(
    df: pd.DataFrame, table_name: str, engine, metadata_obj
):
    """Create a SQL table from a Pandas DataFrame."""
    # Sanitize column names
    sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}
    df = df.rename(columns=sanitized_columns)

    # Define table columns based on DataFrame dtypes
    columns = [
        Column(col, String if dtype == "object" else Integer)
        for col, dtype in zip(df.columns, df.dtypes)
    ]

    table = Table(table_name, metadata_obj, *columns)

    # Create table in database
    metadata_obj.create_all(engine)

    # Insert data into table
    with engine.begin() as conn:
        for _, row in df.iterrows():
            conn.execute(table.insert().values(**row.to_dict()))

@flyte.trace
async def create_table(
    csv_file: File, table_info: TableInfo, database_path: str
) -> str:
    """Safely create a table from CSV if parsing succeeds."""
    df = await read_csv_file(csv_file)
    if df is None:
        return "false"

    print(f"Creating table: {table_info.table_name}")

    engine = create_engine(f"sqlite:///{database_path}")
    metadata_obj = MetaData()

    await create_table_from_dataframe(df, table_info.table_name, engine, metadata_obj)
    return "true"

@flyte.trace
async def llm_structured_predict(
    df_str: str,
    table_names: list[str],
    prompt_tmpl: ChatPromptTemplate,
    feedback: str,
    llm: OpenAI,
) -> TableInfo:
    return llm.structured_predict(
        TableInfo,
        prompt_tmpl,
        feedback=feedback,
        table_str=df_str,
        exclude_table_name_list=str(list(table_names)),
    )

async def generate_unique_table_info(
    df_str: str,
    table_names: list[str],
    prompt_tmpl: ChatPromptTemplate,
    llm: OpenAI,
    tablename_lock: asyncio.Lock,
    retries: int = 3,
) -> TableInfo | None:
    """Process a single CSV file to generate a unique TableInfo."""
    last_table_name = None
    for attempt in range(retries):
        feedback = ""
        if attempt > 0:
            feedback = f"Note: '{last_table_name}' already exists. Please pick a new name not in {table_names}."

        table_info = await llm_structured_predict(
            df_str, table_names, prompt_tmpl, feedback, llm
        )
        last_table_name = table_info.table_name

        async with tablename_lock:
            if table_info.table_name not in table_names:
                table_names.append(table_info.table_name)
                return table_info

        print(f"Table name {table_info.table_name} already exists, retrying...")

    return None

async def process_csv_file(
    csv_file: File,
    table_names: list[str],
    semaphore: asyncio.Semaphore,
    tablename_lock: asyncio.Lock,
    llm: OpenAI,
    prompt_tmpl: ChatPromptTemplate,
) -> TableInfo | None:
    """Process a single CSV file to generate a unique TableInfo."""
    async with semaphore:
        df = await read_csv_file(csv_file, nrows=10)
        if df is None:
            return None
        return await generate_unique_table_info(
            df.to_csv(), table_names, prompt_tmpl, llm, tablename_lock
        )

@env.task
async def extract_table_info(
    data_dir: Dir, model: str, concurrency: int
) -> list[TableInfo | None]:
    """Extract structured table information from CSV files."""
    table_names: list[str] = []
    semaphore = asyncio.Semaphore(concurrency)
    tablename_lock = asyncio.Lock()
    llm = OpenAI(model=model)

    prompt_str = """\
    Provide a JSON object with the following fields:

    - `table_name`: must be unique and descriptive (underscores only, no generic names).
    - `table_summary`: short and concise summary of the table.

    Do NOT use any of these table names: {exclude_table_name_list}

    Table:
    {table_str}

    {feedback}
    """
    prompt_tmpl = ChatPromptTemplate(
        message_templates=[ChatMessage.from_str(prompt_str, role="user")]
    )

    tasks = [
        process_csv_file(
            csv_file, table_names, semaphore, tablename_lock, llm, prompt_tmpl
        )
        async for csv_file in data_dir.walk()
    ]

    return await asyncio.gather(*tasks)

# {{docs-fragment data_ingestion}}
@env.task
async def data_ingestion(
    csv_zip_path: str = "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip",
    search_glob: str = "WikiTableQuestions/csv/200-csv/*.csv",
    concurrency: int = 5,
    model: str = "gpt-4o-mini",
) -> tuple[File, list[TableInfo | None]]:
    """Main data ingestion pipeline: download → extract → analyze → create DB."""
    data_dir = await download_and_extract(csv_zip_path, search_glob)
    table_infos = await extract_table_info(data_dir, model, concurrency)

    database_path = "wiki_table_questions.db"

    i = 0
    async for csv_file in data_dir.walk():
        table_info = table_infos[i]
        if table_info:
            ok = await create_table(csv_file, table_info, database_path)
            if ok == "false":
                table_infos[i] = None
        else:
            print(f"Skipping table creation for {csv_file} due to missing TableInfo.")
        i += 1

    db_file = await File.from_local(database_path)
    return db_file, table_infos

# {{/docs-fragment data_ingestion}}
CODE1
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "sqlalchemy>=2.0.0",
#    "pandas>=2.0.0",
#    "requests>=2.25.0",
#    "pydantic>=2.0.0",
# ]
# main = "text_to_sql"
# params = ""
# ///

import asyncio
from pathlib import Path

import flyte
from data_ingestion import TableInfo, data_ingestion
from flyte.io import Dir, File
from llama_index.core import (
    PromptTemplate,
    SQLDatabase,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
from llama_index.core.llms import ChatResponse
from llama_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema
from llama_index.core.prompts.prompt_type import PromptType
from llama_index.core.retrievers import SQLRetriever
from llama_index.core.schema import TextNode
from llama_index.llms.openai import OpenAI
from sqlalchemy import create_engine, text
from utils import env

# {{docs-fragment index_tables}}
@flyte.trace
async def index_table(table_name: str, table_index_dir: str, database_uri: str) -> str:
    """Index a single table into vector store."""
    path = f"{table_index_dir}/{table_name}"
    engine = create_engine(database_uri)

    def _fetch_rows():
        with engine.connect() as conn:
            cursor = conn.execute(text(f'SELECT * FROM "{table_name}"'))
            return cursor.fetchall()

    result = await asyncio.to_thread(_fetch_rows)
    nodes = [TextNode(text=str(tuple(row))) for row in result]
    index = VectorStoreIndex(nodes)
    index.set_index_id("vector_index")
    index.storage_context.persist(path)

    return path

@env.task
async def index_all_tables(db_file: File) -> Dir:
    """Index all tables concurrently."""
    table_index_dir = "table_indices"
    Path(table_index_dir).mkdir(exist_ok=True)

    await db_file.download(local_path="local_db.sqlite")
    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)

    tasks = [
        index_table(t, table_index_dir, "sqlite:///local_db.sqlite")
        for t in sql_database.get_usable_table_names()
    ]
    await asyncio.gather(*tasks)

    remote_dir = await Dir.from_local(table_index_dir)
    return remote_dir

# {{/docs-fragment index_tables}}

@flyte.trace
async def get_table_schema_context(
    table_schema_obj: SQLTableSchema,
    database_uri: str,
) -> str:
    """Retrieve schema + optional description context for a single table."""
    engine = create_engine(database_uri)
    sql_database = SQLDatabase(engine)

    table_info = sql_database.get_single_table_info(table_schema_obj.table_name)

    if table_schema_obj.context_str:
        table_info += f" The table description is: {table_schema_obj.context_str}"

    return table_info

@flyte.trace
async def get_table_row_context(
    table_schema_obj: SQLTableSchema,
    local_vector_index_dir: str,
    query: str,
) -> str:
    """Retrieve row-level context examples using vector search."""
    storage_context = StorageContext.from_defaults(
        persist_dir=str(f"{local_vector_index_dir}/{table_schema_obj.table_name}")
    )
    vector_index = load_index_from_storage(storage_context, index_id="vector_index")
    vector_retriever = vector_index.as_retriever(similarity_top_k=2)
    relevant_nodes = vector_retriever.retrieve(query)

    if not relevant_nodes:
        return ""

    row_context = "\nHere are some relevant example rows (values in the same order as columns above)\n"
    for node in relevant_nodes:
        row_context += str(node.get_content()) + "\n"

    return row_context

async def process_table(
    table_schema_obj: SQLTableSchema,
    database_uri: str,
    local_vector_index_dir: str,
    query: str,
) -> str:
    """Combine schema + row context for one table."""
    table_info = await get_table_schema_context(table_schema_obj, database_uri)
    row_context = await get_table_row_context(
        table_schema_obj, local_vector_index_dir, query
    )

    full_context = table_info
    if row_context:
        full_context += "\n" + row_context

    print(f"Table Info: {full_context}")
    return full_context

async def get_table_context_and_rows_str(
    query: str,
    database_uri: str,
    table_schema_objs: list[SQLTableSchema],
    vector_index_dir: Dir,
):
    """Get combined schema + row context for all tables."""
    local_vector_index_dir = await vector_index_dir.download()

    # run per-table work concurrently
    context_strs = await asyncio.gather(
        *[
            process_table(t, database_uri, local_vector_index_dir, query)
            for t in table_schema_objs
        ]
    )

    return "\n\n".join(context_strs)

# {{docs-fragment retrieve_tables}}
@env.task
async def retrieve_tables(
    query: str,
    table_infos: list[TableInfo | None],
    db_file: File,
    vector_index_dir: Dir,
) -> str:
    """Retrieve relevant tables and return schema context string."""
    await db_file.download(local_path="local_db.sqlite")
    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)

    table_node_mapping = SQLTableNodeMapping(sql_database)
    table_schema_objs = [
        SQLTableSchema(table_name=t.table_name, context_str=t.table_summary)
        for t in table_infos
        if t is not None
    ]

    obj_index = ObjectIndex.from_objects(
        table_schema_objs,
        table_node_mapping,
        VectorStoreIndex,
    )
    obj_retriever = obj_index.as_retriever(similarity_top_k=3)

    retrieved_schemas = obj_retriever.retrieve(query)
    return await get_table_context_and_rows_str(
        query, "sqlite:///local_db.sqlite", retrieved_schemas, vector_index_dir
    )

# {{/docs-fragment retrieve_tables}}

def parse_response_to_sql(chat_response: ChatResponse) -> str:
    """Extract SQL query from LLM response."""
    response = chat_response.message.content
    sql_query_start = response.find("SQLQuery:")
    if sql_query_start != -1:
        response = response[sql_query_start:]
        if response.startswith("SQLQuery:"):
            response = response[len("SQLQuery:") :]
    sql_result_start = response.find("SQLResult:")
    if sql_result_start != -1:
        response = response[:sql_result_start]
    return response.strip().strip("CODE2

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text_to_sql/text_to_sql.py*

The main `text_to_sql` task orchestrates the pipeline:

- Ingest data
- Build vector indices for each table
- Retrieve relevant tables and rows
- Generate SQL queries with an LLM
- Execute queries and synthesize answers

We use OpenAI GPT models with carefully structured prompts to maximize SQL correctness.

### Vector indexing

We index each table's rows semantically so the model can retrieve relevant examples during SQL generation.

CODE3").strip()

# {{docs-fragment sql_and_response}}
@env.task
async def generate_sql(query: str, table_context: str, model: str, prompt: str) -> str:
    """Generate SQL query from natural language question and table context."""
    llm = OpenAI(model=model)

    fmt_messages = (
        PromptTemplate(
            prompt,
            prompt_type=PromptType.TEXT_TO_SQL,
        )
        .partial_format(dialect="sqlite")
        .format_messages(query_str=query, schema=table_context)
    )

    chat_response = await llm.achat(fmt_messages)
    return parse_response_to_sql(chat_response)

@env.task
async def generate_response(query: str, sql: str, db_file: File, model: str) -> str:
    """Run SQL query on database and synthesize final response."""
    await db_file.download(local_path="local_db.sqlite")

    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)
    sql_retriever = SQLRetriever(sql_database)

    retrieved_rows = sql_retriever.retrieve(sql)

    response_synthesis_prompt = PromptTemplate(
        "Given an input question, synthesize a response from the query results.\n"
        "Query: {query_str}\n"
        "SQL: {sql_query}\n"
        "SQL Response: {context_str}\n"
        "Response: "
    )

    llm = OpenAI(model=model)
    fmt_messages = response_synthesis_prompt.format_messages(
        sql_query=sql,
        context_str=str(retrieved_rows),
        query_str=query,
    )
    chat_response = await llm.achat(fmt_messages)
    return chat_response.message.content

# {{/docs-fragment sql_and_response}}

# {{docs-fragment text_to_sql}}
@env.task
async def text_to_sql(
    system_prompt: str = (
        "Given an input question, first create a syntactically correct {dialect} "
        "query to run, then look at the results of the query and return the answer. "
        "You can order the results by a relevant column to return the most "
        "interesting examples in the database.\n\n"
        "Never query for all the columns from a specific table, only ask for a "
        "few relevant columns given the question.\n\n"
        "Pay attention to use only the column names that you can see in the schema "
        "description. "
        "Be careful to not query for columns that do not exist. "
        "Pay attention to which column is in which table. "
        "Also, qualify column names with the table name when needed. "
        "You are required to use the following format, each taking one line:\n\n"
        "Question: Question here\n"
        "SQLQuery: SQL Query to run\n"
        "SQLResult: Result of the SQLQuery\n"
        "Answer: Final answer here\n\n"
        "Only use tables listed below.\n"
        "{schema}\n\n"
        "Question: {query_str}\n"
        "SQLQuery: "
    ),
    query: str = "What was the year that The Notorious BIG was signed to Bad Boy?",
    model: str = "gpt-4o-mini",
) -> str:
    db_file, table_infos = await data_ingestion()
    vector_index_dir = await index_all_tables(db_file)
    table_context = await retrieve_tables(query, table_infos, db_file, vector_index_dir)
    sql = await generate_sql(query, table_context, model, system_prompt)
    return await generate_response(query, sql, db_file, model)

# {{/docs-fragment text_to_sql}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(text_to_sql)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text_to_sql/text_to_sql.py*

Each row becomes a text node stored in LlamaIndex’s `VectorStoreIndex`. This lets the system pull semantically similar rows when handling queries.

### Table retrieval and context building

We then retrieve the most relevant tables for a given query and build context that combines schema information with sample rows.

CODE4").strip()

# {{docs-fragment sql_and_response}}
@env.task
async def generate_sql(query: str, table_context: str, model: str, prompt: str) -> str:
    """Generate SQL query from natural language question and table context."""
    llm = OpenAI(model=model)

    fmt_messages = (
        PromptTemplate(
            prompt,
            prompt_type=PromptType.TEXT_TO_SQL,
        )
        .partial_format(dialect="sqlite")
        .format_messages(query_str=query, schema=table_context)
    )

    chat_response = await llm.achat(fmt_messages)
    return parse_response_to_sql(chat_response)

@env.task
async def generate_response(query: str, sql: str, db_file: File, model: str) -> str:
    """Run SQL query on database and synthesize final response."""
    await db_file.download(local_path="local_db.sqlite")

    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)
    sql_retriever = SQLRetriever(sql_database)

    retrieved_rows = sql_retriever.retrieve(sql)

    response_synthesis_prompt = PromptTemplate(
        "Given an input question, synthesize a response from the query results.\n"
        "Query: {query_str}\n"
        "SQL: {sql_query}\n"
        "SQL Response: {context_str}\n"
        "Response: "
    )

    llm = OpenAI(model=model)
    fmt_messages = response_synthesis_prompt.format_messages(
        sql_query=sql,
        context_str=str(retrieved_rows),
        query_str=query,
    )
    chat_response = await llm.achat(fmt_messages)
    return chat_response.message.content

# {{/docs-fragment sql_and_response}}

# {{docs-fragment text_to_sql}}
@env.task
async def text_to_sql(
    system_prompt: str = (
        "Given an input question, first create a syntactically correct {dialect} "
        "query to run, then look at the results of the query and return the answer. "
        "You can order the results by a relevant column to return the most "
        "interesting examples in the database.\n\n"
        "Never query for all the columns from a specific table, only ask for a "
        "few relevant columns given the question.\n\n"
        "Pay attention to use only the column names that you can see in the schema "
        "description. "
        "Be careful to not query for columns that do not exist. "
        "Pay attention to which column is in which table. "
        "Also, qualify column names with the table name when needed. "
        "You are required to use the following format, each taking one line:\n\n"
        "Question: Question here\n"
        "SQLQuery: SQL Query to run\n"
        "SQLResult: Result of the SQLQuery\n"
        "Answer: Final answer here\n\n"
        "Only use tables listed below.\n"
        "{schema}\n\n"
        "Question: {query_str}\n"
        "SQLQuery: "
    ),
    query: str = "What was the year that The Notorious BIG was signed to Bad Boy?",
    model: str = "gpt-4o-mini",
) -> str:
    db_file, table_infos = await data_ingestion()
    vector_index_dir = await index_all_tables(db_file)
    table_context = await retrieve_tables(query, table_infos, db_file, vector_index_dir)
    sql = await generate_sql(query, table_context, model, system_prompt)
    return await generate_response(query, sql, db_file, model)

# {{/docs-fragment text_to_sql}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(text_to_sql)
    print(run.url)
    run.wait()
CODE5
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "sqlalchemy>=2.0.0",
#    "pandas>=2.0.0",
#    "requests>=2.25.0",
#    "pydantic>=2.0.0",
# ]
# main = "text_to_sql"
# params = ""
# ///

import asyncio
from pathlib import Path

import flyte
from data_ingestion import TableInfo, data_ingestion
from flyte.io import Dir, File
from llama_index.core import (
    PromptTemplate,
    SQLDatabase,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
from llama_index.core.llms import ChatResponse
from llama_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema
from llama_index.core.prompts.prompt_type import PromptType
from llama_index.core.retrievers import SQLRetriever
from llama_index.core.schema import TextNode
from llama_index.llms.openai import OpenAI
from sqlalchemy import create_engine, text
from utils import env

# {{docs-fragment index_tables}}
@flyte.trace
async def index_table(table_name: str, table_index_dir: str, database_uri: str) -> str:
    """Index a single table into vector store."""
    path = f"{table_index_dir}/{table_name}"
    engine = create_engine(database_uri)

    def _fetch_rows():
        with engine.connect() as conn:
            cursor = conn.execute(text(f'SELECT * FROM "{table_name}"'))
            return cursor.fetchall()

    result = await asyncio.to_thread(_fetch_rows)
    nodes = [TextNode(text=str(tuple(row))) for row in result]
    index = VectorStoreIndex(nodes)
    index.set_index_id("vector_index")
    index.storage_context.persist(path)

    return path

@env.task
async def index_all_tables(db_file: File) -> Dir:
    """Index all tables concurrently."""
    table_index_dir = "table_indices"
    Path(table_index_dir).mkdir(exist_ok=True)

    await db_file.download(local_path="local_db.sqlite")
    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)

    tasks = [
        index_table(t, table_index_dir, "sqlite:///local_db.sqlite")
        for t in sql_database.get_usable_table_names()
    ]
    await asyncio.gather(*tasks)

    remote_dir = await Dir.from_local(table_index_dir)
    return remote_dir

# {{/docs-fragment index_tables}}

@flyte.trace
async def get_table_schema_context(
    table_schema_obj: SQLTableSchema,
    database_uri: str,
) -> str:
    """Retrieve schema + optional description context for a single table."""
    engine = create_engine(database_uri)
    sql_database = SQLDatabase(engine)

    table_info = sql_database.get_single_table_info(table_schema_obj.table_name)

    if table_schema_obj.context_str:
        table_info += f" The table description is: {table_schema_obj.context_str}"

    return table_info

@flyte.trace
async def get_table_row_context(
    table_schema_obj: SQLTableSchema,
    local_vector_index_dir: str,
    query: str,
) -> str:
    """Retrieve row-level context examples using vector search."""
    storage_context = StorageContext.from_defaults(
        persist_dir=str(f"{local_vector_index_dir}/{table_schema_obj.table_name}")
    )
    vector_index = load_index_from_storage(storage_context, index_id="vector_index")
    vector_retriever = vector_index.as_retriever(similarity_top_k=2)
    relevant_nodes = vector_retriever.retrieve(query)

    if not relevant_nodes:
        return ""

    row_context = "\nHere are some relevant example rows (values in the same order as columns above)\n"
    for node in relevant_nodes:
        row_context += str(node.get_content()) + "\n"

    return row_context

async def process_table(
    table_schema_obj: SQLTableSchema,
    database_uri: str,
    local_vector_index_dir: str,
    query: str,
) -> str:
    """Combine schema + row context for one table."""
    table_info = await get_table_schema_context(table_schema_obj, database_uri)
    row_context = await get_table_row_context(
        table_schema_obj, local_vector_index_dir, query
    )

    full_context = table_info
    if row_context:
        full_context += "\n" + row_context

    print(f"Table Info: {full_context}")
    return full_context

async def get_table_context_and_rows_str(
    query: str,
    database_uri: str,
    table_schema_objs: list[SQLTableSchema],
    vector_index_dir: Dir,
):
    """Get combined schema + row context for all tables."""
    local_vector_index_dir = await vector_index_dir.download()

    # run per-table work concurrently
    context_strs = await asyncio.gather(
        *[
            process_table(t, database_uri, local_vector_index_dir, query)
            for t in table_schema_objs
        ]
    )

    return "\n\n".join(context_strs)

# {{docs-fragment retrieve_tables}}
@env.task
async def retrieve_tables(
    query: str,
    table_infos: list[TableInfo | None],
    db_file: File,
    vector_index_dir: Dir,
) -> str:
    """Retrieve relevant tables and return schema context string."""
    await db_file.download(local_path="local_db.sqlite")
    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)

    table_node_mapping = SQLTableNodeMapping(sql_database)
    table_schema_objs = [
        SQLTableSchema(table_name=t.table_name, context_str=t.table_summary)
        for t in table_infos
        if t is not None
    ]

    obj_index = ObjectIndex.from_objects(
        table_schema_objs,
        table_node_mapping,
        VectorStoreIndex,
    )
    obj_retriever = obj_index.as_retriever(similarity_top_k=3)

    retrieved_schemas = obj_retriever.retrieve(query)
    return await get_table_context_and_rows_str(
        query, "sqlite:///local_db.sqlite", retrieved_schemas, vector_index_dir
    )

# {{/docs-fragment retrieve_tables}}

def parse_response_to_sql(chat_response: ChatResponse) -> str:
    """Extract SQL query from LLM response."""
    response = chat_response.message.content
    sql_query_start = response.find("SQLQuery:")
    if sql_query_start != -1:
        response = response[sql_query_start:]
        if response.startswith("SQLQuery:"):
            response = response[len("SQLQuery:") :]
    sql_result_start = response.find("SQLResult:")
    if sql_result_start != -1:
        response = response[:sql_result_start]
    return response.strip().strip("CODE6

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text_to_sql/text_to_sql.py*

The SQL generation prompt includes schema, example rows, and formatting rules. After execution, the system returns a final answer.

At this point, we have an end-to-end Text-to-SQL pipeline: natural language questions go in, SQL queries run, and answers come back. To make this workflow production-ready, we used several Flyte 2 capabilities. Caching ensures that repeated steps, like table ingestion or vector indexing, don’t need to rerun unnecessarily, saving time and compute. Containerization provides consistent, reproducible execution across environments, making it easier to scale and deploy. Observability features let us track every step of the pipeline, monitor performance, and debug issues quickly.

While the pipeline works end-to-end, to get a pulse on how it performs across multiple prompts and to gradually improve performance, we can start experimenting with prompt tuning.

Two things help make this process meaningful:

- **A clean evaluation dataset** - so we can measure accuracy against trusted ground truth.
- **A systematic evaluation loop** - so we can see whether prompt changes or other adjustments actually help.

With these in place, the next step is to build a "golden" QA dataset that will guide iterative prompt optimization.

## Building the QA dataset

> [!NOTE]
> The WikiTableQuestions dataset already includes question-answer pairs, available in its [GitHub repository](https://github.com/ppasupat/WikiTableQuestions/tree/master/data). To use them for this workflow, you'll need to adapt the data into the required format, but the raw material is there for you to build on.

We generate a dataset of natural language questions paired with executable SQL queries. This dataset acts as the benchmark for prompt tuning and evaluation.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "pydantic>=2.0.0",
# ]
# main = "build_eval_dataset"
# params = ""
# ///

import sqlite3

import flyte
import pandas as pd
from data_ingestion import data_ingestion
from flyte.io import File
from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI
from utils import env
from pydantic import BaseModel

class QAItem(BaseModel):
    question: str
    sql: str

class QAList(BaseModel):
    items: list[QAItem]

# {{docs-fragment get_and_split_schema}}
@env.task
async def get_and_split_schema(db_file: File, tables_per_chunk: int) -> list[str]:
    """
    Download the SQLite DB, extract schema info (columns + sample rows),
    then split it into chunks with up to `tables_per_chunk` tables each.
    """
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    tables = cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table';"
    ).fetchall()

    schema_blocks = []
    for table in tables:
        table_name = table[0]

        # columns
        cursor.execute(f"PRAGMA table_info({table_name});")
        columns = [col[1] for col in cursor.fetchall()]
        block = f"Table: {table_name}({', '.join(columns)})"

        # sample rows
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 10;")
        rows = cursor.fetchall()
        if rows:
            block += "\nSample rows:\n"
            for row in rows:
                block += f"{row}\n"

        schema_blocks.append(block)

    conn.close()

    chunks = []
    current_chunk = []
    for block in schema_blocks:
        current_chunk.append(block)
        if len(current_chunk) >= tables_per_chunk:
            chunks.append("\n".join(current_chunk))
            current_chunk = []
    if current_chunk:
        chunks.append("\n".join(current_chunk))

    return chunks

# {{/docs-fragment get_and_split_schema}}

# {{docs-fragment generate_questions_and_sql}}
@flyte.trace
async def generate_questions_and_sql(
    schema: str, num_samples: int, batch_size: int
) -> QAList:
    llm = OpenAI(model="gpt-4.1")

    prompt_tmpl = PromptTemplate(
        """Prompt: You are helping build a Text-to-SQL dataset.

Here is the database schema:
{schema}

Generate {num} natural language questions a user might ask about this database.
For each question, also provide the correct SQL query.

Reasoning process (you must follow this internally):

- Given an input question, first create a syntactically correct {dialect} SQL query.
- Never use SELECT *; only include the relevant columns.
- Use only columns/tables from the schema. Qualify column names when ambiguous.
- You may order results by a meaningful column to make the query more useful.
- Be careful not to add unnecessary columns.
- Use filters, aggregations, joins, grouping, and subqueries when relevant.

Final Output:
Return only a JSON object with one field:

- "items": a list of {num} objects, each with:
    - "question": the natural language question
    - "sql": the corresponding SQL query
"""
    )

    all_items: list[QAItem] = []

    # batch generation
    for start in range(0, num_samples, batch_size):
        current_num = min(batch_size, num_samples - start)
        response = llm.structured_predict(
            QAList,
            prompt_tmpl,
            schema=schema,
            num=current_num,
        )
        all_items.extend(response.items)

    # deduplicate
    seen = set()
    unique_items: list[QAItem] = []
    for item in all_items:
        key = (item.question.strip().lower(), item.sql.strip().lower())
        if key not in seen:
            seen.add(key)
            unique_items.append(item)

    return QAList(items=unique_items[:num_samples])

# {{/docs-fragment generate_questions_and_sql}}

@flyte.trace
async def llm_validate_batch(pairs: list[dict[str, str]]) -> list[str]:
    """Validate a batch of question/sql/result dicts using one LLM call."""
    batch_prompt = """You are validating the correctness of SQL query results against the question.
For each example, answer only "True" (correct) or "False" (incorrect).
Output one answer per line, in the same order as the examples.
---
"""

    for i, pair in enumerate(pairs, start=1):
        batch_prompt += f"""
Example {i}:
Question:
{pair['question']}

SQL:
{pair['sql']}

Result:
{pair['rows']}
---
"""

    llm = OpenAI(model="gpt-4.1")
    resp = await llm.acomplete(batch_prompt)

    # Expect exactly one True/False per example
    results = [
        line.strip()
        for line in resp.text.splitlines()
        if line.strip() in ("True", "False")
    ]
    return results

# {{docs-fragment validate_sql}}
@env.task
async def validate_sql(
    db_file: File, question_sql_pairs: QAList, batch_size: int
) -> list[dict[str, str]]:
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    qa_data = []
    batch = []

    for pair in question_sql_pairs.items:
        q, sql = pair.question, pair.sql
        try:
            cursor.execute(sql)
            rows = cursor.fetchall()
            batch.append({"question": q, "sql": sql, "rows": str(rows)})

            # process when batch is full
            if len(batch) == batch_size:
                results = await llm_validate_batch(batch)
                for pair, is_valid in zip(batch, results):
                    if is_valid == "True":
                        qa_data.append(
                            {
                                "input": pair["question"],
                                "sql": pair["sql"],
                                "target": pair["rows"],
                            }
                        )
                    else:
                        print(f"Filtered out incorrect result for: {pair['question']}")
                batch = []
        except Exception as e:
            print(f"Skipping invalid SQL: {sql} ({e})")

    # process leftover batch
    if batch:
        results = await llm_validate_batch(batch)
        for pair, is_valid in zip(batch, results):
            if is_valid == "True":
                qa_data.append(
                    {
                        "input": pair["question"],
                        "sql": pair["sql"],
                        "target": pair["rows"],
                    }
                )
            else:
                print(f"Filtered out incorrect result for: {pair['question']}")

    conn.close()
    return qa_data

# {{/docs-fragment validate_sql}}

@flyte.trace
async def save_to_csv(qa_data: list[dict]) -> File:
    df = pd.DataFrame(qa_data, columns=["input", "target", "sql"])

    csv_file = "qa_dataset.csv"
    df.to_csv(csv_file, index=False)

    return await File.from_local(csv_file)

# {{docs-fragment build_eval_dataset}}
@env.task
async def build_eval_dataset(
    num_samples: int = 300, batch_size: int = 30, tables_per_chunk: int = 3
) -> File:
    db_file, _ = await data_ingestion()
    schema_chunks = await get_and_split_schema(db_file, tables_per_chunk)

    per_chunk_samples = max(1, num_samples // len(schema_chunks))
    final_qa_data = []

    for chunk in schema_chunks:
        qa_list = await generate_questions_and_sql(
            schema=chunk,
            num_samples=per_chunk_samples,
            batch_size=batch_size,
        )
        qa_data = await validate_sql(db_file, qa_list, batch_size)
        final_qa_data.extend(qa_data)

    csv_file = await save_to_csv(final_qa_data)
    return csv_file

# {{/docs-fragment build_eval_dataset}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(build_eval_dataset)
    print(run.url)
    run.wait()
CODE7
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "pydantic>=2.0.0",
# ]
# main = "build_eval_dataset"
# params = ""
# ///

import sqlite3

import flyte
import pandas as pd
from data_ingestion import data_ingestion
from flyte.io import File
from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI
from utils import env
from pydantic import BaseModel

class QAItem(BaseModel):
    question: str
    sql: str

class QAList(BaseModel):
    items: list[QAItem]

# {{docs-fragment get_and_split_schema}}
@env.task
async def get_and_split_schema(db_file: File, tables_per_chunk: int) -> list[str]:
    """
    Download the SQLite DB, extract schema info (columns + sample rows),
    then split it into chunks with up to `tables_per_chunk` tables each.
    """
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    tables = cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table';"
    ).fetchall()

    schema_blocks = []
    for table in tables:
        table_name = table[0]

        # columns
        cursor.execute(f"PRAGMA table_info({table_name});")
        columns = [col[1] for col in cursor.fetchall()]
        block = f"Table: {table_name}({', '.join(columns)})"

        # sample rows
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 10;")
        rows = cursor.fetchall()
        if rows:
            block += "\nSample rows:\n"
            for row in rows:
                block += f"{row}\n"

        schema_blocks.append(block)

    conn.close()

    chunks = []
    current_chunk = []
    for block in schema_blocks:
        current_chunk.append(block)
        if len(current_chunk) >= tables_per_chunk:
            chunks.append("\n".join(current_chunk))
            current_chunk = []
    if current_chunk:
        chunks.append("\n".join(current_chunk))

    return chunks

# {{/docs-fragment get_and_split_schema}}

# {{docs-fragment generate_questions_and_sql}}
@flyte.trace
async def generate_questions_and_sql(
    schema: str, num_samples: int, batch_size: int
) -> QAList:
    llm = OpenAI(model="gpt-4.1")

    prompt_tmpl = PromptTemplate(
        """Prompt: You are helping build a Text-to-SQL dataset.

Here is the database schema:
{schema}

Generate {num} natural language questions a user might ask about this database.
For each question, also provide the correct SQL query.

Reasoning process (you must follow this internally):

- Given an input question, first create a syntactically correct {dialect} SQL query.
- Never use SELECT *; only include the relevant columns.
- Use only columns/tables from the schema. Qualify column names when ambiguous.
- You may order results by a meaningful column to make the query more useful.
- Be careful not to add unnecessary columns.
- Use filters, aggregations, joins, grouping, and subqueries when relevant.

Final Output:
Return only a JSON object with one field:

- "items": a list of {num} objects, each with:
    - "question": the natural language question
    - "sql": the corresponding SQL query
"""
    )

    all_items: list[QAItem] = []

    # batch generation
    for start in range(0, num_samples, batch_size):
        current_num = min(batch_size, num_samples - start)
        response = llm.structured_predict(
            QAList,
            prompt_tmpl,
            schema=schema,
            num=current_num,
        )
        all_items.extend(response.items)

    # deduplicate
    seen = set()
    unique_items: list[QAItem] = []
    for item in all_items:
        key = (item.question.strip().lower(), item.sql.strip().lower())
        if key not in seen:
            seen.add(key)
            unique_items.append(item)

    return QAList(items=unique_items[:num_samples])

# {{/docs-fragment generate_questions_and_sql}}

@flyte.trace
async def llm_validate_batch(pairs: list[dict[str, str]]) -> list[str]:
    """Validate a batch of question/sql/result dicts using one LLM call."""
    batch_prompt = """You are validating the correctness of SQL query results against the question.
For each example, answer only "True" (correct) or "False" (incorrect).
Output one answer per line, in the same order as the examples.
---
"""

    for i, pair in enumerate(pairs, start=1):
        batch_prompt += f"""
Example {i}:
Question:
{pair['question']}

SQL:
{pair['sql']}

Result:
{pair['rows']}
---
"""

    llm = OpenAI(model="gpt-4.1")
    resp = await llm.acomplete(batch_prompt)

    # Expect exactly one True/False per example
    results = [
        line.strip()
        for line in resp.text.splitlines()
        if line.strip() in ("True", "False")
    ]
    return results

# {{docs-fragment validate_sql}}
@env.task
async def validate_sql(
    db_file: File, question_sql_pairs: QAList, batch_size: int
) -> list[dict[str, str]]:
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    qa_data = []
    batch = []

    for pair in question_sql_pairs.items:
        q, sql = pair.question, pair.sql
        try:
            cursor.execute(sql)
            rows = cursor.fetchall()
            batch.append({"question": q, "sql": sql, "rows": str(rows)})

            # process when batch is full
            if len(batch) == batch_size:
                results = await llm_validate_batch(batch)
                for pair, is_valid in zip(batch, results):
                    if is_valid == "True":
                        qa_data.append(
                            {
                                "input": pair["question"],
                                "sql": pair["sql"],
                                "target": pair["rows"],
                            }
                        )
                    else:
                        print(f"Filtered out incorrect result for: {pair['question']}")
                batch = []
        except Exception as e:
            print(f"Skipping invalid SQL: {sql} ({e})")

    # process leftover batch
    if batch:
        results = await llm_validate_batch(batch)
        for pair, is_valid in zip(batch, results):
            if is_valid == "True":
                qa_data.append(
                    {
                        "input": pair["question"],
                        "sql": pair["sql"],
                        "target": pair["rows"],
                    }
                )
            else:
                print(f"Filtered out incorrect result for: {pair['question']}")

    conn.close()
    return qa_data

# {{/docs-fragment validate_sql}}

@flyte.trace
async def save_to_csv(qa_data: list[dict]) -> File:
    df = pd.DataFrame(qa_data, columns=["input", "target", "sql"])

    csv_file = "qa_dataset.csv"
    df.to_csv(csv_file, index=False)

    return await File.from_local(csv_file)

# {{docs-fragment build_eval_dataset}}
@env.task
async def build_eval_dataset(
    num_samples: int = 300, batch_size: int = 30, tables_per_chunk: int = 3
) -> File:
    db_file, _ = await data_ingestion()
    schema_chunks = await get_and_split_schema(db_file, tables_per_chunk)

    per_chunk_samples = max(1, num_samples // len(schema_chunks))
    final_qa_data = []

    for chunk in schema_chunks:
        qa_list = await generate_questions_and_sql(
            schema=chunk,
            num_samples=per_chunk_samples,
            batch_size=batch_size,
        )
        qa_data = await validate_sql(db_file, qa_list, batch_size)
        final_qa_data.extend(qa_data)

    csv_file = await save_to_csv(final_qa_data)
    return csv_file

# {{/docs-fragment build_eval_dataset}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(build_eval_dataset)
    print(run.url)
    run.wait()
CODE8
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "pydantic>=2.0.0",
# ]
# main = "build_eval_dataset"
# params = ""
# ///

import sqlite3

import flyte
import pandas as pd
from data_ingestion import data_ingestion
from flyte.io import File
from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI
from utils import env
from pydantic import BaseModel

class QAItem(BaseModel):
    question: str
    sql: str

class QAList(BaseModel):
    items: list[QAItem]

# {{docs-fragment get_and_split_schema}}
@env.task
async def get_and_split_schema(db_file: File, tables_per_chunk: int) -> list[str]:
    """
    Download the SQLite DB, extract schema info (columns + sample rows),
    then split it into chunks with up to `tables_per_chunk` tables each.
    """
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    tables = cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table';"
    ).fetchall()

    schema_blocks = []
    for table in tables:
        table_name = table[0]

        # columns
        cursor.execute(f"PRAGMA table_info({table_name});")
        columns = [col[1] for col in cursor.fetchall()]
        block = f"Table: {table_name}({', '.join(columns)})"

        # sample rows
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 10;")
        rows = cursor.fetchall()
        if rows:
            block += "\nSample rows:\n"
            for row in rows:
                block += f"{row}\n"

        schema_blocks.append(block)

    conn.close()

    chunks = []
    current_chunk = []
    for block in schema_blocks:
        current_chunk.append(block)
        if len(current_chunk) >= tables_per_chunk:
            chunks.append("\n".join(current_chunk))
            current_chunk = []
    if current_chunk:
        chunks.append("\n".join(current_chunk))

    return chunks

# {{/docs-fragment get_and_split_schema}}

# {{docs-fragment generate_questions_and_sql}}
@flyte.trace
async def generate_questions_and_sql(
    schema: str, num_samples: int, batch_size: int
) -> QAList:
    llm = OpenAI(model="gpt-4.1")

    prompt_tmpl = PromptTemplate(
        """Prompt: You are helping build a Text-to-SQL dataset.

Here is the database schema:
{schema}

Generate {num} natural language questions a user might ask about this database.
For each question, also provide the correct SQL query.

Reasoning process (you must follow this internally):

- Given an input question, first create a syntactically correct {dialect} SQL query.
- Never use SELECT *; only include the relevant columns.
- Use only columns/tables from the schema. Qualify column names when ambiguous.
- You may order results by a meaningful column to make the query more useful.
- Be careful not to add unnecessary columns.
- Use filters, aggregations, joins, grouping, and subqueries when relevant.

Final Output:
Return only a JSON object with one field:

- "items": a list of {num} objects, each with:
    - "question": the natural language question
    - "sql": the corresponding SQL query
"""
    )

    all_items: list[QAItem] = []

    # batch generation
    for start in range(0, num_samples, batch_size):
        current_num = min(batch_size, num_samples - start)
        response = llm.structured_predict(
            QAList,
            prompt_tmpl,
            schema=schema,
            num=current_num,
        )
        all_items.extend(response.items)

    # deduplicate
    seen = set()
    unique_items: list[QAItem] = []
    for item in all_items:
        key = (item.question.strip().lower(), item.sql.strip().lower())
        if key not in seen:
            seen.add(key)
            unique_items.append(item)

    return QAList(items=unique_items[:num_samples])

# {{/docs-fragment generate_questions_and_sql}}

@flyte.trace
async def llm_validate_batch(pairs: list[dict[str, str]]) -> list[str]:
    """Validate a batch of question/sql/result dicts using one LLM call."""
    batch_prompt = """You are validating the correctness of SQL query results against the question.
For each example, answer only "True" (correct) or "False" (incorrect).
Output one answer per line, in the same order as the examples.
---
"""

    for i, pair in enumerate(pairs, start=1):
        batch_prompt += f"""
Example {i}:
Question:
{pair['question']}

SQL:
{pair['sql']}

Result:
{pair['rows']}
---
"""

    llm = OpenAI(model="gpt-4.1")
    resp = await llm.acomplete(batch_prompt)

    # Expect exactly one True/False per example
    results = [
        line.strip()
        for line in resp.text.splitlines()
        if line.strip() in ("True", "False")
    ]
    return results

# {{docs-fragment validate_sql}}
@env.task
async def validate_sql(
    db_file: File, question_sql_pairs: QAList, batch_size: int
) -> list[dict[str, str]]:
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    qa_data = []
    batch = []

    for pair in question_sql_pairs.items:
        q, sql = pair.question, pair.sql
        try:
            cursor.execute(sql)
            rows = cursor.fetchall()
            batch.append({"question": q, "sql": sql, "rows": str(rows)})

            # process when batch is full
            if len(batch) == batch_size:
                results = await llm_validate_batch(batch)
                for pair, is_valid in zip(batch, results):
                    if is_valid == "True":
                        qa_data.append(
                            {
                                "input": pair["question"],
                                "sql": pair["sql"],
                                "target": pair["rows"],
                            }
                        )
                    else:
                        print(f"Filtered out incorrect result for: {pair['question']}")
                batch = []
        except Exception as e:
            print(f"Skipping invalid SQL: {sql} ({e})")

    # process leftover batch
    if batch:
        results = await llm_validate_batch(batch)
        for pair, is_valid in zip(batch, results):
            if is_valid == "True":
                qa_data.append(
                    {
                        "input": pair["question"],
                        "sql": pair["sql"],
                        "target": pair["rows"],
                    }
                )
            else:
                print(f"Filtered out incorrect result for: {pair['question']}")

    conn.close()
    return qa_data

# {{/docs-fragment validate_sql}}

@flyte.trace
async def save_to_csv(qa_data: list[dict]) -> File:
    df = pd.DataFrame(qa_data, columns=["input", "target", "sql"])

    csv_file = "qa_dataset.csv"
    df.to_csv(csv_file, index=False)

    return await File.from_local(csv_file)

# {{docs-fragment build_eval_dataset}}
@env.task
async def build_eval_dataset(
    num_samples: int = 300, batch_size: int = 30, tables_per_chunk: int = 3
) -> File:
    db_file, _ = await data_ingestion()
    schema_chunks = await get_and_split_schema(db_file, tables_per_chunk)

    per_chunk_samples = max(1, num_samples // len(schema_chunks))
    final_qa_data = []

    for chunk in schema_chunks:
        qa_list = await generate_questions_and_sql(
            schema=chunk,
            num_samples=per_chunk_samples,
            batch_size=batch_size,
        )
        qa_data = await validate_sql(db_file, qa_list, batch_size)
        final_qa_data.extend(qa_data)

    csv_file = await save_to_csv(final_qa_data)
    return csv_file

# {{/docs-fragment build_eval_dataset}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(build_eval_dataset)
    print(run.url)
    run.wait()
CODE9
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
#    "pydantic>=2.0.0",
# ]
# main = "build_eval_dataset"
# params = ""
# ///

import sqlite3

import flyte
import pandas as pd
from data_ingestion import data_ingestion
from flyte.io import File
from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI
from utils import env
from pydantic import BaseModel

class QAItem(BaseModel):
    question: str
    sql: str

class QAList(BaseModel):
    items: list[QAItem]

# {{docs-fragment get_and_split_schema}}
@env.task
async def get_and_split_schema(db_file: File, tables_per_chunk: int) -> list[str]:
    """
    Download the SQLite DB, extract schema info (columns + sample rows),
    then split it into chunks with up to `tables_per_chunk` tables each.
    """
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    tables = cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table';"
    ).fetchall()

    schema_blocks = []
    for table in tables:
        table_name = table[0]

        # columns
        cursor.execute(f"PRAGMA table_info({table_name});")
        columns = [col[1] for col in cursor.fetchall()]
        block = f"Table: {table_name}({', '.join(columns)})"

        # sample rows
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 10;")
        rows = cursor.fetchall()
        if rows:
            block += "\nSample rows:\n"
            for row in rows:
                block += f"{row}\n"

        schema_blocks.append(block)

    conn.close()

    chunks = []
    current_chunk = []
    for block in schema_blocks:
        current_chunk.append(block)
        if len(current_chunk) >= tables_per_chunk:
            chunks.append("\n".join(current_chunk))
            current_chunk = []
    if current_chunk:
        chunks.append("\n".join(current_chunk))

    return chunks

# {{/docs-fragment get_and_split_schema}}

# {{docs-fragment generate_questions_and_sql}}
@flyte.trace
async def generate_questions_and_sql(
    schema: str, num_samples: int, batch_size: int
) -> QAList:
    llm = OpenAI(model="gpt-4.1")

    prompt_tmpl = PromptTemplate(
        """Prompt: You are helping build a Text-to-SQL dataset.

Here is the database schema:
{schema}

Generate {num} natural language questions a user might ask about this database.
For each question, also provide the correct SQL query.

Reasoning process (you must follow this internally):

- Given an input question, first create a syntactically correct {dialect} SQL query.
- Never use SELECT *; only include the relevant columns.
- Use only columns/tables from the schema. Qualify column names when ambiguous.
- You may order results by a meaningful column to make the query more useful.
- Be careful not to add unnecessary columns.
- Use filters, aggregations, joins, grouping, and subqueries when relevant.

Final Output:
Return only a JSON object with one field:

- "items": a list of {num} objects, each with:
    - "question": the natural language question
    - "sql": the corresponding SQL query
"""
    )

    all_items: list[QAItem] = []

    # batch generation
    for start in range(0, num_samples, batch_size):
        current_num = min(batch_size, num_samples - start)
        response = llm.structured_predict(
            QAList,
            prompt_tmpl,
            schema=schema,
            num=current_num,
        )
        all_items.extend(response.items)

    # deduplicate
    seen = set()
    unique_items: list[QAItem] = []
    for item in all_items:
        key = (item.question.strip().lower(), item.sql.strip().lower())
        if key not in seen:
            seen.add(key)
            unique_items.append(item)

    return QAList(items=unique_items[:num_samples])

# {{/docs-fragment generate_questions_and_sql}}

@flyte.trace
async def llm_validate_batch(pairs: list[dict[str, str]]) -> list[str]:
    """Validate a batch of question/sql/result dicts using one LLM call."""
    batch_prompt = """You are validating the correctness of SQL query results against the question.
For each example, answer only "True" (correct) or "False" (incorrect).
Output one answer per line, in the same order as the examples.
---
"""

    for i, pair in enumerate(pairs, start=1):
        batch_prompt += f"""
Example {i}:
Question:
{pair['question']}

SQL:
{pair['sql']}

Result:
{pair['rows']}
---
"""

    llm = OpenAI(model="gpt-4.1")
    resp = await llm.acomplete(batch_prompt)

    # Expect exactly one True/False per example
    results = [
        line.strip()
        for line in resp.text.splitlines()
        if line.strip() in ("True", "False")
    ]
    return results

# {{docs-fragment validate_sql}}
@env.task
async def validate_sql(
    db_file: File, question_sql_pairs: QAList, batch_size: int
) -> list[dict[str, str]]:
    await db_file.download(local_path="local_db.sqlite")
    conn = sqlite3.connect("local_db.sqlite")
    cursor = conn.cursor()

    qa_data = []
    batch = []

    for pair in question_sql_pairs.items:
        q, sql = pair.question, pair.sql
        try:
            cursor.execute(sql)
            rows = cursor.fetchall()
            batch.append({"question": q, "sql": sql, "rows": str(rows)})

            # process when batch is full
            if len(batch) == batch_size:
                results = await llm_validate_batch(batch)
                for pair, is_valid in zip(batch, results):
                    if is_valid == "True":
                        qa_data.append(
                            {
                                "input": pair["question"],
                                "sql": pair["sql"],
                                "target": pair["rows"],
                            }
                        )
                    else:
                        print(f"Filtered out incorrect result for: {pair['question']}")
                batch = []
        except Exception as e:
            print(f"Skipping invalid SQL: {sql} ({e})")

    # process leftover batch
    if batch:
        results = await llm_validate_batch(batch)
        for pair, is_valid in zip(batch, results):
            if is_valid == "True":
                qa_data.append(
                    {
                        "input": pair["question"],
                        "sql": pair["sql"],
                        "target": pair["rows"],
                    }
                )
            else:
                print(f"Filtered out incorrect result for: {pair['question']}")

    conn.close()
    return qa_data

# {{/docs-fragment validate_sql}}

@flyte.trace
async def save_to_csv(qa_data: list[dict]) -> File:
    df = pd.DataFrame(qa_data, columns=["input", "target", "sql"])

    csv_file = "qa_dataset.csv"
    df.to_csv(csv_file, index=False)

    return await File.from_local(csv_file)

# {{docs-fragment build_eval_dataset}}
@env.task
async def build_eval_dataset(
    num_samples: int = 300, batch_size: int = 30, tables_per_chunk: int = 3
) -> File:
    db_file, _ = await data_ingestion()
    schema_chunks = await get_and_split_schema(db_file, tables_per_chunk)

    per_chunk_samples = max(1, num_samples // len(schema_chunks))
    final_qa_data = []

    for chunk in schema_chunks:
        qa_list = await generate_questions_and_sql(
            schema=chunk,
            num_samples=per_chunk_samples,
            batch_size=batch_size,
        )
        qa_data = await validate_sql(db_file, qa_list, batch_size)
        final_qa_data.extend(qa_data)

    csv_file = await save_to_csv(final_qa_data)
    return csv_file

# {{/docs-fragment build_eval_dataset}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(build_eval_dataset)
    print(run.url)
    run.wait()
CODE10
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "sqlalchemy>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from data_ingestion import TableInfo
from flyte.io import Dir, File
from llama_index.core import SQLDatabase
from llama_index.core.retrievers import SQLRetriever
from sqlalchemy import create_engine
from text_to_sql import data_ingestion, generate_sql, index_all_tables, retrieve_tables
from utils import env

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Val/Test split
    df_renamed = df.rename(columns={"input": "question", "target": "answer"})

    n = len(df_renamed)
    split = n // 2

    df_val = df_renamed.iloc[:split]
    df_test = df_renamed.iloc[split:]

    return df_val, df_test

@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

@flyte.trace
async def generate_response(db_file: File, sql: str) -> str:
    await db_file.download(local_path="local_db.sqlite")

    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)
    sql_retriever = SQLRetriever(sql_database)

    retrieved_rows = sql_retriever.retrieve(sql)

    if retrieved_rows:
        # Get the structured result and stringify
        return str(retrieved_rows[0].node.metadata["result"])

    return ""

async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    db_file: File,
    table_infos: list[TableInfo | None],
    vector_index_dir: Dir,
) -> dict:
    # Generate response from target model
    table_context = await retrieve_tables(
        question, table_infos, db_file, vector_index_dir
    )
    sql = await generate_sql(
        question,
        table_context,
        target_model_config.model_name,
        target_model_config.prompt,
    )
    sql = sql.replace("sql\n", "")

    try:
        response = await generate_response(db_file, sql)
    except Exception as e:
        print(f"Failed to generate response for question {question}: {e}")
        response = None

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                query_str=question,
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "sql": sql,
        "is_correct": verdict_clean == "true",
    }

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    sql,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
    db_file,
    table_infos,
    vector_index_dir,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
                db_file,
                table_infos,
                vector_index_dir,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{html.escape(sql)}</td>
                    <td>{result['model_response']}</td>
                    <td>{result['sql']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

@dataclass
class DatabaseConfig:
    csv_zip_path: str
    search_glob: str
    concurrency: int
    model: str

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
    db_config: DatabaseConfig,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Ground Truth Answer</th>
                    <th>Ground Truth SQL</th>
                    <th>Model Response</th>
                    <th>Model SQL</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    db_file, table_infos = await data_ingestion(
        db_config.csv_zip_path,
        db_config.search_glob,
        db_config.concurrency,
        db_config.model,
    )

    vector_index_dir = await index_all_tables(db_file)

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            row.sql,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
            db_file,
            table_infos,
            vector_index_dir,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_val: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
    db_config: DatabaseConfig,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_val,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_val,
                target_model_config,
                review_model_config,
                concurrency,
                db_config,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    ground_truth_csv: File | str = "/root/ground_truth.csv",
    db_config: DatabaseConfig = DatabaseConfig(
        csv_zip_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip",
        search_glob="WikiTableQuestions/csv/200-csv/*.csv",
        concurrency=5,
        model="gpt-4o-mini",
    ),
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""Given an input question, create a syntactically correct {dialect} query to run.

Schema:
{schema}

Question: {query_str}

SQL query to run:
""",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question.
Both answers were generated by running SQL queries on the same database.

- If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True".
- If it adds incorrect or unrelated rows, or omits required values, output "False".

Question:
{query_str}

Ground Truth:
{answer}

Model Response:
{response}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used to translate a natural-language question into a SQL query against a provided database schema.

<EXAMPLE>
<SCHEMA>
artists(id, name)
albums(id, title, artist_id, release_year)
</SCHEMA>
<QUESTION>
How many albums did The Beatles release?
</QUESTION>
<ANSWER>
SELECT COUNT(*) FROM albums a JOIN artists r ON a.artist_id = r.id WHERE r.name = 'The Beatles';
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past.
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past.
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc. for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query_str.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 5,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(ground_truth_csv, str) and os.path.isfile(ground_truth_csv):
        ground_truth_csv = await File.from_local(ground_truth_csv)

    df_val, df_test = await data_prep(ground_truth_csv)

    best_prompt, val_accuracy = await prompt_optimizer(
        df_val,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
        db_config,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

    return {
        "best_prompt": best_prompt,
        "validation_accuracy": val_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
CODE11
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "sqlalchemy>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from data_ingestion import TableInfo
from flyte.io import Dir, File
from llama_index.core import SQLDatabase
from llama_index.core.retrievers import SQLRetriever
from sqlalchemy import create_engine
from text_to_sql import data_ingestion, generate_sql, index_all_tables, retrieve_tables
from utils import env

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Val/Test split
    df_renamed = df.rename(columns={"input": "question", "target": "answer"})

    n = len(df_renamed)
    split = n // 2

    df_val = df_renamed.iloc[:split]
    df_test = df_renamed.iloc[split:]

    return df_val, df_test

@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

@flyte.trace
async def generate_response(db_file: File, sql: str) -> str:
    await db_file.download(local_path="local_db.sqlite")

    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)
    sql_retriever = SQLRetriever(sql_database)

    retrieved_rows = sql_retriever.retrieve(sql)

    if retrieved_rows:
        # Get the structured result and stringify
        return str(retrieved_rows[0].node.metadata["result"])

    return ""

async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    db_file: File,
    table_infos: list[TableInfo | None],
    vector_index_dir: Dir,
) -> dict:
    # Generate response from target model
    table_context = await retrieve_tables(
        question, table_infos, db_file, vector_index_dir
    )
    sql = await generate_sql(
        question,
        table_context,
        target_model_config.model_name,
        target_model_config.prompt,
    )
    sql = sql.replace("sql\n", "")

    try:
        response = await generate_response(db_file, sql)
    except Exception as e:
        print(f"Failed to generate response for question {question}: {e}")
        response = None

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                query_str=question,
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "sql": sql,
        "is_correct": verdict_clean == "true",
    }

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    sql,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
    db_file,
    table_infos,
    vector_index_dir,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
                db_file,
                table_infos,
                vector_index_dir,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{html.escape(sql)}</td>
                    <td>{result['model_response']}</td>
                    <td>{result['sql']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

@dataclass
class DatabaseConfig:
    csv_zip_path: str
    search_glob: str
    concurrency: int
    model: str

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
    db_config: DatabaseConfig,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Ground Truth Answer</th>
                    <th>Ground Truth SQL</th>
                    <th>Model Response</th>
                    <th>Model SQL</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    db_file, table_infos = await data_ingestion(
        db_config.csv_zip_path,
        db_config.search_glob,
        db_config.concurrency,
        db_config.model,
    )

    vector_index_dir = await index_all_tables(db_file)

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            row.sql,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
            db_file,
            table_infos,
            vector_index_dir,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_val: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
    db_config: DatabaseConfig,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_val,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_val,
                target_model_config,
                review_model_config,
                concurrency,
                db_config,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    ground_truth_csv: File | str = "/root/ground_truth.csv",
    db_config: DatabaseConfig = DatabaseConfig(
        csv_zip_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip",
        search_glob="WikiTableQuestions/csv/200-csv/*.csv",
        concurrency=5,
        model="gpt-4o-mini",
    ),
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""Given an input question, create a syntactically correct {dialect} query to run.

Schema:
{schema}

Question: {query_str}

SQL query to run:
""",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question.
Both answers were generated by running SQL queries on the same database.

- If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True".
- If it adds incorrect or unrelated rows, or omits required values, output "False".

Question:
{query_str}

Ground Truth:
{answer}

Model Response:
{response}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used to translate a natural-language question into a SQL query against a provided database schema.

<EXAMPLE>
<SCHEMA>
artists(id, name)
albums(id, title, artist_id, release_year)
</SCHEMA>
<QUESTION>
How many albums did The Beatles release?
</QUESTION>
<ANSWER>
SELECT COUNT(*) FROM albums a JOIN artists r ON a.artist_id = r.id WHERE r.name = 'The Beatles';
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past.
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past.
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc. for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query_str.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 5,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(ground_truth_csv, str) and os.path.isfile(ground_truth_csv):
        ground_truth_csv = await File.from_local(ground_truth_csv)

    df_val, df_test = await data_prep(ground_truth_csv)

    best_prompt, val_accuracy = await prompt_optimizer(
        df_val,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
        db_config,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

    return {
        "best_prompt": best_prompt,
        "validation_accuracy": val_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
CODE12
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas>=2.0.0",
#    "sqlalchemy>=2.0.0",
#    "llama-index-core>=0.11.0",
#    "llama-index-llms-openai>=0.2.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from data_ingestion import TableInfo
from flyte.io import Dir, File
from llama_index.core import SQLDatabase
from llama_index.core.retrievers import SQLRetriever
from sqlalchemy import create_engine
from text_to_sql import data_ingestion, generate_sql, index_all_tables, retrieve_tables
from utils import env

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Val/Test split
    df_renamed = df.rename(columns={"input": "question", "target": "answer"})

    n = len(df_renamed)
    split = n // 2

    df_val = df_renamed.iloc[:split]
    df_test = df_renamed.iloc[split:]

    return df_val, df_test

@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

@flyte.trace
async def generate_response(db_file: File, sql: str) -> str:
    await db_file.download(local_path="local_db.sqlite")

    engine = create_engine("sqlite:///local_db.sqlite")
    sql_database = SQLDatabase(engine)
    sql_retriever = SQLRetriever(sql_database)

    retrieved_rows = sql_retriever.retrieve(sql)

    if retrieved_rows:
        # Get the structured result and stringify
        return str(retrieved_rows[0].node.metadata["result"])

    return ""

async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    db_file: File,
    table_infos: list[TableInfo | None],
    vector_index_dir: Dir,
) -> dict:
    # Generate response from target model
    table_context = await retrieve_tables(
        question, table_infos, db_file, vector_index_dir
    )
    sql = await generate_sql(
        question,
        table_context,
        target_model_config.model_name,
        target_model_config.prompt,
    )
    sql = sql.replace("sql\n", "")

    try:
        response = await generate_response(db_file, sql)
    except Exception as e:
        print(f"Failed to generate response for question {question}: {e}")
        response = None

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                query_str=question,
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "sql": sql,
        "is_correct": verdict_clean == "true",
    }

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    sql,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
    db_file,
    table_infos,
    vector_index_dir,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
                db_file,
                table_infos,
                vector_index_dir,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{html.escape(sql)}</td>
                    <td>{result['model_response']}</td>
                    <td>{result['sql']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

@dataclass
class DatabaseConfig:
    csv_zip_path: str
    search_glob: str
    concurrency: int
    model: str

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
    db_config: DatabaseConfig,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Ground Truth Answer</th>
                    <th>Ground Truth SQL</th>
                    <th>Model Response</th>
                    <th>Model SQL</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    db_file, table_infos = await data_ingestion(
        db_config.csv_zip_path,
        db_config.search_glob,
        db_config.concurrency,
        db_config.model,
    )

    vector_index_dir = await index_all_tables(db_file)

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            row.sql,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
            db_file,
            table_infos,
            vector_index_dir,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_val: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
    db_config: DatabaseConfig,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_val,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_val,
                target_model_config,
                review_model_config,
                concurrency,
                db_config,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    ground_truth_csv: File | str = "/root/ground_truth.csv",
    db_config: DatabaseConfig = DatabaseConfig(
        csv_zip_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip",
        search_glob="WikiTableQuestions/csv/200-csv/*.csv",
        concurrency=5,
        model="gpt-4o-mini",
    ),
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""Given an input question, create a syntactically correct {dialect} query to run.

Schema:
{schema}

Question: {query_str}

SQL query to run:
""",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question.
Both answers were generated by running SQL queries on the same database.

- If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True".
- If it adds incorrect or unrelated rows, or omits required values, output "False".

Question:
{query_str}

Ground Truth:
{answer}

Model Response:
{response}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used to translate a natural-language question into a SQL query against a provided database schema.

<EXAMPLE>
<SCHEMA>
artists(id, name)
albums(id, title, artist_id, release_year)
</SCHEMA>
<QUESTION>
How many albums did The Beatles release?
</QUESTION>
<ANSWER>
SELECT COUNT(*) FROM albums a JOIN artists r ON a.artist_id = r.id WHERE r.name = 'The Beatles';
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past.
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past.
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc. for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query_str.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 5,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(ground_truth_csv, str) and os.path.isfile(ground_truth_csv):
        ground_truth_csv = await File.from_local(ground_truth_csv)

    df_val, df_test = await data_prep(ground_truth_csv)

    best_prompt, val_accuracy = await prompt_optimizer(
        df_val,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
        db_config,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
            db_config,
        )

    return {
        "best_prompt": best_prompt,
        "validation_accuracy": val_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
CODE13
python create_qa_dataset.py
CODE14
python optimizer.py
```

## What we observed

Prompt optimization didn't consistently lift SQL accuracy in this workflow. Accuracy plateaued near the baseline. But the process surfaced valuable lessons about what matters when building LLM-powered systems on real infrastructure.

- **Schema clarity matters**: CSV ingestion produced tables with overlapping names, creating ambiguity. This showed how schema design and metadata hygiene directly affect downstream evaluation.
- **Ground truth needs trust**: Because the dataset came from LLM outputs, noise remained even after filtering. Human review proved essential. Golden datasets need deliberate curation, not just automation.
- **Optimization needs context**: The optimizer couldn't “see” which examples failed, limiting its ability to improve. Feeding failures directly risks overfitting. A structured way to capture and reuse evaluation signals is the right long-term path.

Sometimes prompt tweaks alone can lift accuracy, but other times the real bottleneck lives in the data, the schema, or the evaluation loop. The lesson isn't "prompt optimization doesn't work", but that its impact depends on the system around it. Accuracy improves most reliably when prompts evolve alongside clean data, trusted evaluation, and observable feedback loops.

## The bigger lesson

Evaluation and optimization aren’t one-off experiments; they’re continuous processes. What makes them sustainable isn't a clever prompt, it’s the platform around it.

Systems succeed when they:

- **Observe** failures with clarity: track exactly what failed and why.
- **Remain durable** across iterations: run pipelines that are stable, reproducible, and comparable over time.

That's where Flyte 2 comes in. Prompt optimization is one lever, but it becomes powerful only when combined with:

- Clean, human-validated evaluation datasets.
- Systematic reporting and feedback loops.

**The real takeaway: improving LLM pipelines isn't about chasing the perfect prompt. It's about designing workflows with observability and durability at the core, so that every experiment compounds into long-term progress.**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/auto_prompt_engineering ===

# Automatic prompt engineering

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/auto_prompt_engineering).

When building with LLMs and agents, the first prompt almost never works. We usually need several iterations before results are useful. Doing this manually is slow, inconsistent, and hard to reproduce.

Flyte turns prompt engineering into a systematic process. With Flyte we can:

- Generate candidate prompts automatically.
- Run evaluations in parallel.
- Track results in real time with built-in observability.
- Recover from failures without losing progress.
- Trace the lineage of every experiment for reproducibility.

And we're not limited to prompts. Just like [hyperparameter optimization](../../model-training/hpo/_index) in ML, we can tune model temperature, retrieval strategies, tool usage, and more. Over time, this grows into full agentic evaluations, tracking not only prompts but also how agents behave, make decisions, and interact with their environment.

In this tutorial, we'll build an automated prompt engineering pipeline with Flyte, step by step.

## Set up the environment

First, let's configure our task environment.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE0
flyte create secret openai_api_key <YOUR_OPENAI_API_KEY>
```

We also define CSS styles for live HTML reports that track prompt optimization in real time:

![Results](../../../_static/images/tutorials/prompt_engineering/results.gif)

## Prepare the evaluation dataset

Next, we define our golden dataset, a set of prompts with known outputs. This dataset is used to evaluate the quality of generated prompts.

For this tutorial, we use a small geometric shapes dataset. To keep it portable, the data prep task takes a CSV file (as a Flyte `File` or a string for files available remotely) and splits it into train and test subsets.

If you already have prompts and outputs in Google Sheets, simply export them as CSV with two columns: `input` and `target`.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE1
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto_prompt_engineering/optimizer.py*

Then we define a Flyte `trace` to call the model. Unlike a task, a trace runs within the same runtime as the parent process. Since the model is hosted externally, this keeps the call lightweight but still observable.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE2
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto_prompt_engineering/optimizer.py*

## Evaluate prompts

We now define the evaluation process.

Each prompt in the dataset is tested in parallel, but we use a semaphore to control concurrency. A helper function ties together the `generate_and_review` task with an HTML report template. Using `asyncio.gather`, we evaluate multiple prompts at once.

The function measures accuracy as the fraction of responses that match the ground truth. Flyte streams these results to the UI, so you can watch evaluations happen live.

```
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE3
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE4
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE5
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pandas==2.3.1",
#    "pyarrow==21.0.0",
#    "litellm==1.75.0",
# ]
# main = "auto_prompt_engineering"
# params = ""
# ///

# {{docs-fragment env}}
import asyncio
import html
import os
import re
from dataclasses import dataclass
from typing import Optional, Union

import flyte
import flyte.report
import pandas as pd
from flyte.io._file import File

env = flyte.TaskEnvironment(
    name="auto-prompt-engineering",
    image=flyte.Image.from_uv_script(
        __file__, name="auto-prompt-engineering", pre=True
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    resources=flyte.Resources(cpu=1),
)

CSS = """
<style>
    body {
        font-family: 'Segoe UI', Roboto, Arial, sans-serif;
    }
    .results-table {
        border-collapse: collapse;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        font-size: 14px;
    }
    .results-table th {
        background: linear-gradient(135deg, #4CAF50, #2E7D32);
        color: white;
        padding: 10px;
        text-align: left;
    }
    .results-table td {
        border: 1px solid #ddd;
        padding: 8px;
        vertical-align: top;
    }
    .results-table tr:nth-child(even) {background-color: #f9f9f9;}
    .results-table tr:hover {background-color: #f1f1f1;}
    .correct {color: #2E7D32; font-weight: bold;}
    .incorrect {color: #C62828; font-weight: bold;}
    .summary-card {
        background: #f9fbfd;
        padding: 14px 18px;
        border-radius: 8px;
        box-shadow: 0 1px 4px rgba(0,0,0,0.05);
        max-width: 800px;
        margin-top: 12px;
    }
    .summary-card h3 {
        margin-top: 0;
        color: #1e88e5;
        font-size: 16px;
    }
</style>
"""

# {{/docs-fragment env}}

# {{docs-fragment data_prep}}
@env.task
async def data_prep(csv_file: File | str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames.
    The sheet should have columns: 'input' and 'target'.
    """
    df = pd.read_csv(
        await csv_file.download() if isinstance(csv_file, File) else csv_file
    )

    if "input" not in df.columns or "target" not in df.columns:
        raise ValueError("Sheet must contain 'input' and 'target' columns.")

    # Shuffle rows
    df = df.sample(frac=1, random_state=1234).reset_index(drop=True)

    # Train/Test split
    df_train = df.iloc[:150].rename(columns={"input": "question", "target": "answer"})
    df_test = df.iloc[150:250].rename(columns={"input": "question", "target": "answer"})

    return df_train, df_test

# {{/docs-fragment data_prep}}

# {{docs-fragment model_config}}
@dataclass
class ModelConfig:
    model_name: str
    hosted_model_uri: Optional[str] = None
    temperature: float = 0.0
    max_tokens: Optional[int] = 1000
    timeout: int = 600
    prompt: str = ""

# {{/docs-fragment model_config}}

# {{docs-fragment call_model}}
@flyte.trace
async def call_model(
    model_config: ModelConfig,
    messages: list[dict[str, str]],
) -> str:
    from litellm import acompletion

    response = await acompletion(
        model=model_config.model_name,
        api_base=model_config.hosted_model_uri,
        messages=messages,
        temperature=model_config.temperature,
        timeout=model_config.timeout,
        max_tokens=model_config.max_tokens,
    )
    return response.choices[0].message["content"]

# {{/docs-fragment call_model}}

# {{docs-fragment generate_and_review}}
async def generate_and_review(
    index: int,
    question: str,
    answer: str,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
) -> dict:
    # Generate response from target model
    response = await call_model(
        target_model_config,
        [
            {"role": "system", "content": target_model_config.prompt},
            {"role": "user", "content": question},
        ],
    )

    # Format review prompt with response + answer
    review_messages = [
        {
            "role": "system",
            "content": review_model_config.prompt.format(
                response=response,
                answer=answer,
            ),
        }
    ]
    verdict = await call_model(review_model_config, review_messages)

    # Normalize verdict
    verdict_clean = verdict.strip().lower()
    if verdict_clean not in {"true", "false"}:
        verdict_clean = "not sure"

    return {
        "index": index,
        "model_response": response,
        "is_correct": verdict_clean == "true",
    }

# {{/docs-fragment generate_and_review}}

async def run_grouped_task(
    i,
    index,
    question,
    answer,
    semaphore,
    target_model_config,
    review_model_config,
    counter,
    counter_lock,
):
    async with semaphore:
        with flyte.group(name=f"row-{i}"):
            result = await generate_and_review(
                index,
                question,
                answer,
                target_model_config,
                review_model_config,
            )

            async with counter_lock:
                # Update counters
                counter["processed"] += 1
                if result["is_correct"]:
                    counter["correct"] += 1
                    correct_html = "<span class='correct'>✔ Yes</span>"
                else:
                    correct_html = "<span class='incorrect'>✘ No</span>"

                # Calculate accuracy
                accuracy_pct = (counter["correct"] / counter["processed"]) * 100

            # Update chart
            await flyte.report.log.aio(
                f"<script>updateAccuracy({accuracy_pct});</script>",
                do_flush=True,
            )

            # Add row to table
            await flyte.report.log.aio(
                f"""
                <tr>
                    <td>{html.escape(question)}</td>
                    <td>{html.escape(answer)}</td>
                    <td>{result['model_response']}</td>
                    <td>{correct_html}</td>
                </tr>
                """,
                do_flush=True,
            )

            return result

# {{docs-fragment evaluate_prompt}}
@env.task(report=True)
async def evaluate_prompt(
    df: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    concurrency: int,
) -> float:
    semaphore = asyncio.Semaphore(concurrency)
    counter = {"correct": 0, "processed": 0}
    counter_lock = asyncio.Lock()

    # Write initial HTML structure
    await flyte.report.log.aio(
        CSS
        + """
        <script>
            function updateAccuracy(percent) {
                const bar = document.getElementById('acc-bar');
                const label = document.getElementById('acc-label');
                bar.setAttribute('width', percent * 3);
                label.textContent = `Accuracy: ${percent.toFixed(1)}%`;
            }
        </script>

        <h2 style="margin-top:0;">Model Evaluation Results</h2>
        <h3>Live Accuracy</h3>
        <svg width="320" height="30" id="accuracy-chart">
            <defs>
                <linearGradient id="acc-gradient" x1="0" x2="1" y1="0" y2="0">
                    <stop offset="0%" stop-color="#66bb6a"/>
                    <stop offset="100%" stop-color="#2e7d32"/>
                </linearGradient>
            </defs>
            <rect width="300" height="20" fill="#ddd" rx="5" ry="5"></rect>
            <rect id="acc-bar" width="0" height="20" fill="url(#acc-gradient)" rx="5" ry="5"></rect>
            <text id="acc-label" x="150" y="15" font-size="12" font-weight="bold" text-anchor="middle" fill="#000">
                Accuracy: 0.0%
            </text>
        </svg>

        <table class="results-table">
            <thead>
                <tr>
                    <th>Question</th>
                    <th>Answer</th>
                    <th>Model Response</th>
                    <th>Correct?</th>
                </tr>
            </thead>
            <tbody>
        """,
        do_flush=True,
    )

    # Launch tasks concurrently
    tasks = [
        run_grouped_task(
            i,
            row.Index,
            row.question,
            row.answer,
            semaphore,
            target_model_config,
            review_model_config,
            counter,
            counter_lock,
        )
        for i, row in enumerate(df.itertuples(index=True))
    ]
    await asyncio.gather(*tasks)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    async with counter_lock:
        return (
            (counter["correct"] / counter["processed"]) if counter["processed"] else 0.0
        )

# {{/docs-fragment evaluate_prompt}}

@dataclass
class PromptResult:
    prompt: str
    accuracy: float

# {{docs-fragment prompt_optimizer}}
@env.task(report=True)
async def prompt_optimizer(
    df_train: pd.DataFrame,
    target_model_config: ModelConfig,
    review_model_config: ModelConfig,
    optimizer_model_config: ModelConfig,
    max_iterations: int,
    concurrency: int,
) -> tuple[str, float]:
    prompt_accuracies: list[PromptResult] = []

    # Send styling + table header immediately
    await flyte.report.log.aio(
        CSS
        + """
    <h2 style="margin-bottom:6px;">📊 Prompt Accuracy Comparison</h2>
    <table class="results-table">
        <thead>
            <tr>
                <th>Prompt</th>
                <th>Accuracy</th>
            </tr>
        </thead>
    <tbody>
    """,
        do_flush=True,
    )

    # Step 1: Evaluate starting prompt and stream row
    with flyte.group(name="baseline_evaluation"):
        starting_accuracy = await evaluate_prompt(
            df_train,
            target_model_config,
            review_model_config,
            concurrency,
        )
        prompt_accuracies.append(
            PromptResult(prompt=target_model_config.prompt, accuracy=starting_accuracy)
        )

        await _log_prompt_row(target_model_config.prompt, starting_accuracy)

    # Step 2: Optimize prompts one by one, streaming after each
    while len(prompt_accuracies) <= max_iterations:
        with flyte.group(name=f"prompt_optimization_step_{len(prompt_accuracies)}"):
            # Prepare prompt scores string for optimizer
            prompt_scores_str = "\n".join(
                f"{result.prompt}: {result.accuracy:.2f}"
                for result in sorted(prompt_accuracies, key=lambda x: x.accuracy)
            )

            optimizer_model_prompt = optimizer_model_config.prompt.format(
                prompt_scores_str=prompt_scores_str
            )
            response = await call_model(
                optimizer_model_config,
                [{"role": "system", "content": optimizer_model_prompt}],
            )
            response = response.strip()

            match = re.search(r"\[\[(.*?)\]\]", response, re.DOTALL)
            if not match:
                print("No new prompt found. Skipping.")
                continue

            new_prompt = match.group(1)
            target_model_config.prompt = new_prompt
            accuracy = await evaluate_prompt(
                df_train,
                target_model_config,
                review_model_config,
                concurrency,
            )
            prompt_accuracies.append(PromptResult(prompt=new_prompt, accuracy=accuracy))

            # Log this new prompt row immediately
            await _log_prompt_row(new_prompt, accuracy)

    # Close table
    await flyte.report.log.aio("</tbody></table>", do_flush=True)

    # Find best
    best_result = max(prompt_accuracies, key=lambda x: x.accuracy)
    improvement = best_result.accuracy - starting_accuracy

    # Summary
    await flyte.report.log.aio(
        f"""
    <div class="summary-card">
        <h3>🏆 Summary</h3>
        <p><strong>Best Prompt:</strong> {html.escape(best_result.prompt)}</p>
        <p><strong>Best Accuracy:</strong> {best_result.accuracy*100:.2f}%</p>
        <p><strong>Improvement Over Baseline:</strong> {improvement*100:.2f}%</p>
    </div>
    """,
        do_flush=True,
    )

    return best_result.prompt, best_result.accuracy

# {{/docs-fragment prompt_optimizer}}

async def _log_prompt_row(prompt: str, accuracy: float):
    """Helper to log a single prompt/accuracy row to Flyte report."""
    pct = accuracy * 100
    if pct > 80:
        color = "linear-gradient(90deg, #4CAF50, #81C784)"
    elif pct > 60:
        color = "linear-gradient(90deg, #FFC107, #FFD54F)"
    else:
        color = "linear-gradient(90deg, #F44336, #E57373)"

    await flyte.report.log.aio(
        f"""
        <tr>
            <td>{html.escape(prompt)}</td>
            <td>
                {pct:.1f}%
                <div class="accuracy-bar-container">
                    <div class="accuracy-bar" style="width:{pct*1.6}px; background:{color};"></div>
                </div>
            </td>
        </tr>
        """,
        do_flush=True,
    )

# {{docs-fragment auto_prompt_engineering}}
@env.task
async def auto_prompt_engineering(
    csv_file: File | str = "https://dub.sh/geometric-shapes",
    target_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="Solve the given problem about geometric shapes. Think step by step.",
        max_tokens=10000,
    ),
    review_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1-mini",
        hosted_model_uri=None,
        prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem.
The response may contain detailed steps and explanations, but the final answer is the key point.
Please determine if the final answer provided in the response is correct based on the ground truth number.
Respond with 'True' if the final answer is correct and 'False' if it is not.
Only respond with 'True' or 'False', nothing else.

Model Response:
{response}

Ground Truth:
{answer}
""",
    ),
    optimizer_model_config: ModelConfig = ModelConfig(
        model_name="gpt-4.1",
        hosted_model_uri=None,
        temperature=0.7,
        max_tokens=None,
        prompt="""
<EXPLANATION>
I have some prompts along with their corresponding accuracies.
The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality.
</EXPLANATION>

<PROMPTS>
{prompt_scores_str}
</PROMPTS>

Each prompt was used together with a problem statement around geometric shapes.

<EXAMPLE>
<QUESTION>
This SVG path element <path d="M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L 45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69"/> draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle
</QUESTION>
<ANSWER>
(B)
</ANSWER>
</EXAMPLE>

<TASK>
Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones.
</TASK>

<RULES>
- It is very important that the new prompt is distinct from ALL the old ones!
- Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past
- Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past
- Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt.
- Use all available information like prompt length, formal/informal use of language, etc for your analysis.
- Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy.
- You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task.
- Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc).
</RULES>
""",
    ),
    max_iterations: int = 3,
    concurrency: int = 10,
) -> dict[str, Union[str, float]]:
    if isinstance(csv_file, str) and os.path.isfile(csv_file):
        csv_file = await File.from_local(csv_file)

    df_train, df_test = await data_prep(csv_file)

    best_prompt, training_accuracy = await prompt_optimizer(
        df_train,
        target_model_config,
        review_model_config,
        optimizer_model_config,
        max_iterations,
        concurrency,
    )

    with flyte.group(name="test_data_evaluation"):
        baseline_test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

        target_model_config.prompt = best_prompt
        test_accuracy = await evaluate_prompt(
            df_test,
            target_model_config,
            review_model_config,
            concurrency,
        )

    return {
        "best_prompt": best_prompt,
        "training_accuracy": training_accuracy,
        "baseline_test_accuracy": baseline_test_accuracy,
        "test_accuracy": test_accuracy,
    }

# {{/docs-fragment auto_prompt_engineering}}

# {{docs-fragment main}}
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(auto_prompt_engineering)
    print(run.url)
    run.wait()
# {{/docs-fragment main}}
CODE6
uv run optimizer.py
```

![Execution](../../../_static/images/tutorials/prompt_engineering/execution.gif)

## Why this matters

Most prompt engineering pipelines start as quick scripts or notebooks. They're fine for experimenting, but they're difficult to scale, reproduce, or debug when things go wrong.

With Flyte 2, we get a more reliable setup:

- Run many evaluations in parallel with **Migration > From Flyte 1 to 2 > Migration overview > Asynchronous model > True parallelism for all workloads** or **Migration > From Flyte 1 to 2 > Migration overview > Asynchronous model > The `flyte.map` function: Familiar patterns**.
- Watch accuracy improve in real time and link results back to the exact dataset, prompt, and model config used.
- Resume cleanly after failures without rerunning everything from scratch.
- Reuse the same pattern to tune other parameters like temperature, retrieval depth, or agent strategies, not just prompts.

## Next steps

You now have a working automated prompt engineering pipeline. Here’s how you can take it further:

- **Optimize beyond prompts**: Tune temperature, retrieval strategies, or tool usage just like prompts.
- **Expand evaluation metrics**: Add latency, cost, robustness, or diversity alongside accuracy.
- **Move toward agentic evaluation**: Instead of single prompts, test how agents plan, use tools, and recover from failures in long-horizon tasks.

With this foundation, prompt engineering becomes repeatable, observable, and scalable, ready for production-grade LLM and agent systems.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training ===

# Model training

Tutorials for training, fine-tuning, and hyperparameter optimization of models at scale.

### **Model training > Hyperparameter optimization**

Run large-scale HPO experiments with zero manual tracking, deterministic results, and automatic recovery.

### **Model training > LLM fine-tuning with LoRA and QLoRA**

Fine-tune a language model for SQL generation using full, LoRA, or QLoRA methods in one Flyte pipeline.

### **Model training > BERT emotion classification**

Fine-tune ModernBERT on Twitter emotion labels with confusion-matrix evaluation and attention visualizations.

## Subpages
- **Model training > LLM fine-tuning with LoRA and QLoRA**
- **Model training > BERT emotion classification**
- **Model training > Hyperparameter optimization**


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training/llm-fine-tuning-lora-qlora ===

# LLM fine-tuning with LoRA and QLoRA

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/llm_fine_tuning_lora_qlora).

This tutorial fine-tunes a language model for SQL generation using three methods in one workflow: **full** fine-tuning, **LoRA** adapters, and **QLoRA** (4-bit quantized base + LoRA). The pipeline prepares an instruction dataset from HuggingFace, trains with [TRL](https://huggingface.co/docs/trl) `SFTTrainer`, evaluates against a base-model baseline, and streams training charts into Flyte reports.

Flyte provides:

- **GPU training** with live loss and learning-rate charts via `report=True`.
- **Method switching** through a single `method` parameter (`full`, `lora`, or `qlora`).
- **Cached dataset preparation** for fast iteration on hyperparameters.

## Define the task environments

The GPU environment declares a HuggingFace token secret for gated models.

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "peft>=0.13.0",
#    "trl>=0.12.0",
#    "datasets>=3.0.0",
#    "bitsandbytes>=0.44.0",
#    "accelerate>=0.34.0",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import json
import logging
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
import os

main_img = flyte.Image.from_uv_script(__file__, name="llm-fine-tuning-lora-qlora", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

cpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)

HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}

from report_helpers import make_bar_chart, make_line_chart, pipeline_step_indicator, wrap_report

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Task 1: Prepare dataset
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def prepare_data(
    dataset_name: str = "b-mc2/sql-create-context",
    max_train_samples: int = 5000,
    max_eval_samples: int = 500,
) -> flyte.io.Dir:
    """Download dataset from HuggingFace and format for instruction fine-tuning."""
    from datasets import DatasetDict, load_dataset

    log.info(f"Loading dataset: {dataset_name}")
    ds = load_dataset(dataset_name, split="train")

    def format_example(ex):
        return {
            "text": (
                "### Task: Generate a SQL query to answer the question.\n"
                f"### Schema:\n{ex['context']}\n"
                f"### Question:\n{ex['question']}\n"
                f"### SQL:\n{ex['answer']}\n<|endoftext|>"
            )
        }

    ds = ds.map(format_example)

    # Split into train and eval
    total = len(ds)
    train_end = min(max_train_samples, total - max_eval_samples)
    eval_start = train_end
    eval_end = min(eval_start + max_eval_samples, total)

    processed = DatasetDict({
        "train": ds.select(range(train_end)),
        "eval": ds.select(range(eval_start, eval_end)),
    })

    output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
    processed.save_to_disk(output_dir)
    log.info(f"Dataset ready: {len(processed['train'])} train, {len(processed['eval'])} eval")

    return await flyte.io.Dir.from_local(output_dir)

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    method: str = "lora",
    epochs: int = 3,
    lr: float = 2e-4,
    batch_size: int = 4,
    lora_r: int = 16,
    lora_alpha: int = 32,
) -> flyte.io.Dir:
    """Fine-tune a model using full, LoRA, or QLoRA method."""
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
    from trl import SFTConfig, SFTTrainer

    log.info(f"Training: model={model_name}, method={method}")

    # -- Load data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)

    # -- Load tokenizer --
    token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
    tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # -- Initial report: loading model --
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Loading Model...</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="card">'
            f"<p><b>Method:</b> <span class=\"badge badge-info\">{method.upper()}</span></p>"
            f"<p><b>Dataset:</b> {len(dataset['train']):,} train / {len(dataset['eval']):,} eval</p>"
            f"</div>"
        ),
        do_flush=True,
    )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    if method == "qlora":
        from transformers import BitsAndBytesConfig

        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            **token_kwargs,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=dtype,
                bnb_4bit_use_double_quant=True,
            ),
            dtype=dtype,
            device_map="auto",
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            **token_kwargs,
            dtype=dtype,
            device_map="auto",
        )

    # -- Apply LoRA adapters --
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())

    if method in ("lora", "qlora"):
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

        if method == "qlora":
            model = prepare_model_for_kbit_training(model)

        lora_config = LoraConfig(
            r=lora_r,              # Rank — size of the low-rank matrices. Higher = more capacity but more params
            lora_alpha=lora_alpha,  # Scaling factor — controls adapter impact. Effective scale = alpha/r
            # Attention layers — LoRA adapters inject low-rank updates here:
            #   q_proj (Query)     — what to look for in context
            #   k_proj (Key)       — what each token offers to match against
            #   v_proj (Value)     — what information to extract once matched
            #   o_proj (Output)    — combines multi-head attention results
            # MLP layers — LoRA adapters also update the feed-forward network:
            #   gate_proj (Gate)   — controls how much information flows through (SwiGLU activation)
            #   up_proj (Up)       — projects to a higher dimension for richer representations
            #   down_proj (Down)   — projects back down to the model's hidden size
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            lora_dropout=0.05,     # Dropout on adapter weights — light regularization to prevent overfitting
            bias="none",           # Don't train bias terms — keeps adapter small and stable
            task_type="CAUSAL_LM", # Tells PEFT this is a text generation model (vs classification, etc.)
        )
        model = get_peft_model(model, lora_config)
        trainable_params, total_params = model.get_nb_trainable_parameters()
        log.info(f"Trainable params: {trainable_params:,} / {total_params:,} ({trainable_params / total_params * 100:.1f}%)")

    # -- Live training report state --
    training_log: list[dict] = []
    loop = asyncio.get_running_loop()

    method_badge = f'<span class="badge badge-info">{method.upper()}</span>'
    if method == "qlora":
        method_badge = f'<span class="badge badge-success">QLoRA (4-bit)</span>'
    elif method == "full":
        method_badge = f'<span class="badge badge-danger">Full Fine-Tune</span>'

    def _build_training_report(max_steps: int) -> str:
        """Build the live training report HTML from current training_log."""
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>
          <div class="stat"><div class="value">{len(dataset['train']):,}</div><div class="label">Train Examples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
        </div>
        <p>Method: {method_badge} | Total params: {total_params:,} | Trainable: {trainable_params:,}</p>
        """

        charts_html = ""
        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs} |
              Loss: <span class="highlight">{current['loss']:.4f}</span>
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_chart = make_line_chart(
                data=training_log,
                x_key="epoch",
                y_keys=["loss"],
                title="Training Loss",
                x_label="Epoch",
                y_label="Loss",
                colors=["#5a7db5"],
            )
            charts_html += f'<div class="chart-container">{loss_chart}</div>'

            if "lr" in training_log[0]:
                lr_chart = make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["lr"],
                    title="Learning Rate Schedule",
                    x_label="Epoch",
                    y_label="LR",
                    colors=["#0f3460"],
                )
                charts_html += f'<div class="chart-container">{lr_chart}</div>'

            if "grad_norm" in training_log[0]:
                grad_chart = make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["grad_norm"],
                    title="Gradient Norm",
                    x_label="Epoch",
                    y_label="Grad Norm",
                    colors=["#06d6a0"],
                )
                charts_html += f'<div class="chart-container">{grad_chart}</div>'

        return wrap_report(stats_html + charts_html)

    # -- Metrics callback with live report updates --
    class MetricsCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs or "loss" not in logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
                "loss": round(logs["loss"], 4),
            }
            if "learning_rate" in logs:
                entry["lr"] = logs["learning_rate"]
            if "grad_norm" in logs:
                entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
            training_log.append(entry)
            log.info(
                f"step={state.global_step}/{state.max_steps} "
                f"epoch={entry['epoch']:.2f} "
                f"loss={entry['loss']:.4f}"
            )

            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

    # -- Train --
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=lr,
        logging_steps=10,
        save_strategy="epoch",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        gradient_accumulation_steps=4,
        warmup_steps=10,
        report_to="none",
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["eval"],
        processing_class=tokenizer,
        callbacks=[MetricsCallback()],
    )

    log.info("Starting training...")
    await asyncio.to_thread(trainer.train)
    log.info("Training complete.")

    # -- Merge LoRA weights and save --
    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")

    if method in ("lora", "qlora"):
        log.info("Merging LoRA weights into base model...")
        model = model.merge_and_unload()

    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Final training report --
    final_loss = training_log[-1]["loss"] if training_log else "N/A"

    loss_chart = make_line_chart(
        data=training_log,
        x_key="epoch",
        y_keys=["loss"],
        title="Training Loss",
        x_label="Epoch",
        y_label="Loss",
        colors=["#5a7db5"],
    ) if training_log else ""

    lr_chart = ""
    if training_log and "lr" in training_log[0]:
        lr_chart = make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["lr"],
            title="Learning Rate Schedule",
            x_label="Epoch",
            y_label="LR",
            colors=["#0f3460"],
        )

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Training Complete</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>'
            f'  <div class="stat"><div class="value">{final_loss}</div><div class="label">Final Loss</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Trainable Params</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">% Trainable</div></div>'
            f'</div>'
            f'<div class="chart-container">{loss_chart}</div>'
            f'{f"""<div class="chart-container">{lr_chart}</div>""" if lr_chart else ""}'
        ),
        do_flush=True,
    )

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Task 3: Evaluate — before/after comparison
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    model_name: str,
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 50,
) -> str:
    """Compare base model vs fine-tuned model on test examples."""
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(
        wrap_report(
            "<h2>Evaluation</h2>"
            '<div class="card"><p>Loading models and running inference...</p></div>'
        ),
        do_flush=True,
    )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    # Load eval data
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"].select(range(min(num_examples, len(dataset["eval"]))))

    # Load tokenizer
    token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
    tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def generate_sql(model, prompt, max_new_tokens=128):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()

    def normalize_sql(sql):
        """Extract the first SQL statement and normalize for comparison."""
        # Truncate at first ### or newline to isolate the SQL
        for stop in ["###", "\n"]:
            if stop in sql:
                sql = sql[:sql.index(stop)]
        return " ".join(sql.lower().split()).strip().rstrip(";")

    def build_prompt(example):
        return (
            "### Task: Generate a SQL query to answer the question.\n"
            f"### Schema:\n{example['context']}\n"
            f"### Question:\n{example['question']}\n"
            "### SQL:\n"
        )

    # -- Run base model --
    log.info(f"Loading base model: {model_name}")
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{len(eval_ds)}</div><div class="label">Eval Examples</div></div>'
            f'  <div class="stat"><div class="value">1/2</div><div class="label">Phase</div></div>'
            f'</div>'
            f'<div class="card"><p>Running <b>base model</b> inference...</p>'
            f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
            f'<div style="background:#adb5bd;width:25%;height:100%;border-radius:4px;"></div>'
            f'</div></div>'
        ),
        do_flush=True,
    )

    base_model = AutoModelForCausalLM.from_pretrained(
        model_name, **token_kwargs, dtype=dtype, device_map="auto",
    )

    base_results = []
    for i, example in enumerate(eval_ds):
        prompt = build_prompt(example)
        generated = generate_sql(base_model, prompt)
        base_results.append(generated)
        if (i + 1) % 10 == 0:
            log.info(f"Base model: {i + 1}/{len(eval_ds)}")
            pct = (i + 1) / len(eval_ds) * 50
            await flyte.report.replace.aio(
                wrap_report(
                    f"<h2>Evaluation</h2>"
                    f'<div class="card"><p>Running <b>base model</b> inference... {i + 1}/{len(eval_ds)}</p>'
                    f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                    f'<div style="background:#adb5bd;width:{pct:.0f}%;height:100%;border-radius:4px;"></div>'
                    f'</div></div>'
                ),
                do_flush=True,
            )

    del base_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Run fine-tuned model --
    log.info("Loading fine-tuned model...")
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation</h2>"
            f'<div class="card"><p>Running <b>fine-tuned model</b> inference...</p>'
            f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
            f'<div style="background:#0f3460;width:50%;height:100%;border-radius:4px;"></div>'
            f'</div></div>'
        ),
        do_flush=True,
    )

    ft_path = await finetuned_dir.download()
    ft_model = AutoModelForCausalLM.from_pretrained(
        ft_path, dtype=dtype, device_map="auto",
    )

    ft_results = []
    for i, example in enumerate(eval_ds):
        prompt = build_prompt(example)
        generated = generate_sql(ft_model, prompt)
        ft_results.append(generated)
        if (i + 1) % 10 == 0:
            log.info(f"Fine-tuned model: {i + 1}/{len(eval_ds)}")
            pct = 50 + (i + 1) / len(eval_ds) * 50
            await flyte.report.replace.aio(
                wrap_report(
                    f"<h2>Evaluation</h2>"
                    f'<div class="card"><p>Running <b>fine-tuned model</b> inference... {i + 1}/{len(eval_ds)}</p>'
                    f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                    f'<div style="background:#0f3460;width:{pct:.0f}%;height:100%;border-radius:4px;"></div>'
                    f'</div></div>'
                ),
                do_flush=True,
            )

    del ft_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Score --
    base_correct = 0
    ft_correct = 0
    comparisons = []

    for i, example in enumerate(eval_ds):
        expected = example["answer"]
        base_gen = base_results[i]
        ft_gen = ft_results[i]

        base_match = normalize_sql(base_gen) == normalize_sql(expected)
        ft_match = normalize_sql(ft_gen) == normalize_sql(expected)

        if base_match:
            base_correct += 1
        if ft_match:
            ft_correct += 1

        comparisons.append({
            "question": example["question"],
            "schema": example["context"],
            "expected": expected,
            "base": base_gen,
            "finetuned": ft_gen,
            "base_correct": base_match,
            "ft_correct": ft_match,
        })

    total = len(eval_ds)
    base_acc = base_correct / total * 100
    ft_acc = ft_correct / total * 100
    improvement = ft_acc - base_acc

    log.info(f"Base model accuracy: {base_acc:.1f}% ({base_correct}/{total})")
    log.info(f"Fine-tuned accuracy: {ft_acc:.1f}% ({ft_correct}/{total})")

    # -- Build final eval report --
    improvement_badge = (
        f'<span class="badge badge-success">+{improvement:.1f}pp</span>'
        if improvement > 0
        else f'<span class="badge badge-danger">{improvement:.1f}pp</span>'
    )

    bar_chart = make_bar_chart(
        labels=["Exact Match Accuracy"],
        series={
            "Base Model": [base_acc],
            "Fine-Tuned": [ft_acc],
        },
        title="Base vs Fine-Tuned Accuracy",
        colors=["#adb5bd", "#0f3460"],
        y_max_cap=100.0,
    )

    examples_html = ""
    for c in comparisons[:10]:
        base_badge = '<span class="badge badge-success">correct</span>' if c["base_correct"] else '<span class="badge badge-danger">wrong</span>'
        ft_badge = '<span class="badge badge-success">correct</span>' if c["ft_correct"] else '<span class="badge badge-danger">wrong</span>'
        examples_html += f"""
        <div class="card">
          <p><b>Q:</b> {c['question']}</p>
          <p style="font-size:0.85em; color:#6c757d;"><b>Schema:</b> {c['schema'][:200]}...</p>
          <table>
            <tr><th>Source</th><th>SQL</th><th>Result</th></tr>
            <tr><td>Expected</td><td><code>{c['expected']}</code></td><td></td></tr>
            <tr><td>Base</td><td><code>{c['base'][:200]}</code></td><td>{base_badge}</td></tr>
            <tr><td>Fine-tuned</td><td><code>{c['finetuned'][:200]}</code></td><td>{ft_badge}</td></tr>
          </table>
        </div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation Results</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{base_acc:.1f}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{ft_acc:.1f}%</div><div class="label">Fine-Tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{improvement:+.1f}pp</div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{total}</div><div class="label">Eval Examples</div></div>'
            f'</div>'
            f'<div class="chart-container">{bar_chart}</div>'
            f'<h3>Example Comparisons {improvement_badge}</h3>'
            f'{examples_html}'
            f'<div class="note">'
            f'<b>Note:</b> Exact match accuracy compares normalized SQL output. '
            f'The fine-tuned model may generate semantically correct queries that differ in formatting.'
            f'</div>'
        ),
        do_flush=True,
    )

    return json.dumps({
        "base_accuracy": round(base_acc, 1),
        "finetuned_accuracy": round(ft_acc, 1),
        "improvement": round(ft_acc - base_acc, 1),
        "num_examples": total,
        "comparisons": comparisons[:10],
    })

# ------------------------------------------------------------------
# Pipeline: orchestrate everything
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "HuggingFaceTB/SmolLM2-135M",
    dataset_name: str = "b-mc2/sql-create-context",
    method: str = "lora",
    epochs: int = 3,
    lr: float = 2e-4,
    batch_size: int = 4,
    max_train_samples: int = 5000,
    max_eval_samples: int = 500,
    num_eval_examples: int = 50,
    lora_r: int = 16,
    lora_alpha: int = 32,
) -> flyte.io.Dir:
    """
    End-to-end LLM fine-tuning pipeline.

    1. Download and format dataset
    2. Fine-tune model (full / LoRA / QLoRA)
    3. Evaluate: before/after comparison on test set

    Returns the fine-tuned model directory so it can be served directly.
    """
    log.info(f"Pipeline: {model_name} | method={method} | dataset={dataset_name}")
    steps = ["Prepare Data", "Train", "Evaluate"]

    method_badge = f'<span class="badge badge-info">{method.upper()}</span>'

    # Step 1: Prepare data
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(0, steps)}'
            f'<div class="card"><p>Downloading and formatting dataset: <b>{dataset_name}</b>...</p></div>'
        ),
        do_flush=True,
    )

    data_dir = await prepare_data(dataset_name, max_train_samples, max_eval_samples)

    # Step 2: Train
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(1, steps)}'
            f'<div class="card"><p>Training in progress... check the <b>train</b> task report for live charts.</p></div>'
        ),
        do_flush=True,
    )

    finetuned_dir = await train(
        model_name, data_dir, method, epochs, lr, batch_size, lora_r, lora_alpha,
    )

    # Step 3: Evaluate
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(2, steps)}'
            f'<div class="card"><p>Evaluating base vs fine-tuned model...</p></div>'
        ),
        do_flush=True,
    )

    result = await evaluate(model_name, finetuned_dir, data_dir, num_eval_examples)
    metrics = json.loads(result)

    # Final pipeline report
    improvement = metrics["improvement"]
    improvement_badge = (
        f'<span class="badge badge-success">+{improvement:.1f}pp</span>'
        if improvement > 0
        else f'<span class="badge badge-danger">{improvement:.1f}pp</span>'
    )

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Pipeline Complete</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(3, steps)}'
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{metrics["base_accuracy"]}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{metrics["finetuned_accuracy"]}%</div><div class="label">Fine-Tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{improvement:+.1f}pp</div><div class="label">Improvement {improvement_badge}</div></div>'
            f'  <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{metrics["num_examples"]}</div><div class="label">Eval Examples</div></div>'
            f'</div>'
            f'<div class="note">'
            f'Check the <b>train</b> task report for training loss/LR charts, '
            f'and the <b>evaluate</b> task report for detailed example comparisons.'
            f'</div>'
        ),
        do_flush=True,
    )

    log.info(f"Pipeline complete. Improvement: {metrics['improvement']:+.1f}pp")
    return finetuned_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/llm_fine_tuning_lora_qlora/llm_fine_tuning_lora_qlora.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "peft>=0.13.0",
#    "trl>=0.12.0",
#    "bitsandbytes>=0.44.0",
#    ...
# ]
# ///
```

## Orchestrate the pipeline

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "peft>=0.13.0",
#    "trl>=0.12.0",
#    "datasets>=3.0.0",
#    "bitsandbytes>=0.44.0",
#    "accelerate>=0.34.0",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import json
import logging
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
import os

main_img = flyte.Image.from_uv_script(__file__, name="llm-fine-tuning-lora-qlora", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

cpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)

HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}

from report_helpers import make_bar_chart, make_line_chart, pipeline_step_indicator, wrap_report

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

# ------------------------------------------------------------------
# Task 1: Prepare dataset
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def prepare_data(
    dataset_name: str = "b-mc2/sql-create-context",
    max_train_samples: int = 5000,
    max_eval_samples: int = 500,
) -> flyte.io.Dir:
    """Download dataset from HuggingFace and format for instruction fine-tuning."""
    from datasets import DatasetDict, load_dataset

    log.info(f"Loading dataset: {dataset_name}")
    ds = load_dataset(dataset_name, split="train")

    def format_example(ex):
        return {
            "text": (
                "### Task: Generate a SQL query to answer the question.\n"
                f"### Schema:\n{ex['context']}\n"
                f"### Question:\n{ex['question']}\n"
                f"### SQL:\n{ex['answer']}\n<|endoftext|>"
            )
        }

    ds = ds.map(format_example)

    # Split into train and eval
    total = len(ds)
    train_end = min(max_train_samples, total - max_eval_samples)
    eval_start = train_end
    eval_end = min(eval_start + max_eval_samples, total)

    processed = DatasetDict({
        "train": ds.select(range(train_end)),
        "eval": ds.select(range(eval_start, eval_end)),
    })

    output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
    processed.save_to_disk(output_dir)
    log.info(f"Dataset ready: {len(processed['train'])} train, {len(processed['eval'])} eval")

    return await flyte.io.Dir.from_local(output_dir)

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    method: str = "lora",
    epochs: int = 3,
    lr: float = 2e-4,
    batch_size: int = 4,
    lora_r: int = 16,
    lora_alpha: int = 32,
) -> flyte.io.Dir:
    """Fine-tune a model using full, LoRA, or QLoRA method."""
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
    from trl import SFTConfig, SFTTrainer

    log.info(f"Training: model={model_name}, method={method}")

    # -- Load data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)

    # -- Load tokenizer --
    token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
    tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # -- Initial report: loading model --
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Loading Model...</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="card">'
            f"<p><b>Method:</b> <span class=\"badge badge-info\">{method.upper()}</span></p>"
            f"<p><b>Dataset:</b> {len(dataset['train']):,} train / {len(dataset['eval']):,} eval</p>"
            f"</div>"
        ),
        do_flush=True,
    )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    if method == "qlora":
        from transformers import BitsAndBytesConfig

        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            **token_kwargs,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=dtype,
                bnb_4bit_use_double_quant=True,
            ),
            dtype=dtype,
            device_map="auto",
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            **token_kwargs,
            dtype=dtype,
            device_map="auto",
        )

    # -- Apply LoRA adapters --
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())

    if method in ("lora", "qlora"):
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

        if method == "qlora":
            model = prepare_model_for_kbit_training(model)

        lora_config = LoraConfig(
            r=lora_r,              # Rank — size of the low-rank matrices. Higher = more capacity but more params
            lora_alpha=lora_alpha,  # Scaling factor — controls adapter impact. Effective scale = alpha/r
            # Attention layers — LoRA adapters inject low-rank updates here:
            #   q_proj (Query)     — what to look for in context
            #   k_proj (Key)       — what each token offers to match against
            #   v_proj (Value)     — what information to extract once matched
            #   o_proj (Output)    — combines multi-head attention results
            # MLP layers — LoRA adapters also update the feed-forward network:
            #   gate_proj (Gate)   — controls how much information flows through (SwiGLU activation)
            #   up_proj (Up)       — projects to a higher dimension for richer representations
            #   down_proj (Down)   — projects back down to the model's hidden size
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            lora_dropout=0.05,     # Dropout on adapter weights — light regularization to prevent overfitting
            bias="none",           # Don't train bias terms — keeps adapter small and stable
            task_type="CAUSAL_LM", # Tells PEFT this is a text generation model (vs classification, etc.)
        )
        model = get_peft_model(model, lora_config)
        trainable_params, total_params = model.get_nb_trainable_parameters()
        log.info(f"Trainable params: {trainable_params:,} / {total_params:,} ({trainable_params / total_params * 100:.1f}%)")

    # -- Live training report state --
    training_log: list[dict] = []
    loop = asyncio.get_running_loop()

    method_badge = f'<span class="badge badge-info">{method.upper()}</span>'
    if method == "qlora":
        method_badge = f'<span class="badge badge-success">QLoRA (4-bit)</span>'
    elif method == "full":
        method_badge = f'<span class="badge badge-danger">Full Fine-Tune</span>'

    def _build_training_report(max_steps: int) -> str:
        """Build the live training report HTML from current training_log."""
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>
          <div class="stat"><div class="value">{len(dataset['train']):,}</div><div class="label">Train Examples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">Trainable</div></div>
        </div>
        <p>Method: {method_badge} | Total params: {total_params:,} | Trainable: {trainable_params:,}</p>
        """

        charts_html = ""
        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs} |
              Loss: <span class="highlight">{current['loss']:.4f}</span>
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_chart = make_line_chart(
                data=training_log,
                x_key="epoch",
                y_keys=["loss"],
                title="Training Loss",
                x_label="Epoch",
                y_label="Loss",
                colors=["#5a7db5"],
            )
            charts_html += f'<div class="chart-container">{loss_chart}</div>'

            if "lr" in training_log[0]:
                lr_chart = make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["lr"],
                    title="Learning Rate Schedule",
                    x_label="Epoch",
                    y_label="LR",
                    colors=["#0f3460"],
                )
                charts_html += f'<div class="chart-container">{lr_chart}</div>'

            if "grad_norm" in training_log[0]:
                grad_chart = make_line_chart(
                    data=training_log,
                    x_key="epoch",
                    y_keys=["grad_norm"],
                    title="Gradient Norm",
                    x_label="Epoch",
                    y_label="Grad Norm",
                    colors=["#06d6a0"],
                )
                charts_html += f'<div class="chart-container">{grad_chart}</div>'

        return wrap_report(stats_html + charts_html)

    # -- Metrics callback with live report updates --
    class MetricsCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs or "loss" not in logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
                "loss": round(logs["loss"], 4),
            }
            if "learning_rate" in logs:
                entry["lr"] = logs["learning_rate"]
            if "grad_norm" in logs:
                entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
            training_log.append(entry)
            log.info(
                f"step={state.global_step}/{state.max_steps} "
                f"epoch={entry['epoch']:.2f} "
                f"loss={entry['loss']:.4f}"
            )

            asyncio.run_coroutine_threadsafe(
                flyte.report.replace.aio(
                    _build_training_report(state.max_steps),
                    do_flush=True,
                ),
                loop,
            )

    # -- Train --
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=lr,
        logging_steps=10,
        save_strategy="epoch",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        gradient_accumulation_steps=4,
        warmup_steps=10,
        report_to="none",
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["eval"],
        processing_class=tokenizer,
        callbacks=[MetricsCallback()],
    )

    log.info("Starting training...")
    await asyncio.to_thread(trainer.train)
    log.info("Training complete.")

    # -- Merge LoRA weights and save --
    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")

    if method in ("lora", "qlora"):
        log.info("Merging LoRA weights into base model...")
        model = model.merge_and_unload()

    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Final training report --
    final_loss = training_log[-1]["loss"] if training_log else "N/A"

    loss_chart = make_line_chart(
        data=training_log,
        x_key="epoch",
        y_keys=["loss"],
        title="Training Loss",
        x_label="Epoch",
        y_label="Loss",
        colors=["#5a7db5"],
    ) if training_log else ""

    lr_chart = ""
    if training_log and "lr" in training_log[0]:
        lr_chart = make_line_chart(
            data=training_log,
            x_key="epoch",
            y_keys=["lr"],
            title="Learning Rate Schedule",
            x_label="Epoch",
            y_label="LR",
            colors=["#0f3460"],
        )

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Training Complete</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>'
            f'  <div class="stat"><div class="value">{final_loss}</div><div class="label">Final Loss</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{total_params:,}</div><div class="label">Total Params</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Trainable Params</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params / total_params * 100:.1f}%</div><div class="label">% Trainable</div></div>'
            f'</div>'
            f'<div class="chart-container">{loss_chart}</div>'
            f'{f"""<div class="chart-container">{lr_chart}</div>""" if lr_chart else ""}'
        ),
        do_flush=True,
    )

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Task 3: Evaluate — before/after comparison
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    model_name: str,
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 50,
) -> str:
    """Compare base model vs fine-tuned model on test examples."""
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForCausalLM, AutoTokenizer

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(
        wrap_report(
            "<h2>Evaluation</h2>"
            '<div class="card"><p>Loading models and running inference...</p></div>'
        ),
        do_flush=True,
    )

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    # Load eval data
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"].select(range(min(num_examples, len(dataset["eval"]))))

    # Load tokenizer
    token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
    tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def generate_sql(model, prompt, max_new_tokens=128):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()

    def normalize_sql(sql):
        """Extract the first SQL statement and normalize for comparison."""
        # Truncate at first ### or newline to isolate the SQL
        for stop in ["###", "\n"]:
            if stop in sql:
                sql = sql[:sql.index(stop)]
        return " ".join(sql.lower().split()).strip().rstrip(";")

    def build_prompt(example):
        return (
            "### Task: Generate a SQL query to answer the question.\n"
            f"### Schema:\n{example['context']}\n"
            f"### Question:\n{example['question']}\n"
            "### SQL:\n"
        )

    # -- Run base model --
    log.info(f"Loading base model: {model_name}")
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{len(eval_ds)}</div><div class="label">Eval Examples</div></div>'
            f'  <div class="stat"><div class="value">1/2</div><div class="label">Phase</div></div>'
            f'</div>'
            f'<div class="card"><p>Running <b>base model</b> inference...</p>'
            f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
            f'<div style="background:#adb5bd;width:25%;height:100%;border-radius:4px;"></div>'
            f'</div></div>'
        ),
        do_flush=True,
    )

    base_model = AutoModelForCausalLM.from_pretrained(
        model_name, **token_kwargs, dtype=dtype, device_map="auto",
    )

    base_results = []
    for i, example in enumerate(eval_ds):
        prompt = build_prompt(example)
        generated = generate_sql(base_model, prompt)
        base_results.append(generated)
        if (i + 1) % 10 == 0:
            log.info(f"Base model: {i + 1}/{len(eval_ds)}")
            pct = (i + 1) / len(eval_ds) * 50
            await flyte.report.replace.aio(
                wrap_report(
                    f"<h2>Evaluation</h2>"
                    f'<div class="card"><p>Running <b>base model</b> inference... {i + 1}/{len(eval_ds)}</p>'
                    f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                    f'<div style="background:#adb5bd;width:{pct:.0f}%;height:100%;border-radius:4px;"></div>'
                    f'</div></div>'
                ),
                do_flush=True,
            )

    del base_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Run fine-tuned model --
    log.info("Loading fine-tuned model...")
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation</h2>"
            f'<div class="card"><p>Running <b>fine-tuned model</b> inference...</p>'
            f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
            f'<div style="background:#0f3460;width:50%;height:100%;border-radius:4px;"></div>'
            f'</div></div>'
        ),
        do_flush=True,
    )

    ft_path = await finetuned_dir.download()
    ft_model = AutoModelForCausalLM.from_pretrained(
        ft_path, dtype=dtype, device_map="auto",
    )

    ft_results = []
    for i, example in enumerate(eval_ds):
        prompt = build_prompt(example)
        generated = generate_sql(ft_model, prompt)
        ft_results.append(generated)
        if (i + 1) % 10 == 0:
            log.info(f"Fine-tuned model: {i + 1}/{len(eval_ds)}")
            pct = 50 + (i + 1) / len(eval_ds) * 50
            await flyte.report.replace.aio(
                wrap_report(
                    f"<h2>Evaluation</h2>"
                    f'<div class="card"><p>Running <b>fine-tuned model</b> inference... {i + 1}/{len(eval_ds)}</p>'
                    f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                    f'<div style="background:#0f3460;width:{pct:.0f}%;height:100%;border-radius:4px;"></div>'
                    f'</div></div>'
                ),
                do_flush=True,
            )

    del ft_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Score --
    base_correct = 0
    ft_correct = 0
    comparisons = []

    for i, example in enumerate(eval_ds):
        expected = example["answer"]
        base_gen = base_results[i]
        ft_gen = ft_results[i]

        base_match = normalize_sql(base_gen) == normalize_sql(expected)
        ft_match = normalize_sql(ft_gen) == normalize_sql(expected)

        if base_match:
            base_correct += 1
        if ft_match:
            ft_correct += 1

        comparisons.append({
            "question": example["question"],
            "schema": example["context"],
            "expected": expected,
            "base": base_gen,
            "finetuned": ft_gen,
            "base_correct": base_match,
            "ft_correct": ft_match,
        })

    total = len(eval_ds)
    base_acc = base_correct / total * 100
    ft_acc = ft_correct / total * 100
    improvement = ft_acc - base_acc

    log.info(f"Base model accuracy: {base_acc:.1f}% ({base_correct}/{total})")
    log.info(f"Fine-tuned accuracy: {ft_acc:.1f}% ({ft_correct}/{total})")

    # -- Build final eval report --
    improvement_badge = (
        f'<span class="badge badge-success">+{improvement:.1f}pp</span>'
        if improvement > 0
        else f'<span class="badge badge-danger">{improvement:.1f}pp</span>'
    )

    bar_chart = make_bar_chart(
        labels=["Exact Match Accuracy"],
        series={
            "Base Model": [base_acc],
            "Fine-Tuned": [ft_acc],
        },
        title="Base vs Fine-Tuned Accuracy",
        colors=["#adb5bd", "#0f3460"],
        y_max_cap=100.0,
    )

    examples_html = ""
    for c in comparisons[:10]:
        base_badge = '<span class="badge badge-success">correct</span>' if c["base_correct"] else '<span class="badge badge-danger">wrong</span>'
        ft_badge = '<span class="badge badge-success">correct</span>' if c["ft_correct"] else '<span class="badge badge-danger">wrong</span>'
        examples_html += f"""
        <div class="card">
          <p><b>Q:</b> {c['question']}</p>
          <p style="font-size:0.85em; color:#6c757d;"><b>Schema:</b> {c['schema'][:200]}...</p>
          <table>
            <tr><th>Source</th><th>SQL</th><th>Result</th></tr>
            <tr><td>Expected</td><td><code>{c['expected']}</code></td><td></td></tr>
            <tr><td>Base</td><td><code>{c['base'][:200]}</code></td><td>{base_badge}</td></tr>
            <tr><td>Fine-tuned</td><td><code>{c['finetuned'][:200]}</code></td><td>{ft_badge}</td></tr>
          </table>
        </div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation Results</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{base_acc:.1f}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{ft_acc:.1f}%</div><div class="label">Fine-Tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{improvement:+.1f}pp</div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{total}</div><div class="label">Eval Examples</div></div>'
            f'</div>'
            f'<div class="chart-container">{bar_chart}</div>'
            f'<h3>Example Comparisons {improvement_badge}</h3>'
            f'{examples_html}'
            f'<div class="note">'
            f'<b>Note:</b> Exact match accuracy compares normalized SQL output. '
            f'The fine-tuned model may generate semantically correct queries that differ in formatting.'
            f'</div>'
        ),
        do_flush=True,
    )

    return json.dumps({
        "base_accuracy": round(base_acc, 1),
        "finetuned_accuracy": round(ft_acc, 1),
        "improvement": round(ft_acc - base_acc, 1),
        "num_examples": total,
        "comparisons": comparisons[:10],
    })

# ------------------------------------------------------------------
# Pipeline: orchestrate everything
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "HuggingFaceTB/SmolLM2-135M",
    dataset_name: str = "b-mc2/sql-create-context",
    method: str = "lora",
    epochs: int = 3,
    lr: float = 2e-4,
    batch_size: int = 4,
    max_train_samples: int = 5000,
    max_eval_samples: int = 500,
    num_eval_examples: int = 50,
    lora_r: int = 16,
    lora_alpha: int = 32,
) -> flyte.io.Dir:
    """
    End-to-end LLM fine-tuning pipeline.

    1. Download and format dataset
    2. Fine-tune model (full / LoRA / QLoRA)
    3. Evaluate: before/after comparison on test set

    Returns the fine-tuned model directory so it can be served directly.
    """
    log.info(f"Pipeline: {model_name} | method={method} | dataset={dataset_name}")
    steps = ["Prepare Data", "Train", "Evaluate"]

    method_badge = f'<span class="badge badge-info">{method.upper()}</span>'

    # Step 1: Prepare data
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(0, steps)}'
            f'<div class="card"><p>Downloading and formatting dataset: <b>{dataset_name}</b>...</p></div>'
        ),
        do_flush=True,
    )

    data_dir = await prepare_data(dataset_name, max_train_samples, max_eval_samples)

    # Step 2: Train
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(1, steps)}'
            f'<div class="card"><p>Training in progress... check the <b>train</b> task report for live charts.</p></div>'
        ),
        do_flush=True,
    )

    finetuned_dir = await train(
        model_name, data_dir, method, epochs, lr, batch_size, lora_r, lora_alpha,
    )

    # Step 3: Evaluate
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>LLM Fine-Tuning Pipeline</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(2, steps)}'
            f'<div class="card"><p>Evaluating base vs fine-tuned model...</p></div>'
        ),
        do_flush=True,
    )

    result = await evaluate(model_name, finetuned_dir, data_dir, num_eval_examples)
    metrics = json.loads(result)

    # Final pipeline report
    improvement = metrics["improvement"]
    improvement_badge = (
        f'<span class="badge badge-success">+{improvement:.1f}pp</span>'
        if improvement > 0
        else f'<span class="badge badge-danger">{improvement:.1f}pp</span>'
    )

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Pipeline Complete</h2>"
            f"<h3>{model_name} {method_badge}</h3>"
            f'{pipeline_step_indicator(3, steps)}'
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{metrics["base_accuracy"]}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{metrics["finetuned_accuracy"]}%</div><div class="label">Fine-Tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{improvement:+.1f}pp</div><div class="label">Improvement {improvement_badge}</div></div>'
            f'  <div class="stat"><div class="value">{method.upper()}</div><div class="label">Method</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{metrics["num_examples"]}</div><div class="label">Eval Examples</div></div>'
            f'</div>'
            f'<div class="note">'
            f'Check the <b>train</b> task report for training loss/LR charts, '
            f'and the <b>evaluate</b> task report for detailed example comparisons.'
            f'</div>'
        ),
        do_flush=True,
    )

    log.info(f"Pipeline complete. Improvement: {metrics['improvement']:+.1f}pp")
    return finetuned_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/llm_fine_tuning_lora_qlora/llm_fine_tuning_lora_qlora.py*

## Run the workflow

Create a HuggingFace token secret if you use a gated base model:

```
flyte create secret huggingface-token <YOUR_HF_TOKEN>
```

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/llm_fine_tuning_lora_qlora):

```
cd v2/tutorials/llm_fine_tuning_lora_qlora
uv run --script llm_fine_tuning_lora_qlora.py
```

Try QLoRA on a GPU:

```
flyte run llm_fine_tuning_lora_qlora.py pipeline --method qlora --epochs 3
```

QLoRA requires CUDA; LoRA and full fine-tuning follow the same entry point with different memory requirements.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training/bert-fine-tuning-emotion ===

# BERT emotion classification

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/bert_fine_tuning_emotion).

This tutorial fine-tunes a BERT-style model (ModernBERT by default) on the [dair-ai/emotion](https://huggingface.co/datasets/dair-ai/emotion) Twitter dataset for six-way emotion classification: sadness, joy, love, anger, fear, and surprise. The pipeline trains the classifier, evaluates with a confusion matrix and per-class F1, and explores inference with attention and token-importance visualizations in Flyte reports.

Flyte provides:

- **GPU fine-tuning** with live training loss charts.
- **Rich evaluation reports** including confusion matrices and confidence bars.
- **Cached dataset loading** for repeatable experiments.

## Define the task environments

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "datasets>=3.0.0",
#    "accelerate>=0.34.0",
#    "scikit-learn",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
import os

main_img = flyte.Image.from_uv_script(__file__, name="bert-fine-tuning-emotion", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="bert-fine-tuning-emotion-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu=1),
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

cpu_env = flyte.TaskEnvironment(
    name="bert-fine-tuning-emotion-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)

HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}

from report_helpers import (
    make_attention_text,
    make_bar_chart,
    make_confidence_bars,
    make_confusion_matrix,
    make_line_chart,
    make_token_importance_text,
    pipeline_step_indicator,
    wrap_report,
)

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

EMOTION_LABELS = ["sadness", "joy", "love", "anger", "fear", "surprise"]
EMOTION_DATASET = "dair-ai/emotion"

# ------------------------------------------------------------------
# Task 1: Get data
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def get_data(
    max_train_samples: int = 10000,
    max_eval_samples: int = 2000,
) -> flyte.io.Dir:
    """Download the emotion dataset and save train/eval splits.

    The dair-ai/emotion dataset contains ~20k English Twitter messages labeled
    with one of 6 emotions: sadness, joy, love, anger, fear, surprise.
    """
    from datasets import DatasetDict, load_dataset

    log.info("Loading emotion dataset...")
    ds = load_dataset(EMOTION_DATASET)

    train_ds = ds["train"].shuffle(seed=42).select(range(min(max_train_samples, len(ds["train"]))))
    eval_ds = ds["test"].shuffle(seed=42).select(range(min(max_eval_samples, len(ds["test"]))))

    processed = DatasetDict({"train": train_ds, "eval": eval_ds})

    output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
    processed.save_to_disk(output_dir)
    log.info(f"Dataset ready: {len(train_ds)} train, {len(eval_ds)} eval")

    return await flyte.io.Dir.from_local(output_dir)

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    epochs: int = 3,
    lr: float = 2e-5,
    batch_size: int = 16,
    warmup_steps: int = 100,
) -> flyte.io.Dir:
    """Fine-tune a BERT-style model for 6-class emotion classification."""
    import numpy as np
    import torch
    from datasets import load_from_disk
    from sklearn.metrics import accuracy_score, f1_score
    from transformers import (
        AutoModelForSequenceClassification,
        AutoTokenizer,
        Trainer,
        TrainerCallback,
        TrainingArguments,
    )

    log.info(f"Training: model={model_name}")

    id2label = {i: l for i, l in enumerate(EMOTION_LABELS)}
    label2id = {l: i for i, l in enumerate(EMOTION_LABELS)}

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Loading Model...</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="card"><p>Preparing for emotion classification training...</p></div>'
        ),
        do_flush=True,
    )

    # -- Load data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)

    # -- Tokenize --
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)

    def tokenize(examples):
        return tokenizer(examples["text"], truncation=True, max_length=128, padding="max_length")

    dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])

    # -- Load model --
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        token=HF_TOKEN,
        num_labels=6,
        id2label=id2label,
        label2id=label2id,
    )

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log.info(f"Parameters: {trainable_params:,} / {total_params:,}")

    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        log.info(f"GPU: {gpu_name} ({gpu_mem:.1f} GB)")

    # -- Metrics tracking for live report --
    training_log: list[dict] = []
    eval_log: list[dict] = []

    def _build_training_report(max_steps: int) -> str:
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{len(dataset['train']):,}</div><div class="label">Train Samples</div></div>
          <div class="stat"><div class="value">{len(dataset['eval']):,}</div><div class="label">Eval Samples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Parameters</div></div>
        </div>
        """

        charts_html = ""

        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            loss_display = f"Loss: <span class=\"highlight\">{current['loss']:.4f}</span>" if current.get("loss") else ""
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs}
              {f' | {loss_display}' if loss_display else ''}
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_entries = [e for e in training_log if "loss" in e]
            if len(loss_entries) >= 2:
                loss_chart = make_line_chart(
                    data=loss_entries,
                    x_key="epoch",
                    y_keys=["loss"],
                    title="Training Loss",
                    x_label="Epoch",
                    y_label="Loss",
                    colors=["#5a7db5"],
                )
                charts_html += f'<div class="chart-container">{loss_chart}</div>'

        if eval_log:
            latest_eval = eval_log[-1]
            best_acc = max(e.get("accuracy", 0) for e in eval_log)
            best_f1 = max(e.get("f1", 0) for e in eval_log)
            charts_html += f"""
            <div class="stat-grid" style="margin-top:16px;">
              <div class="stat"><div class="value">{latest_eval.get('accuracy', 0):.1%}</div><div class="label">Eval Accuracy</div></div>
              <div class="stat"><div class="value">{latest_eval.get('f1', 0):.1%}</div><div class="label">Eval F1</div></div>
              <div class="stat"><div class="value">{best_acc:.1%}</div><div class="label">Best Accuracy</div></div>
              <div class="stat"><div class="value">{latest_eval.get('eval_loss', 0):.4f}</div><div class="label">Eval Loss</div></div>
            </div>
            """

            if len(eval_log) >= 2:
                eval_chart = make_line_chart(
                    data=eval_log,
                    x_key="epoch",
                    y_keys=["accuracy", "f1"],
                    title="Eval Metrics Over Training",
                    x_label="Epoch",
                    y_label="Score",
                    colors=["#0f3460", "#06d6a0"],
                    y_max_cap=1.05,
                    y_display_names={"accuracy": "Accuracy", "f1": "Weighted F1"},
                )
                charts_html += f'<div class="chart-container">{eval_chart}</div>'

                eval_loss_chart = make_line_chart(
                    data=[e for e in eval_log if "eval_loss" in e],
                    x_key="epoch",
                    y_keys=["eval_loss"],
                    title="Eval Loss",
                    x_label="Epoch",
                    y_label="Loss",
                    colors=["#e63946"],
                )
                if any("eval_loss" in e for e in eval_log):
                    charts_html += f'<div class="chart-container">{eval_loss_chart}</div>'

        return wrap_report(stats_html + charts_html)

    # -- Callbacks --
    class ReportCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
            }
            if "loss" in logs:
                entry["loss"] = round(logs["loss"], 4)
            if "eval_accuracy" in logs:
                eval_log.append({
                    "epoch": entry["epoch"],
                    "accuracy": logs["eval_accuracy"],
                    "f1": logs.get("eval_f1", 0),
                    "eval_loss": logs.get("eval_loss", 0),
                })
            if "loss" in entry:
                training_log.append(entry)

            flyte.report.replace(
                _build_training_report(state.max_steps),
                do_flush=True,
            )

    # -- Compute metrics --
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        return {
            "accuracy": accuracy_score(labels, preds),
            "f1": f1_score(labels, preds, average="weighted"),
        }

    # -- Training --
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        learning_rate=lr,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        warmup_steps=warmup_steps,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["eval"],
        processing_class=tokenizer,
        compute_metrics=compute_metrics,
        callbacks=[ReportCallback()],
    )

    log.info("Starting training...")
    await flyte.report.replace.aio(
        _build_training_report(0),
        do_flush=True,
    )

    trainer.train()
    log.info("Training complete.")

    # -- Save model --
    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
    trainer.save_model(save_dir)
    tokenizer.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Final eval + report --
    metrics = trainer.evaluate()
    final_acc = metrics.get("eval_accuracy", 0)
    final_f1 = metrics.get("eval_f1", 0)

    final_charts = ""
    loss_entries = [e for e in training_log if "loss" in e]
    if len(loss_entries) >= 2:
        loss_chart = make_line_chart(
            data=loss_entries,
            x_key="epoch",
            y_keys=["loss"],
            title="Training Loss",
            x_label="Epoch",
            y_label="Loss",
            colors=["#5a7db5"],
        )
        final_charts += f'<div class="chart-container">{loss_chart}</div>'

    if len(eval_log) >= 2:
        eval_chart = make_line_chart(
            data=eval_log,
            x_key="epoch",
            y_keys=["accuracy", "f1"],
            title="Eval Metrics Over Training",
            x_label="Epoch",
            y_label="Score",
            colors=["#0f3460", "#06d6a0"],
            y_max_cap=1.05,
            y_display_names={"accuracy": "Accuracy", "f1": "Weighted F1"},
        )
        final_charts += f'<div class="chart-container">{eval_chart}</div>'

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Training Complete</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{final_acc:.1%}</div><div class="label">Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{final_f1:.1%}</div><div class="label">Weighted F1</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Parameters</div></div>'
            f'</div>'
            f"{final_charts}"
        ),
        do_flush=True,
    )

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Task 3: Evaluate
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    model_name: str,
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 200,
) -> str:
    """Compare base model (random head) vs fine-tuned on emotion classification.

    Produces confusion matrix, per-class precision/recall/F1, and overall metrics.
    """
    import numpy as np
    import torch
    from datasets import load_from_disk
    from sklearn.metrics import (
        accuracy_score,
        classification_report,
        confusion_matrix as sk_confusion_matrix,
        f1_score,
    )
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Loading models...</p>"),
        do_flush=True,
    )

    # -- Load eval data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"].select(range(min(num_examples, len(dataset["eval"]))))
    texts = eval_ds["text"]
    labels = eval_ds["label"]

    def predict_batch(model, tokenizer, texts, batch_size=32):
        preds = []
        probs_all = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            inputs = tokenizer(batch, truncation=True, max_length=128, padding=True, return_tensors="pt")
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = model(**inputs)
            batch_probs = torch.softmax(outputs.logits, dim=-1).cpu()
            batch_preds = torch.argmax(batch_probs, dim=-1).tolist()
            preds.extend(batch_preds)
            probs_all.extend(batch_probs.tolist())
        return preds, probs_all

    # -- Base model --
    log.info(f"Loading base model: {model_name}")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Running base model (random classifier head)...</p>"),
        do_flush=True,
    )

    base_tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
    base_model = AutoModelForSequenceClassification.from_pretrained(
        model_name, token=HF_TOKEN, num_labels=6,
    )
    base_model.eval()
    if torch.cuda.is_available():
        base_model = base_model.cuda()

    base_preds, base_probs = predict_batch(base_model, base_tokenizer, texts)
    del base_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Fine-tuned model --
    log.info("Loading fine-tuned model...")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Running fine-tuned model...</p>"),
        do_flush=True,
    )

    ft_path = await finetuned_dir.download()
    ft_tokenizer = AutoTokenizer.from_pretrained(ft_path)
    ft_model = AutoModelForSequenceClassification.from_pretrained(ft_path)
    ft_model.eval()
    if torch.cuda.is_available():
        ft_model = ft_model.cuda()

    ft_preds, ft_probs = predict_batch(ft_model, ft_tokenizer, texts)
    del ft_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Compute metrics --
    base_acc = accuracy_score(labels, base_preds) * 100
    base_f1 = f1_score(labels, base_preds, average="weighted") * 100
    ft_acc = accuracy_score(labels, ft_preds) * 100
    ft_f1 = f1_score(labels, ft_preds, average="weighted") * 100

    log.info(f"Base:      Accuracy={base_acc:.1f}%, F1={base_f1:.1f}%")
    log.info(f"Fine-tuned: Accuracy={ft_acc:.1f}%, F1={ft_f1:.1f}%")

    # -- Confusion matrix --
    ft_cm = sk_confusion_matrix(labels, ft_preds, labels=list(range(6)))
    cm_list = ft_cm.tolist()
    cm_svg = make_confusion_matrix(cm_list, EMOTION_LABELS, title="Fine-tuned Model — Confusion Matrix")

    # -- Per-class metrics --
    report_dict = classification_report(
        labels, ft_preds, labels=list(range(6)), target_names=EMOTION_LABELS,
        output_dict=True, zero_division=0,
    )
    per_class_html = "<table><tr><th>Emotion</th><th>Precision</th><th>Recall</th><th>F1</th><th>Support</th></tr>"
    for label_name in EMOTION_LABELS:
        if label_name in report_dict:
            m = report_dict[label_name]
            per_class_html += (
                f"<tr><td><b>{label_name}</b></td>"
                f"<td>{m['precision']:.1%}</td>"
                f"<td>{m['recall']:.1%}</td>"
                f"<td>{m['f1-score']:.1%}</td>"
                f"<td>{int(m['support'])}</td></tr>"
            )
    per_class_html += "</table>"

    # -- Bar chart: base vs fine-tuned --
    per_class_base_acc = []
    per_class_ft_acc = []
    for cls_idx in range(6):
        cls_mask = [i for i, l in enumerate(labels) if l == cls_idx]
        if cls_mask:
            base_cls_acc = sum(1 for i in cls_mask if base_preds[i] == cls_idx) / len(cls_mask) * 100
            ft_cls_acc = sum(1 for i in cls_mask if ft_preds[i] == cls_idx) / len(cls_mask) * 100
        else:
            base_cls_acc = 0
            ft_cls_acc = 0
        per_class_base_acc.append(base_cls_acc)
        per_class_ft_acc.append(ft_cls_acc)

    bar_chart = make_bar_chart(
        labels=EMOTION_LABELS,
        series={"Base": per_class_base_acc, "Fine-tuned": per_class_ft_acc},
        title="Per-Class Accuracy — Base vs Fine-tuned",
        colors=["#adb5bd", "#0f3460"],
        y_max_cap=105.0,
    )

    # -- Example predictions --
    improvement = ft_acc - base_acc
    imp_badge = "badge-success" if improvement > 0 else "badge-danger" if improvement < 0 else "badge-info"

    examples_html = ""
    for i in range(min(10, len(texts))):
        true_label = EMOTION_LABELS[labels[i]]
        ft_label = EMOTION_LABELS[ft_preds[i]]
        base_label = EMOTION_LABELS[base_preds[i]]
        ft_correct = ft_preds[i] == labels[i]
        base_correct = base_preds[i] == labels[i]
        text_preview = texts[i][:200]

        ft_badge = "badge-success" if ft_correct else "badge-danger"
        base_badge = "badge-success" if base_correct else "badge-danger"

        examples_html += f"""
<div class="card">
  <p style="font-size:0.95em;">"{text_preview}"</p>
  <p>True: <b>{true_label}</b> |
  Base: <span class="badge {base_badge}">{base_label}</span> |
  Fine-tuned: <span class="badge {ft_badge}">{ft_label}</span></p>
</div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation Results — Emotion Classification</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{base_acc:.1f}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{ft_acc:.1f}%</div><div class="label">Fine-tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value"><span class="badge {imp_badge}">{improvement:+.1f}pp</span></div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{ft_f1:.1f}%</div><div class="label">Fine-tuned F1</div></div>'
            f'</div>'
            f'<div class="chart-container">{bar_chart}</div>'
            f'<div class="chart-container">{cm_svg}</div>'
            f"<h3>Per-Class Metrics (Fine-tuned)</h3>"
            f"{per_class_html}"
            f"<h3>Example Predictions</h3>"
            f"{examples_html}"
        ),
        do_flush=True,
    )

    return json.dumps({
        "base_accuracy": round(base_acc, 1),
        "base_f1": round(base_f1, 1),
        "finetuned_accuracy": round(ft_acc, 1),
        "finetuned_f1": round(ft_f1, 1),
        "improvement": round(improvement, 1),
        "num_examples": len(texts),
        "confusion_matrix": cm_list,
        "per_class": {k: report_dict[k] for k in EMOTION_LABELS if k in report_dict},
    })

# ------------------------------------------------------------------
# Task 4: Explore inference
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def explore_inference(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 8,
) -> str:
    """Deep-dive into model behavior with attention and token importance.

    For a set of examples, this task produces:
    1. Predictions with full confidence distribution across all 6 emotions
    2. Attention heatmaps — which tokens the model focuses on for classification
       (CLS token attention from the last layer, averaged across heads)
    3. Token importance via gradient-based attribution — which tokens most
       influence the predicted class (gradient x embedding norm)
    4. Misclassification analysis — confident wrong predictions with explanations
    """
    import numpy as np
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    log.info("Starting explore_inference...")
    await flyte.report.replace.aio(
        wrap_report(
            "<h2>Explore Inference</h2>"
            "<p>Loading model for attention and attribution analysis...</p>"
        ),
        do_flush=True,
    )

    # -- Load model (with eager attention for weight extraction) --
    ft_path = await finetuned_dir.download()
    tokenizer = AutoTokenizer.from_pretrained(ft_path)

    # Need eager attention to extract attention weights (flash attention doesn't return them)
    model = AutoModelForSequenceClassification.from_pretrained(
        ft_path,
        output_attentions=True,
        attn_implementation="eager",
    )
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # -- Load eval data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"]

    # Pick a diverse set of examples — try to get some from each class
    examples_per_class = max(1, num_examples // 6)
    selected_indices = []
    for cls_idx in range(6):
        cls_indices = [i for i in range(len(eval_ds)) if eval_ds[i]["label"] == cls_idx]
        selected_indices.extend(cls_indices[:examples_per_class])
    # Fill remaining with random
    remaining = num_examples - len(selected_indices)
    if remaining > 0:
        other_indices = [i for i in range(len(eval_ds)) if i not in selected_indices]
        selected_indices.extend(other_indices[:remaining])
    selected_indices = selected_indices[:num_examples]

    # -- Analyze each example --
    analyses = []
    for idx_num, ds_idx in enumerate(selected_indices):
        text = eval_ds[ds_idx]["text"]
        true_label = eval_ds[ds_idx]["label"]

        await flyte.report.replace.aio(
            wrap_report(
                f"<h2>Explore Inference</h2>"
                f"<p>Analyzing example {idx_num + 1}/{len(selected_indices)}...</p>"
                f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                f'<div style="background:#0f3460;width:{(idx_num + 1) / len(selected_indices) * 100:.1f}%;height:100%;border-radius:4px;"></div>'
                f'</div>'
            ),
            do_flush=True,
        )

        # Tokenize
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        token_ids = inputs["input_ids"][0]
        tokens = tokenizer.convert_ids_to_tokens(token_ids)

        # Forward pass with attention
        with torch.no_grad():
            outputs = model(**inputs)

        logits = outputs.logits[0]
        probs = torch.softmax(logits, dim=-1).cpu().tolist()
        pred_idx = int(torch.argmax(logits).item())

        # -- Attention: CLS token attention from last layer --
        # attentions shape: (num_layers, batch, num_heads, seq_len, seq_len)
        last_layer_attention = outputs.attentions[-1][0]  # (num_heads, seq_len, seq_len)
        # Average across heads, take CLS row (index 0)
        cls_attention = last_layer_attention.mean(dim=0)[0].cpu().numpy()  # (seq_len,)

        # Remove [CLS] and [SEP] and padding from visualization
        real_token_mask = []
        clean_tokens = []
        clean_attention = []
        for i, tok in enumerate(tokens):
            if tok in ("[CLS]", "[SEP]", "<s>", "</s>", "[PAD]", "<pad>"):
                continue
            if tok == tokenizer.pad_token:
                continue
            clean_tokens.append(tok)
            clean_attention.append(float(cls_attention[i]))
            real_token_mask.append(i)

        # -- Token importance via gradient attribution --
        # Re-run with gradients enabled on embeddings
        embedding_layer = None
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Embedding) and "word" in name.lower():
                embedding_layer = module
                break
        if embedding_layer is None:
            # Fallback: find the first large embedding
            for name, module in model.named_modules():
                if isinstance(module, torch.nn.Embedding) and module.weight.shape[0] > 1000:
                    embedding_layer = module
                    break

        importance_scores = [0.0] * len(clean_tokens)
        if embedding_layer is not None:
            inputs_grad = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
            inputs_grad = {k: v.to(device) for k, v in inputs_grad.items()}

            embeddings = embedding_layer(inputs_grad["input_ids"])
            embeddings.retain_grad()

            # Run model with embeddings instead of input_ids
            # We need to hook into the model to replace the embedding output
            embedding_output = [None]

            def hook_fn(module, input, output):
                embedding_output[0] = output
                return embeddings.requires_grad_(True)

            handle = embedding_layer.register_forward_hook(hook_fn)

            outputs_grad = model(**inputs_grad)
            handle.remove()

            # Gradient of predicted class w.r.t. embeddings
            pred_score = outputs_grad.logits[0, pred_idx]
            pred_score.backward()

            if embeddings.grad is not None:
                # Token importance = L2 norm of (gradient * embedding) per token
                token_importance = (embeddings.grad[0] * embeddings[0]).norm(dim=-1).detach().cpu().numpy()
                for clean_idx, orig_idx in enumerate(real_token_mask):
                    if orig_idx < len(token_importance):
                        importance_scores[clean_idx] = float(token_importance[orig_idx])

            model.zero_grad()

        analyses.append({
            "text": text,
            "true_label": true_label,
            "pred_idx": pred_idx,
            "probs": probs,
            "tokens": clean_tokens,
            "attention": clean_attention,
            "importance": importance_scores,
            "correct": pred_idx == true_label,
        })

    # -- Build report --
    log.info("Building explore_inference report...")

    # Overall summary
    correct = sum(1 for a in analyses if a["correct"])
    total = len(analyses)

    # Separate correct vs wrong
    correct_analyses = [a for a in analyses if a["correct"]]
    wrong_analyses = [a for a in analyses if not a["correct"]]

    # -- Build example cards --
    examples_html = ""
    for a in analyses:
        true_name = EMOTION_LABELS[a["true_label"]]
        pred_name = EMOTION_LABELS[a["pred_idx"]]
        status_badge = "badge-success" if a["correct"] else "badge-danger"
        status_text = "Correct" if a["correct"] else "Wrong"

        # Confidence bars
        conf_bars = make_confidence_bars(
            labels=EMOTION_LABELS,
            probabilities=a["probs"],
            predicted_idx=a["pred_idx"],
            true_idx=a["true_label"],
        )

        # Attention heatmap
        attention_viz = make_attention_text(
            tokens=a["tokens"],
            weights=a["attention"],
            title="Attention (what the model looks at for its prediction — darker = more attention)",
        )

        # Token importance
        importance_viz = make_token_importance_text(
            tokens=a["tokens"],
            importance=a["importance"],
            title="Token importance (gradient attribution — green = supports prediction, red = opposes)",
        )

        text_preview = a["text"][:300]
        examples_html += f"""
<div class="card">
  <p style="font-size:1em;"><b>"{text_preview}"</b></p>
  <p>True: <b>{true_name}</b> | Predicted: <span class="badge {status_badge}">{pred_name} ({status_text})</span>
     | Confidence: <b>{a['probs'][a['pred_idx']]:.1%}</b></p>
  <div style="margin:12px 0;">{conf_bars}</div>
  <div style="margin:12px 0;">{attention_viz}</div>
  <div style="margin:12px 0;">{importance_viz}</div>
</div>"""

    # -- Misclassification spotlight --
    misclass_html = ""
    if wrong_analyses:
        # Sort by confidence (most confident wrong first)
        wrong_sorted = sorted(wrong_analyses, key=lambda a: a["probs"][a["pred_idx"]], reverse=True)

        misclass_html = "<h3>Misclassification Spotlight</h3>"
        misclass_html += '<div class="note">These are the model\'s most confident wrong predictions — cases where the model is sure but incorrect. These reveal the model\'s blind spots.</div>'

        for a in wrong_sorted[:3]:
            true_name = EMOTION_LABELS[a["true_label"]]
            pred_name = EMOTION_LABELS[a["pred_idx"]]
            conf = a["probs"][a["pred_idx"]]
            true_conf = a["probs"][a["true_label"]]

            misclass_html += f"""
<div class="card" style="border-left:4px solid #e63946;">
  <p><b>"{a['text'][:200]}"</b></p>
  <p>Predicted <span class="badge badge-danger">{pred_name}</span> ({conf:.1%})
     but true label is <span class="badge badge-info">{true_name}</span> ({true_conf:.1%})</p>
  <p style="font-size:0.85em;color:#6c757d;">
     The model assigned {conf:.1%} confidence to {pred_name} vs {true_conf:.1%} to {true_name}.
     {"The model was very sure here — this is a genuine blind spot." if conf > 0.7 else "The model was uncertain — the true class was a close second."}
  </p>
</div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Explore Inference — Attention &amp; Attribution</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{correct}/{total}</div><div class="label">Correct</div></div>'
            f'  <div class="stat"><div class="value">{correct/total:.0%}</div><div class="label">Accuracy (sample)</div></div>'
            f'  <div class="stat"><div class="value">{len(wrong_analyses)}</div><div class="label">Errors to Analyze</div></div>'
            f'</div>'
            f'<div class="note">'
            f'<b>How to read the visualizations below:</b><br/>'
            f'<b>Attention heatmap:</b> Shows which tokens the [CLS] token attends to in the final layer '
            f'(averaged across all attention heads). Darker = more attention. This reveals what the model "looks at" when making its classification decision.<br/>'
            f'<b>Token importance:</b> Gradient-based attribution showing which tokens most influence the prediction. '
            f'Green = supports the prediction, Red = opposes it. Computed as gradient &times; embedding norm.'
            f'</div>'
            f"<h3>Example Analysis</h3>"
            f"{examples_html}"
            f"{misclass_html}"
        ),
        do_flush=True,
    )

    return json.dumps({
        "num_examples": total,
        "correct": correct,
        "accuracy": round(correct / total * 100, 1),
        "num_misclassifications": len(wrong_analyses),
        "analyses": [
            {
                "text": a["text"][:200],
                "true_label": EMOTION_LABELS[a["true_label"]],
                "predicted": EMOTION_LABELS[a["pred_idx"]],
                "confidence": round(a["probs"][a["pred_idx"]], 3),
                "correct": a["correct"],
            }
            for a in analyses
        ],
    })

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "answerdotai/ModernBERT-base",
    epochs: int = 3,
    lr: float = 2e-5,
    batch_size: int = 16,
    warmup_steps: int = 100,
    max_train_samples: int = 10000,
    max_eval_samples: int = 2000,
    num_eval_examples: int = 200,
    num_explore_examples: int = 12,
) -> flyte.io.Dir:
    """
    ModernBERT emotion classification pipeline.

    Returns the fine-tuned model directory (used by serve.py for deployment).

    1. Download emotion dataset (6 classes from Twitter text)
    2. Fine-tune ModernBERT for sequence classification
    3. Evaluate: base vs fine-tuned with confusion matrix
    4. Explore inference: attention heatmaps + token importance

    Args:
        model_name: HuggingFace encoder model to fine-tune.
        num_explore_examples: Number of examples for attention/attribution analysis.
    """
    log.info(f"Pipeline: {model_name} | emotion classification")
    steps = ["Get Data", "Train", "Evaluate", "Explore Inference"]

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(0, steps)}"
            f'<div class="card"><p>Downloading emotion dataset...</p></div>'
        ),
        do_flush=True,
    )

    # Step 1: Get data
    data_dir = await get_data(max_train_samples, max_eval_samples)

    # Step 2: Train
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(1, steps)}"
            f'<div class="card"><p>Fine-tuning for emotion classification...</p></div>'
        ),
        do_flush=True,
    )

    finetuned_dir = await train(model_name, data_dir, epochs, lr, batch_size, warmup_steps)

    # Step 3: Evaluate
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(2, steps)}"
            f'<div class="card"><p>Evaluating base vs fine-tuned model...</p></div>'
        ),
        do_flush=True,
    )

    eval_result = await evaluate(model_name, finetuned_dir, data_dir, num_eval_examples)
    eval_metrics = json.loads(eval_result)

    # Step 4: Explore inference
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(3, steps)}"
            f'<div class="card"><p>Analyzing attention patterns and token importance...</p></div>'
        ),
        do_flush=True,
    )

    explore_result = await explore_inference(finetuned_dir, data_dir, num_explore_examples)

    # -- Final report --
    improvement = eval_metrics["improvement"]
    imp_badge = "badge-success" if improvement > 0 else "badge-danger" if improvement < 0 else "badge-info"

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline Complete</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(4, steps)}"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{eval_metrics["base_accuracy"]}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{eval_metrics["finetuned_accuracy"]}%</div><div class="label">Fine-tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value"><span class="badge {imp_badge}">{improvement:+.1f}pp</span></div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{eval_metrics["finetuned_f1"]}%</div><div class="label">Weighted F1</div></div>'
            f'</div>'
        ),
        do_flush=True,
    )

    log.info(f"Pipeline complete. Accuracy improvement: {improvement:+.1f}pp")
    return finetuned_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/bert_fine_tuning_emotion/bert_fine_tuning_emotion.py*

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "datasets>=3.0.0",
#    "scikit-learn",
#    ...
# ]
# ///
```

## Orchestrate the pipeline

```
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.4.0",
#    "torch>=2.1.0",
#    "transformers>=4.45.0",
#    "datasets>=3.0.0",
#    "accelerate>=0.34.0",
#    "scikit-learn",
#    "numpy",
# ]
# main = "pipeline"
# params = ""
# ///
import json
import logging
import os
import tempfile

import flyte
import flyte.io
import flyte.report

# {{docs-fragment env}}
import os

main_img = flyte.Image.from_uv_script(__file__, name="bert-fine-tuning-emotion", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="bert-fine-tuning-emotion-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu=1),
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

cpu_env = flyte.TaskEnvironment(
    name="bert-fine-tuning-emotion-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)

HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}

from report_helpers import (
    make_attention_text,
    make_bar_chart,
    make_confidence_bars,
    make_confusion_matrix,
    make_line_chart,
    make_token_importance_text,
    pipeline_step_indicator,
    wrap_report,
)

logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

EMOTION_LABELS = ["sadness", "joy", "love", "anger", "fear", "surprise"]
EMOTION_DATASET = "dair-ai/emotion"

# ------------------------------------------------------------------
# Task 1: Get data
# ------------------------------------------------------------------

@cpu_env.task(cache="auto")
async def get_data(
    max_train_samples: int = 10000,
    max_eval_samples: int = 2000,
) -> flyte.io.Dir:
    """Download the emotion dataset and save train/eval splits.

    The dair-ai/emotion dataset contains ~20k English Twitter messages labeled
    with one of 6 emotions: sadness, joy, love, anger, fear, surprise.
    """
    from datasets import DatasetDict, load_dataset

    log.info("Loading emotion dataset...")
    ds = load_dataset(EMOTION_DATASET)

    train_ds = ds["train"].shuffle(seed=42).select(range(min(max_train_samples, len(ds["train"]))))
    eval_ds = ds["test"].shuffle(seed=42).select(range(min(max_eval_samples, len(ds["test"]))))

    processed = DatasetDict({"train": train_ds, "eval": eval_ds})

    output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
    processed.save_to_disk(output_dir)
    log.info(f"Dataset ready: {len(train_ds)} train, {len(eval_ds)} eval")

    return await flyte.io.Dir.from_local(output_dir)

# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def train(
    model_name: str,
    data_dir: flyte.io.Dir,
    epochs: int = 3,
    lr: float = 2e-5,
    batch_size: int = 16,
    warmup_steps: int = 100,
) -> flyte.io.Dir:
    """Fine-tune a BERT-style model for 6-class emotion classification."""
    import numpy as np
    import torch
    from datasets import load_from_disk
    from sklearn.metrics import accuracy_score, f1_score
    from transformers import (
        AutoModelForSequenceClassification,
        AutoTokenizer,
        Trainer,
        TrainerCallback,
        TrainingArguments,
    )

    log.info(f"Training: model={model_name}")

    id2label = {i: l for i, l in enumerate(EMOTION_LABELS)}
    label2id = {l: i for i, l in enumerate(EMOTION_LABELS)}

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Loading Model...</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="card"><p>Preparing for emotion classification training...</p></div>'
        ),
        do_flush=True,
    )

    # -- Load data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)

    # -- Tokenize --
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)

    def tokenize(examples):
        return tokenizer(examples["text"], truncation=True, max_length=128, padding="max_length")

    dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])

    # -- Load model --
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        token=HF_TOKEN,
        num_labels=6,
        id2label=id2label,
        label2id=label2id,
    )

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log.info(f"Parameters: {trainable_params:,} / {total_params:,}")

    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        log.info(f"GPU: {gpu_name} ({gpu_mem:.1f} GB)")

    # -- Metrics tracking for live report --
    training_log: list[dict] = []
    eval_log: list[dict] = []

    def _build_training_report(max_steps: int) -> str:
        stats_html = f"""
        <h2>Training in Progress...</h2>
        <h3>{model_name}</h3>
        <div class="stat-grid">
          <div class="stat"><div class="value">{len(dataset['train']):,}</div><div class="label">Train Samples</div></div>
          <div class="stat"><div class="value">{len(dataset['eval']):,}</div><div class="label">Eval Samples</div></div>
          <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>
          <div class="stat"><div class="value">{lr}</div><div class="label">Learning Rate</div></div>
          <div class="stat"><div class="value">{batch_size}</div><div class="label">Batch Size</div></div>
          <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Parameters</div></div>
        </div>
        """

        charts_html = ""

        if training_log:
            current = training_log[-1]
            progress_pct = current["step"] / max_steps * 100 if max_steps else 0
            loss_display = f"Loss: <span class=\"highlight\">{current['loss']:.4f}</span>" if current.get("loss") else ""
            charts_html += f"""
            <div class="card">
              <b>Step {current['step']}/{max_steps}</b>
              ({progress_pct:.0f}%) |
              Epoch {current['epoch']:.2f}/{epochs}
              {f' | {loss_display}' if loss_display else ''}
              <div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">
                <div style="background:#0f3460;width:{progress_pct:.1f}%;height:100%;border-radius:4px;"></div>
              </div>
            </div>
            """

            loss_entries = [e for e in training_log if "loss" in e]
            if len(loss_entries) >= 2:
                loss_chart = make_line_chart(
                    data=loss_entries,
                    x_key="epoch",
                    y_keys=["loss"],
                    title="Training Loss",
                    x_label="Epoch",
                    y_label="Loss",
                    colors=["#5a7db5"],
                )
                charts_html += f'<div class="chart-container">{loss_chart}</div>'

        if eval_log:
            latest_eval = eval_log[-1]
            best_acc = max(e.get("accuracy", 0) for e in eval_log)
            best_f1 = max(e.get("f1", 0) for e in eval_log)
            charts_html += f"""
            <div class="stat-grid" style="margin-top:16px;">
              <div class="stat"><div class="value">{latest_eval.get('accuracy', 0):.1%}</div><div class="label">Eval Accuracy</div></div>
              <div class="stat"><div class="value">{latest_eval.get('f1', 0):.1%}</div><div class="label">Eval F1</div></div>
              <div class="stat"><div class="value">{best_acc:.1%}</div><div class="label">Best Accuracy</div></div>
              <div class="stat"><div class="value">{latest_eval.get('eval_loss', 0):.4f}</div><div class="label">Eval Loss</div></div>
            </div>
            """

            if len(eval_log) >= 2:
                eval_chart = make_line_chart(
                    data=eval_log,
                    x_key="epoch",
                    y_keys=["accuracy", "f1"],
                    title="Eval Metrics Over Training",
                    x_label="Epoch",
                    y_label="Score",
                    colors=["#0f3460", "#06d6a0"],
                    y_max_cap=1.05,
                    y_display_names={"accuracy": "Accuracy", "f1": "Weighted F1"},
                )
                charts_html += f'<div class="chart-container">{eval_chart}</div>'

                eval_loss_chart = make_line_chart(
                    data=[e for e in eval_log if "eval_loss" in e],
                    x_key="epoch",
                    y_keys=["eval_loss"],
                    title="Eval Loss",
                    x_label="Epoch",
                    y_label="Loss",
                    colors=["#e63946"],
                )
                if any("eval_loss" in e for e in eval_log):
                    charts_html += f'<div class="chart-container">{eval_loss_chart}</div>'

        return wrap_report(stats_html + charts_html)

    # -- Callbacks --
    class ReportCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs:
                return
            entry = {
                "step": state.global_step,
                "epoch": round(logs.get("epoch", 0), 2),
            }
            if "loss" in logs:
                entry["loss"] = round(logs["loss"], 4)
            if "eval_accuracy" in logs:
                eval_log.append({
                    "epoch": entry["epoch"],
                    "accuracy": logs["eval_accuracy"],
                    "f1": logs.get("eval_f1", 0),
                    "eval_loss": logs.get("eval_loss", 0),
                })
            if "loss" in entry:
                training_log.append(entry)

            flyte.report.replace(
                _build_training_report(state.max_steps),
                do_flush=True,
            )

    # -- Compute metrics --
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        return {
            "accuracy": accuracy_score(labels, preds),
            "f1": f1_score(labels, preds, average="weighted"),
        }

    # -- Training --
    output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        learning_rate=lr,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        bf16=use_bf16,
        fp16=not use_bf16 and torch.cuda.is_available(),
        warmup_steps=warmup_steps,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["eval"],
        processing_class=tokenizer,
        compute_metrics=compute_metrics,
        callbacks=[ReportCallback()],
    )

    log.info("Starting training...")
    await flyte.report.replace.aio(
        _build_training_report(0),
        do_flush=True,
    )

    trainer.train()
    log.info("Training complete.")

    # -- Save model --
    save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
    trainer.save_model(save_dir)
    tokenizer.save_pretrained(save_dir)
    log.info(f"Model saved to {save_dir}")

    # -- Final eval + report --
    metrics = trainer.evaluate()
    final_acc = metrics.get("eval_accuracy", 0)
    final_f1 = metrics.get("eval_f1", 0)

    final_charts = ""
    loss_entries = [e for e in training_log if "loss" in e]
    if len(loss_entries) >= 2:
        loss_chart = make_line_chart(
            data=loss_entries,
            x_key="epoch",
            y_keys=["loss"],
            title="Training Loss",
            x_label="Epoch",
            y_label="Loss",
            colors=["#5a7db5"],
        )
        final_charts += f'<div class="chart-container">{loss_chart}</div>'

    if len(eval_log) >= 2:
        eval_chart = make_line_chart(
            data=eval_log,
            x_key="epoch",
            y_keys=["accuracy", "f1"],
            title="Eval Metrics Over Training",
            x_label="Epoch",
            y_label="Score",
            colors=["#0f3460", "#06d6a0"],
            y_max_cap=1.05,
            y_display_names={"accuracy": "Accuracy", "f1": "Weighted F1"},
        )
        final_charts += f'<div class="chart-container">{eval_chart}</div>'

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Training Complete</h2>"
            f"<h3>{model_name}</h3>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{final_acc:.1%}</div><div class="label">Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{final_f1:.1%}</div><div class="label">Weighted F1</div></div>'
            f'  <div class="stat"><div class="value">{epochs}</div><div class="label">Epochs</div></div>'
            f'  <div class="stat"><div class="value">{trainable_params:,}</div><div class="label">Parameters</div></div>'
            f'</div>'
            f"{final_charts}"
        ),
        do_flush=True,
    )

    return await flyte.io.Dir.from_local(save_dir)

# ------------------------------------------------------------------
# Task 3: Evaluate
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def evaluate(
    model_name: str,
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 200,
) -> str:
    """Compare base model (random head) vs fine-tuned on emotion classification.

    Produces confusion matrix, per-class precision/recall/F1, and overall metrics.
    """
    import numpy as np
    import torch
    from datasets import load_from_disk
    from sklearn.metrics import (
        accuracy_score,
        classification_report,
        confusion_matrix as sk_confusion_matrix,
        f1_score,
    )
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    log.info("Starting evaluation...")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Loading models...</p>"),
        do_flush=True,
    )

    # -- Load eval data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"].select(range(min(num_examples, len(dataset["eval"]))))
    texts = eval_ds["text"]
    labels = eval_ds["label"]

    def predict_batch(model, tokenizer, texts, batch_size=32):
        preds = []
        probs_all = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            inputs = tokenizer(batch, truncation=True, max_length=128, padding=True, return_tensors="pt")
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = model(**inputs)
            batch_probs = torch.softmax(outputs.logits, dim=-1).cpu()
            batch_preds = torch.argmax(batch_probs, dim=-1).tolist()
            preds.extend(batch_preds)
            probs_all.extend(batch_probs.tolist())
        return preds, probs_all

    # -- Base model --
    log.info(f"Loading base model: {model_name}")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Running base model (random classifier head)...</p>"),
        do_flush=True,
    )

    base_tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
    base_model = AutoModelForSequenceClassification.from_pretrained(
        model_name, token=HF_TOKEN, num_labels=6,
    )
    base_model.eval()
    if torch.cuda.is_available():
        base_model = base_model.cuda()

    base_preds, base_probs = predict_batch(base_model, base_tokenizer, texts)
    del base_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Fine-tuned model --
    log.info("Loading fine-tuned model...")
    await flyte.report.replace.aio(
        wrap_report("<h2>Evaluation</h2><p>Running fine-tuned model...</p>"),
        do_flush=True,
    )

    ft_path = await finetuned_dir.download()
    ft_tokenizer = AutoTokenizer.from_pretrained(ft_path)
    ft_model = AutoModelForSequenceClassification.from_pretrained(ft_path)
    ft_model.eval()
    if torch.cuda.is_available():
        ft_model = ft_model.cuda()

    ft_preds, ft_probs = predict_batch(ft_model, ft_tokenizer, texts)
    del ft_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # -- Compute metrics --
    base_acc = accuracy_score(labels, base_preds) * 100
    base_f1 = f1_score(labels, base_preds, average="weighted") * 100
    ft_acc = accuracy_score(labels, ft_preds) * 100
    ft_f1 = f1_score(labels, ft_preds, average="weighted") * 100

    log.info(f"Base:      Accuracy={base_acc:.1f}%, F1={base_f1:.1f}%")
    log.info(f"Fine-tuned: Accuracy={ft_acc:.1f}%, F1={ft_f1:.1f}%")

    # -- Confusion matrix --
    ft_cm = sk_confusion_matrix(labels, ft_preds, labels=list(range(6)))
    cm_list = ft_cm.tolist()
    cm_svg = make_confusion_matrix(cm_list, EMOTION_LABELS, title="Fine-tuned Model — Confusion Matrix")

    # -- Per-class metrics --
    report_dict = classification_report(
        labels, ft_preds, labels=list(range(6)), target_names=EMOTION_LABELS,
        output_dict=True, zero_division=0,
    )
    per_class_html = "<table><tr><th>Emotion</th><th>Precision</th><th>Recall</th><th>F1</th><th>Support</th></tr>"
    for label_name in EMOTION_LABELS:
        if label_name in report_dict:
            m = report_dict[label_name]
            per_class_html += (
                f"<tr><td><b>{label_name}</b></td>"
                f"<td>{m['precision']:.1%}</td>"
                f"<td>{m['recall']:.1%}</td>"
                f"<td>{m['f1-score']:.1%}</td>"
                f"<td>{int(m['support'])}</td></tr>"
            )
    per_class_html += "</table>"

    # -- Bar chart: base vs fine-tuned --
    per_class_base_acc = []
    per_class_ft_acc = []
    for cls_idx in range(6):
        cls_mask = [i for i, l in enumerate(labels) if l == cls_idx]
        if cls_mask:
            base_cls_acc = sum(1 for i in cls_mask if base_preds[i] == cls_idx) / len(cls_mask) * 100
            ft_cls_acc = sum(1 for i in cls_mask if ft_preds[i] == cls_idx) / len(cls_mask) * 100
        else:
            base_cls_acc = 0
            ft_cls_acc = 0
        per_class_base_acc.append(base_cls_acc)
        per_class_ft_acc.append(ft_cls_acc)

    bar_chart = make_bar_chart(
        labels=EMOTION_LABELS,
        series={"Base": per_class_base_acc, "Fine-tuned": per_class_ft_acc},
        title="Per-Class Accuracy — Base vs Fine-tuned",
        colors=["#adb5bd", "#0f3460"],
        y_max_cap=105.0,
    )

    # -- Example predictions --
    improvement = ft_acc - base_acc
    imp_badge = "badge-success" if improvement > 0 else "badge-danger" if improvement < 0 else "badge-info"

    examples_html = ""
    for i in range(min(10, len(texts))):
        true_label = EMOTION_LABELS[labels[i]]
        ft_label = EMOTION_LABELS[ft_preds[i]]
        base_label = EMOTION_LABELS[base_preds[i]]
        ft_correct = ft_preds[i] == labels[i]
        base_correct = base_preds[i] == labels[i]
        text_preview = texts[i][:200]

        ft_badge = "badge-success" if ft_correct else "badge-danger"
        base_badge = "badge-success" if base_correct else "badge-danger"

        examples_html += f"""
<div class="card">
  <p style="font-size:0.95em;">"{text_preview}"</p>
  <p>True: <b>{true_label}</b> |
  Base: <span class="badge {base_badge}">{base_label}</span> |
  Fine-tuned: <span class="badge {ft_badge}">{ft_label}</span></p>
</div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Evaluation Results — Emotion Classification</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{base_acc:.1f}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{ft_acc:.1f}%</div><div class="label">Fine-tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value"><span class="badge {imp_badge}">{improvement:+.1f}pp</span></div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{ft_f1:.1f}%</div><div class="label">Fine-tuned F1</div></div>'
            f'</div>'
            f'<div class="chart-container">{bar_chart}</div>'
            f'<div class="chart-container">{cm_svg}</div>'
            f"<h3>Per-Class Metrics (Fine-tuned)</h3>"
            f"{per_class_html}"
            f"<h3>Example Predictions</h3>"
            f"{examples_html}"
        ),
        do_flush=True,
    )

    return json.dumps({
        "base_accuracy": round(base_acc, 1),
        "base_f1": round(base_f1, 1),
        "finetuned_accuracy": round(ft_acc, 1),
        "finetuned_f1": round(ft_f1, 1),
        "improvement": round(improvement, 1),
        "num_examples": len(texts),
        "confusion_matrix": cm_list,
        "per_class": {k: report_dict[k] for k in EMOTION_LABELS if k in report_dict},
    })

# ------------------------------------------------------------------
# Task 4: Explore inference
# ------------------------------------------------------------------

@gpu_env.task(report=True)
async def explore_inference(
    finetuned_dir: flyte.io.Dir,
    data_dir: flyte.io.Dir,
    num_examples: int = 8,
) -> str:
    """Deep-dive into model behavior with attention and token importance.

    For a set of examples, this task produces:
    1. Predictions with full confidence distribution across all 6 emotions
    2. Attention heatmaps — which tokens the model focuses on for classification
       (CLS token attention from the last layer, averaged across heads)
    3. Token importance via gradient-based attribution — which tokens most
       influence the predicted class (gradient x embedding norm)
    4. Misclassification analysis — confident wrong predictions with explanations
    """
    import numpy as np
    import torch
    from datasets import load_from_disk
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    log.info("Starting explore_inference...")
    await flyte.report.replace.aio(
        wrap_report(
            "<h2>Explore Inference</h2>"
            "<p>Loading model for attention and attribution analysis...</p>"
        ),
        do_flush=True,
    )

    # -- Load model (with eager attention for weight extraction) --
    ft_path = await finetuned_dir.download()
    tokenizer = AutoTokenizer.from_pretrained(ft_path)

    # Need eager attention to extract attention weights (flash attention doesn't return them)
    model = AutoModelForSequenceClassification.from_pretrained(
        ft_path,
        output_attentions=True,
        attn_implementation="eager",
    )
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # -- Load eval data --
    data_path = await data_dir.download()
    dataset = load_from_disk(data_path)
    eval_ds = dataset["eval"]

    # Pick a diverse set of examples — try to get some from each class
    examples_per_class = max(1, num_examples // 6)
    selected_indices = []
    for cls_idx in range(6):
        cls_indices = [i for i in range(len(eval_ds)) if eval_ds[i]["label"] == cls_idx]
        selected_indices.extend(cls_indices[:examples_per_class])
    # Fill remaining with random
    remaining = num_examples - len(selected_indices)
    if remaining > 0:
        other_indices = [i for i in range(len(eval_ds)) if i not in selected_indices]
        selected_indices.extend(other_indices[:remaining])
    selected_indices = selected_indices[:num_examples]

    # -- Analyze each example --
    analyses = []
    for idx_num, ds_idx in enumerate(selected_indices):
        text = eval_ds[ds_idx]["text"]
        true_label = eval_ds[ds_idx]["label"]

        await flyte.report.replace.aio(
            wrap_report(
                f"<h2>Explore Inference</h2>"
                f"<p>Analyzing example {idx_num + 1}/{len(selected_indices)}...</p>"
                f'<div style="background:#e9ecef;border-radius:4px;height:8px;margin-top:8px;">'
                f'<div style="background:#0f3460;width:{(idx_num + 1) / len(selected_indices) * 100:.1f}%;height:100%;border-radius:4px;"></div>'
                f'</div>'
            ),
            do_flush=True,
        )

        # Tokenize
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        token_ids = inputs["input_ids"][0]
        tokens = tokenizer.convert_ids_to_tokens(token_ids)

        # Forward pass with attention
        with torch.no_grad():
            outputs = model(**inputs)

        logits = outputs.logits[0]
        probs = torch.softmax(logits, dim=-1).cpu().tolist()
        pred_idx = int(torch.argmax(logits).item())

        # -- Attention: CLS token attention from last layer --
        # attentions shape: (num_layers, batch, num_heads, seq_len, seq_len)
        last_layer_attention = outputs.attentions[-1][0]  # (num_heads, seq_len, seq_len)
        # Average across heads, take CLS row (index 0)
        cls_attention = last_layer_attention.mean(dim=0)[0].cpu().numpy()  # (seq_len,)

        # Remove [CLS] and [SEP] and padding from visualization
        real_token_mask = []
        clean_tokens = []
        clean_attention = []
        for i, tok in enumerate(tokens):
            if tok in ("[CLS]", "[SEP]", "<s>", "</s>", "[PAD]", "<pad>"):
                continue
            if tok == tokenizer.pad_token:
                continue
            clean_tokens.append(tok)
            clean_attention.append(float(cls_attention[i]))
            real_token_mask.append(i)

        # -- Token importance via gradient attribution --
        # Re-run with gradients enabled on embeddings
        embedding_layer = None
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Embedding) and "word" in name.lower():
                embedding_layer = module
                break
        if embedding_layer is None:
            # Fallback: find the first large embedding
            for name, module in model.named_modules():
                if isinstance(module, torch.nn.Embedding) and module.weight.shape[0] > 1000:
                    embedding_layer = module
                    break

        importance_scores = [0.0] * len(clean_tokens)
        if embedding_layer is not None:
            inputs_grad = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
            inputs_grad = {k: v.to(device) for k, v in inputs_grad.items()}

            embeddings = embedding_layer(inputs_grad["input_ids"])
            embeddings.retain_grad()

            # Run model with embeddings instead of input_ids
            # We need to hook into the model to replace the embedding output
            embedding_output = [None]

            def hook_fn(module, input, output):
                embedding_output[0] = output
                return embeddings.requires_grad_(True)

            handle = embedding_layer.register_forward_hook(hook_fn)

            outputs_grad = model(**inputs_grad)
            handle.remove()

            # Gradient of predicted class w.r.t. embeddings
            pred_score = outputs_grad.logits[0, pred_idx]
            pred_score.backward()

            if embeddings.grad is not None:
                # Token importance = L2 norm of (gradient * embedding) per token
                token_importance = (embeddings.grad[0] * embeddings[0]).norm(dim=-1).detach().cpu().numpy()
                for clean_idx, orig_idx in enumerate(real_token_mask):
                    if orig_idx < len(token_importance):
                        importance_scores[clean_idx] = float(token_importance[orig_idx])

            model.zero_grad()

        analyses.append({
            "text": text,
            "true_label": true_label,
            "pred_idx": pred_idx,
            "probs": probs,
            "tokens": clean_tokens,
            "attention": clean_attention,
            "importance": importance_scores,
            "correct": pred_idx == true_label,
        })

    # -- Build report --
    log.info("Building explore_inference report...")

    # Overall summary
    correct = sum(1 for a in analyses if a["correct"])
    total = len(analyses)

    # Separate correct vs wrong
    correct_analyses = [a for a in analyses if a["correct"]]
    wrong_analyses = [a for a in analyses if not a["correct"]]

    # -- Build example cards --
    examples_html = ""
    for a in analyses:
        true_name = EMOTION_LABELS[a["true_label"]]
        pred_name = EMOTION_LABELS[a["pred_idx"]]
        status_badge = "badge-success" if a["correct"] else "badge-danger"
        status_text = "Correct" if a["correct"] else "Wrong"

        # Confidence bars
        conf_bars = make_confidence_bars(
            labels=EMOTION_LABELS,
            probabilities=a["probs"],
            predicted_idx=a["pred_idx"],
            true_idx=a["true_label"],
        )

        # Attention heatmap
        attention_viz = make_attention_text(
            tokens=a["tokens"],
            weights=a["attention"],
            title="Attention (what the model looks at for its prediction — darker = more attention)",
        )

        # Token importance
        importance_viz = make_token_importance_text(
            tokens=a["tokens"],
            importance=a["importance"],
            title="Token importance (gradient attribution — green = supports prediction, red = opposes)",
        )

        text_preview = a["text"][:300]
        examples_html += f"""
<div class="card">
  <p style="font-size:1em;"><b>"{text_preview}"</b></p>
  <p>True: <b>{true_name}</b> | Predicted: <span class="badge {status_badge}">{pred_name} ({status_text})</span>
     | Confidence: <b>{a['probs'][a['pred_idx']]:.1%}</b></p>
  <div style="margin:12px 0;">{conf_bars}</div>
  <div style="margin:12px 0;">{attention_viz}</div>
  <div style="margin:12px 0;">{importance_viz}</div>
</div>"""

    # -- Misclassification spotlight --
    misclass_html = ""
    if wrong_analyses:
        # Sort by confidence (most confident wrong first)
        wrong_sorted = sorted(wrong_analyses, key=lambda a: a["probs"][a["pred_idx"]], reverse=True)

        misclass_html = "<h3>Misclassification Spotlight</h3>"
        misclass_html += '<div class="note">These are the model\'s most confident wrong predictions — cases where the model is sure but incorrect. These reveal the model\'s blind spots.</div>'

        for a in wrong_sorted[:3]:
            true_name = EMOTION_LABELS[a["true_label"]]
            pred_name = EMOTION_LABELS[a["pred_idx"]]
            conf = a["probs"][a["pred_idx"]]
            true_conf = a["probs"][a["true_label"]]

            misclass_html += f"""
<div class="card" style="border-left:4px solid #e63946;">
  <p><b>"{a['text'][:200]}"</b></p>
  <p>Predicted <span class="badge badge-danger">{pred_name}</span> ({conf:.1%})
     but true label is <span class="badge badge-info">{true_name}</span> ({true_conf:.1%})</p>
  <p style="font-size:0.85em;color:#6c757d;">
     The model assigned {conf:.1%} confidence to {pred_name} vs {true_conf:.1%} to {true_name}.
     {"The model was very sure here — this is a genuine blind spot." if conf > 0.7 else "The model was uncertain — the true class was a close second."}
  </p>
</div>"""

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Explore Inference — Attention &amp; Attribution</h2>"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{correct}/{total}</div><div class="label">Correct</div></div>'
            f'  <div class="stat"><div class="value">{correct/total:.0%}</div><div class="label">Accuracy (sample)</div></div>'
            f'  <div class="stat"><div class="value">{len(wrong_analyses)}</div><div class="label">Errors to Analyze</div></div>'
            f'</div>'
            f'<div class="note">'
            f'<b>How to read the visualizations below:</b><br/>'
            f'<b>Attention heatmap:</b> Shows which tokens the [CLS] token attends to in the final layer '
            f'(averaged across all attention heads). Darker = more attention. This reveals what the model "looks at" when making its classification decision.<br/>'
            f'<b>Token importance:</b> Gradient-based attribution showing which tokens most influence the prediction. '
            f'Green = supports the prediction, Red = opposes it. Computed as gradient &times; embedding norm.'
            f'</div>'
            f"<h3>Example Analysis</h3>"
            f"{examples_html}"
            f"{misclass_html}"
        ),
        do_flush=True,
    )

    return json.dumps({
        "num_examples": total,
        "correct": correct,
        "accuracy": round(correct / total * 100, 1),
        "num_misclassifications": len(wrong_analyses),
        "analyses": [
            {
                "text": a["text"][:200],
                "true_label": EMOTION_LABELS[a["true_label"]],
                "predicted": EMOTION_LABELS[a["pred_idx"]],
                "confidence": round(a["probs"][a["pred_idx"]], 3),
                "correct": a["correct"],
            }
            for a in analyses
        ],
    })

# ------------------------------------------------------------------
# Pipeline
# ------------------------------------------------------------------

# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
    model_name: str = "answerdotai/ModernBERT-base",
    epochs: int = 3,
    lr: float = 2e-5,
    batch_size: int = 16,
    warmup_steps: int = 100,
    max_train_samples: int = 10000,
    max_eval_samples: int = 2000,
    num_eval_examples: int = 200,
    num_explore_examples: int = 12,
) -> flyte.io.Dir:
    """
    ModernBERT emotion classification pipeline.

    Returns the fine-tuned model directory (used by serve.py for deployment).

    1. Download emotion dataset (6 classes from Twitter text)
    2. Fine-tune ModernBERT for sequence classification
    3. Evaluate: base vs fine-tuned with confusion matrix
    4. Explore inference: attention heatmaps + token importance

    Args:
        model_name: HuggingFace encoder model to fine-tune.
        num_explore_examples: Number of examples for attention/attribution analysis.
    """
    log.info(f"Pipeline: {model_name} | emotion classification")
    steps = ["Get Data", "Train", "Evaluate", "Explore Inference"]

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(0, steps)}"
            f'<div class="card"><p>Downloading emotion dataset...</p></div>'
        ),
        do_flush=True,
    )

    # Step 1: Get data
    data_dir = await get_data(max_train_samples, max_eval_samples)

    # Step 2: Train
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(1, steps)}"
            f'<div class="card"><p>Fine-tuning for emotion classification...</p></div>'
        ),
        do_flush=True,
    )

    finetuned_dir = await train(model_name, data_dir, epochs, lr, batch_size, warmup_steps)

    # Step 3: Evaluate
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(2, steps)}"
            f'<div class="card"><p>Evaluating base vs fine-tuned model...</p></div>'
        ),
        do_flush=True,
    )

    eval_result = await evaluate(model_name, finetuned_dir, data_dir, num_eval_examples)
    eval_metrics = json.loads(eval_result)

    # Step 4: Explore inference
    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(3, steps)}"
            f'<div class="card"><p>Analyzing attention patterns and token importance...</p></div>'
        ),
        do_flush=True,
    )

    explore_result = await explore_inference(finetuned_dir, data_dir, num_explore_examples)

    # -- Final report --
    improvement = eval_metrics["improvement"]
    imp_badge = "badge-success" if improvement > 0 else "badge-danger" if improvement < 0 else "badge-info"

    await flyte.report.replace.aio(
        wrap_report(
            f"<h2>Emotion Classification Pipeline Complete</h2>"
            f"<h3>{model_name}</h3>"
            f"{pipeline_step_indicator(4, steps)}"
            f'<div class="stat-grid">'
            f'  <div class="stat"><div class="value">{eval_metrics["base_accuracy"]}%</div><div class="label">Base Accuracy</div></div>'
            f'  <div class="stat"><div class="value">{eval_metrics["finetuned_accuracy"]}%</div><div class="label">Fine-tuned Accuracy</div></div>'
            f'  <div class="stat"><div class="value"><span class="badge {imp_badge}">{improvement:+.1f}pp</span></div><div class="label">Improvement</div></div>'
            f'  <div class="stat"><div class="value">{eval_metrics["finetuned_f1"]}%</div><div class="label">Weighted F1</div></div>'
            f'</div>'
        ),
        do_flush=True,
    )

    log.info(f"Pipeline complete. Accuracy improvement: {improvement:+.1f}pp")
    return finetuned_dir

# {{/docs-fragment pipeline}}

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(pipeline)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/bert_fine_tuning_emotion/bert_fine_tuning_emotion.py*

## Run the workflow

From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/bert_fine_tuning_emotion):

```
cd v2/tutorials/bert_fine_tuning_emotion
uv run --script bert_fine_tuning_emotion.py
```

Quick smoke test with a small sample:

```
flyte run bert_fine_tuning_emotion.py pipeline --max_train_samples 200 --max_eval_samples 50 --epochs 1
```

Open the **evaluate** and **explore_inference** task reports for confusion matrices and attention visualizations.


=== PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training/hpo ===

# Hyperparameter optimization

> [!NOTE]
> Code available [on GitHub](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/ml/optimizer.py).

Hyperparameter Optimization (HPO) is a critical step in the machine learning (ML) lifecycle. Hyperparameters are the knobs and dials of a model: values such as learning rates, tree depths, or dropout rates that significantly impact performance but cannot be learned during training. Instead, we must select them manually or optimize them through guided search.

Model developers often enjoy the flexibility of choosing from a wide variety of model types, whether gradient boosted machines (GBMs), generalized linear models (GLMs), deep learning architectures, or dozens of others. A common challenge across all these options is the need to systematically explore model performance across hyperparameter configurations tailored to the specific dataset and task.

Thankfully, this exploration can be automated. Frameworks like [Optuna](https://optuna.org/), [Hyperopt](https://hyperopt.github.io/hyperopt/), and [Ray Tune](https://docs.ray.io/en/latest/tune/index.html) use advanced sampling algorithms to efficiently search the hyperparameter space and identify optimal configurations. HPO may be executed in two distinct ways:

- **Serial HPO** runs one trial at a time, which is easy to set up but can be painfully slow.
- **Parallel HPO** distributes trials across multiple processes. It typically follows a pattern with two parameters: **_N_**, the total number of trials to run, and **_C_**, the maximum number of trials that can run concurrently. Trials are executed asynchronously, and new ones are scheduled based on the results and status of completed or in-progress ones.

However, parallel HPO introduces a new complexity: the need for a centralized state that tracks:

- All past trials (successes and failures)
- All ongoing trials

This state is essential so that the optimization algorithm can make informed decisions about which hyperparameters to try next.

## A better way to run HPO

This is where Flyte shines.

- There's no need to manage a separate centralized database for state tracking, as every objective run is **cached**, **recorded**, and **recoverable** via Flyte's execution engine.
- The entire HPO process is observable in the UI with full lineage and metadata for each trial.
- Each objective is seeded for reproducibility, enabling deterministic trial results.
- If the main optimization task crashes or is terminated, **Flyte can resume from the last successful or failed trial, making the experiment highly fault-tolerant**.
- Trial functions can be strongly typed, enabling rich, flexible hyperparameter spaces while maintaining strict type safety across trials.

In this example, we combine Flyte with Optuna to optimize a `RandomForestClassifier` on the Iris dataset. Each trial runs in an isolated task, and the optimization process is orchestrated asynchronously, with Flyte handling the underlying scheduling, retries, and caching.

## Declare dependencies

We start by declaring a Python environment using Python 3.13 and specifying our runtime dependencies.

```
# /// script
requires-python = "==3.13"
dependencies = [
   "optuna>=4.0.0,<5.0.0",
   "flyte>=2.0.0b0",
   "scikit-learn==1.7.0",
]
# ///
```

With the environment defined, we begin by importing standard library and third-party modules necessary for both the ML task and distributed execution.

```
import asyncio
import typing
from collections import Counter
from typing import Optional, Union
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

These standard library imports are essential for asynchronous execution (`asyncio`), type annotations (`typing`, `Optional`, `Union`), and aggregating trial state counts (`Counter`).

```
import optuna
from optuna import Trial
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
from sklearn.utils import shuffle
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

We use Optuna for hyperparameter optimization and several utilities from scikit-learn to prepare data (`load_iris`), define the model (`RandomForestClassifier`), evaluate it (`cross_val_score`), and shuffle the dataset for randomness (`shuffle`).

```
import flyte
import flyte.errors
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

Flyte is our orchestration framework. We use it to define tasks, manage resources, and recover from execution errors.

## Define the task environment

We define a Flyte task environment called `driver`, which encapsulates metadata, compute resources, the container image context needed for remote execution, and caching behavior.

```
driver = flyte.TaskEnvironment(
    name="driver",
    resources=flyte.Resources(cpu=1, memory="250Mi"),
    image=flyte.Image.from_uv_script(__file__, name="optimizer"),
    cache="auto",
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

This environment specifies that the tasks will run with 1 CPU and 250Mi of memory, the image is built using the current script (`__file__`), and caching is enabled.

## Define the optimizer

Next, we define an `Optimizer` class that handles parallel execution of Optuna trials using async coroutines. This class abstracts the full optimization loop and supports concurrent trial execution with live logging.

```
class Optimizer:
    def __init__(
        self,
        objective: callable,
        n_trials: int,
        concurrency: int = 1,
        delay: float = 0.1,
        study: Optional[optuna.Study] = None,
        log_delay: float = 0.1,
    ):
        self.n_trials: int = n_trials
        self.concurrency: int = concurrency
        self.objective: typing.Callable = objective
        self.delay: float = delay
        self.log_delay = log_delay

        self.study = study if study else optuna.create_study()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

We pass the `objective` function, number of trials to run (`n_trials`), and maximum parallel trials (`concurrency`). The optional delay throttles execution between trials, while `log_delay` controls how often logging runs. If no existing Optuna Study is provided, a new one is created automatically.

```
    async def log(self):
        while True:
            await asyncio.sleep(self.log_delay)

            counter = Counter()

            for trial in self.study.trials:
                counter[trial.state.name.lower()] += 1

            counts = dict(counter, queued=self.n_trials - len(self))

            # print items in dictionary in a readable format
            formatted = [f"{name}: {count}" for name, count in counts.items()]
            print(f"{'    '.join(formatted)}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

This method periodically prints the number of trials in each state (e.g., running, complete, fail). It keeps users informed of ongoing optimization progress and is invoked as a background task when logging is enabled.

![Optuna logging](../../../_static/images/tutorials/hpo/logging.png)
_Logs are streamed live as the execution progresses._

```
    async def spawn(self, semaphore: asyncio.Semaphore):
        async with semaphore:
            trial: Trial = self.study.ask()

            try:
                print("Starting trial", trial.number)

                params = {
                    "n_estimators": trial.suggest_int("n_estimators", 10, 200),
                    "max_depth": trial.suggest_int("max_depth", 2, 20),
                    "min_samples_split": trial.suggest_float(
                        "min_samples_split", 0.1, 1.0
                    ),
                }

                output = await self.objective(params)

                self.study.tell(trial, output, state=optuna.trial.TrialState.COMPLETE)
            except flyte.errors.RuntimeUserError as e:
                print(f"Trial {trial.number} failed: {e}")

                self.study.tell(trial, state=optuna.trial.TrialState.FAIL)

            await asyncio.sleep(self.delay)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

Each call to `spawn` runs a single Optuna trial. The `semaphore` ensures that only a fixed number of concurrent trials are active at once, respecting the `concurrency` parameter. We first ask Optuna for a new trial and generate a parameter dictionary by querying the trial object for suggested hyperparameters. The trial is then evaluated by the objective function. If successful, we mark it as `COMPLETE`. If the trial fails due to a `RuntimeUserError` from Flyte, we log and record the failure in the Optuna study.

```
    async def __call__(self):
        # create semaphore to manage concurrency
        semaphore = asyncio.Semaphore(self.concurrency)

        # create list of async trials
        trials = [self.spawn(semaphore) for _ in range(self.n_trials)]

        logger: Optional[asyncio.Task] = None
        if self.log_delay:
            logger = asyncio.create_task(self.log())

        # await all trials to complete
        await asyncio.gather(*trials)

        if self.log_delay and logger:
            logger.cancel()
            try:
                await logger
            except asyncio.CancelledError:
                pass
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

The `__call__` method defines the overall async optimization routine. It creates the semaphore, spawns `n_trials` coroutines, and optionally starts the background logging task. All trials are awaited with `asyncio.gather`.

```
    def __len__(self) -> int:
        """Return the number of trials in history."""
        return len(self.study.trials)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

This method simply allows us to query the number of trials already associated with the study.

## Define the objective function

The objective task defines how we evaluate a particular set of hyperparameters. It's an async task, allowing for caching, tracking, and recoverability across executions.

```
@driver.task
async def objective(params: dict[str, Union[int, float]]) -> float:
    data = load_iris()
    X, y = shuffle(data.data, data.target, random_state=42)

    clf = RandomForestClassifier(
        n_estimators=params["n_estimators"],
        max_depth=params["max_depth"],
        min_samples_split=params["min_samples_split"],
        random_state=42,
        n_jobs=-1,
    )

    # Use cross-validation to evaluate performance
    score = cross_val_score(clf, X, y, cv=3, scoring="accuracy").mean()

    return score.item()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

We use the Iris dataset as a toy classification problem. The input params dictionary contains the trial's hyperparameters, which we unpack into a `RandomForestClassifier`. We shuffle the dataset for randomness, and compute a 3-fold cross-validation accuracy.

## Define the main optimization loop

The optimize task is the main driver of our optimization experiment. It creates the `Optimizer` instance and invokes it.

```
@driver.task
async def optimize(
    n_trials: int = 20,
    concurrency: int = 5,
    delay: float = 0.05,
    log_delay: float = 0.1,
) -> dict[str, Union[int, float]]:
    optimizer = Optimizer(
        objective=objective,
        n_trials=n_trials,
        concurrency=concurrency,
        delay=delay,
        log_delay=log_delay,
        study=optuna.create_study(
            direction="maximize", sampler=optuna.samplers.TPESampler(seed=42)
        ),
    )

    await optimizer()

    best = optimizer.study.best_trial

    print("✅ Best Trial")
    print("  Number :", best.number)
    print("  Params :", best.params)
    print("  Score  :", best.value)

    return best.params
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

We configure a `TPESampler` for Optuna and `seed` it for determinism. After running all trials, we extract the best-performing trial and print its parameters and score. Returning the best params allows downstream tasks or clients to use the tuned model.

## Run the experiment

Finally, we include an executable entry point to run this optimization using `flyte.run`.

```
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(optimize, 100, 10)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py*

We load Flyte config from `config.yaml`, launch the optimize task with 100 trials and concurrency of 10, and print a link to view the execution in the Flyte UI.

![HPO execution](../../../_static/images/tutorials/hpo/execution.png)
_Each objective run is cached, recorded, and recoverable. With concurrency set to 10, only 10 trials execute in parallel at any given time._


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations ===

# Integrations

Flyte 2 is designed to be extensible by default. While the core platform covers the most common orchestration needs, many production workloads require specialized infrastructure, external services or execution semantics that go beyond the core runtime.

Flyte 2 exposes these capabilities through integrations.

Under the hood, integrations are implemented using Flyte 2's plugin system, which provides a consistent way to extend the platform without modifying core execution logic.

An integration allows you to declaratively enable new capabilities such as distributed compute frameworks or third-party services without manually managing infrastructure. You specify what you need, and Flyte takes care of how it is provisioned, used and cleaned up.

This page covers:

- The types of integrations Flyte 2 supports today
- How integrations fit into Flyte 2's execution model
- How to use integrations in your tasks
- The integrations available out of the box

If you need functionality that doesn't exist yet, Flyte 2's plugin system is intentionally open-ended. You can build and register your own integrations using the same architecture described here.

## Integration categories

Flyte 2 integrations fall into the following categories:

1. **Distributed compute**: Provision transient compute clusters to run tasks across multiple nodes, with automatic lifecycle management.
2. **Agentic AI**: Support for various common aspects of agentic AI applications.
3. **Configuration**: Compose and pass hierarchical configuration objects between tasks, with type-safe schemas and CLI/YAML composition.
4. **Experiment tracking**: Integrate with experiment tracking platforms for logging metrics, parameters, and artifacts.
5. **Data validation**: Enforce schema contracts on dataframes flowing between tasks, with automatic validation reports.
6. **Data types**: Add native support for additional file and dataframe types as task inputs and outputs.
7. **Connectors**: Stateless, long-running services that receive execution requests via gRPC and then submit work to external (or internal) systems.
8. **LLM Serving**: Deploy and serve large language models with an OpenAI-compatible API.
9. **Notebook execution**: Run parameterized Jupyter notebooks as typed Flyte tasks with cell-level reports.
10. **Observability**: Export task and agent telemetry to external tracing and observability backends.

## Distributed compute

Distributed compute integrations allow tasks to run on dynamically provisioned clusters. These clusters are created just-in-time, scoped to the task execution and torn down automatically when the task completes.

This enables large-scale parallelism without requiring users to operate or maintain long-running infrastructure.

### Supported distributed compute integrations

| Plugin                      | Description                                      | Common use cases                                       |
| --------------------------- | ------------------------------------------------ | ------------------------------------------------------ |
| [Ray](./ray/_index)         | Provisions Ray clusters via KubeRay              | Distributed Python, ML training, hyperparameter tuning |
| [Spark](./spark/_index)     | Provisions Spark clusters via Spark Operator     | Large-scale data processing, ETL pipelines             |
| [Dask](./dask/_index)       | Provisions Dask clusters via Dask Operator       | Parallel Python workloads, dataframe operations        |
| [PyTorch](./pytorch/_index) | Distributed PyTorch training with elastic launch | Single-node and multi-node training                    |

Each plugin encapsulates:

- Cluster provisioning
- Resource configuration
- Networking and service discovery
- Lifecycle management and teardown

From the task author's perspective, these details are abstracted away.

### How the plugin system works

At a high level, Flyte 2's distributed compute plugin architecture follows a simple and consistent pattern.

#### 1. Registration

Each plugin registers itself with Flyte 2's core plugin registry:

- **`TaskPluginRegistry`**: The central registry for all distributed compute plugins
- Each plugin declares:
  - Its configuration schema
  - How that configuration maps to execution behavior

This registration step makes the plugin discoverable by the runtime.

#### 2. Task environments and plugin configuration

Integrations are activated through a `TaskEnvironment`.

A `TaskEnvironment` bundles:

- A container image
- Execution settings
- A plugin configuration object enabled with `plugin_config`

The plugin configuration describes _what_ infrastructure or integration the task requires.

#### 3. Automatic provisioning and execution

When a task associated with a `TaskEnvironment` runs:

1. Flyte inspects the environment's plugin configuration
2. The plugin provisions the required infrastructure or integration
3. The task executes with access to that capability
4. Flyte cleans up all transient resources after completion

### Example: Using the Dask plugin

Below is a complete example showing how a task gains access to a Dask cluster simply by running inside an environment configured with the Dask plugin.

```python
from flyteplugins.dask import Dask, WorkerGroup
import flyte

# Define the Dask cluster configuration
dask_config = Dask(
    workers=WorkerGroup(number_of_workers=4)
)

# Create a task environment that enables Dask
env = flyte.TaskEnvironment(
    name="dask_env",
    plugin_config=dask_config,
    image=image,
)

# Any task in this environment has access to the Dask cluster
@env.task
async def process_data(data: list) -> list:
    from distributed import Client

    client = Client()  # Automatically connects to the provisioned cluster
    futures = client.map(transform, data)
    return client.gather(futures)
```

When `process_data` executes, Flyte performs the following steps:

1. Provisions a Dask cluster with 4 workers
2. Executes the task with network access to the cluster
3. Tears down the cluster once the task completes

No cluster management logic appears in the task code. The task only expresses intent.

### Key design principle

All distributed compute integrations follow the same mental model:

- You declare the required capability via configuration
- You attach that configuration to a task environment
- Tasks decorated with that environment automatically gain access to the capability

This makes it easy to swap execution backends or introduce distributed compute incrementally without rewriting workflows.

## Agentic AI

Agentic AI integrations let you run agents written in a third-party framework as durable Flyte tasks. You keep the framework's own idioms; Flyte supplies the runtime underneath, so tool calls become containerized child actions with their own resources, retries and caching, completed model turns replay instead of re-billing, and conversations persist across runs.

### Supported agentic AI integrations

| Plugin                              | Description                                                                                                              | Common use cases                                 |
| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------ |
| [Agent frameworks](./agents/_index) | Adapters for ten agent SDKs, including OpenAI, Claude, Google ADK, Mistral, LangChain, LangGraph, CrewAI and Pydantic AI | Durable agents, tools as tasks, cross-run memory |
| [Code generation](./codegen/_index) | LLM-driven code generation with automatic testing in sandboxes                                                           | Data processing, ETL, analysis pipelines         |

## Experiment tracking

Experiment tracking integrations let you log metrics, parameters, and artifacts to external tracking platforms during Flyte task execution.

### Supported experiment tracking integrations

| Plugin                               | Description                  | Common use cases                                 |
| ------------------------------------ | ---------------------------- | ------------------------------------------------ |
| [MLflow](./mlflow/_index)            | MLflow experiment tracking   | Experiment tracking, autologging, model registry |
| [Weights and Biases](./wandb/_index) | Weights & Biases integration | Experiment tracking and hyperparameter tuning    |

## Configuration

Configuration integrations let you compose and pass hierarchical configuration objects between Flyte tasks, with type-safe schemas and CLI/YAML composition.

### Supported configuration integrations

| Plugin                          | Description                                                       | Common use cases                                                                  |
| ------------------------------- | ----------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [OmegaConf](./omegaconf/_index) | `DictConfig` / `ListConfig` as native task input and output types | Passing composed configs between tasks, structured configs, YAML-driven pipelines |
| [Hydra](./hydra/_index)         | Hydra config composition and sweep submission for Flyte tasks     | YAML-driven experiment composition, grid and Bayesian sweeps, hardware presets    |

## Data validation

Data validation integrations enforce schema contracts on the dataframes flowing between tasks. They validate data at task boundaries, catch type and constraint violations early, and produce HTML reports visible in the Flyte UI.

### Supported data validation integrations

| Plugin                      | Description                                                | Common use cases                                            |
| --------------------------- | ---------------------------------------------------------- | ----------------------------------------------------------- |
| [Pandera](./pandera/_index) | Validates dataframes with pandera `DataFrameModel` schemas | Schema enforcement, data quality checks, validation reports |

## Data types

Data type integrations add native support for additional file and dataframe types as task inputs and outputs. They register typed encoders and decoders with Flyte's type engine, so you can annotate task signatures with the type directly and let Flyte handle serialization.

### Supported data type integrations

| Plugin                    | Description                                                  | Common use cases                                            |
| ------------------------- | ------------------------------------------------------------ | ----------------------------------------------------------- |
| [JSONL](./jsonl/_index)   | Typed `JsonlFile` / `JsonlDir` for streaming JSON Lines data | LLM dataset pipelines, event logs, large line-delimited I/O |
| [Polars](./polars/_index) | Native `pl.DataFrame` / `pl.LazyFrame` support via Parquet   | High-performance dataframe ETL, feature engineering         |

## Connectors

Connectors are stateless, long-running services that receive execution requests via gRPC and then submit work to external (or internal) systems. Each connector runs as its own Kubernetes deployment, and is triggered when a Flyte task of the matching type is executed.

Although they normally run inside the data plane, you can also run connectors locally as long as the required secrets/credentials are present locally. This is useful because connectors are just Python services that can be spawned in-process.

Connectors are designed to scale horizontally and reduce load on the core Flyte backend because they execute _outside_ the core system. This decoupling makes connectors efficient, resilient, and easy to iterate on. You can even test them locally without modifying backend configuration, which reduces friction during development.

### Supported connectors

| Connector                         | Description                                 | Common use cases                         |
| --------------------------------- | ------------------------------------------- | ---------------------------------------- |
| [Snowflake](./snowflake/_index)   | Run SQL queries on Snowflake asynchronously | Data warehousing, ETL, analytics queries |
| [BigQuery](./bigquery/_index)     | Run SQL queries on Google BigQuery          | Data warehousing, ETL, analytics queries |
| [Databricks](./databricks/_index) | Run PySpark jobs on Databricks clusters     | Large-scale data processing, Spark ETL   |

### Creating a new connector

If none of the existing connectors meet your needs, you can build your own.

> [!NOTE]
> Connectors communicate via Protobuf, so in theory they can be implemented in any language.
> Today, only **Python** connectors are supported.

### Async connector interface

To implement a new async connector, extend `AsyncConnector` and implement the following methods, all of which must be idempotent:

| Method     | Purpose                                                     |
| ---------- | ----------------------------------------------------------- |
| `create`   | Launch the external job (via REST, gRPC, SDK, or other API) |
| `get`      | Fetch current job state (return job status or output)       |
| `delete`   | Delete / cancel the external job                            |
| `get_logs` | Stream paginated log lines to the Flyte UI                  |

To test the connector locally, the connector task should inherit from
[AsyncConnectorExecutorMixin](https://github.com/flyteorg/flyte-sdk/blob/1d49299294cd5e15385fe8c48089b3454b7a4cd1/src/flyte/connectors/_connector.py#L206). This mixin simulates how the Flyte 2 system executes asynchronous connector tasks, making it easier to validate your connector implementation before deploying it.

### Example: Batch job connector

The following example implements a connector that simulates submitting and polling an external batch job. Replace the mock logic with real API calls for your use case.

**Connector** (`my_connector/connector.py`):

```
import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Optional

from flyteidl2.connector.connector_pb2 import (
    GetTaskLogsResponse,
    GetTaskLogsResponseBody,
    GetTaskLogsResponseHeader,
)
from flyteidl2.core.execution_pb2 import TaskExecution
from flyteidl2.logs.dataplane.payload_pb2 import LogLine, LogLineOriginator
from google.protobuf.timestamp_pb2 import Timestamp

from flyte import logger
from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta

@dataclass
class BatchJobMetadata(ResourceMeta):
    job_id: str
    created_at: float

class BatchJobConnector(AsyncConnector):
    name = "Batch Job Connector"
    task_type_name = "batch_job"
    metadata_type = BatchJobMetadata

    async def create(self, task_template, inputs: Optional[Dict[str, Any]] = None, **kwargs) -> BatchJobMetadata:
        job_id = str(uuid.uuid4())[:8]
        logger.info(f"Submitted batch job {job_id}")
        return BatchJobMetadata(job_id=job_id, created_at=time.time())

    async def get(self, resource_meta: BatchJobMetadata, **kwargs) -> Resource:
        elapsed = time.time() - resource_meta.created_at
        if elapsed < 5:
            return Resource(phase=TaskExecution.RUNNING, message="Job in progress")
        return Resource(
            phase=TaskExecution.SUCCEEDED,
            message="Job completed",
            outputs={"result": f"output-from-{resource_meta.job_id}"},
        )

    async def delete(self, resource_meta: BatchJobMetadata, **kwargs):
        logger.info(f"Cancelled job {resource_meta.job_id}")

    async def get_logs(self, resource_meta: BatchJobMetadata, token: str = "", **kwargs):
        def line(message: str, ts: float) -> LogLine:
            t = Timestamp()
            t.FromSeconds(int(ts))
            return LogLine(timestamp=t, message=message, originator=LogLineOriginator.USER)

        start = resource_meta.created_at
        job_id = resource_meta.job_id
        pages = {
            "": GetTaskLogsResponseBody(lines=[
                line(f"[INFO] Job {job_id} submitted", start),
                line(f"[INFO] Job {job_id} started", start + 1),
            ]),
            "page-2": GetTaskLogsResponseBody(lines=[
                line(f"[INFO] Job {job_id} finished", start + 5),
            ]),
        }
        next_tokens = {"": "page-2", "page-2": ""}
        yield GetTaskLogsResponse(body=pages.get(token, GetTaskLogsResponseBody(lines=[])))
        next_token = next_tokens.get(token, "")
        if next_token:
            yield GetTaskLogsResponse(header=GetTaskLogsResponseHeader(token=next_token))

ConnectorRegistry.register(BatchJobConnector())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch_job/connector.py*

**Task plugin** (`my_connector/task.py`):

```
from dataclasses import dataclass
from typing import Any, Dict, Optional, Type

from flyte.connectors import AsyncConnectorExecutorMixin
from flyte.extend import TaskTemplate
from flyte.models import NativeInterface, SerializationContext

@dataclass
class BatchJobConfig:
    timeout_seconds: int = 300

class BatchJobTask(AsyncConnectorExecutorMixin, TaskTemplate):
    _TASK_TYPE = "batch_job"

    def __init__(self, name: str, plugin_config: BatchJobConfig,
                 inputs: Optional[Dict[str, Type]] = None,
                 outputs: Optional[Dict[str, Type]] = None, **kwargs):
        super().__init__(
            name=name,
            interface=NativeInterface(
                {k: (v, None) for k, v in inputs.items()} if inputs else {},
                outputs or {},
            ),
            task_type=self._TASK_TYPE,
            image=None,
            **kwargs,
        )
        self.plugin_config = plugin_config

    def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]:
        return {"timeout_seconds": self.plugin_config.timeout_seconds}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch_job/task.py*

**Usage**:

```python
import flyte
from my_connector.task import BatchJobConfig, BatchJobTask

batch_job = BatchJobTask(
    name="my_batch_job",
    plugin_config=BatchJobConfig(timeout_seconds=60),
    inputs={"name": str},
    outputs={"result": str},
)

flyte.TaskEnvironment.from_task("batch-job-env", batch_job)
```

### Connector-level secrets

If your connector needs credentials (API keys, tokens) shared across all tasks, pass them as environment variables into the connector process.

Set environment variables on the connector Kubernetes deployment:

```bash
kubectl set env deployment/<connector-deployment> MY_API_KEY=<value> -n <flyte-namespace>
```

Inside the connector, read the secret from the environment:

```python
import os

api_key = os.environ["MY_API_KEY"]
```

See **Tasks > Configure tasks > Secrets** for how to store and manage secrets.

### Deploy a custom connector

Deploying a connector requires two steps: building a Docker image that contains your connector code and then patching the connector Kubernetes deployment to use it.

**Step 1: Build the connector image**

```python
import asyncio
from flyte import Image
from flyte.extend import ImageBuildEngine

async def build_connector_image(registry: str, name: str, builder: str = "local"):
    image = Image.from_debian_base(
        registry=registry, name=name
    ).with_pip_packages("flyte[connector]", "my-connector-package")
    await ImageBuildEngine.build(image, builder=builder)

if __name__ == "__main__":
    asyncio.run(
        build_connector_image(
            registry="<YOUR_REGISTRY>", name="my-connector", builder="local"
        )
    )
```

**Step 2: Override the connector deployment image**

Once the image is pushed, patch the connector Kubernetes deployment to use it:

```bash
kubectl set image deployment/<connector-deployment-name> \
    connector=<YOUR_REGISTRY>/my-connector:<TAG> \
    -n <flyte-namespace>
```

Replace `<connector-deployment-name>` with the name of your connector deployment (e.g. `flyte-connector`), and `<flyte-namespace>` with the namespace where Flyte is installed (typically `flyte`).

## LLM serving

LLM serving integrations let you deploy and serve large language models as Flyte apps with an OpenAI-compatible API. They handle model loading, GPU management, and autoscaling.

### Supported LLM serving integrations

| Plugin                                                          | Description                                         | Common use cases             |
| --------------------------------------------------------------- | --------------------------------------------------- | ---------------------------- |
| **Apps > Native app integrations > SGLang app** | Deploy models with SGLang's high-throughput runtime | LLM inference, model serving |
| **Apps > Native app integrations > vLLM app**     | Deploy models with vLLM's PagedAttention engine     | LLM inference, model serving |

For full setup instructions including multi-GPU deployment, model prefetching, and autoscaling, see the **Apps > Native app integrations > SGLang app** and **Apps > Native app integrations > vLLM app** pages.

## Notebook execution

Notebook execution integrations let you run Jupyter notebooks as first-class Flyte tasks with typed inputs and outputs, HTML reports surfaced in the Flyte UI, and the ability to call other Flyte tasks from within the notebook.

### Supported notebook execution integrations

| Plugin                          | Description                                                                                | Common use cases                                                                                     |
| ------------------------------- | ------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------- |
| [Papermill](./papermill/_index) | Parameterize and execute `.ipynb` files via [papermill](https://papermill.readthedocs.io/) | Productionizing exploratory notebooks, cell-by-cell HTML reports, notebook-driven analysis pipelines |

## Observability

Observability integrations export telemetry from a Flyte run to an external backend. They understand that a durable run is several processes over time, so a run that crashes and resumes arrives as one trace rather than several, and steps replayed from the durable log are still recorded.

### Supported observability integrations

| Plugin                                                              | Description                                                                                 | Common use cases                                                     |
| ------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [OpenTelemetry](./opentelemetry/_index)                             | Records tasks and traced steps as OpenTelemetry spans and exports them over OTLP            | Distributed tracing, debugging cross-service latency, durable traces |
| [Grafana Agent Observability](./grafana-agent-observability/_index) | Sends agent generations, tool calls, token usage, and cost to Grafana, grouped by Flyte run | LLM cost tracking, prompt iteration, agent debugging                 |

Both carry trace context across task boundaries using Flyte's **Tasks > Build tasks > Custom context** primitive, so a run submitted from inside a caller's span joins that caller's trace.

## Subpages
- **Agent frameworks**
- **BigQuery**
- **Code generation**
- **Dask**
- **Databricks**
- **Grafana Agent Observability**
- **Hydra**
- **JSONL**
- **MLflow**
- **OmegaConf**
- **OpenTelemetry**
- **Pandera**
- **Papermill**
- **Polars**
- **PyTorch**
- **Ray**
- **Snowflake**
- **Spark**
- **Weights & Biases**


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents ===

# Agent frameworks

Agent frameworks are good at deciding what an agent should do next. They are less good at what happens when the worker running the agent dies on turn seven, when one tool needs a GPU and another needs 200 MB of RAM, or when you need to explain to someone what the agent actually did last Tuesday.

The Flyte agent plugins cover that half. You keep writing agents in your framework's own idioms. Flyte becomes the runtime underneath: completed model turns replay instead of re-billing, every tool call is a containerized child action with its own resources and cache, conversations persist across runs, and the whole thing renders as a timeline in the task report.

Ten frameworks are supported, each as a separate package on a shared core. The call shape is identical across all of them, so switching frameworks is mostly a change of import.

## Supported frameworks

| Framework | Page | Package |
|---|---|---|
| [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/) | **Agent frameworks > OpenAI Agents SDK** | `flyteplugins-agents-openai` |
| [Claude Agent SDK](https://github.com/anthropics/claude-agent-sdk-python) | **Agent frameworks > Claude Agent SDK** | `flyteplugins-agents-claude` |
| [Google ADK](https://github.com/google/adk-python) | **Agent frameworks > Google ADK** | `flyteplugins-agents-google` |
| [Mistral Agents](https://docs.mistral.ai/agents/agents_introduction/) | **Agent frameworks > Mistral Agents** | `flyteplugins-agents-mistral` |
| [LangChain](https://docs.langchain.com/oss/python/langchain/agents) | **Agent frameworks > LangChain** | `flyteplugins-agents-langchain` |
| [LangGraph](https://langchain-ai.github.io/langgraph/) | **Agent frameworks > LangGraph** | `flyteplugins-agents-langgraph` |
| [Deep Agents](https://docs.langchain.com/oss/python/deepagents/overview) | **Agent frameworks > Deep Agents** | `flyteplugins-agents-deepagents` |
| [CrewAI](https://docs.crewai.com/) | **Agent frameworks > CrewAI** | `flyteplugins-agents-crewai` |
| [Pydantic AI](https://ai.pydantic.dev/) | **Agent frameworks > Pydantic AI** | `flyteplugins-agents-pydantic-ai` |
| [Hermes](https://pypi.org/project/hermes-agent/) | **Agent frameworks > Hermes** | `flyteplugins-agents-hermes` |

Install the one you need. Each package pulls in `flyteplugins-agents-core` and the underlying SDK.

```bash
pip install flyteplugins-agents-openai
```

## Two decorators

Every adapter exports the same two things: `tool` and `run_agent`.

`tool` stacks on top of `@env.task`. The result is simultaneously a normal Flyte task and a tool your framework recognizes, so when the model calls it, the call becomes a durable child action rather than a function call inside the agent process.

`run_agent` drives the framework's own agent loop from inside a Flyte task. That task is the durable parent: give it `retries=` for self-healing and `report=True` for the timeline.

```python
import flyte
from flyteplugins.agents.openai import run_agent, tool

env = flyte.TaskEnvironment("agent")

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(question, tools=[get_weather], model="gpt-4.1")
```

Swap the import line for `flyteplugins.agents.crewai` or `flyteplugins.agents.langchain` and the rest of the file stays as it is, apart from the model name.

## Quick start

A complete, runnable agent. The API key is read from the environment, so wire it as a Flyte secret rather than passing it as a task input.

```python{hl_lines=[2, 6, 14, 21, "30-35"]}
import flyte
from flyteplugins.agents.openai import run_agent, tool

env = flyte.TaskEnvironment(
    "city-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-agents-openai",
    ),
    resources=flyte.Resources(cpu=1),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@tool
@env.task(cache="auto", retries=3)
async def get_population(city: str) -> int:
    """Get the population of a city."""
    return {"Paris": 2102650, "Tokyo": 13929286}.get(city, 1_000_000)

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[get_weather, get_population],
        instructions="You are a concise city-facts assistant. Use the tools to answer.",
        model="gpt-4.1",
    )

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(city_agent, question="What's the weather and population of Paris?")
    print(run.url)
```

Run it:

```bash
flyte run city_agent.py city_agent --question "What's the weather and population of Paris?"
```

Add `--local` right after `run` to execute on your machine instead. The durability, memory and observability layers become transparent no-ops outside a task context, so the same file runs unchanged.

![City agent](../../_static/images/integrations/agents/city_agent_index.png)

## What Flyte adds

| Capability | What it means |
|---|---|
| Tools as child actions | Each tool call runs in its own container with its own resources, retries and cache. A retrieval tool can hold a GPU while the agent task holds one CPU. |
| Model-turn replay | Completed turns are recorded. When the parent task is retried, they replay from the record instead of calling the model again. |
| Self-healing | `retries=` on the agent task, combined with per-turn and per-tool replay, means a transient failure resumes rather than restarting. |
| Cross-run memory | A `memory_key` continues a conversation across separate runs, workers and restarts, backed by object storage. |
| Observability | Turns, tool calls, results and token usage render into the task report. |
| Human in the loop | A tool can suspend on a Flyte condition and wait for a human. The run survives restarts while it waits. |

**Agent frameworks > How it works** covers each of these in detail, including where the durability seam sits and why.

## Capability matrix

The adapters share a contract but the underlying SDKs differ, so durability lands in different places.

| Framework | Model-turn durability | Tool type | What memory persists | Python |
|---|---|---|---|---|
| **Agent frameworks > OpenAI Agents SDK** | Per turn | `FunctionTool` | Conversation transcript | 3.10+ |
| **Agent frameworks > Claude Agent SDK** | Per session, via resume | In-process MCP tool | Conversation transcript | 3.10+ |
| **Agent frameworks > Google ADK** | Per turn | Plain callable | Session events | 3.10+ |
| **Agent frameworks > Mistral Agents** | Per turn | Plain callable | Server-side conversation ID | 3.10+ |
| **Agent frameworks > LangChain** | Per turn, built agents | `StructuredTool` | Conversation transcript | 3.10+ |
| **Agent frameworks > LangGraph** | Per turn, via `ai_node` | `StructuredTool` | Conversation transcript | 3.10+ |
| **Agent frameworks > Deep Agents** | Per turn, built agents | `StructuredTool` | Transcript and virtual filesystem | 3.11+ |
| **Agent frameworks > CrewAI** | Per turn, built agents | `BaseTool` | Conversation transcript | 3.10+ |
| **Agent frameworks > Pydantic AI** | Per turn | Plain callable | Message history | 3.10+ |
| **Agent frameworks > Hermes** | Not available | Registry tool | Conversation transcript | 3.11+ |

"Built agents" means durability applies when `run_agent` constructs the agent for you. If you hand it a fully pre-built agent, Flyte cannot reach inside to wrap the model, so you wrap it yourself. Each page says exactly how.

Tool calls are durable in every case, including Hermes, regardless of the `durable` setting.

## Choosing a framework

The plugins do not have an opinion here. Pick the framework you would have picked anyway. The two things worth knowing:

- If you want per-turn replay and you are starting fresh, everything except Hermes gives it to you on the builder path.
- If you already own a compiled graph or a configured agent object, check the framework's page for how durability is applied on the pre-built path. LangGraph is designed around this case: you build the `StateGraph`, and `ai_node` and `tool_node` supply the durable pieces.

## Next steps

- **Agent frameworks > How it works**: the runtime model, from the durable parent down to the trace leaf.
- Pick a framework page above for SDK-specific setup, options and limitations.
- [Build an agent](../../user-guide/agents/build-agent/_index): Flyte's own agent harness, if you would rather not bring a framework at all.

## Subpages
- **Agent frameworks > How it works**
- **Agent frameworks > OpenAI Agents SDK**
- **Agent frameworks > Claude Agent SDK**
- **Agent frameworks > Google ADK**
- **Agent frameworks > Mistral Agents**
- **Agent frameworks > LangChain**
- **Agent frameworks > LangGraph**
- **Agent frameworks > Deep Agents**
- **Agent frameworks > CrewAI**
- **Agent frameworks > Pydantic AI**
- **Agent frameworks > Hermes**


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/how-it-works ===

# How it works

Every adapter follows the same division of labor. Understanding it once means you can read any of the framework pages quickly, and it explains why some capabilities land differently depending on the SDK.

## The division of labor

Three levels, each mapping to a Flyte primitive:

| Level | Flyte primitive | What it gives you |
|---|---|---|
| The agent run | An `@env.task` (the durable parent) | Retries, timeout, resources, the report |
| Each model turn | A `flyte.trace` leaf | Replay on retry, no repeat billing |
| Each tool call | A child action | Own container, own resources, retries, caching |

The framework still owns the loop. Nothing here reimplements tool-calling or turn management. `run_agent` starts the SDK's own runner inside your task and instruments the seams around it.

```python
@env.task(report=True, retries=3)      # the durable parent
async def city_agent(question: str) -> str:
    return await run_agent(            # the SDK's loop runs in here
        question,
        tools=[get_weather],           # each call is a child action
        model="gpt-4.1",
    )
```

```mermaid
flowchart TB
    user(["user question"]) --> t1

    subgraph parent["city_agent · @env.task · the durable parent (retries · timeout · report)"]
        direction TB
        t1["Model turn 1<br/>flyte.trace leaf · replays on retry"]

        subgraph c1["own container"]
            w["get_weather<br/>child action · own resources · retries · cache"]
        end
        subgraph c2["own container"]
            p["get_population<br/>child action · own resources · retries · cache"]
        end

        t2["Model turn 2<br/>flyte.trace leaf · replays on retry"]

        t1 --> w
        t1 --> p
        w --> t2
        p --> t2
    end

    t2 --> answer(["final answer"])
```

The parent task is the box. The model turns inside it are `flyte.trace` leaves that replay on retry. Each tool call is a child action in its own container, sized and cached on its own terms.

## Tools are Flyte tasks

Stacking `tool` on `@env.task` produces one object that is both things at once. The framework sees whatever tool type it expects, and Flyte sees the task.

```python
@tool
@env.task(cache="auto", retries=3, resources=flyte.Resources(gpu="T4"))
async def embed_documents(query: str) -> list[float]:
    """Embed a query for semantic search."""
    ...
```

When the model calls `embed_documents`, Flyte submits a child action. That action runs in its own container with a T4, retries on failure, and hits the cache on identical inputs. The agent task itself keeps whatever modest resources you gave it.

This is the part that is hard to get any other way. A tool in a normal agent process is a function call: same machine, same memory limit, same failure domain. Here each tool is sized and cached on its own terms, and a tool crash does not take down the conversation.

The tool's schema, name and description come from the task. The docstring becomes the description the model sees, so write it for the model.

> [!NOTE] Docstrings are prompts
> The first line of the docstring is what the model reads when deciding whether to call the tool. Vague docstrings produce vague tool selection.

### Passing tools

`tools=` accepts `tool`-wrapped tasks. It also accepts bare `@env.task` templates, which are wrapped for you:

```python
await run_agent(question, tools=[get_weather])       # tool-wrapped, or
await run_agent(question, tools=[some_plain_task])   # bare task, wrapped on the fly
```

Wrap explicitly with `@tool` when you want the tool object at module scope, for example to attach it to a pre-built agent or a subagent.

### Renaming a tool

```python
search = tool(query_warehouse, name="search", description="Search the product catalog.")
```

The OpenAI adapter forwards to the SDK's own kwargs (`name_override`, `description_override`) instead.

## Durable model turns

A retried task normally starts from scratch. For an agent that means paying for every completed turn a second time, and getting different answers the second time around.

Instead, each model turn is recorded as a `flyte.trace` leaf keyed by a fingerprint of the request. On a retry, a turn whose fingerprint already has a record returns that record without calling the model.

The recording happens at the seam below the framework's loop, not around it. For the OpenAI adapter that is a `ModelProvider`; for Google ADK it is `BaseLlm.generate_content_async`; for the LangChain family it is the chat model itself; for Mistral it is the two HTTP methods the runner uses per turn. Different seam, same mechanism. The loop above it is untouched, so handoffs, guardrails, structured output and everything else the SDK does keep working.

Turn durability is on by default. Switch it off with `durable=False`.

```python
await run_agent(question, tools=[get_weather], model="gpt-4.1", durable=False)
```

### What replay means in practice

A task that crashes partway through an agent run, with `retries=3`, resumes like this:

1. Completed model turns return from their trace records. No model calls, no tokens.
2. Completed tool calls return from cache, if the task was declared with `cache="auto"`.
3. Execution continues from the first step that never finished.

Transient model failures such as 429s and 5xx are a separate matter. Those are retried in place by the provider's own client, below the durable wrapper, so a turn is only recorded once it has actually succeeded.

Here is that recovery in the Flyte report. It is one crash-resume run viewed at each attempt, using the `openai_crash_resume.py` example: the task runs the agent for real, crashes on its first attempt, and Flyte retries it. The `Attempt` selector at the top right switches between the two views.

**Attempt 1, the first run.** The agent does the full job. Both model turns are live calls with real token usage (103 and 154 input tokens), and both tools execute as child actions, `get_weather` in 5.1 s and `get_population` in 16.6 s. The agent timeline totals 20.3 s. The task then crashes.

![Attempt 1](../../_static/images/integrations/agents/attempt_1.png)

**Attempt 2, the retry.** The two model `response` rows are gone. The turns replayed from their `flyte.trace` records, so the model was never called and no tokens were spent. The tool calls are cache hits, `get_weather` in 59 ms and `get_population` in 102 ms. Same answer, with the agent timeline down from 20.3 s to 0.44 s.

![Attempt 2](../../_static/images/integrations/agents/attempt_2.png)

The absence of those model rows on the retry is the replay. The second attempt re-drives the agent loop, but every completed turn comes back from its record and every tool from cache, so no work is repeated and nothing is re-billed.

### Where durability does not reach

Two cases, both called out on the relevant framework pages:

**Pre-built agents:** If you construct the agent object yourself and pass it as `run_agent(agent=...)`, Flyte often cannot reach the model inside it to wrap it. The LangChain family exposes `DurableChatModel` for this; Pydantic AI applies the wrapper through `Agent.override`. Tool calls stay durable either way.

**Subprocess loops:** The Claude Agent SDK runs its loop in the Claude Code runtime, a subprocess Flyte does not intercept, so a turn cannot be a trace leaf. That adapter uses the SDK's own session resume against a `flyte.Checkpoint` instead. It is coarser, whole-session rather than per-turn, but it is real. Hermes exposes no per-turn hook at all, so `durable=` is accepted and ignored there.

## Cross-run memory

Pass a `memory_key` and the conversation continues across separate runs, separate workers and restarts:

```python
@env.task(report=True, retries=3)
async def chat(message: str, memory_key: str) -> str:
    return await run_agent(message, model="gpt-4.1", memory_key=memory_key)
```

```bash
flyte run chat.py chat --message "Hi, I'm Alice and I love hiking." --memory_key user-alice
flyte run chat.py chat --message "What do I like?" --memory_key user-alice
```

The second run answers correctly. It is a separate run on possibly a different worker, and the transcript came from object storage.

The backing store is Flyte's keyed `MemoryStore`, which resolves a deterministic remote path from the key and the run context. Two runs sharing a key share one store. It carries a message transcript and a path-addressed key-value space with audit and version history, so the same key covers both conversation history and durable named facts.

What each adapter actually persists differs, because the SDKs represent conversation state differently. Mistral keeps transcripts server-side, so only the conversation ID is stored. Google ADK persists its event list. Deep Agents persists the virtual filesystem alongside the transcript. The [capability matrix](./_index) has the full list.

`memory_key` should be a single segment, such as a user or thread ID. Memory is best-effort by design: if no durable store can be resolved, the adapter logs a warning and the run continues without memory rather than failing.

> [!NOTE] Memory needs a configured context
> The store path is derived from the active org, project and domain. Run with `flyte.init_from_config()` or against a backend. Local runs without a context skip memory silently.

## Observability

With `report=True` on the agent task, the run renders as a timeline in the report tab: assistant turns, tool calls with their arguments, tool results, errors, and a token usage summary.

```python
@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(question, tools=[get_weather], model="gpt-4.1")
```

Turn it off with `observability=False`.

Token accounting is honest about replay. On a retried run, turns served from their durable records are counted as cached rather than presented as fresh spend.

## Human in the loop

A tool is a Flyte task, and a Flyte task can suspend on a condition. That gives you an approval gate no agent SDK has an equivalent for, because the run genuinely suspends rather than blocking a thread, and survives a restart while it waits.

```python
@tool
@env.task(retries=3)
async def issue_refund(account_id: str, amount_usd: float) -> str:
    """Issue a refund. Requires human approval before it runs."""
    condition = await flyte.new_condition.aio(
        f"approve_refund_{account_id}",
        prompt=f"Approve a ${amount_usd:.2f} refund to account {account_id}?",
        data_type=bool,
    )
    if not await condition.wait.aio():
        return f"Refund to {account_id} was declined by a human reviewer."
    return f"refunded ${amount_usd:.2f} to account {account_id}"
```

The model decides whether to call the tool. You decide what happens when it does. The agent sees the decline as an ordinary tool result and carries on.

![Approval](../../_static/images/integrations/agents/agent_approval.png)

## Multi-agent orchestration

Handoffs and subagents inside a single `run_agent` work as the framework defines them. Flyte adds a layer above: each agent can be its own task, composed with ordinary control flow.

```python
@env.task(retries=3)
async def research(subtopic: str) -> str:
    return await run_agent(f"Research: {subtopic}", tools=[search_web], model="gpt-4.1")

@env.task(report=True, retries=3)
async def pipeline(topic: str) -> str:
    subtopics = await plan(topic)
    with flyte.group("parallel-research"):
        findings = await asyncio.gather(*(research(s) for s in subtopics))
    return await synthesize(topic, list(findings))
```

Each researcher is a separate durable action with its own retries, cache and report. The fan-out is real distributed parallelism across workers, not asyncio inside one process.

## Sync and async

`run_agent` is a coroutine function. Await it from an async task. From a sync task, call `run_agent_sync`, which every adapter also exports with the same signature.

```python
@env.task(report=True)
async def async_agent(q: str) -> str:
    return await run_agent(q, tools=[get_weather], model="gpt-4.1")

@env.task(report=True)
def sync_agent(q: str) -> str:
    return run_agent_sync(q, tools=[get_weather], model="gpt-4.1")
```

## Running locally

Call `run_agent` from inside an `@env.task`. That task is what makes the durability, memory and observability layers real.

Outside a task context they are transparent no-ops: `flyte.trace` passes through, memory resolves to nothing, the report is not rendered. The same file runs locally unchanged, which is what you want for iteration, but it also means a local run tells you nothing about whether replay works. Run on a backend to see that.

## The shared contract

Every adapter exports `tool`, `run_agent` and `run_agent_sync`, and every `run_agent` accepts `tools`, `model`, `instructions`, `durable`, `observability` and `memory_key`. This is enforced in CI by a conformance check that each adapter runs as a one-line test, so the surface cannot drift between packages.

Adapters add their own keyword arguments on top where the SDK calls for it, such as `run_config` for OpenAI, `options` for Claude, `agent_id` for Mistral, and `subagents` for Deep Agents. Those are documented on each framework's page.

The shared machinery lives in `flyteplugins-agents-core`, which every adapter depends on. It has no agent SDK dependency of its own. If you are writing an adapter for a framework that is not listed, that package is the contract to implement.

## Next steps

- [Agent frameworks](./_index): the supported list and the capability matrix.
- [Secrets](../../user-guide/tasks/task-configuration/secrets): how to wire provider API keys.
- [Caching](../../user-guide/tasks/task-configuration/caching): what `cache="auto"` does on a tool task.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/openai ===

# OpenAI Agents SDK

Run [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/) agents on Flyte. The SDK's `Runner` still drives the loop, including handoffs, guardrails and structured output. Flyte supplies the runtime: tools become durable child actions, model turns replay on retry, and the SDK's trace telemetry renders into the task report instead of being shipped to OpenAI's traces dashboard.

## Installation

```bash
pip install flyteplugins-agents-openai
```

Requires Python 3.10 or later.

## Quick start

```python{hl_lines=[2, 6, 11, "20-25"]}
import flyte
from flyteplugins.agents.openai import run_agent, tool

env = flyte.TaskEnvironment(
    "openai-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-openai"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[get_weather],
        instructions="You are a concise assistant. Use the tools to answer.",
        model="gpt-4.1",
    )
```

The API key is read from the environment, so it never lands in task inputs. Wire it as a Flyte secret.

## How it maps to Flyte

**Tools:** `tool` turns an `@env.task` into an `agents.FunctionTool`. The SDK derives the JSON schema, name and description from the task signature, so strict tool calling works unchanged. When the agent invokes the tool, the call dispatches to the task instead of running inline.

Applied to a plain function or a `@flyte.trace` helper, `tool` forwards to the SDK's native `function_tool`, so you can mix durable and inline tools in one agent.

**Model turns:** `FlyteModelProvider` wraps whatever `ModelProvider` the run is configured with and records each turn through `flyte.trace`. That is the seam directly below the loop, so the `Runner` above it is untouched.

**Tracing:** `install_flyte_tracing()` registers a trace processor that forwards turns, tool calls, handoffs and token usage into the Flyte report. It runs with `exclusive=True` by default, which replaces the SDK's trace processors so the run's trace telemetry is not exported to OpenAI's traces dashboard. Pass `exclusive=False` to keep the SDK's default exporter and render into the report alongside it.

This applies to the observability spans, not the inference. The model calls themselves still go to OpenAI whenever you use an OpenAI model, carrying the prompts, tool schemas, tool-call arguments and completions, because that is how the model runs. To keep prompt data off OpenAI entirely, point the run at a self-hosted or OpenAI-compatible endpoint with a custom `RunConfig(model_provider=...)`, which is a separate choice from tracing.

## Bring your own agent

If you already have an `agents.Agent` with handoffs and guardrails configured, pass it through. Durability spans the handoff: a crash mid-chain replays both agents' turns.

```python{hl_lines=[3, 8]}
from agents import Agent

triage = Agent(name="triage", handoffs=[billing, technical], input_guardrails=[...])

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    return await run_agent(request, agent=triage)
```

`agent` and `tools` are mutually exclusive. A pre-built agent carries its own tools.

## Custom run configuration

Pass a `RunConfig` to control the client, model settings or provider. The `model_provider` you set is wrapped for durability unless you pass `durable=False`.

```python{hl_lines=[1, 12]}
from agents import OpenAIProvider, RunConfig
from openai import AsyncOpenAI

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    client = AsyncOpenAI(max_retries=5, timeout=30)
    return await run_agent(
        question,
        tools=[get_weather],
        model="gpt-4.1",
        run_config=RunConfig(model_provider=OpenAIProvider(openai_client=client)),
    )
```

Client-level retries sit below the durable wrapper, so a 429 is retried in place and the turn is recorded only once it succeeds.

## Memory

```python
await run_agent(message, model="gpt-4.1", memory_key="user-alice")
```

This backs the SDK's `Session` with a durable, keyed `MemoryStore` on object storage. The SDK's default session is local SQLite, which does not survive a distributed backend. The same store also holds path-addressed facts for long-term recall.

See [How it works](./how-it-works) for the full memory model.

## Building blocks

`run_agent` wires three independently usable pieces together. Reach for them directly when you want to drive `Runner.run` yourself.

| Export | Purpose |
|---|---|
| `tool` | Turn a Flyte task into an Agents SDK tool |
| `FunctionTool` | The task-backed `FunctionTool` subclass `tool` produces |
| `FlyteModelProvider` | Set on `RunConfig.model_provider` for durable turns |
| `FlyteModel` | The per-model durable wrapper the provider hands out |
| `FlyteSession` | The `MemoryStore`-backed `Session` implementation |
| `install_flyte_tracing` | Register the Flyte trace processor |
| `FlyteTracingProcessor` | The processor itself, if you want to configure it |

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str \| list` | required | The user prompt, or a list of input items |
| `agent` | `Agent \| None` | `None` | A pre-built `agents.Agent`. Mutually exclusive with `tools` |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str` | `"gpt-4.1"` | Model name, when `agent` is not given |
| `instructions` | `str \| None` | `None` | System instructions, when `agent` is not given |
| `name` | `str` | `"flyte-agent"` | Agent name, when `agent` is not given |
| `max_turns` | `int` | `10` | Maximum model-to-tool turns before the SDK raises |
| `durable` | `bool` | `True` | Record and replay each model turn |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `run_config` | `RunConfig \| None` | `None` | Custom run configuration. Its `model_provider` is wrapped unless `durable=False` |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |

Returns the agent's final output as a string. Use `run_agent_sync` with the same signature from a sync task.

## Notes

- Streamed runs via `Runner.run_streamed` are not memoized per turn in this version. Tool calls remain durable.
- `max_turns` counts model-to-tool turns, not model calls. To bound the whole run in wall-clock terms, set `timeout=` on the enclosing task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/openai/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/openai/examples):

- `openai_durable_agent.py`: a single durable agent, with both the async and sync call forms.
- `openai_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `openai_handoffs.py`: handoffs plus a human approval gate on a sensitive refund tool.
- `openai_crash_resume.py`: the task crashes on its first attempt and finishes on retry without re-calling the model.
- `openai_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/claude-agent-sdk ===

# Claude Agent SDK

Run [Claude Agent SDK](https://github.com/anthropics/claude-agent-sdk-python) agents on Flyte. Tools you expose become durable Flyte child actions, the run streams into the task report as a timeline, and a crashed attempt resumes the conversation instead of restarting it.

This adapter differs from the others in one respect worth knowing up front: the Claude SDK runs its agent loop inside the Claude Code runtime, a subprocess Flyte does not intercept. Durability is therefore whole-session rather than per-turn. The section on **Agent frameworks > Claude Agent SDK > Durability** explains what that means in practice.

## Installation

```bash
pip install flyteplugins-agents-claude
```

Requires Python 3.10 or later.

The `claude-agent-sdk` wheel bundles the native `claude` CLI as a per-platform binary, including the `manylinux` build. It is around 250 MB, and it means the runtime image needs no separate Node.js install. A pip install and an Anthropic API key is the whole setup.

## Quick start

```python{hl_lines=[2, 6, 11, 20]}
import flyte
from flyteplugins.agents.claude import run_agent, tool

env = flyte.TaskEnvironment(
    "claude-agent",
    secrets=[flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-claude"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(question, tools=[get_weather], model="claude-sonnet-4-5")
```

## How it maps to Flyte

**Tools:** The Claude SDK exposes custom tools as in-process MCP tools. `tool` wraps an `@env.task` as an `SdkMcpTool` whose handler dispatches to the task, so a tool call becomes a durable child action with its own container, resources, retries and cache. The input schema is derived through the Flyte type engine, which handles `Literal` enums, `File`, `Dir`, `DataFrame` and dataclasses correctly.

`run_agent` builds an in-process MCP server from the tools you pass, registers it under `server_name`, and adds the tool names to `allowed_tools`.

**The loop:** `run_agent` runs the SDK's loop inside your task, streams the messages, and renders assistant turns, tool calls, cost and token usage into the report.

## Durability

**Tool calls** are durable Flyte child actions always, regardless of the `durable` setting. Their retries and caching behave exactly as they do for any Flyte task.

**The conversation** survives a crash through session resume. With `durable=True`, `run_agent` wires the SDK's own session mirror onto a `flyte.Checkpoint`. A deterministic `session_id`, derived from the task's action so it is stable across retries, is pinned on the first attempt. On a retry, the prior attempt's transcript is restored from the checkpoint and the run resumes.

The reason for delegating to the SDK here is structural. A model turn cannot be a `flyte.trace` leaf when the loop that produces it runs in a subprocess. Session resume is the coarser-grained equivalent: whole-session rather than per-turn. It no-ops cleanly when there is no checkpoint context, such as a local run.

## Observability

With `report=True`, the timeline shows assistant turns from the streamed messages plus each tool's outcome. The message stream does not surface tool results, so `run_agent` installs `PostToolUse` and `PostToolUseFailure` hooks to capture them.

If you pass your own `ClaudeAgentOptions(hooks=...)`, the Flyte hooks are merged into yours rather than replacing them. They observe only and return an empty decision, so they never affect the agent's behavior.

The result row carries the turn count, wall-clock duration, the SDK's cost estimate, and a token breakdown covering input, output, cache reads and cache writes. Those are the counts that drive the dollar figure, so you can check it at a glance.

![Timeline](../../_static/images/integrations/agents/claude.png)

## Bring your own options

Pass a fully-built `ClaudeAgentOptions` to keep SDK-native configuration such as subagents, permissions, hooks and session settings. The `tools`, `model`, `instructions` and `max_turns` arguments are layered on top of it.

```python{hl_lines=[1, 9]}
from claude_agent_sdk import ClaudeAgentOptions

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    return await run_agent(
        request,
        tools=[lookup_account],
        options=ClaudeAgentOptions(agents={"billing": {...}}),
        model="claude-sonnet-4-5",
    )
```

## Human in the loop

A tool that pauses for human approval is a durable gate the SDK has no equivalent for. The run genuinely suspends and survives restarts while it waits.

```python{hl_lines=[5, 10]}
@tool
@env.task(retries=3)
async def issue_refund(account_id: str, amount_usd: float) -> str:
    """Issue a refund. Requires human approval before it runs."""
    condition = await flyte.new_condition.aio(
        f"approve_refund_{account_id}",
        prompt=f"Approve a ${amount_usd:.2f} refund to account {account_id}?",
        data_type=bool,
    )
    if not await condition.wait.aio():
        return f"Refund to {account_id} was declined by a human reviewer."
    return f"refunded ${amount_usd:.2f} to account {account_id}"
```

## Memory

```python
await run_agent(message, model="claude-sonnet-4-5", memory_key="user-alice")
```

The transcript is persisted to a durable, keyed `MemoryStore` and resumed through the SDK's session mirror on the next run with the same key.

Memory takes precedence over the per-run `durable` checkpoint, because it covers crash resume as well. When `memory_key` is set, the checkpoint path is not used.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str \| None` | `"claude-sonnet-4-5"` | Model name |
| `instructions` | `str \| None` | `None` | System prompt |
| `max_turns` | `int \| None` | `None` | Maximum turns. `None` uses the SDK default |
| `durable` | `bool` | `True` | Wire session resume onto a `flyte.Checkpoint` |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `options` | `ClaudeAgentOptions \| None` | `None` | SDK-native configuration, layered under the arguments above |
| `server_name` | `str` | `"flyte_tools"` | Name of the in-process MCP server holding the tools |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |

Returns the final text. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/claude/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/claude/examples):

- `claude_durable_agent.py`: a single durable agent with tool outcomes in the report.
- `claude_crash_resume.py`: the task crashes on its first attempt and resumes the conversation from the checkpoint on retry.
- `claude_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `claude_hitl.py`: a refund tool gated on human approval.
- `claude_memory.py`: two separate runs sharing a `memory_key`.
- `claude_handoffs.py`: native subagent delegation, with the whole run durable on Flyte.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/google-adk ===

# Google ADK

Run [Google ADK](https://github.com/google/adk-python) (Agent Development Kit) agents on Flyte. ADK's `Runner` drives the loop and yields events. Flyte supplies the runtime: tools become durable child actions, each model turn is recorded for replay, and the event stream renders into the task report.

## Installation

```bash
pip install flyteplugins-agents-google
```

Requires Python 3.10 or later and `google-adk` 2.0 or later.

## Quick start

```python{hl_lines=[2, 6, 11, 20]}
import flyte
from flyteplugins.agents.google import run_agent, tool

env = flyte.TaskEnvironment(
    "google-agent",
    secrets=[flyte.Secret(key="google_api_key", as_env_var="GOOGLE_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-google"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(question, tools=[get_weather], model="gemini-2.0-flash")
```

Credentials come from the environment, whether that is `GOOGLE_API_KEY` for Gemini or your Vertex AI configuration. Wire them as Flyte secrets so they cannot leak into task inputs.

## How it maps to Flyte

**Tools:** ADK accepts plain Python callables and derives the tool declaration from the signature. `tool` produces one whose body dispatches to `task.aio()`, so each call runs as a durable child action.

**Model turns:** With `durable=True`, the agent's model is wrapped in `FlyteLlm`, which records each pass through `BaseLlm.generate_content_async` via `flyte.trace`. That method is the seam directly below the loop, the ADK equivalent of swapping OpenAI's `ModelProvider`. On a retry, completed turns replay from their recorded `LlmResponse` and tool calls come back from cache.

**Observability:** Turns and tool calls render into the report, followed by a usage row summarizing model turns, prompt tokens, completion tokens and total tokens. Gemini's context-cache tokens appear as `cached`, and thinking tokens as `thinking` on models that report them.

## Bring your own agent

Pass a pre-built `LlmAgent` or any `BaseAgent`, including a tree with sub-agent transfers.

```python{hl_lines=[4]}
from google.adk.agents import LlmAgent
from flyteplugins.agents.google import durable_model

triage = LlmAgent(
    name="triage",
    model=durable_model("gemini-2.0-flash"),
    instruction="Route the request to the right specialist.",
    sub_agents=[billing, technical],
)

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    return await run_agent(request, agent=triage)
```

`run_agent` cannot reach inside a pre-built tree to wrap the models, so wrap them yourself with `durable_model` when you want per-turn replay on that path. Tool calls stay durable regardless.

`agent` and `tools` are mutually exclusive.

## Memory

```python
await run_agent(message, model="gemini-2.0-flash", memory_key="user-alice")
```

ADK keeps the conversation as a list of `Event` objects on the session. Those events are persisted to a durable, keyed `MemoryStore` and restored into a fresh session on the next run with the same key.

## Bounding a run

`max_llm_calls` caps model calls before ADK raises `LlmCallsLimitExceededError`, its runaway-loop guard. It counts LLM calls rather than conversational turns, so a single tool round is roughly two calls. Leaving it at `None` uses ADK's default of 500.

For a wall-clock bound on the whole run, including tool calls, set `timeout=` on the enclosing task instead.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `agent` | `Any` | `None` | A pre-built ADK agent. Mutually exclusive with `tools` |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str` | `"gemini-2.0-flash"` | Model name, when `agent` is not given |
| `instructions` | `str \| None` | `None` | System instruction, when `agent` is not given |
| `name` | `str` | `"assistant"` | Agent name. Must be a valid Python identifier |
| `max_llm_calls` | `int \| None` | `None` | Cap on model calls. `None` uses ADK's default of 500 |
| `durable` | `bool` | `True` | Record and replay each model turn |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `app_name` | `str` | `"flyte-agent"` | ADK app name, used for namespacing |
| `user_id` | `str` | `"flyte-user"` | ADK user ID |

Returns the final text. Use `run_agent_sync` with the same signature from a sync task.

> [!NOTE] `name` is visible to the model
> ADK injects the agent name into the system prompt as the model's internal name, so it can surface in replies. Keep it natural. An internal or brand-heavy label will show up in the conversation.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/google/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/google/examples):

- `google_durable_agent.py`: a single durable agent with traced model turns.
- `google_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `google_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `google_memory.py`: two separate runs sharing a `memory_key`.
- `google_handoffs.py`: native agent transfer to a specialist sub-agent, which can pause on a Flyte condition for a human to supply details mid-conversation.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/mistral ===

# Mistral Agents

Run [Mistral Agents](https://docs.mistral.ai/agents/agents_introduction/) on Flyte, built on the Conversations API in `mistralai` 2.x. Mistral's own runner drives the loop and executes tools. Flyte registers task-backed tools with it, records each conversation turn for replay, and renders the run into the task report.

Mistral is server-side, which changes two things relative to the other adapters: the transcript lives on Mistral's side, and you can drive a pre-created agent by ID instead of an inline model.

## Installation

```bash
pip install flyteplugins-agents-mistral
```

Requires Python 3.10 or later and `mistralai[agents]` 2.0 or later.

## Quick start

```python{hl_lines=[2, 6, 11, 20]}
import flyte
from flyteplugins.agents.mistral import run_agent, tool

env = flyte.TaskEnvironment(
    "mistral-agent",
    secrets=[flyte.Secret(key="mistral_api_key", as_env_var="MISTRAL_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-mistral"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(question, tools=[get_weather], model="mistral-large-latest")
```

`run_agent` raises with a clear message if the API key is missing, naming the environment variable it looked in.

## How it maps to Flyte

**Tools:** Mistral's `RunContext` takes plain Python functions and registers them with the runner. `tool` produces one whose body dispatches to `task.aio()`, so each call runs as a durable child action.

**Model turns:** The runner makes each turn by calling `conversations.start_async` or `conversations.append_async`, both in-process HTTP calls. With `durable=True` those two methods are wrapped, and each turn is recorded via `flyte.trace`. The `ConversationResponse` round-trips through pydantic JSON, including the polymorphic output entries.

That is the seam directly below the loop. The SDK still owns the loop; on a retry, completed turns replay from their records and completed tool calls come back from cache.

**Observability:** The turns, tool calls and final answer render into the report. The SDK's `RunResult` exposes no token usage, but each turn's `ConversationResponse` does, so the same wrapper tallies it and adds a usage row.

Replayed turns are counted as cached rather than fresh spend, so a retried run does not present a free replay as though it cost money.

## Driving a pre-created agent

Mistral agents can be created server-side and referenced by ID. Pass `agent_id` instead of `model` and the tool calls still run as durable Flyte actions.

```python{hl_lines=[6]}
@env.task(report=True, retries=3)
async def support(request: str) -> str:
    return await run_agent(
        request,
        tools=[lookup_account],
        agent_id="ag_01jd...",
    )
```

Native handoffs work on this path too. A triage agent can hand the conversation to a billing or technical agent by ID, with the whole multi-agent run durable on Flyte.

## Memory

```python
await run_agent(message, model="mistral-large-latest", memory_key="user-alice")
```

Mistral keeps the transcript server-side, so there is nothing to copy. Flyte persists the thread's `conversation_id` in a keyed `MemoryStore` and continues that conversation when the key recurs.

## Bounding a run

`timeout_ms` is a per-turn request timeout that the SDK applies to each model call inside its loop. It bounds a single hung turn. It is not a whole-run cap, and Mistral exposes no turn-count limit.

To bound the entire agent run, including every turn and tool call, set `timeout=` on the enclosing task.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str \| None` | `"mistral-large-latest"` | Model for an inline run, when `agent_id` is not given |
| `instructions` | `str \| None` | `None` | System instructions |
| `timeout_ms` | `int \| None` | `None` | Per-turn request timeout in milliseconds. `None` uses the SDK default |
| `durable` | `bool` | `True` | Record and replay each conversation turn |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `agent_id` | `str \| None` | `None` | Drive an existing server-side agent instead of `model` |
| `api_key_env_var` | `str` | `"MISTRAL_API_KEY"` | Environment variable holding the API key |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |

Returns the final text. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/mistral/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/mistral/examples):

- `mistral_durable_agent.py`: a single durable agent with per-turn tracing.
- `mistral_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `mistral_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `mistral_agent_id.py`: driving a pre-created server-side agent by ID.
- `mistral_memory.py`: two separate runs sharing a `memory_key`.
- `mistral_handoffs.py`: native handoffs to a specialist agent, which can pause on a Flyte condition for a human detail.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/langchain ===

# LangChain

Run [LangChain agents](https://docs.langchain.com/oss/python/langchain/agents) on Flyte. In LangChain 1.x an agent is built with `create_agent(model, tools, system_prompt=...)`, which returns a compiled graph. `run_agent` drives that graph inside your task, with tools running as durable child actions and model turns recorded for replay.

## Installation

```bash
pip install flyteplugins-agents-langchain
```

Requires Python 3.10 or later. Install the provider integration you need alongside it, for example `langchain-openai` or `langchain-anthropic`.

## Quick start

```python{hl_lines=[2, 6, 13, "24-29"]}
import flyte
from flyteplugins.agents.langchain import run_agent, tool

env = flyte.TaskEnvironment(
    "langchain-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-agents-langchain", "langchain-openai",
    ),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    from langchain_openai import ChatOpenAI

    return await run_agent(
        question,
        tools=[get_weather],
        model=ChatOpenAI(model="gpt-4o"),
        instructions="You are a concise assistant. Use the tools to answer.",
    )
```

Build the chat model inside the task, where the provider API key is available.

## How it maps to Flyte

**Tools:** `tool` turns an `@env.task` into a LangChain `StructuredTool`, a real `BaseTool` that drops straight into `create_agent(model, tools=[...])`. The args schema is built as a pydantic model from the task's typed signature, with annotations and defaults preserved, rather than being inferred from the wrapper. When the agent calls the tool, the coroutine dispatches to `task.aio()` and the call becomes a durable child action.

**Model turns:** With `durable=True`, the chat model is wrapped in `DurableChatModel`, which records each turn via `flyte.trace`. On a retry, completed turns replay from their records and tool calls come back from cache.

**Observability:** The run timeline renders into the task report.

## Pass a model instance, not a string

Durability is applied by wrapping a `BaseChatModel` instance. `create_agent` also accepts a `provider:model` string, and that will run, but a string is passed straight through unwrapped, so you lose per-turn replay.

```python
model=ChatOpenAI(model="gpt-4o")   # wrapped, turns are durable
model="openai:gpt-4o"              # runs, but turns are not recorded
```

Tool calls stay durable either way. If you want the string form with durability, use the [LangGraph](./langgraph) or [Deep Agents](./deepagents) adapter, both of which resolve the string before wrapping.

## Bring your own agent

Pass a compiled `create_agent` graph as `agent=`.

```python{hl_lines=["9-14"]}
from langchain.agents import create_agent
from flyteplugins.agents.langchain import DurableChatModel

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    from langchain_openai import ChatOpenAI

    graph = create_agent(
        DurableChatModel(inner=ChatOpenAI(model="gpt-4o")),
        [lookup_account],
        system_prompt="You are a billing support agent.",
    )
    return await run_agent(request, agent=graph)
```

A fully compiled graph owns its own model and cannot be rewrapped from outside, so wrap the model yourself with `DurableChatModel` when building it. Tool calls remain durable regardless.

`agent` and `tools` are mutually exclusive.

## Memory

```python
await run_agent(message, model=ChatOpenAI(model="gpt-4o"), memory_key="user-alice")
```

The conversation transcript is persisted to a durable, keyed `MemoryStore`. On the next run with the same key, prior messages are loaded and prepended to the new user turn, and the full transcript is saved back afterwards.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `Any` | `None` | A LangChain chat model. Required when `agent` is not given |
| `instructions` | `str \| None` | `None` | System prompt for the built agent |
| `agent` | `Any` | `None` | A pre-built compiled `create_agent` graph. Mutually exclusive with `tools` |
| `name` | `str` | `"langchain-agent"` | Agent name, used for debugging and observability |
| `durable` | `bool` | `True` | Record and replay each model turn. Applies on the builder path |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `**agent_kwargs` | | | Forwarded to `create_agent` |

Returns the final text, taken from the content of the last message. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/langchain/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/langchain/examples):

- `langchain_durable_agent.py`: a single durable agent with traced model turns.
- `langchain_custom_agent.py`: building the agent yourself and passing it as `agent=`.
- `langchain_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `langchain_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `langchain_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/langgraph ===

# LangGraph

Run [LangGraph](https://langchain-ai.github.io/langgraph/) graphs on Flyte. This adapter is shaped differently from the others: LangGraph is a framework for building your own control flow, so rather than hiding the graph behind `run_agent`, it gives you durable node factories and expects you to wire them up yourself.

You build the `StateGraph`. `ai_node` and `tool_node` are the pieces Flyte makes durable and observable.

## Installation

```bash
pip install flyteplugins-agents-langgraph
```

Requires Python 3.10 or later. Install the provider integration you need alongside it, for example `langchain-openai`.

## Quick start

Build the graph with the two node factories, compile it, and hand the compiled graph to `run_agent`.

```python{hl_lines=[2, 6, 13, "38-41"]}
import flyte
from flyteplugins.agents.langgraph import ai_node, run_agent, tool, tool_node

env = flyte.TaskEnvironment(
    "langgraph-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-agents-langgraph", "langchain-openai",
    ),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

def build_city_graph():
    """A standard tool-calling loop: ai, then tools, then back to ai."""
    from langchain_openai import ChatOpenAI
    from langgraph.graph import START, MessagesState, StateGraph
    from langgraph.prebuilt import tools_condition

    tools = [get_weather]
    builder = StateGraph(MessagesState)
    builder.add_node("ai", ai_node(ChatOpenAI(model="gpt-4o"), tools))
    builder.add_node("tools", tool_node(tools))
    builder.add_edge(START, "ai")
    builder.add_conditional_edges("ai", tools_condition)
    builder.add_edge("tools", "ai")
    return builder.compile()

@env.task(report=True, retries=3)
async def city_agent(city: str) -> str:
    return await run_agent(
        f"What's the weather in {city}?",
        agent=build_city_graph(),
    )
```

Build the graph inside the task, where the provider API key is available.

## The node factories

**`ai_node(model, tools, *, name="ai", durable=True, observability=True)`**

The model-calling node. It binds the tools to your chat model and runs one turn over `state["messages"]`, appending the response. With `durable=True`, each turn is recorded as a `flyte.trace` leaf keyed by a fingerprint of the message list, so a retry replays the recorded response instead of calling the model again.

Returns an async node with the signature `state -> {"messages": [ai_message]}`.

**`tool_node(tools, *, name="tools", observability=True)`**

The tool-executing node. It reads the tool calls off the last message and runs each one, appending a `ToolMessage` per call. Tools wrapped with `tool` run as durable Flyte child actions; anything else runs as the tool defines. Tool errors are caught and surfaced back to the model as the tool result rather than failing the node.

Returns an async node with the signature `state -> {"messages": [tool_message, ...]}`.

Both render their activity into the task report.

## How it maps to Flyte

**Tools:** `tool` turns an `@env.task` into a LangChain `StructuredTool`. It is a first-class LangGraph tool, so it works with `model.bind_tools(...)`, with `tool_node`, and with LangGraph's own `ToolNode`. The args schema comes from the task's typed signature.

**Model turns:** Durability lives in `ai_node`, not in a model wrapper. That means any chat model works, including a `provider:model` string, and durability applies uniformly.

## Skipping the graph

If you do not need a custom topology, pass `tools` and a model instead of `agent`, and `run_agent` assembles the same tool-calling loop from the same two factories.

```python{hl_lines=[7, 8]}
@env.task(report=True, retries=3)
async def quick_city_agent(city: str) -> str:
    from langchain_openai import ChatOpenAI

    return await run_agent(
        f"What's the weather in {city}?",
        tools=[get_weather],
        model=ChatOpenAI(model="gpt-4o"),
        instructions="You are a concise assistant. Use the tools to answer.",
    )
```

`model` accepts a chat model instance or a `provider:model` string. The string form is resolved through `init_chat_model`, which requires the `langchain` package.

`agent` and `tools` are mutually exclusive.

## Custom state

`input` accepts a full graph input state as a dict, not just a prompt string, so a graph with a state schema beyond `MessagesState` works.

```python
await run_agent({"messages": [...], "budget": 3}, agent=graph)
```

When memory is in play, prior messages are merged into the `messages` key of whatever state you pass.

## Memory

```python
await run_agent(question, agent=graph, memory_key="user-alice")
```

The conversation transcript is persisted to a durable, keyed `MemoryStore` and prepended to the graph's messages on the next run with the same key. On a resumed run the system prompt is not re-added, since it already lives in the prior transcript.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str \| dict` | required | The user prompt, or a full graph input state |
| `tools` | `Sequence` | `()` | Tools for the default graph. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `Any` | `None` | A chat model instance or `provider:model` string. Required when building the graph |
| `instructions` | `str \| None` | `None` | System prompt prepended to a built graph's messages |
| `agent` | `Any` | `None` | A pre-built compiled graph. Mutually exclusive with `tools` |
| `name` | `str` | `"langgraph-agent"` | Graph name, used for debugging and observability |
| `durable` | `bool` | `True` | Record each model turn. Applies to built graphs |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `**run_kwargs` | | | Forwarded to the graph's `ainvoke` |

Returns the final assistant message as a string. Use `run_agent_sync` with the same signature from a sync task.

> [!NOTE] `durable` applies to graphs `run_agent` builds
> When you build the graph yourself, durability is whatever you configured on `ai_node`. Passing `durable=False` alongside `agent=` does not turn off a node you already built with `durable=True`.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/langgraph/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/langgraph/examples):

- `langgraph_custom_agent.py`: building a `StateGraph` from `ai_node` and `tool_node`, plus the default-graph shortcut.
- `langgraph_durable_agent.py`: a single durable agent with traced model turns.
- `langgraph_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `langgraph_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `langgraph_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/deepagents ===

# Deep Agents

Run [Deep Agents](https://docs.langchain.com/oss/python/deepagents/overview) on Flyte. Deep Agents is LangChain's agent harness, with built-in planning through todos, a virtual filesystem, and subagents. `create_deep_agent` returns a compiled LangGraph graph; `run_agent` drives it inside your task.

The virtual filesystem is the part worth calling out. It is agent state that outlives a single turn, and `memory_key` persists it alongside the conversation, so a later run picks up both the transcript and whatever files the agent wrote.

## Installation

```bash
pip install flyteplugins-agents-deepagents
```

Requires Python 3.11 or later.

## Quick start

```python{hl_lines=[2, 6, 11, "20-30"]}
import flyte
from flyteplugins.agents.deepagents import run_agent, tool

env = flyte.TaskEnvironment(
    "deep-agent",
    secrets=[flyte.Secret(key="anthropic_api_key", as_env_var="ANTHROPIC_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-deepagents"),
)

@tool
@env.task(cache="auto", retries=3)
async def search_web(query: str) -> str:
    """Search the web for a query."""
    ...

@env.task(report=True, retries=3)
async def research_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[search_web],
        instructions="You are an expert researcher.",
        model="anthropic:claude-sonnet-4-6",
        subagents=[{
            "name": "critic",
            "description": "Critiques draft answers.",
            "system_prompt": "You are a ruthless critic.",
        }],
    )
```

Deep-Agents-specific options such as `subagents`, `skills`, `backend` and `interrupt_on` pass straight through as keyword arguments.

## How it maps to Flyte

**Tools:** `tool` turns an `@env.task` into a LangChain `StructuredTool`. It attaches to the main agent through `create_deep_agent(tools=[...])` and equally to a subagent's tool list, so a subagent's tool calls are durable child actions too.

**Model turns:** `model` accepts a chat model instance or a `provider:model` string. A string is resolved through `init_chat_model` first, then wrapped in `DurableChatModel`, so both forms get per-turn replay.

**Observability:** The run timeline renders into the task report.

## Bring your own agent

Pass a compiled `create_deep_agent` graph as `agent=`. Wrap the model in `DurableChatModel` when you build it, since a compiled graph cannot be rewrapped from outside.

```python{hl_lines=[1, "9-14"]}
from deepagents import create_deep_agent
from flyteplugins.agents.deepagents import DurableChatModel

@env.task(report=True, retries=3)
async def research_agent(question: str) -> str:
    from langchain_anthropic import ChatAnthropic

    graph = create_deep_agent(
        model=DurableChatModel(inner=ChatAnthropic(model="claude-sonnet-4-6")),
        tools=[search_web],
        system_prompt="You are an expert researcher.",
    )
    return await run_agent(question, agent=graph)
```

Tool calls remain durable regardless.

`agent` and `tools` are mutually exclusive.

## Memory

```python
await run_agent(message, model="anthropic:claude-sonnet-4-6", memory_key="user-alice")
```

Both the conversation and the agent's virtual filesystem are persisted to a durable, keyed `MemoryStore`. On the next run with the same key, prior messages are prepended to the new turn and the `files` state is restored, so an agent that wrote notes in one run can read them back in the next.

## Composing with Flyte

Deep agents plan internally and can spawn their own subagents. That composes with Flyte's orchestration rather than competing with it: Flyte fans out the team, and each member is a full deep agent with its own internal planning.

```python{hl_lines=[13, 14]}
@env.task(retries=3)
async def research(subtopic: str) -> str:
    return await run_agent(
        f"Research this subtopic:\n{subtopic}",
        tools=[search_web],
        model="anthropic:claude-sonnet-4-6",
    )

@env.task(report=True, retries=3)
async def pipeline(topic: str) -> str:
    subtopics = await plan(topic)
    with flyte.group("parallel-research"):
        findings = await asyncio.gather(*(research(s) for s in subtopics))
    return await synthesize(topic, list(findings))
```

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `Any` | `None` | A chat model instance or `provider:model` string. Required when `agent` is not given |
| `instructions` | `str \| None` | `None` | System prompt for the built agent |
| `agent` | `Any` | `None` | A pre-built compiled `create_deep_agent` graph. Mutually exclusive with `tools` |
| `name` | `str` | `"deep-agent"` | Agent name, used for debugging and observability |
| `durable` | `bool` | `True` | Record and replay each model turn. Applies on the builder path |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for the conversation and the virtual filesystem |
| `**agent_kwargs` | | | Forwarded to `create_deep_agent`, including `subagents`, `skills` and `backend` |

Returns the final text, taken from the content of the last message. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/deepagents/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/deepagents/examples):

- `deepagents_durable_agent.py`: a single durable deep agent with traced model turns.
- `deepagents_custom_agent.py`: building the graph yourself with `create_deep_agent`.
- `deepagents_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `deepagents_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `deepagents_memory.py`: two separate runs sharing a `memory_key`, carrying the virtual filesystem across.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/crewai ===

# CrewAI

Run [CrewAI](https://docs.crewai.com/) agents on Flyte. CrewAI drives the loop through `Agent.kickoff_async`. Flyte supplies the runtime: tools become durable child actions, model turns are recorded for replay, and the run renders into the task report.

## Installation

```bash
pip install flyteplugins-agents-crewai
```

Requires Python 3.10 or later.

## Quick start

```python{hl_lines=[2, 6, 11, "20-25"]}
import flyte
from flyteplugins.agents.crewai import run_agent, tool

env = flyte.TaskEnvironment(
    "crewai-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-crewai"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[get_weather],
        model="gpt-4o",
        instructions="You are a concise assistant. Use the tools to answer.",
    )
```

`model` is required on the builder path. The adapter is provider agnostic and assumes no default.

## How it maps to Flyte

**Tools:** CrewAI requires tools attached to `Agent(tools=[...])` to be `crewai.tools.BaseTool` instances; plain callables are rejected by pydantic validation. `tool` therefore produces a real `BaseTool` subclass whose execution dispatches to `task.aio()`. The input schema comes from the Flyte type engine.

CrewAI invokes tools synchronously, which is awkward inside an already-running event loop. The adapter handles this by making the synchronous `_run` path bridge to the task through a dedicated background-thread loop, and by awaiting the task directly on CrewAI's native async path. You do not have to think about it, but it explains why the tool object is a class rather than a function.

**Model turns:** On the builder path the agent is driven by a durable `LLM`, so each turn is recorded via `flyte.trace` and replayed on retry.

**Observability:** The run timeline renders into the task report.

## Bring your own agent

Pass a pre-built CrewAI `Agent` with its tools already attached.

```python{hl_lines=["6-12"]}
from crewai import Agent

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    agent = Agent(
        role="Billing specialist",
        goal="Resolve billing questions accurately.",
        backstory="You have handled billing escalations for years.",
        tools=[lookup_account],
        llm="gpt-4o",
    )
    return await run_agent(request, agent=agent)
```

> [!WARNING] Pre-built agents keep their own model
> Model-turn durability is applied only when `run_agent` builds the agent, because the builder is what sets the durable `llm`. A pre-built agent keeps whatever `llm` you gave it and is not rewrapped, so its turns are not recorded. Tool calls remain durable either way.

`agent` and `tools` are mutually exclusive. A pre-built agent carries its own tools.

## Instructions and the built agent

On the builder path, `run_agent` constructs an agent with the role `Assistant` and a goal of answering accurately and concisely. `instructions` is folded into the backstory rather than replacing the whole persona. If you need full control over role, goal and backstory, build the agent yourself and pass it as `agent=`.

## Memory

```python
await run_agent(message, model="gpt-4o", memory_key="user-alice")
```

The conversation transcript is persisted to a durable, keyed `MemoryStore`. On the next run with the same key, the prior transcript is loaded and passed to `kickoff_async` as a message list so the agent continues the conversation.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str \| None` | `None` | Model name, for example `"gpt-4o"`. Required when `agent` is not given |
| `instructions` | `str \| None` | `None` | Extra guidance folded into the built agent's backstory |
| `agent` | `Any` | `None` | A pre-built CrewAI `Agent`. Mutually exclusive with `tools` |
| `name` | `str` | `"crewai-agent"` | Agent name, used for debugging and observability |
| `durable` | `bool` | `True` | Record and replay each model turn. Builder path only |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `**run_kwargs` | | | Forwarded to `Agent.kickoff_async` |

Returns the final text, taken from the result's `raw` field. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/crewai/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/crewai/examples):

- `crewai_durable_agent.py`: a single durable agent with traced model turns.
- `crewai_custom_agent.py`: building the `Agent` yourself and passing it as `agent=`.
- `crewai_sync_agent.py`: driving the same agent from a sync task with `run_agent_sync`.
- `crewai_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `crewai_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `crewai_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/pydantic-ai ===

# Pydantic AI

Run [Pydantic AI](https://ai.pydantic.dev/) agents on Flyte. Pydantic AI owns the loop through `Agent.run`. Flyte supplies the runtime: tools become durable child actions, model turns are recorded for replay, and the run renders into the task report.

This is the one adapter that applies model-turn durability on both the builder path and the pre-built path, because `Agent.override` gives it a clean way in.

## Installation

```bash
pip install flyteplugins-agents-pydantic-ai
```

Requires Python 3.10 or later and `pydantic-ai` 2.x.

## Quick start

```python{hl_lines=[2, 6, 11, "20-25"]}
import flyte
from flyteplugins.agents.pydantic_ai import run_agent, tool

env = flyte.TaskEnvironment(
    "pydantic-ai-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-pydantic-ai"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[get_weather],
        model="openai:gpt-4o",
        instructions="You are a concise assistant. Use the tools to answer.",
    )
```

Note the import path uses an underscore, `flyteplugins.agents.pydantic_ai`, while the package on PyPI is `flyteplugins-agents-pydantic-ai`.

`model` is required on the builder path. The adapter is provider agnostic and assumes no default.

## How it maps to Flyte

**Tools:** Pydantic AI accepts plain async callables in `Agent(tools=[...])` and infers each tool's schema from the signature. `tool` here is the shared core wrapper, which preserves the signature through `functools.wraps` and dispatches to `task.aio()`, so schema inference works unchanged and every call is a durable child action.

**Model turns:** On the builder path the model is resolved through `infer_model` and wrapped in `FlyteModel`. On the pre-built path the wrapper is applied through `Agent.override(model=...)`, scoped to the run.

Both paths are best-effort. If the model cannot be inferred or the agent exposes no accessible `Model`, a warning is logged and the run proceeds without per-turn durability rather than failing. Tool calls stay durable regardless.

**Observability:** The run timeline renders into the task report.

## Bring your own agent

Tools are attached at construction in Pydantic AI. `Agent.run` takes no `tools` argument, so a pre-built agent carries its own.

```python{hl_lines=[1, "6-11"]}
from pydantic_ai import Agent

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    agent = Agent(
        "openai:gpt-4o",
        system_prompt="You are a billing support agent.",
        tools=[lookup_account],
    )
    return await run_agent(request, agent=agent)
```

Durability is applied through `override` on this path, so you do not have to wrap the model yourself.

`agent` and `tools` are mutually exclusive.

## Memory

```python
await run_agent(message, model="openai:gpt-4o", memory_key="user-alice")
```

Prior conversation history is loaded from a durable, keyed `MemoryStore` and passed as `message_history=`. After the run, the full history is saved back, so a later run with the same key continues the conversation.

An explicit `message_history=` in `**run_kwargs` takes precedence over loaded memory.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `Any` | `None` | Model name such as `"openai:gpt-4o"`, or a `Model` instance. Required when `agent` is not given |
| `instructions` | `str \| None` | `None` | System prompt for the built agent |
| `agent` | `Any` | `None` | A pre-built Pydantic AI `Agent` with tools attached. Mutually exclusive with `tools` |
| `name` | `str` | `"pydantic-ai-agent"` | Agent name, used for debugging and observability |
| `durable` | `bool` | `True` | Record and replay each model turn. Applies on both paths |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `**run_kwargs` | | | Forwarded to `agent.run`, including an explicit `message_history=` |

Returns the final output as a string, taken from `result.output`. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/pydantic_ai/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/pydantic_ai/examples):

- `pydantic_ai_durable_agent.py`: a single durable agent with traced model turns.
- `pydantic_ai_custom_agent.py`: building the `Agent` yourself and passing it as `agent=`.
- `pydantic_ai_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `pydantic_ai_crash_resume.py`: the task crashes on its first attempt and replays completed turns on retry.
- `pydantic_ai_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/agents/hermes ===

# Hermes

Run [Hermes](https://pypi.org/project/hermes-agent/) agents on Flyte. Hermes, from Nous Research, drives the loop through `AIAgent.run_conversation`. Flyte supplies the runtime: tools become durable child actions, the run renders into the task report, and `memory_key` carries the conversation across runs.

Hermes is the one adapter without model-turn replay. The package exposes no per-turn hook, so `durable=` is accepted for contract consistency and does nothing. Tool calls are durable regardless, so a retried task still self-heals at tool granularity.

## Installation

```bash
pip install flyteplugins-agents-hermes
```

Requires Python 3.11 or later.

## Quick start

```python{hl_lines=[2, 6, 11, "20-25"]}
import flyte
from flyteplugins.agents.hermes import run_agent, tool

env = flyte.TaskEnvironment(
    "hermes-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-agents-hermes"),
)

@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."

@env.task(report=True, retries=3)
async def city_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[get_weather],
        model="gpt-4o",
        instructions="You are a concise assistant. Use the tools to answer.",
    )
```

`model` is required on the builder path. There is no default.

## Credentials

Hermes normally reads credentials from its own `hermes setup` configuration, which a fresh container does not have. To make the common case work, `run_agent` fills in the gap: when none of `api_key`, `base_url` or `provider` are passed and `OPENAI_API_KEY` is set in the environment, the built agent is pointed at OpenAI with that key.

For any other provider, pass the credentials explicitly. They go through `**agent_kwargs` to the `AIAgent` constructor.

```python{hl_lines=[5, 6]}
await run_agent(
    question,
    tools=[get_weather],
    model="Hermes-4-405B",
    api_key=os.environ["NOUS_API_KEY"],
    base_url="https://inference-api.nousresearch.com/v1",
)
```

## How it maps to Flyte

**Tools:** Hermes does not accept tool callables on the agent object. Tools live in a process-global registry keyed by name and grouped into toolsets, and an `AIAgent` exposes whatever its `enabled_toolsets` resolve to.

`tool` therefore does two things: it wraps the `@env.task` so a call dispatches to `task.aio()` as a durable child action, and it registers that wrapper in the Hermes registry under the `FLYTE_TOOLSET` toolset, with an OpenAI-format schema derived through the Flyte type engine.

**Toolset scoping:** Every `tool` registers under the same shared toolset. To keep two agents in one process from seeing each other's tools, `run_agent` creates a scoped toolset per built agent, named from the agent's `name`, holding exactly the tools you passed.

**The loop:** `run_conversation` is synchronous. The adapter runs it off the event loop through `asyncio.to_thread`, which propagates the Flyte task context into the worker thread.

## Bring your own agent

Pass a pre-configured `AIAgent`. It needs `FLYTE_TOOLSET` in its `enabled_toolsets` to see Flyte-backed tools.

```python{hl_lines=[1, 2, 7, 9]}
from run_agent import AIAgent
from flyteplugins.agents.hermes import FLYTE_TOOLSET

@env.task(report=True, retries=3)
async def support(request: str) -> str:
    agent = AIAgent(
        model="gpt-4o",
        enabled_toolsets=[FLYTE_TOOLSET],
        quiet_mode=True,
    )
    return await run_agent(request, agent=agent, instructions="Be concise.")
```

On this path, `instructions` is passed as the run's `system_message` rather than replacing the agent's own prompt, and `**agent_kwargs` is rejected, since those configure a built agent.

`agent` and `tools` are mutually exclusive.

> [!NOTE] The `AIAgent` import path
> `hermes-agent` exposes `AIAgent` from a top-level module named `run_agent`, which is easy to confuse with this adapter's `run_agent` function. The `from run_agent import AIAgent` form above does not bind the name `run_agent`, so the two coexist, but a bare `import run_agent` would shadow the function.

## Memory

```python
await run_agent(message, model="gpt-4o", memory_key="user-alice")
```

The transcript is persisted to a durable, keyed `MemoryStore` and passed back to Hermes as `conversation_history` on the next run with the same key.

## `run_agent` parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | `str` | required | The user prompt |
| `tools` | `Sequence` | `()` | Tools to expose. Accepts `tool`-wrapped tools or bare `@env.task` templates |
| `model` | `str \| None` | `None` | Model name. Required when `agent` is not given |
| `instructions` | `str \| None` | `None` | System prompt. Becomes `ephemeral_system_prompt` on the builder path, or the run's `system_message` with a pre-built agent |
| `agent` | `Any` | `None` | A pre-built Hermes `AIAgent`. Mutually exclusive with `tools` |
| `name` | `str` | `"hermes-agent"` | Agent name. Also names the scoped toolset |
| `durable` | `bool` | `True` | Accepted for contract consistency. No effect on Hermes |
| `observability` | `bool` | `True` | Render the timeline into the task report |
| `memory_key` | `str \| None` | `None` | Stable user or thread ID for cross-run memory |
| `**agent_kwargs` | | | Forwarded to the built `AIAgent`, including `api_key`, `base_url`, `provider` and `max_iterations`. Builder path only |

Returns the final text from the result's `final_response` field. Use `run_agent_sync` with the same signature from a sync task.

## Examples

Full runnable examples live in the SDK repository under [`plugins/agents/hermes/examples`](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agents/hermes/examples):

- `hermes_durable_agent.py`: a single agent with durable tool calls.
- `hermes_custom_agent.py`: building the `AIAgent` yourself and passing it as `agent=`.
- `hermes_multi_agent.py`: a planner, parallel researchers and an editor, each its own durable action.
- `hermes_crash_resume.py`: the task crashes on its first attempt and completed tool calls are cache hits on retry.
- `hermes_memory.py`: two separate runs sharing a `memory_key`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/bigquery ===

# BigQuery

The BigQuery connector lets you run SQL queries against [Google BigQuery](https://cloud.google.com/bigquery) directly from Flyte tasks. Queries are submitted asynchronously via the BigQuery Jobs API and polled for completion, so they don't block a worker while waiting for results.

The connector supports:

- Parameterized SQL queries with typed inputs
- Google Cloud service account authentication
- Returns query results as DataFrames
- Query cancellation on task abort

## Installation

```bash
pip install flyteplugins-bigquery
```

This installs the Google Cloud BigQuery client libraries.

## Quick start

Here's a minimal example that runs a SQL query on BigQuery:

```python
from flyte.io import DataFrame
from flyteplugins.bigquery import BigQueryConfig, BigQueryTask

config = BigQueryConfig(
    ProjectID="my-gcp-project",
    Location="US",
)

count_users = BigQueryTask(
    name="count_users",
    query_template="SELECT COUNT(*) FROM dataset.users",
    plugin_config=config,
    output_dataframe_type=DataFrame,
)
```

This defines a task called `count_users` that runs the query on the configured BigQuery instance. When executed, the connector:

1. Connects to BigQuery using the provided configuration
2. Submits the query asynchronously via the Jobs API
3. Polls until the query completes or fails

To run the task, create a `TaskEnvironment` from it and execute it locally or remotely:

```python
import flyte

bigquery_env = flyte.TaskEnvironment.from_task("bigquery_env", count_users)

if __name__ == "__main__":
    flyte.init_from_config()

    # Run locally (connector runs in-process, requires credentials locally)
    run = flyte.with_runcontext(mode="local").run(count_users)

    # Run remotely (connector runs as a service in your data plane)
    run = flyte.with_runcontext(mode="remote").run(count_users)

    print(run.url)
```

> [!NOTE]
> The `TaskEnvironment` created by `from_task` does not need an image or pip packages. BigQuery tasks are connector tasks, which means the query executes on the connector service, not in your task container. In `local` mode, the connector runs in-process and requires `flyteplugins-bigquery` and credentials to be available on your machine.

## Configuration

### `BigQueryConfig` parameters

| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `ProjectID` | `str` | Yes | GCP project ID |
| `Location` | `str` | No | BigQuery region (e.g., `"US"`, `"EU"`) |
| `QueryJobConfig` | `bigquery.QueryJobConfig` | No | Native BigQuery [QueryJobConfig](https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJobConfig) object for advanced settings |

### `BigQueryTask` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `name` | `str` | Unique task name |
| `query_template` | `str` | SQL query (whitespace is normalized before execution) |
| `plugin_config` | `BigQueryConfig` | Connection configuration |
| `inputs` | `Dict[str, Type]` | Named typed inputs bound as query parameters |
| `output_dataframe_type` | `Type[DataFrame]` | If set, query results are returned as a `DataFrame` |
| `google_application_credentials` | `str` | Name of the Flyte secret containing the GCP service account JSON key |

## Authentication

Pass the name of a Flyte secret containing your GCP service account JSON key:

```python
query = BigQueryTask(
    name="secure_query",
    query_template="SELECT * FROM dataset.sensitive_data",
    plugin_config=config,
    google_application_credentials="my-gcp-sa-key",
)
```

## Query templating

Use the `inputs` parameter to define typed inputs for your query. Input values are bound as BigQuery `ScalarQueryParameter` values.

### Supported input types

| Python type | BigQuery type |
|-------------|---------------|
| `int` | `INT64` |
| `float` | `FLOAT64` |
| `str` | `STRING` |
| `bool` | `BOOL` |
| `bytes` | `BYTES` |
| `datetime` | `DATETIME` |
| `list` | `ARRAY` |

### Parameterized query example

```python
from flyte.io import DataFrame

events_by_region = BigQueryTask(
    name="events_by_region",
    query_template="SELECT * FROM dataset.events WHERE region = @region AND score > @min_score",
    plugin_config=config,
    inputs={"region": str, "min_score": float},
    output_dataframe_type=DataFrame,
)
```

> [!NOTE]
> The query template is normalized before execution: newlines and tabs are replaced with spaces and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution.

## Retrieving query results

Set `output_dataframe_type` to capture results as a DataFrame:

```python
from flyte.io import DataFrame

top_customers = BigQueryTask(
    name="top_customers",
    query_template="""
        SELECT customer_id, SUM(amount) AS total_spend
        FROM dataset.orders
        GROUP BY customer_id
        ORDER BY total_spend DESC
        LIMIT 100
    """,
    plugin_config=config,
    output_dataframe_type=DataFrame,
)
```

If you don't need query results (for example, DDL statements or INSERT queries), omit `output_dataframe_type`.

## API reference

See the [BigQuery API reference](../../api-reference/integrations/bigquery/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/codegen ===

# Code generation

The code generation plugin turns natural-language prompts into tested, production-ready Python code.

You describe what the code should do, along with sample data, schema definitions, constraints, and typed inputs/outputs, and the plugin handles the rest: generating code, writing tests, building an isolated [code sandbox](https://www.union.ai/docs/v2/union/user-guide/sandboxing/code-sandboxing) with the right dependencies, running the tests, diagnosing failures, and iterating until everything passes. The result is a validated script you can execute against real data or deploy as a reusable Flyte task.

## Installation

```bash
pip install flyteplugins-codegen

# For Agent mode (Claude-only)
pip install flyteplugins-codegen[agent]
```

## Quick start

```python{hl_lines=[3, 4, 6, 12, 14, "20-25"]}
import flyte
from flyte.io import File
from flyte.sandbox import sandbox_environment
from flyteplugins.codegen import AutoCoderAgent

agent = AutoCoderAgent(model="gpt-4.1", name="summarize-sales")

env = flyte.TaskEnvironment(
    name="my-env",
    secrets=[flyte.Secret(key="openai_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-codegen",
    ),
    depends_on=[sandbox_environment],
)

@env.task
async def process_data(csv_file: File) -> tuple[float, int, int]:
    result = await agent.generate.aio(
        prompt="Read the CSV and compute total_revenue, total_units and row_count.",
        samples={"sales": csv_file},
        outputs={"total_revenue": float, "total_units": int, "row_count": int},
    )
    return await result.run.aio()
```

The `depends_on=[sandbox_environment]` declaration is required. It ensures the sandbox runtime is available when dynamically-created sandboxes execute.

![Sandbox](../../_static/images/integrations/codegen/sandbox.png)

## Two execution backends

The plugin supports two backends for generating and validating code. Both share the same `AutoCoderAgent` interface and produce the same `CodeGenEvalResult`.

### LiteLLM (default)

Uses structured-output LLM calls to generate code, detect packages, build sandbox images, run tests, diagnose failures, and iterate. Works with any model that supports structured outputs (GPT-4, Claude, Gemini, etc. via LiteLLM).

```python{hl_lines=[1, 3]}
agent = AutoCoderAgent(
    name="my-task",
    model="gpt-4.1",
    max_iterations=10,
)
```

The LiteLLM backend follows a fixed pipeline:

```mermaid
flowchart TD
    A["prompt + samples"] --> B["generate_plan"]
    B --> C["generate_code"]
    C --> D["detect_packages"]
    D --> E["build_image"]
    E --> F{skip_tests?}
    F -- yes --> G["return result"]
    F -- no --> H["generate_tests"]
    H --> I["execute_tests"]
    I --> J{pass?}
    J -- yes --> G
    J -- no --> K["diagnose_error"]
    K --> L{error type?}
    L -- "logic error" --> M["regenerate code"]
    L -- "environment error" --> N["add packages, rebuild image"]
    L -- "test error" --> O["fix test expectations"]
    M --> I
    N --> I
    O --> I
```

The loop continues until tests pass or `max_iterations` is reached.

![LiteLLM](../../_static/images/integrations/codegen/litellm.png)

### Agent (Claude)

Uses the Claude Agent SDK to autonomously generate, test, and fix code. The agent has access to `Bash`, `Read`, `Write`, and `Edit` tools and decides what to do at each step. Test execution commands (`pytest`) are intercepted and run inside isolated sandboxes.

```python{hl_lines=["3-4"]}
agent = AutoCoderAgent(
    name="my-task",
    model="claude-sonnet-4-5-20250929",
    backend="claude",
)
```

> [!NOTE]
> Agent mode requires `ANTHROPIC_API_KEY` as a Flyte secret and is Claude-only.

**Key differences from LiteLLM:**

|                       | LiteLLM                           | Agent                                          |
| --------------------- | --------------------------------- | ---------------------------------------------- |
| **Execution**         | Fixed generate-test-fix pipeline  | Autonomous agent decides actions               |
| **Model support**     | Any model with structured outputs | Claude only                                    |
| **Iteration control** | `max_iterations`                  | `agent_max_turns`                              |
| **Test execution**    | Direct sandbox execution          | `pytest` commands intercepted via hooks        |
| **Tool safety**       | N/A                               | Commands classified as safe/denied/intercepted |
| **Observability**     | Logs + token counts               | Full tool call tracing in Flyte UI             |

In Agent mode, Bash commands are classified before execution:

- **Safe** (`ls`, `cat`, `grep`, `head`, etc.): allowed to run directly
- **Intercepted** (`pytest`): routed to sandbox execution
- **Denied** (`apt`, `pip install`, `curl`, etc.): blocked for safety

## Providing data

### Sample data

Pass sample data via `samples` as `File` objects or pandas `DataFrame`s. The plugin automatically:

1. Converts DataFrames to CSV files
2. Infers [Pandera](https://pandera.readthedocs.io/) schemas from the data: column types, nullability
3. Parses natural-language `constraints` into Pandera checks (e.g., `"quantity must be positive"` becomes `pa.Check.gt(0)`)
4. Extracts data context: column statistics, distributions, patterns, sample rows
5. Injects all of this into the LLM prompt so the generated code is aware of the exact data structure

Pandera is used purely for prompt enrichment, not runtime validation. The generated code does not import Pandera; it benefits from the LLM knowing the precise data structure. The generated schemas are stored on `result.generated_schemas` for inspection.

```python{hl_lines=[3]}
result = await agent.generate.aio(
    prompt="Clean and validate the data, remove duplicates",
    samples={"orders": orders_df, "products": products_file},
    constraints=["quantity must be positive", "price between 0 and 10000"],
    outputs={"cleaned_orders": File},
)
```

### Schema and constraints

Use `schema` to provide free-form context about data formats or target structures (e.g., a database schema). Use `constraints` to declare business rules that the generated code must respect:

```python{hl_lines=["4-17"]}
result = await agent.generate.aio(
    prompt=prompt,
    samples={"readings": sensor_df},
    schema="""Output JSON schema for report_json:
    {
        "sensor_id": str,
        "avg_temp": float,
        "min_temp": float,
        "max_temp": float,
        "avg_humidity": float,
    }
    """,
    constraints=[
        "Temperature values must be between -40 and 60 Celsius",
        "Humidity values must be between 0 and 100 percent",
        "Output report must have one row per unique sensor_id",
    ],
    outputs={
        "report_json": str,
        "total_anomalies": int,
    },
)
```

![Pandera Constraints](../../_static/images/integrations/codegen/pandera_constraints.png)

### Inputs and outputs

Declare `inputs` for non-sample arguments (e.g., thresholds, flags) and `outputs` for the expected result types.

Supported output types: `str`, `int`, `float`, `bool`, `datetime.datetime`, `datetime.timedelta`, `File`.

Sample entries are automatically added as `File` inputs; you do not need to redeclare them.

```python{hl_lines=[4, 5]}
result = await agent.generate.aio(
    prompt="Filter transactions above the threshold",
    samples={"transactions": tx_file},
    inputs={"threshold": float, "include_pending": bool},
    outputs={"filtered": File, "count": int},
)
```

## Running generated code

`agent.generate()` returns a `CodeGenEvalResult`. If `result.success` is `True`, the generated code passed all tests and you can execute it against real data. If `max_iterations` (LiteLLM) or `agent_max_turns` (Agent) is reached without tests passing, `result.success` is `False` and `result.error` contains the failure details.

Both `run()` and `as_task()` return output values as a tuple in the order declared in `outputs`. If there is a single output, the value is returned directly (not wrapped in a tuple).

### One-shot execution with `result.run()`

Runs the generated code in a sandbox. If samples were provided during `generate()`, they are used as default inputs.

```python
# Use sample data as defaults
total_revenue, total_units, count = await result.run.aio()

# Override specific inputs
total_revenue, total_units, count = await result.run.aio(threshold=0.5)

# Sync version
total_revenue, total_units, count = result.run()
```

`result.run()` accepts optional configuration:

```python{hl_lines=["4-6"]}
total_revenue, total_units, count = await result.run.aio(
    name="execute-on-data",
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    retries=2,
    timeout=600,
    cache="auto",
)
```

### Reusable task with `result.as_task()`

Creates a callable sandbox task from the generated code. Useful when you want to run the same generated code against different data.

```python{hl_lines=[1, "6-7", "9-10"]}
task = result.as_task(
    name="run-sensor-analysis",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

# Call with sample defaults
report, total_anomalies = await task.aio()

# Call with different data
report, total_anomalies = await task.aio(readings=new_data_file)
```

## Error diagnosis

The LiteLLM backend classifies test failures into three categories and applies targeted fixes:

| Error type    | Meaning                       | Action                                           |
| ------------- | ----------------------------- | ------------------------------------------------ |
| `logic`       | Bug in the generated code     | Regenerate code with specific patch instructions |
| `environment` | Missing package or dependency | Add the package and rebuild the sandbox image    |
| `test_error`  | Bug in the generated test     | Fix the test expectations                        |

If the same error persists after a fix, the plugin reclassifies it (e.g., `logic` to `test_error`) to try the other approach.

In Agent mode, the agent diagnoses and fixes issues autonomously based on error output.

## Durable execution

Code generation is expensive: it involves multiple LLM calls, image builds, and sandbox executions. Without durability, a transient failure in the pipeline (network blip, OOM, downstream service error) would force the entire process to restart from scratch: regenerating code, rebuilding images, re-running sandboxes, making additional LLM calls.

Flyte solves this through two complementary mechanisms: **replay logs** and **caching**.

### Replay logs

Flyte maintains a replay log that records every trace and task execution within a run. When a task crashes and retries, the system replays the log from the previous attempt rather than recomputing everything:

- No additional model calls
- No code regeneration
- No sandbox re-execution
- No container rebuilds

The workflow breezes through the earlier steps and resumes from the failure point. This applies as long as the traces and tasks execute in the same order and use the same inputs as the first attempt.

### Caching

Separately, Flyte can cache task results across runs. With `cache="auto"`, sandbox executions (image builds, test runs, code execution) are cached. This is useful when you re-run the same pipeline, not just when recovering from a crash, but across entirely separate invocations with the same inputs.

Together, replay logs handle crash recovery within a run, and caching avoids redundant work across runs.

### Non-determinism in agent mode

One challenge with agents is that they are inherently non-deterministic: the sequence of actions can vary between runs, which could break replay.

In practice, the codegen agent follows a predictable pattern (write code, generate tests, run tests, inspect results), which works in replay's favor. The plugin also embeds logic that instructs the agent not to regenerate or re-execute steps that already completed successfully in the first run. This acts as an additional safety check alongside the replay log to account for non-determinism.

![Agent](../../_static/images/integrations/codegen/agent.png)

On the first attempt, the full pipeline runs. If a transient failure occurs, the system instantly replays the traces (which track model calls) and sandbox executions, allowing the pipeline to resume from the point of failure.

![Durability](../../_static/images/integrations/codegen/durability.png)

## Observability

### LiteLLM backend

- Logs every iteration with attempt count, error type, and package changes
- Tracks total input/output tokens across all LLM calls (available on `result.total_input_tokens` and `result.total_output_tokens`)
- Results include full conversation history for debugging (`result.conversation_history`)

### Agent backend

- Traces each tool call (name + input) via `PostToolUse` hooks
- Traces tool failures via `PostToolUseFailure` hooks
- Traces a summary when the agent finishes (total tool calls, tool distribution, final image/packages)
- Classifies Bash commands as safe, denied, or intercepted (for sandbox execution)
- All traces appear in the Flyte UI

## Examples

### Processing CSVs with different schemas

Generate code that handles varying CSV formats, then run on real data:

```python{hl_lines=[1, 3, 14, 16, 27]}
from flyteplugins.codegen import AutoCoderAgent

agent = AutoCoderAgent(
    name="sales-processor",
    model="gpt-4.1",
    max_iterations=5,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    litellm_params={"temperature": 0.2, "max_tokens": 4096},
)

@env.task
async def process_sales(csv_file: File) -> dict[str, float | int]:
    result = await agent.generate.aio(
        prompt="Read the CSV and compute total_revenue, total_units, and transaction_count.",
        samples={"csv_data": csv_file},
        outputs={
            "total_revenue": float,
            "total_units": int,
            "transaction_count": int,
        },
    )

    if not result.success:
        raise RuntimeError(f"Code generation failed: {result.error}")

    total_revenue, total_units, transaction_count = await result.run.aio()

    return {
        "total_revenue": total_revenue,
        "total_units": total_units,
        "transaction_count": transaction_count,
    }
```

### DataFrame analysis with constraints

Pass DataFrames directly and enforce business rules with constraints:

```python{hl_lines=[10, "15-19"]}
agent = AutoCoderAgent(
    model="gpt-4.1",
    name="sensor-analysis",
    base_packages=["numpy"],
    max_sample_rows=30,
)

@env.task
async def analyze_sensors(sensor_df: pd.DataFrame) -> tuple[File, int]:
    result = await agent.generate.aio(
        prompt="""Analyze IoT sensor data. For each sensor, calculate mean/min/max
temperature, mean humidity, and count warnings. Output a summary CSV.""",
        samples={"readings": sensor_df},
        constraints=[
            "Temperature values must be between -40 and 60 Celsius",
            "Humidity values must be between 0 and 100 percent",
            "Output report must have one row per unique sensor_id",
        ],
        outputs={
            "report": File,
            "total_anomalies": int,
        },
    )

    if not result.success:
        raise RuntimeError(f"Code generation failed: {result.error}")

    task = result.as_task(
        name="run-sensor-analysis",
        resources=flyte.Resources(cpu=1, memory="512Mi"),
    )

    return await task.aio(readings=result.original_samples["readings"])
```

### Agent mode

The same task using Claude as an autonomous agent:

```python{hl_lines=[3]}
agent = AutoCoderAgent(
    name="sales-agent",
    backend="claude",
    model="claude-sonnet-4-5-20250929",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

@env.task
async def process_sales_with_agent(csv_file: File) -> dict[str, float | int]:
    result = await agent.generate.aio(
        prompt="Read the CSV and compute total_revenue, total_units, and transaction_count.",
        samples={"csv_data": csv_file},
        outputs={
            "total_revenue": float,
            "total_units": int,
            "transaction_count": int,
        },
    )

    if not result.success:
        raise RuntimeError(f"Agent code generation failed: {result.error}")

    total_revenue, total_units, transaction_count = await result.run.aio()

    return {
        "total_revenue": total_revenue,
        "total_units": total_units,
        "transaction_count": transaction_count,
    }
```

## Configuration

### LiteLLM parameters

Tune model behavior with `litellm_params`:

```python{hl_lines=["5-8"]}
agent = AutoCoderAgent(
    name="my-task",
    model="anthropic/claude-sonnet-4-20250514",
    api_key="ANTHROPIC_API_KEY",
    litellm_params={
        "temperature": 0.3,
        "max_tokens": 4000,
    },
)
```

### Image configuration

Control the registry and Python version for sandbox images:

```python{hl_lines=["6-10"]}
from flyte.sandbox import ImageConfig

agent = AutoCoderAgent(
    name="my-task",
    model="gpt-4.1",
    image_config=ImageConfig(
        registry="my-registry.io",
        registry_secret="registry-creds",
        python_version=(3, 12),
    ),
)
```

### Skipping tests

Set `skip_tests=True` to skip test generation and execution. The agent still generates code, detects packages, and builds the sandbox image, but does not generate or run tests.

```python{hl_lines=[4]}
agent = AutoCoderAgent(
    name="my-task",
    model="gpt-4.1",
    skip_tests=True,
)
```

> [!NOTE]
> `skip_tests` only applies to LiteLLM mode. In Agent mode, the agent autonomously decides when to test.

### Base packages

Ensure specific packages are always installed in every sandbox:

```python{hl_lines=[4]}
agent = AutoCoderAgent(
    name="my-task",
    model="gpt-4.1",
    base_packages=["numpy", "pandas"],
)
```

## Best practices

- **One agent per task.** Each `generate()` call builds its own sandbox image and manages its own package state. Running multiple agents in the same task can cause resource contention and makes failures harder to diagnose.
- **Keep `cache="auto"` (the default).** Caching flows to all internal sandboxes, making retries near-instant. Use `"disable"` during development if you want fresh executions, or `"override"` to force re-execution and update the cached result.
- **Set `max_iterations` conservatively.** Start with 5-10 iterations. If the model cannot produce correct code in that budget, the prompt or constraints likely need refinement.
- **Provide constraints for data-heavy tasks.** Explicit constraints (e.g., `"quantity must be positive"`) produce better schemas and better generated code.
- **Inspect `result.generated_schemas`.** Review the inferred Pandera schemas to verify the model understood your data structure correctly.

## API reference

### `AutoCoderAgent` constructor

| Parameter         | Type              | Default        | Description                                                                            |
| ----------------- | ----------------- | -------------- | -------------------------------------------------------------------------------------- |
| `name`            | `str`             | `"auto-coder"` | Unique name for tracking and image naming                                              |
| `model`           | `str`             | `"gpt-4.1"`    | LiteLLM model identifier                                                               |
| `backend`         | `str`             | `"litellm"`    | Execution backend: `"litellm"` or `"claude"`                                           |
| `system_prompt`   | `str`             | `None`         | Custom system prompt override                                                          |
| `api_key`         | `str`             | `None`         | Name of the environment variable containing the LLM API key (e.g., `"OPENAI_API_KEY"`) |
| `api_base`        | `str`             | `None`         | Custom API base URL                                                                    |
| `litellm_params`  | `dict`            | `None`         | Extra LiteLLM params (temperature, max_tokens, etc.)                                   |
| `base_packages`   | `list[str]`       | `None`         | Always-install pip packages                                                            |
| `resources`       | `flyte.Resources` | `None`         | Resources for sandbox execution (default: 1 CPU, 1Gi)                                  |
| `image_config`    | `ImageConfig`     | `None`         | Registry, secret, and Python version                                                   |
| `max_iterations`  | `int`             | `10`           | Max generate-test-fix iterations (LiteLLM mode)                                        |
| `max_sample_rows` | `int`             | `100`          | Rows to sample from data for LLM context                                               |
| `skip_tests`      | `bool`            | `False`        | Skip test generation and execution (LiteLLM mode)                                      |
| `sandbox_retries` | `int`             | `0`            | Flyte task-level retries for each sandbox execution                                    |
| `timeout`         | `int`             | `None`         | Timeout in seconds for sandboxes                                                       |
| `env_vars`        | `dict[str, str]`  | `None`         | Environment variables for sandboxes                                                    |
| `secrets`         | `list[Secret]`    | `None`         | Flyte secrets for sandboxes                                                            |
| `cache`           | `str`             | `"auto"`       | Cache behavior: `"auto"`, `"override"`, or `"disable"`                                 |
| `agent_max_turns` | `int`             | `50`           | Max turns when `backend="claude"`                                                      |

### `generate()` parameters

| Parameter     | Type                           | Default  | Description                                                                             |
| ------------- | ------------------------------ | -------- | --------------------------------------------------------------------------------------- |
| `prompt`      | `str`                          | required | Natural-language task description                                                       |
| `schema`      | `str`                          | `None`   | Free-form context about data formats or target structures                               |
| `constraints` | `list[str]`                    | `None`   | Natural-language constraints (e.g., `"quantity must be positive"`)                      |
| `samples`     | `dict[str, File \| DataFrame]` | `None`   | Sample data. DataFrames are auto-converted to CSV files.                                |
| `inputs`      | `dict[str, type]`              | `None`   | Non-sample input types (e.g., `{"threshold": float}`)                                   |
| `outputs`     | `dict[str, type]`              | `None`   | Output types. Supported: `str`, `int`, `float`, `bool`, `datetime`, `timedelta`, `File` |

### `CodeGenEvalResult` fields

| Field                      | Type                      | Description                                               |
| -------------------------- | ------------------------- | --------------------------------------------------------- |
| `success`                  | `bool`                    | Whether tests passed                                      |
| `solution`                 | `CodeSolution`            | Generated code (`.code`, `.language`, `.system_packages`) |
| `tests`                    | `str`                     | Generated test code                                       |
| `output`                   | `str`                     | Test output                                               |
| `exit_code`                | `int`                     | Test exit code                                            |
| `error`                    | `str \| None`             | Error message if failed                                   |
| `attempts`                 | `int`                     | Number of iterations used                                 |
| `image`                    | `str`                     | Built sandbox image with all dependencies                 |
| `detected_packages`        | `list[str]`               | Pip packages detected                                     |
| `detected_system_packages` | `list[str]`               | Apt packages detected                                     |
| `generated_schemas`        | `dict[str, str] \| None`  | Pandera schemas as Python code strings                    |
| `data_context`             | `str \| None`             | Extracted data context                                    |
| `original_samples`         | `dict[str, File] \| None` | Sample data as Files (defaults for `run()`/`as_task()`)   |
| `total_input_tokens`       | `int`                     | Total input tokens across all LLM calls                   |
| `total_output_tokens`      | `int`                     | Total output tokens across all LLM calls                  |
| `conversation_history`     | `list[dict]`              | Full LLM conversation history for debugging               |

### `CodeGenEvalResult` methods

| Method                              | Description                                                        |
| ----------------------------------- | ------------------------------------------------------------------ |
| `result.run(**overrides)`           | Execute generated code in a sandbox. Sample data used as defaults. |
| `await result.run.aio(**overrides)` | Async version of `run()`.                                          |
| `result.as_task(name, ...)`         | Create a reusable callable sandbox task from the generated code.   |

Both `run()` and `as_task()` accept optional `name`, `resources`, `retries`, `timeout`, `env_vars`, `secrets`, and `cache` parameters.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/dask ===

# Dask

The Dask plugin lets you run [Dask](https://www.dask.org/) jobs natively on Kubernetes. Flyte provisions a transient Dask cluster for each task execution using the [Dask Kubernetes Operator](https://kubernetes.dask.org/en/latest/operator.html) and tears it down on completion.

## When to use this plugin

- Parallel Python workloads that outgrow a single machine
- Distributed DataFrame operations on large datasets
- Workloads that use Dask's task scheduler for arbitrary computation graphs
- Jobs that need to scale NumPy, pandas, or scikit-learn workflows across multiple nodes

## Installation

```bash
pip install flyteplugins-dask
```

Your task image must also include the Dask distributed scheduler:

```python
image = flyte.Image.from_debian_base(name="dask").with_pip_packages("flyteplugins-dask")
```

## Configuration

Create a `Dask` configuration and pass it as `plugin_config` to a `TaskEnvironment`:

```python
from flyteplugins.dask import Dask, Scheduler, WorkerGroup

dask_config = Dask(
    scheduler=Scheduler(),
    workers=WorkerGroup(number_of_workers=4),
)

dask_env = flyte.TaskEnvironment(
    name="dask_env",
    plugin_config=dask_config,
    image=image,
)
```

### `Dask` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `scheduler` | `Scheduler` | Scheduler pod configuration (defaults to `Scheduler()`) |
| `workers` | `WorkerGroup` | Worker group configuration (defaults to `WorkerGroup()`) |

### `Scheduler` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `image` | `str` | Custom scheduler image (must include `dask[distributed]`) |
| `resources` | `Resources` | Resource requests for the scheduler pod |

### `WorkerGroup` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `number_of_workers` | `int` | Number of worker pods (default: `1`) |
| `image` | `str` | Custom worker image (must include `dask[distributed]`) |
| `resources` | `Resources` | Resource requests per worker pod |

> [!NOTE]
> The scheduler and all workers should use the same Python environment to avoid serialization issues.

### Accessing the Dask client

Inside a Dask task, create a `distributed.Client()` with no arguments. It automatically connects to the provisioned cluster:

```python
from distributed import Client

@dask_env.task
async def my_dask_task(n: int) -> list:
    client = Client()
    futures = client.map(lambda x: x + 1, range(n))
    return client.gather(futures)
```

## Example

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-dask",
#    "distributed"
# ]
# main = "hello_dask_nested"
# params = ""
# ///

import asyncio
import typing

from distributed import Client
from flyteplugins.dask import Dask, Scheduler, WorkerGroup

import flyte.remote
import flyte.storage
from flyte import Resources

image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("flyteplugins-dask")

dask_config = Dask(
    scheduler=Scheduler(),
    workers=WorkerGroup(number_of_workers=4),
)

task_env = flyte.TaskEnvironment(
    name="hello_dask", resources=Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
dask_env = flyte.TaskEnvironment(
    name="dask_env",
    plugin_config=dask_config,
    image=image,
    resources=Resources(cpu="1", memory="1Gi"),
    depends_on=[task_env],
)

@task_env.task()
async def hello_dask():
    await asyncio.sleep(5)
    print("Hello from the Dask task!")

@dask_env.task
async def hello_dask_nested(n: int = 3) -> typing.List[int]:
    print("running dask task")
    t = asyncio.create_task(hello_dask())
    client = Client()
    futures = client.map(lambda x: x + 1, range(n))
    res = client.gather(futures)
    await t
    return res

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(hello_dask_nested)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/dask/dask_example.py*

## API reference

See the [Dask API reference](../../api-reference/integrations/dask/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/databricks ===

# Databricks

The Databricks plugin lets you run PySpark jobs on [Databricks](https://www.databricks.com/) clusters directly from Flyte tasks. You write normal PySpark code in a Flyte task, and the plugin submits it to Databricks via the [Jobs API 2.1](https://docs.databricks.com/api/workspace/jobs/submit). The connector handles job submission, polling, and cancellation.

The plugin supports:

- Running PySpark tasks on new or existing Databricks clusters
- Full Spark configuration (driver/executor memory, cores, instances)
- Databricks cluster auto-scaling
- API token-based authentication

## Installation

```bash
pip install flyteplugins-databricks
```

This also installs `flyteplugins-spark` as a dependency, since the Databricks plugin extends the Spark plugin.

## Quick start

Create a `Databricks` configuration and pass it as `plugin_config` to a `TaskEnvironment`:

```python
from flyteplugins.databricks import Databricks
import flyte

image = (
    flyte.Image.from_base("databricksruntime/standard:16.4-LTS")
    .clone(name="spark", registry="ghcr.io/flyteorg", extendable=True)
    .with_env_vars({"UV_PYTHON": "/databricks/python3/bin/python"})
    .with_pip_packages("flyteplugins-databricks", pre=True)
)

databricks_conf = Databricks(
    spark_conf={
        "spark.driver.memory": "2000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
    },
    executor_path="/databricks/python3/bin/python",
    databricks_conf={
        "run_name": "flyte databricks plugin",
        "new_cluster": {
            "spark_version": "13.3.x-scala2.12",
            "node_type_id": "m6i.large",
            "autoscale": {"min_workers": 1, "max_workers": 2},
        },
        "timeout_seconds": 3600,
        "max_retries": 1,
    },
    databricks_instance="myaccount.cloud.databricks.com",
    databricks_token="DATABRICKS_TOKEN",
)

databricks_env = flyte.TaskEnvironment(
    name="databricks_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")),
    plugin_config=databricks_conf,
    image=image,
)
```

Then use the environment to decorate your task:

```python
@databricks_env.task
async def hello_databricks() -> float:
    spark = flyte.ctx().data["spark_session"]
    # Use spark as a normal SparkSession
    count = spark.sparkContext.parallelize(range(100)).count()
    return float(count)
```

## Configuration

The `Databricks` config extends the [Spark](../spark/_index) config with Databricks-specific fields.

### Spark fields (inherited)

| Parameter | Type | Description |
|-----------|------|-------------|
| `spark_conf` | `Dict[str, str]` | Spark configuration key-value pairs |
| `hadoop_conf` | `Dict[str, str]` | Hadoop configuration key-value pairs |
| `executor_path` | `str` | Path to the Python binary on the Databricks cluster (e.g., `/databricks/python3/bin/python`) |
| `applications_path` | `str` | Path to the main application file |

### Databricks-specific fields

| Parameter | Type | Description |
|-----------|------|-------------|
| `databricks_conf` | `Dict[str, Union[str, dict]]` | Databricks [run-submit](https://docs.databricks.com/api/workspace/jobs/submit) job configuration. Must contain either `existing_cluster_id` or `new_cluster` |
| `databricks_instance` | `str` | Your workspace domain (e.g., `myaccount.cloud.databricks.com`). Can also be set via the `FLYTE_DATABRICKS_INSTANCE` env var on the connector |
| `databricks_token` | `str` | Name of the Flyte secret containing the Databricks API token |

### `databricks_conf` structure

The `databricks_conf` dict maps to the Databricks run-submit API payload. Key fields:

| Field | Description |
|-------|-------------|
| `new_cluster` | Cluster spec with `spark_version`, `node_type_id`, `autoscale`, etc. |
| `existing_cluster_id` | ID of an existing cluster to use instead of creating a new one |
| `run_name` | Display name in the Databricks UI |
| `timeout_seconds` | Maximum job duration |
| `max_retries` | Number of retries before marking the job as failed |

The connector automatically injects the Docker image, Spark configuration, and environment variables from the task container into the cluster spec.

## Authentication

Store your Databricks API token as a Flyte secret. The `databricks_token` parameter specifies the secret name:

```python
databricks_conf = Databricks(
    # ...
    databricks_token="DATABRICKS_TOKEN",
)
```

## Accessing the Spark session

Inside a Databricks task, the `SparkSession` is available through the task context, just like the [Spark plugin](../spark/_index):

```python
@databricks_env.task
async def my_databricks_task() -> float:
    spark = flyte.ctx().data["spark_session"]
    df = spark.read.parquet("s3://my-bucket/data.parquet")
    return float(df.count())
```

## API reference

See the [Databricks API reference](../../api-reference/integrations/databricks/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/grafana-agent-observability ===

# Grafana Agent Observability

`flyteplugins-agento11y` sends your agent's generations, tool calls, token usage and cost to [Grafana Agent Observability](https://grafana.com/docs/grafana-cloud/observe-and-act/agent-observability/), nested inside the Flyte task span and grouped by Flyte run.

It instruments agents built with the [agent framework plugins](../agents/_index), the `flyteplugins-agents-*` adapters that run a framework's agent loop inside a Flyte task, with each model turn as a durable traced step and each tool as a child action.

One call at module scope is the whole integration. Your agent code does not change:

```python{hl_lines=[3,7]}
import flyte
from flyteplugins.agents.openai import run_agent, tool
from flyteplugins.agento11y import init

# Module scope, not inside a task. The task span opens before the task body runs,
# and the Flyte identity binding rides on that span.
init(service_name="my-agent")

env = flyte.TaskEnvironment(
    name="agent_env",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-agents-openai",
        "flyteplugins-agento11y[openai]",
    ),
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
)

@env.task
async def lookup_order(order_id: str) -> str:
    """A durable Flyte child action, and a tool call in Grafana."""
    return f"Order {order_id} shipped on 2026-07-20."

@env.task
async def support_agent(question: str) -> str:
    return await run_agent(
        question,
        tools=[tool(lookup_order)],
        model="gpt-4.1",
        instructions="You are a support agent. Use the tools to answer.",
    )
```

That covers instrumentation. **Grafana Agent Observability > Configuration** covers the credentials that decide where the generations actually go.

![Agent Observability conversation view showing the agent's two model calls, its tool call, and the prompt and answer](../../_static/images/integrations/grafana-agent-observability/openai_agent.png)

*The agent above, in Agent Observability. The flow on the left is the two model calls with the `lookup_order` tool call between them; the thread on the right is the prompt and the answer. Call count, token usage, and cost sit in the header.*

This plugin builds on [`flyteplugins-otel`](../opentelemetry/_index), which it initializes for you. Read that page first if you want to understand where the spans come from.

## Installation

```bash
pip install "flyteplugins-agento11y[openai]"
```

The extra is what makes your framework's instrumentor available. Install the one matching the agent adapter you use:

| Extra         | Agent adapter                     | agento11y integration package |
| ------------- | --------------------------------- | ----------------------------- |
| `langchain`   | `flyteplugins-agents-langchain`   | `agento11y-langchain`         |
| `langgraph`   | `flyteplugins-agents-langgraph`   | `agento11y-langgraph`         |
| `openai`      | `flyteplugins-agents-openai`      | `agento11y-openai-agents`     |
| `claude`      | `flyteplugins-agents-claude`      | `agento11y-claude-agent-sdk`  |
| `google`      | `flyteplugins-agents-google`      | `agento11y-google-adk`        |
| `pydantic-ai` | `flyteplugins-agents-pydantic-ai` | `agento11y-pydantic-ai`       |

Nothing else has to be configured: `init()` registers an instrumentor for every framework whose integration package it finds. `instrumented_frameworks()` returns the ones that were registered, which is the quickest way to confirm the extra actually installed.

```python
from flyteplugins.agento11y import instrumented_frameworks

print(instrumented_frameworks())   # ('openai',)
```

## What you get beyond agento11y on its own

agento11y works inside a Flyte task without any of this and Grafana's dashboards will light up because they are driven by generation records rather than by trace structure. What is missing is everything Flyte knows and agento11y cannot.

Without the plugin, three model calls in a task become three unrelated root traces. There is no task boundary, nothing tying a generation to the run that produced it, and on a resume the replayed steps produce nothing at all.

### Generations nest inside the task span

One run is one trace, and generation records carry that trace ID. That ID is the link from a generation in Grafana back to the Flyte run that produced it.

![Tempo trace with each generation span nested inside the Flyte step that produced it, across three attempts](../../_static/images/integrations/grafana-agent-observability/agent_trace.png)

*A durable agent's trace. Each `generateText` span sits inside the `flyte.trace` step that produced it, which sits inside the task span. In the second attempt the steps are microsecond replays with no `generateText` child at all: the resume did not call the model again.*

![Expanded generation span listing gen_ai attributes bound to Flyte's run, task, and version](../../_static/images/integrations/grafana-agent-observability/trace_id.png)

*Expanding one generation shows the binding described below: `gen_ai.conversation.id` is the Flyte run name, `gen_ai.agent.name` the task, and `gen_ai.agent.version` the task version.*

### Flyte identity is bound onto agento11y's context

Nothing has to be restated by hand:

| agento11y concept | Flyte value  |
| ----------------- | ------------ |
| Conversation ID   | Run name     |
| Agent name        | Task name    |
| Agent version     | Task version |

A Flyte run therefore shows up in Grafana as one conversation, and a redeploy shows up as a new agent version, so the before and after of a prompt change is directly comparable.

![Agent Observability conversations list with one row per Flyte run, agents named after Flyte tasks](../../_static/images/integrations/grafana-agent-observability/conversations.png)

*The conversations list, one row per Flyte run. **Conversation** holds run names and **Agents** holds task names, so runs driving different frameworks and models line up in a single view.*

Both bindings are switchable because both assume something that is not always true:

- `bind_conversation=False` keeps your own conversation IDs, for a product where a conversation spans more than one run.
- `bind_agent_name=False` lets each framework name its own agents, for a task that drives several: a planner and a worker would otherwise both report as the task.

### Durability is preserved end to end

A crashed and resumed run stays a single trace, because **OpenTelemetry > Traces across crashes and resumes > One run, one trace** rather than generated per process.

Steps the resumed run replayed from its durable log appear marked `flyte.replayed`, so the trace has no holes where durability did its job. And those steps do not call the model again, so a resume does not pay for the generations the first attempt already bought.

## Framework coverage

All six frameworks in the table above capture generations and tool calls. Two capture more:

| Framework   | Also captures                                                                         |
| ----------- | ------------------------------------------------------------------------------------- |
| `langgraph` | Workflow steps, so non-LLM nodes (routing, retrieval) appear too                      |
| `claude`    | Model turns read off the SDK's message stream, via a call wrapper rather than options |

Adapters without an agento11y integration package, crewai and mistral among them, still work: their runs are traced and their tasks and tool calls appear as spans. Only the generations are not captured automatically. **Grafana Agent Observability > Recording generations by hand**.

## Configuration

With no arguments, `init()` reads the standard `AGENTO11Y_*` variables, which is how the Grafana documentation configures it.

| Variable                   | What it is                                                                 |
| -------------------------- | -------------------------------------------------------------------------- |
| `AGENTO11Y_ENDPOINT`       | Generation export endpoint, for example `https://<your-stack>.grafana.net` |
| `AGENTO11Y_AUTH_MODE`      | `none` (the agento11y default) or `basic`                                  |
| `AGENTO11Y_AUTH_TOKEN`     | The token or password                                                      |
| `AGENTO11Y_AUTH_TENANT_ID` | Basic-auth username. On Grafana Cloud this is your instance ID             |

Supply the credentials as a `flyte.Secret` rather than hardcoding them:

```python{hl_lines=[4,"6-8"]}
env = flyte.TaskEnvironment(
    name="agent_env",
    image=image,
    env_vars={"AGENTO11Y_AUTH_MODE": "basic"},
    secrets=[
        flyte.Secret(key="agento11y_endpoint", as_env_var="AGENTO11Y_ENDPOINT"),
        flyte.Secret(key="agento11y_token", as_env_var="AGENTO11Y_AUTH_TOKEN"),
        flyte.Secret(key="agento11y_tenant_id", as_env_var="AGENTO11Y_AUTH_TENANT_ID"),
        # Spans go to Tempo over OTLP; generations go to Agent Observability over
        # their own channel. Both are needed for the two UI links to resolve.
        flyte.Secret(key="otlp_endpoint", as_env_var="OTEL_EXPORTER_OTLP_ENDPOINT"),
        flyte.Secret(key="otlp_headers", as_env_var="OTEL_EXPORTER_OTLP_HEADERS"),
    ],
)
```

> [!WARNING] Set the auth mode explicitly on Grafana Cloud
> agento11y defaults `AGENTO11Y_AUTH_MODE` to `none`, so a token on its own is never sent and
> the export comes back `401`. Grafana Cloud uses Basic auth with the instance ID as the
> username, which agento11y fills from `AGENTO11Y_AUTH_TENANT_ID` when the mode is `basic`.

Generations and spans travel over two different channels. `AGENTO11Y_ENDPOINT` decides where generations go; `OTEL_EXPORTER_OTLP_ENDPOINT` decides where spans go. Configuring one does not configure the other.

### `init()` parameters

| Parameter           | Default              | What it does                                                                                                     |
| ------------------- | -------------------- | ---------------------------------------------------------------------------------------------------------------- |
| `service_name`      | None                 | Value for `service.name` on the OpenTelemetry side                                                               |
| `endpoint`          | `AGENTO11Y_ENDPOINT` | Generation export endpoint                                                                                       |
| `client`            | None                 | Use an agento11y client you built yourself. It is left alone and not shut down                                   |
| `client_options`    | None                 | Extra `ClientConfig` fields: auth mode, protocol, content capture, a custom generation exporter                  |
| `bind_conversation` | `True`               | Bind the Flyte run name as the conversation ID                                                                   |
| `bind_agent_name`   | `True`               | Bind the Flyte task name as the agent name                                                                       |
| `trace`             | `True`               | Also initialize `flyteplugins-otel`. Turn off if you configure tracing yourself, or if you only want generations |

Anything else is forwarded to `flyteplugins.otel.init()`, including `tracer_provider` for an **OpenTelemetry > Exporters and configuration > Adopting a tracer provider you already have**, and `exporter`, `headers` and `disable_batch`.

`init()` returns the agento11y client and `get_client()` returns it later for recording generations directly.

## Linking back from Grafana

`GrafanaAgentObservability` links a Flyte action to its conversation in Agent Observability, rendered on the action in the Flyte UI. It works precisely because this plugin binds the run name as the conversation ID.

```python{hl_lines=["6-7"]}
from flyteplugins.agento11y import GrafanaAgentObservability
from flyteplugins.otel.grafana import GrafanaTrace

@env.task(links=(
    GrafanaAgentObservability(host="https://myorg.grafana.net"),
    GrafanaTrace(host="https://myorg.grafana.net", datasource_uid="<tempo-uid>"),
))
async def support_agent(question: str) -> str:
    ...
```

![Flyte UI action summary with both Grafana links highlighted in its Links section](../../_static/images/integrations/grafana-agent-observability/ui_links.png)

*Both links on the same action. **Grafana Agent Observability** opens this run's conversation; **Grafana trace** opens its spans in Tempo.*

The two answer different questions. The first goes to the generations, prompts and cost. The second goes to the distributed trace in Tempo; it lives in **OpenTelemetry > Exporters and configuration > Linking back from Grafana** because it needs nothing from this package.

The conversation link opens the conversation itself rather than the filtered list, and fills the app's back navigation with the list scoped to the same run.

| Parameter           | Default                                     | What it does                                             |
| ------------------- | ------------------------------------------- | -------------------------------------------------------- |
| `host`              | Required                                    | Stack URL, for example `https://myorg.grafana.net`       |
| `name`              | `"Grafana Agent Observability"`             | Label shown in the Flyte UI                              |
| `app_id`            | `"grafana-agento11y-app"`                   | Grafana app plugin ID                                    |
| `conversation_path` | `"conversations/{conversation_id}/explore"` | Path template within the app                             |
| `list_path`         | `"conversations"`                           | Path of the conversations list, used for back navigation |
| `return_to`         | `True`                                      | Include the back-navigation parameter                    |
| `by_run`            | `True`                                      | Address the conversation by the Flyte run                |

Set `by_run=False` when something other than Flyte owns the conversation ID, typically alongside `bind_conversation=False`. The link then lands on the conversations list rather than on a URL that resolves to nothing.

> [!NOTE]
> The Grafana app moved from `grafana-sigil-app` to `grafana-agento11y-app`. The old ID still
> resolves but is deprecated, which is why the ID and both path templates are settable.

## Recording generations by hand

The client is available whether or not a framework integration is installed, so you can record generations explicitly. They still land inside the Flyte task span and still carry the run's identity because neither of those depends on a framework integration.

This is the path for an agent written against a provider SDK directly or for an adapter that has no agento11y package yet.

```python{hl_lines=[5,8]}
from agento11y import GenerationStart, ModelRef, assistant_text_message, user_text_message
from flyteplugins.agento11y import get_client

@flyte.trace
async def ask(question: str) -> str:
    """A durable model turn, recorded as a generation."""
    client = get_client()
    with client.start_generation(GenerationStart(model=ModelRef(provider="openai", name="gpt-4o"))) as rec:
        answer = await call_the_model(question)
        rec.set_result(
            input=[user_text_message(question)],
            output=[assistant_text_message(answer)],
        )
    return answer
```

Putting the call inside a **Tasks > Build tasks > Traces** step is what makes it durable: a resumed run replays the recorded result instead of calling the model again.

## Content capture

agento11y sends metadata by default (model, token usage, tool names, timing) and keeps prompts and responses local unless you opt in.

That is an agento11y setting rather than a Flyte one. `client_options` is the passthrough for it: every key becomes a field on agento11y's own `ClientConfig`, so content capture is switched on exactly as it would be outside Flyte. Check the [agento11y documentation](https://grafana.com/docs/grafana-cloud/monitor-applications/agent-observability/) for the current field names and defaults; they belong to that library, not to this plugin.

The same passthrough covers anything else `init()` does not surface: auth mode and token, protocol or a custom generation exporter.

```python
init(service_name="my-agent", client_options={"generation_exporter": MyExporter()})
```

`init()` sets three `ClientConfig` fields itself: `tracer` (so generations nest inside the Flyte task span), `generation_export_endpoint` (from `endpoint=`), and `generation_exporter` (a no-op when no endpoint is configured). Anything you put in `client_options` wins over all three.

## Instrumenting a different backend

`flyteplugins-agents-core` exposes two registries that let an out-of-tree package instrument the frameworks the adapters drive. They are how this plugin attaches agento11y's handlers to calls the adapter owns rather than you, and neither knows anything about Grafana. If you maintain instrumentation for a different vendor, register against the same hooks.

Use `register_instrumentor` when the framework accepts a handler in its run payload. The adapter offers you the framework-native payload and uses whatever you return:

```python
from flyteplugins.agents.core import register_instrumentor

def add_my_handler(config):
    config = dict(config or {})
    config.setdefault("callbacks", []).append(MyHandler())
    return config

register_instrumentor("langgraph", add_my_handler)
```

Use `register_call_wrapper` when the SDK cannot be instrumented by handing it an object, and the only way in is to wrap the call itself. That is the case for the Claude Agent SDK, whose model turns arrive as messages on the stream returned by `query`:

```python
from flyteplugins.agents.core import register_call_wrapper

def wrap(call):
    def instrumented(*args, **kwargs):
        return my_recording_query(_query_fn=call, **kwargs)

    return instrumented

register_call_wrapper("claude", wrap)
```

Framework names match the adapter directory: `langchain`, `langgraph`, `openai`, `claude`, `google`, `pydantic_ai`.

Both registries are best-effort by construction. If your instrumentor or wrapper raises, the adapter logs at debug level and runs the agent uninstrumented, because observing an agent must never be the reason it stops working. The flip side is that a handler which never attaches fails quietly, so confirm registration with `flyteplugins.agents.core.instrumented_frameworks()` rather than assuming it.

## Limitations

**`init()` must be called at module scope:** The task span opens before the task body runs, so initializing from inside the body means that task's span and the identity binding that rides on it have already been missed. `flyteplugins-otel` logs a warning when it detects this.

**Not every adapter has an integration:** crewai and mistral have Flyte adapters but no agento11y package, so their generations are not captured automatically.

**Short tasks need the exit flush:** agento11y batches generations and flushes on an interval and unlike OpenTelemetry's tracer provider it registers no exit hook of its own. The plugin registers one for a client it created, so a task that finishes inside the flush window does not lose its generations. A client you pass in with `client=` is yours to flush.

**With no endpoint configured, generations are dropped:** Without `AGENTO11Y_ENDPOINT` or `endpoint=`, the plugin installs a no-op exporter and logs a warning once. OpenTelemetry spans are unaffected and follow their own exporter settings.

## Related

- **[OpenTelemetry](../opentelemetry/_index)**: the tracing layer this plugin builds on.
- ****OpenTelemetry > Traces across crashes and resumes****: why a durable run needs more than a stock OpenTelemetry setup.
- **[Agent frameworks](../agents/_index)**: the `flyteplugins-agents-*` adapters this instruments.
- **[Build an agent](../../user-guide/agents/build-agent/_index)**: building the agent in the first place.

> [!NOTE] Runnable examples
> The plugin ships [worked examples](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/agento11y/examples)
> for OpenAI Agents, LangGraph, Claude, Google ADK, PydanticAI, manual generations and a
> crash-and-resume agent whose trace stays intact.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/hydra ===

# Hydra

[Hydra](https://hydra.cc) is a framework for composing and overriding configuration trees from YAML files, dataclasses and the command line. The `flyteplugins-hydra` plugin makes Hydra a first-class submission layer for Flyte, so you can compose a config exactly as you would in any other Hydra app and have each composed run executed as a Flyte task, locally or as a remote execution on a Flyte cluster.

The plugin offers three complementary entry points that share a single launcher implementation:

| Entry point                                    | Use it when                                                                                                                              |
| ---------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
| `hydra/launcher=flyte` (Hydra Launcher plugin) | You already have a `@hydra.main` script and want standard Hydra CLI ergonomics, including `--multirun` and custom sweepers.              |
| `flyte hydra run` (Flyte CLI extension)        | You want a Flyte-style CLI that imports a task from a Python file and composes a Hydra config without requiring a `@hydra.main` wrapper. |
| `hydra_run` / `hydra_sweep` (Python SDK)       | You want to submit runs directly from Python -- notebooks, tests, examples or another orchestration script.                              |

All three paths converge on the same `FlyteLauncher`.

## Installation

```bash
pip install flyteplugins-hydra
```

The plugin depends on `flyteplugins-omegaconf`, which is installed automatically and provides the `DictConfig`/`ListConfig` type transformers that allow Hydra-composed configs to flow into Flyte tasks. Both packages must be available in the same environment as `flyte`.

If you call `apply_task_env` for child tasks (see **Hydra > Task environment overrides**), include `flyteplugins-hydra` in the task image as well.

## Requirements on tasks

Every task launched through this plugin must accept an OmegaConf `DictConfig` input. Any other parameters are passed through as ordinary task arguments.

```python{hl_lines=[1, 5]}
from omegaconf import DictConfig

@env.task
async def pipeline(cfg: DictConfig, dataset: str) -> float:
    ...
```

The plugin auto-detects the `DictConfig` parameter name. If your parameter is `cfg`, app-level overrides are passed through `--cfg` on the CLI; if it is `config`, they are passed through `--config`; and so on.

## A walkthrough config

The examples in this page assume a small project layout:

```
project/
├── train.py
└── conf/
    ├── training.yaml
    ├── model/
    │   ├── resnet.yaml
    │   └── vit.yaml
    ├── optimizer/
    │   ├── adam.yaml
    │   └── sgd.yaml
    └── task_env/
        ├── a100.yaml
        └── prebuilt_image.yaml
```

`conf/training.yaml`:

```yaml
defaults:
  - optimizer: adam
  - model: resnet
  - _self_

data:
  path: s3://my-bucket/imagenet
  dataset: imagenet

training:
  epochs: 30
  batch_size: 64
```

`train.py` (abbreviated):

```python
import flyte
from omegaconf import DictConfig
from flyteplugins.hydra import apply_task_env

env = flyte.TaskEnvironment(name="training", image=...)

@env.task
async def preprocess(cfg: DictConfig) -> flyte.io.Dir: ...

@env.task
async def train_model(cfg: DictConfig, data: flyte.io.Dir) -> tuple[flyte.io.Dir, float]: ...

@env.task
async def pipeline(cfg: DictConfig, dataset: str) -> float:
    data = await preprocess(cfg)
    train_task = apply_task_env(train_model, cfg)
    _, val_loss = await train_task(cfg, data)
    return val_loss
```

The same `pipeline` task is the target of every example below.

> **📝 Note**
>
> `config_path` is resolved relative to the current working directory. If you submit runs from a directory other than `project/`, pass an absolute path (or an absolute path on the CLI via `--config-path /abs/path/to/conf`). For structured-config-only setups (no YAML files), omit `config_path` / `--config-path` entirely.

## Execution mode

Remote execution is the default. Every entry point exposes an explicit knob:

| Surface                | Local                       | Remote                                 |
| ---------------------- | --------------------------- | -------------------------------------- |
| `@hydra.main` launcher | `hydra.launcher.mode=local` | `hydra.launcher.mode=remote` (default) |
| `flyte hydra run`      | `--local`                   | `--mode remote` (default)              |
| Python SDK             | `mode="local"`              | `mode="remote"` (default)              |

For the `@hydra.main` launcher, the default applies as soon as `hydra/launcher=flyte` is selected.

Remote runs print the Flyte run URL immediately after submission, before any waiting. By default the plugin then waits for every submitted run to reach a terminal phase, capped at 32 worker threads. To tune or disable waiting:

| Surface                | Tune wait threads                    | Fire and forget             |
| ---------------------- | ------------------------------------ | --------------------------- |
| `@hydra.main` launcher | `hydra.launcher.wait_max_workers=64` | `hydra.launcher.wait=false` |
| `flyte hydra run`      | `--wait-max-workers 64`              | `--no-wait`                 |
| Python SDK             | `wait_max_workers=64`                | `wait=False`                |

For a sweep, every job is submitted first, and then the plugin waits for all runs concurrently. Submission is not blocked by earlier runs reaching a terminal phase.

## Hydra launcher (`@hydra.main` scripts)

Use this path when your script already has a `@hydra.main` entry point. Selecting `hydra/launcher=flyte` swaps Hydra's built-in `BasicLauncher` for `FlyteLauncher`.

Single remote run:

```bash
python train.py hydra/launcher=flyte hydra.launcher.mode=remote
```

Single local run:

```bash
python train.py hydra/launcher=flyte hydra.launcher.mode=local
```

Remote grid sweep submission: Each comma-separated value expands into a separate Flyte execution; six executions in this example:

```bash{hl_lines=[4]}
python train.py --multirun \
  hydra/launcher=flyte hydra.launcher.mode=remote \
  hydra.launcher.wait_max_workers=64 \
  optimizer.lr=0.001,0.01,0.1 training.epochs=10,20
```

Fire-and-forget sweep submission:

```bash{hl_lines=[2]}
python train.py --multirun \
  hydra/launcher=flyte hydra.launcher.wait=false \
  optimizer.lr=0.001,0.01,0.1
```

Custom sweepers (Optuna) work exactly as they do with the BasicLauncher. Selecting `hydra/sweeper=...` activates the sweeper and `FlyteLauncher` runs each trial as a Flyte execution:

```bash{hl_lines=["3-5"]}
python train.py --multirun \
  hydra/launcher=flyte hydra.launcher.mode=remote \
  hydra/sweeper=optuna hydra.sweeper.n_trials=20 \
  hydra.sweeper.n_jobs=4 \
  "optimizer.lr=interval(1e-4,1e-1)"
```

Inside `@hydra.main`, the standard pattern is:

```python{hl_lines=[7]}
import flyte
import hydra
from omegaconf import DictConfig
from flyteplugins.hydra import apply_task_env

@hydra.main(version_base=None, config_path="conf", config_name="training")
def main(cfg: DictConfig):
    flyte.init_from_config()
    entry_task = apply_task_env(pipeline, cfg)
    return flyte.run(entry_task, cfg=cfg, dataset=cfg.data.dataset)

if __name__ == "__main__":
    main()
```

## Python SDK

`hydra_run` composes one config and runs the task once. `hydra_sweep` expands sweep overrides and runs the task once per combination.

### Single run

```python{hl_lines=[1, 3, 7]}
from flyteplugins.hydra import hydra_run

run = hydra_run(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.01"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
    wait=True,
    wait_max_workers=64,
)
```

For a remote run with `wait=True`, the return value is a wrapper exposing both `run.url` and `run.value` (the resolved task output). The wrapper is `float()`-castable so Hydra sweepers such as Optuna can consume scalar objectives directly. With `wait=False`, the return value is the underlying `flyte.remote.Run`.

### Grid sweep

```python{hl_lines=[7]}
from flyteplugins.hydra import hydra_sweep

runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.001,0.01,0.1", "training.epochs=10,20"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
)
```

Six executions are submitted (3 × 2). `runs` is a list aligned with the Cartesian-product order Hydra's `BasicSweeper` produces.

### Custom sweepers

Custom sweeper plugins are activated by passing their selection in `overrides`:

```python{hl_lines=["5-10"]}
runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=[
        "hydra/sweeper=optuna",
        "hydra.sweeper.n_trials=20",
        "hydra.sweeper.n_jobs=4",
        "optimizer.lr=interval(1e-4,1e-1)",
    ],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
)
```

Whenever an override starts with `hydra/`, the plugin invokes the full Hydra runtime so plugin discovery (sweepers, launchers, callbacks) can run. Pure value overrides on the `hydra.*` namespace (for example `hydra.run.dir=...`) do not need the full runtime and are applied per-job by the launcher directly.

### Forwarding `flyte.with_runcontext` options

Use `run_options` to pass Flyte runtime options through to every job:

```python{hl_lines=["8-14"]}
runs = hydra_sweep(
    pipeline,
    config_path="conf",
    config_name="training",
    overrides=["optimizer.lr=0.001,0.01,0.1"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
    run_options={
        "name": "my-training-sweep",
        "service_account": "default",
        "copy_style": "all",
        "raw_data_path": "s3://my-bucket/raw-data",
        "debug": True,
    },
)
```

## Flyte CLI (`flyte hydra run`)

`flyte hydra run` is registered through the `flyte.plugins.cli.commands` entry point. It loads a task from a Python file, composes a Hydra config, and runs the task without requiring the script to have its own `@hydra.main` function. It also inherits the relevant flags from `flyte run` (`--project`, `--domain`, `--image`, `--name`, `--service-account`, `--raw-data-path`, `--copy-style`, `--debug`, `--local`, `--follow`).

### Single run

Remote (default):

```bash
flyte hydra run --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet
```

Forced local:

```bash{hl_lines=[1]}
flyte hydra run --local --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet
```

### Grid sweep

```bash{hl_lines=[4]}
flyte hydra run --multirun --config-path conf --config-name training \
  --wait-max-workers 64 \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --cfg "optimizer.lr=0.001,0.01,0.1" --cfg "training.epochs=10,20"
```

### App-level vs Hydra-namespace overrides

The CLI keeps app-level overrides separate from Hydra runtime overrides so they do not collide with ordinary Flyte task arguments.

App-level overrides target the composed config and are passed through the **task's `DictConfig` parameter name**. For `pipeline(cfg: DictConfig, ...)`, use `--cfg`. For `pipeline_with_config(config: DictConfig, ...)`, use `--config`:

```bash{hl_lines=["3-4", 8]}
flyte hydra run --config-path conf --config-name training \
  train.py pipeline \
  --cfg optimizer.lr=0.01 \
  --cfg training.epochs=20

flyte hydra run --config-path conf --config-name training \
  train.py pipeline_with_config \
  --config optimizer.lr=0.01
```

Hydra runtime overrides: Anything in the `hydra.*` or `hydra/*` namespace go through `--hydra-override`:

```bash{hl_lines=[3, 4]}
flyte hydra run --config-path conf --config-name training \
  train.py pipeline \
  --hydra-override hydra.run.dir=./outputs/exp1 \
  --hydra-override hydra/launcher=flyte
```

Custom sweepers combine the two:

```bash{hl_lines=["3-7"]}
flyte hydra run --multirun --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --hydra-override hydra/sweeper=optuna \
  --hydra-override hydra.sweeper.n_trials=20 \
  --hydra-override hydra.sweeper.n_jobs=4 \
  --cfg "optimizer.lr=interval(1e-4,1e-1)" \
  --cfg "training.epochs=choice(10,20,50)"
```

### `--follow` and `--no-wait`

`--follow` streams logs from the launched run after submission; it implies waiting and cannot be combined with `--no-wait`. `--no-wait` returns immediately after submission and skips log streaming.

### Shell completion

Install Click's completion hook for the `flyte` executable. For zsh:

```zsh
eval "$(_FLYTE_COMPLETE=zsh_source flyte)"
```

For bash:

```bash
eval "$(_FLYTE_COMPLETE=bash_source flyte)"
```

Once installed, `flyte hydra run` adds Hydra-aware completion after `SCRIPT TASK_NAME`. The command imports the script, inspects the task signature, and suggests:

- The app override flag matching the task's `DictConfig` parameter (`--cfg`, `--config`, ...).
- Override values for that flag and `--hydra-override` via Hydra's own completion engine, including config keys, config-group selections and sweep functions.

```bash{hl_lines=["2-3", "6-7"]}
flyte hydra run --config-path conf --config-name training \
  train.py pipeline --cfg optimizer.<TAB>
# suggests optimizer.lr=, optimizer.weight_decay=, ...

flyte hydra run --config-path conf --config-name training \
  train.py pipeline --hydra-override hydra/launcher=<TAB>
# suggests hydra launcher choices
```

Because completion has to import the target script, keep task definitions and `ConfigStore` registration import-safe, and avoid expensive top-level work in scripts you reach via `flyte hydra run`.

![Auto Completion](../../_static/images/integrations/hydra/auto_complete.gif)

## Override grammar

The override grammar is identical to standard Hydra; what differs is only how you pass the strings (positional in `python train.py ...`, list entries in `overrides=[...]`, repeated `--cfg`/`--hydra-override` on the Flyte CLI).

| Form                               | Meaning                                                                                  |
| ---------------------------------- | ---------------------------------------------------------------------------------------- |
| `optimizer.lr=0.01`                | Set an existing key.                                                                     |
| `optimizer=sgd`                    | Select a config group (replaces the `optimizer` subtree with `conf/optimizer/sgd.yaml`). |
| `+task_env=a100`                   | Append a config group whose key is not currently in the config.                          |
| `+training.grad_clip=1.0`          | Append a key that does not exist.                                                        |
| `++optimizer.lr=0.05`              | Force-set a key, creating it if missing and overriding strict-schema errors.             |
| `~training.warmup_steps`           | Delete a key from the composed config.                                                   |
| `optimizer.lr=0.001,0.01,0.1`      | Sweep value (with `--multirun`); expanded into one job per element.                      |
| `optimizer.lr=interval(1e-4,1e-1)` | Continuous sweep range; consumed by samplers like Optuna.                                |
| `optimizer=choice(adam,sgd)`       | Categorical sweep; consumed by samplers.                                                 |
| `hydra.run.dir=./outputs/exp1`     | Hydra-namespace value override (single run output dir).                                  |
| `hydra.sweep.dir=./outputs/sweep1` | Hydra-namespace sweep output dir.                                                        |
| `hydra/sweeper=optuna`             | Hydra-namespace config group selection (activates the Optuna sweeper plugin).            |

## Sweeps

### Grid sweeps (BasicSweeper)

Comma-separated overrides expand into a Cartesian product. The plugin uses Hydra's `BasicSweeper` to expand them, then submits one Flyte execution per combination.

```python{hl_lines=[1, 4, 7]}
from flyteplugins.hydra import hydra_sweep

runs = hydra_sweep(
    pipeline,
    config_path="conf", config_name="training",
    overrides=["model=resnet,vit", "optimizer.lr=0.001,0.01,0.1"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
)  # 6 executions
```

```bash{hl_lines=[3]}
flyte hydra run --multirun --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --cfg "model=resnet,vit" --cfg "optimizer.lr=0.001,0.01,0.1"
```

Hardware presets can sweep alongside hyperparameters:

```bash{hl_lines=[3]}
flyte hydra run --multirun --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --cfg "+task_env=a10g,a100" --cfg "optimizer.lr=0.001,0.01,0.1"
```

### Bayesian / TPE sweeps (Optuna)

Install the sweeper, then activate it via `hydra/sweeper=optuna`. Continuous parameters use `interval(...)`; categorical parameters use `choice(...)`.

```bash
pip install hydra-optuna-sweeper
```

```bash{hl_lines=["3-8"]}
flyte hydra run --multirun --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --hydra-override "hydra/sweeper=optuna" \
  --hydra-override "hydra.sweeper.n_trials=30" \
  --hydra-override "hydra.sweeper.n_jobs=5" \
  --cfg "optimizer.lr=interval(1e-4,1e-1)" \
  --cfg "optimizer.weight_decay=interval(1e-6,1e-2)" \
  --cfg "model=choice(resnet,vit)"
```

When `wait=True`, each remote run's wrapped result exposes the task output as a float (via `__float__`), so Optuna can use it directly as the trial objective. With `wait=False`, the sweeper sees the run URL but cannot read objective values; use this only for fire-and-forget submission.

Other sweepers that respect Hydra's plugin protocol are activated the same way: install the package, select `hydra/sweeper=<name>`, and set the sweeper's parameters under `hydra.sweeper.*`.

### Sweep output directories

Hydra-namespace overrides redirect where Hydra writes per-job logs and config snapshots:

```bash{hl_lines=[3, 4]}
flyte hydra run --multirun --config-path conf --config-name training \
  train.py pipeline --dataset s3://my-bucket/imagenet \
  --hydra-override "hydra.sweep.dir=./outputs/sweep1" \
  --hydra-override "hydra.sweep.subdir=\${hydra.job.num}" \
  --cfg "optimizer.lr=0.001,0.01,0.1"
```

## Task environment overrides

Hydra is good at composing flat YAML; Flyte tasks need richer settings such as resources and container images. The plugin reserves a config key named `task_env` by default that maps task names to `task.override` kwargs.

```yaml
task_env:
  pipeline:
    resources:
      cpu: "2"
      memory: 8Gi
  train_model:
    resources:
      cpu: "16"
      memory: 64Gi
      gpu: "A100:1"
```

When the plugin launches a task, it looks up `task_env[<entry-task-name>]` (`pipeline` in this example) and applies the values via `task.override(...)`. Resource mappings are converted into `flyte.Resources(**values)` automatically.

### Prebuilt images

To run a task in a prebuilt container image, set `image` (and optionally `primary_container_name`):

```yaml{hl_lines=[3]}
task_env:
  pipeline:
    image: ghcr.io/acme/flyte-training:latest
    primary_container_name: main
    resources:
      cpu: "4"
      memory: 16Gi
```

`task.override` does not accept `image` directly. The task image is part of the task definition. Instead, the plugin lowers the override to a `flyte.PodTemplate` whose primary container uses the requested image:

- If the task has no inline pod template, a new one is created.
- If the task already has an inline `flyte.PodTemplate`, the plugin deep-copies it and sets only the image on the primary container.
- If the task references a pod template by name (a string), the plugin raises an error. You must patch a string-named template by editing it in cluster config rather than at submission time.

### Applying overrides to child tasks

The launcher only controls the entry task it submits. Child tasks called from within the entry task are not patched automatically. Use `apply_task_env` to apply the same `resources`/`image` handling to a child task before invoking it:

```python{hl_lines=[1, 7]}
from flyteplugins.hydra import apply_task_env

@env.task
async def pipeline(cfg: DictConfig, dataset: str) -> float:
    data = await preprocess(cfg)
    train_task = apply_task_env(train_model, cfg)
    _, val_loss = await train_task(cfg, data)
    return val_loss
```

This keeps the override knobs in YAML/CLI surfaces while leaving each task in control of which children it patches.

### Renaming the task-env key

If your config uses a different name for the task-env subtree, pass it explicitly:

```python
hydra_run(..., task_env_key="task_environment")
```

```bash
flyte hydra run --task-env-key task_environment ...
```

### What `task_env` should not model

The YAML schema intentionally omits the full Kubernetes `V1PodSpec`. Keep advanced pod configuration (volumes, init containers, node selectors, etc.) in Python task/environment code where you have a real type. Use Hydra `task_env` presets for the common knobs only: image, primary container name and resources.

## Structured configs (without YAML)

Structured configs work with this plugin as long as they are registered before the launcher composes the config. `flyte hydra run` imports the script first, so top-level `ConfigStore.instance().store(...)` calls run before composition.

```python{hl_lines=[17]}
from dataclasses import dataclass, field
from hydra.core.config_store import ConfigStore
from omegaconf import DictConfig

@dataclass
class TrainingConf:
    epochs: int = 30
    batch_size: int = 64

@dataclass
class RootConf:
    training: TrainingConf = field(default_factory=TrainingConf)

ConfigStore.instance().store(name="structured_training", node=RootConf)
```

Run a fully-structured config without YAML:

```bash{hl_lines=[1]}
flyte hydra run --config-name structured_training \
  train.py pipeline --dataset s3://my-bucket/imagenet
```

The same config also works through `@hydra.main`:

```bash
python train.py --config-name structured_training
```

If the structured config still references YAML config groups, keep `--config-path conf`. If everything is registered in `ConfigStore`, omit `--config-path`.

> **⚠️ Warning**
>
> Do not register structured configs only inside `if __name__ == "__main__":` or inside the `@hydra.main` function body. `flyte hydra run` and shell completion inspect the script at import time, before either of those blocks runs, and registrations placed there will not be visible.

Structured configs sweep just like YAML configs:

```python{hl_lines=[4, 5]}
runs = hydra_sweep(
    pipeline,
    config_path=None,
    config_name="structured_training",
    overrides=["training.epochs=10,20", "training.batch_size=32,64"],
    dataset="s3://my-bucket/imagenet",
    mode="remote",
)
```


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/jsonl ===

# JSONL

The JSONL plugin adds two typed I/O types for working with [JSON Lines](https://jsonlines.org/) data as task inputs and outputs: `flyteplugins.jsonl.JsonlFile` for a single JSONL file and `flyteplugins.jsonl.JsonlDir` for a directory of sharded JSONL files. Both are backed by [`orjson`](https://github.com/ijl/orjson) for fast serialization and stream records one at a time, so you can process datasets that don't fit in memory.

`JsonlFile` and `JsonlDir` extend the built-in `flyte.io.File` and `flyte.io.Dir` types, so they inherit remote-storage, upload/download, and caching behavior. They simply add JSONL-aware streaming readers and writers on top. Every read/write method has a synchronous `_sync` counterpart (`writer_sync()`, `iter_records_sync()`) for use in non-`async` tasks.

## When to use this plugin

- Passing line-delimited JSON datasets (LLM training/eval sets, event logs, model outputs) between tasks
- Streaming records without loading an entire file into memory
- Writing large outputs as automatically rotated, sharded directories
- Working with compressed JSONL (`.jsonl.zst`) transparently

## Installation

```bash
pip install flyteplugins-jsonl
```

Add the plugin to your task image. Installing it registers `JsonlFile` and `JsonlDir` with the Flyte type engine automatically. No explicit registration call is needed:

```
import flyte
from flyteplugins.jsonl import JsonlDir, JsonlFile

env = flyte.TaskEnvironment(
    name="jsonl-examples",
    image=flyte.Image.from_debian_base(name="jsonl").with_pip_packages(
        "flyteplugins-jsonl"
    ),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## Working with `JsonlFile`

Create a writable file reference with `JsonlFile.new_remote()`, then stream records through the `writer()` context manager without holding the whole dataset in memory:

```
@env.task
async def write_records() -> JsonlFile:
    """Write records to a single JSONL file."""
    out = JsonlFile.new_remote("results.jsonl")
    async with out.writer() as writer:
        for i in range(500_000):
            await writer.write({"id": i, "score": i * 0.1})
    return out
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

Reading is equally streaming. `iter_records()` yields one parsed `dict` per line:

```
@env.task
async def read_records(data: JsonlFile) -> int:
    """Read records from a JsonlFile and return the count."""
    count = 0
    async for record in data.iter_records():
        count += 1
    return count
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## Working with `JsonlDir`

`JsonlDir` writes a directory of shard files (`part-00000.jsonl`, `part-00001.jsonl`, …) and reads them back transparently in sorted order. Pass `max_records_per_shard` (or `max_bytes_per_shard`) to control shard rotation:

```
@env.task
async def write_large_dataset() -> JsonlDir:
    """Write a large dataset to a sharded JsonlDir.

    JsonlDir automatically rotates to a new shard file once the
    current shard reaches the record or byte limit. Shards are named
    part-00000.jsonl, part-00001.jsonl, etc.
    """
    out = JsonlDir.new_remote("dataset/")
    async with out.writer(
        max_records_per_shard=100_000,
        max_bytes_per_shard=256 * 1024 * 1024,  # 256 MB
    ) as writer:
        for i in range(500_000):
            await writer.write({"index": i, "value": i * i})
    return out
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

Reading iterates across all shards transparently, prefetching the next shard in the background to overlap network I/O with processing:

```
@env.task
async def sum_values(dataset: JsonlDir) -> int:
    """Read all records across all shards and compute a sum.

    Iteration is transparent across shards and handles mixed
    compressed/uncompressed shards automatically. The next shard is
    prefetched in the background for higher throughput.
    """
    total = 0
    async for record in dataset.iter_records():
        total += record["value"]
    return total
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

For bulk processing, `iter_batches()` yields lists of records at a time; `JsonlDir` also inherits all `flyte.io.Dir` capabilities (`walk()`, `list_files()`, `download()`):

```
@env.task
async def process_in_batches(dataset: JsonlDir) -> int:
    """Process records in batches of dicts for bulk operations."""
    total = 0
    async for batch in dataset.iter_batches(batch_size=1000):
        # Each batch is a list[dict]
        total += len(batch)
    return total
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## Configuration and options

### Compression

Give the file a `.jsonl.zst` (or `.jsonl.zstd`) extension and records are zstd-compressed transparently on write and decompressed on read. Tune the level via the writer:

```
@env.task
async def write_compressed() -> JsonlFile:
    """Write a zstd-compressed JSONL file.

    Compression is activated by using a .jsonl.zst extension.
    Both reading and writing handle compression transparently.
    """
    out = JsonlFile.new_remote("results.jsonl.zst")
    async with out.writer(compression_level=3) as writer:
        for i in range(100_000):
            await writer.write({"id": i, "compressed": True})
    return out
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

For `JsonlDir`, set `shard_extension=".jsonl.zst"` on `writer()`. Mixed compressed and uncompressed shards within a directory are supported on read:

```
@env.task
async def write_compressed_dir() -> JsonlDir:
    """Write zstd-compressed shards by specifying the shard extension."""
    out = JsonlDir.new_remote("compressed_dataset/")
    async with out.writer(
        shard_extension=".jsonl.zst",
        max_records_per_shard=50_000,
    ) as writer:
        for i in range(200_000):
            await writer.write({"id": i, "data": f"payload-{i}"})
    return out
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

### Error handling on read

The record iterators accept an `on_error` argument: `"raise"` (default), `"skip"` to drop malformed lines, or a callable `(line_number, raw_line, exception) -> None` for custom handling:

```
@env.task
async def read_with_error_handling(data: JsonlFile) -> int:
    """Read records, skipping any corrupt lines instead of raising."""
    count = 0
    async for record in data.iter_records(on_error="skip"):
        count += 1
    return count

@env.task
async def read_with_custom_handler(data: JsonlFile) -> int:
    """Read records with a custom error handler that collects errors."""
    errors: list[dict] = []

    def on_error(line_number: int, raw_line: bytes, exc: Exception) -> None:
        errors.append({"line": line_number, "error": str(exc)})

    count = 0
    async for record in data.iter_records(on_error=on_error):
        count += 1
    print(f"{count} valid records, {len(errors)} errors")
    return count
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

### Arrow batches

To hand JSONL data to columnar tooling, stream it as Arrow `RecordBatch`es with `iter_arrow_batches(batch_size=...)`. Memory usage stays bounded by the batch size. Arrow iteration requires the optional `pyarrow` dependency. Install it with `pip install 'flyteplugins-jsonl[arrow]'`:

```
arrow_env = flyte.TaskEnvironment(
    name="jsonl-arrow",
    image=flyte.Image.from_debian_base(name="jsonl-arrow").with_pip_packages(
        "flyteplugins-jsonl[arrow]"
    ),
)

@arrow_env.task
async def analyze_with_arrow(dataset: JsonlDir) -> float:
    """Stream records as Arrow RecordBatches for analytics.

    Memory usage is bounded by batch_size — the full dataset is
    never loaded into memory at once.
    """
    import pyarrow as pa

    batches = []
    async for batch in dataset.iter_arrow_batches(batch_size=65_536):
        batches.append(batch)

    table = pa.Table.from_batches(batches)
    mean_value = table.column("value").to_pylist()
    return sum(mean_value) / len(mean_value)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## Common use cases

- **LLM dataset pipelines**: stream prompt/completion or eval records between preprocessing, generation, and scoring tasks.
- **Event and log processing**: read large line-delimited logs shard by shard without buffering the whole file.
- **Fan-out writes**: produce a `JsonlDir` of rotated shards from a task that emits millions of records, then consume it downstream.

## API reference

See the [JSONL API reference](../../api-reference/integrations/jsonl/_index) for the full `JsonlFile` and `JsonlDir` method listings.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/mlflow ===

# MLflow

The MLflow plugin integrates [MLflow](https://mlflow.org/) experiment tracking with Flyte. It provides a `@mlflow_run` decorator that automatically manages MLflow runs within Flyte tasks, with support for autologging, parent-child run sharing, distributed training, and auto-generated UI links.

The decorator works with both sync and async tasks.

## Installation

```bash
pip install flyteplugins-mlflow
```

Requires `mlflow` and `flyte`.

## Quick start

```python{hl_lines=[3, 9, "13-16", 22]}
import flyte
import mlflow
from flyteplugins.mlflow import mlflow_run, get_mlflow_run

env = flyte.TaskEnvironment(
    name="mlflow-tracking",
    resources=flyte.Resources(cpu=1, memory="500Mi"),
    image=flyte.Image.from_debian_base(name="mlflow_example").with_pip_packages(
        "flyteplugins-mlflow"
    ),
)

@mlflow_run(
    tracking_uri="http://localhost:5000",
    experiment_name="my-experiment",
)
@env.task
async def train_model(learning_rate: float) -> str:
    mlflow.log_param("lr", learning_rate)
    mlflow.log_metric("loss", 0.42)

    run = get_mlflow_run()
    return run.info.run_id
```

![Link](../../_static/images/integrations/mlflow/link.png)

![Mlflow UI](../../_static/images/integrations/mlflow/mlflow_dashboard.png)

> [!NOTE]
> `@mlflow_run` must be the outermost decorator, before `@env.task`:
>
> ```python{hl_lines=["1-2"]}
> @mlflow_run          # outermost
> @env.task            # innermost
> async def my_task(): ...
> ```

## Autologging

Enable MLflow's autologging to automatically capture parameters, metrics, and models without manual `mlflow.log_*` calls.

### Generic autologging

```python{hl_lines=[1]}
@mlflow_run(autolog=True)
@env.task
async def train():
    from sklearn.linear_model import LogisticRegression

    model = LogisticRegression()
    model.fit(X, y)  # Parameters, metrics, and model are logged automatically
```

### Framework-specific autologging

Pass `framework` to use a framework-specific autolog implementation:

```python{hl_lines=[3]}
@mlflow_run(
    autolog=True,
    framework="sklearn",
    log_models=True,
    log_datasets=False,
)
@env.task
async def train_sklearn():
    from sklearn.ensemble import RandomForestClassifier

    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)
```

Supported frameworks include any framework with an `mlflow.{framework}.autolog()` function. You can find the [full list of supported frameworks](https://mlflow.org/docs/latest/ml/tracking/autolog/#supported-libraries) in the MLflow documentation.

You can pass additional autolog parameters via `autolog_kwargs`:

```python{hl_lines=[4]}
@mlflow_run(
    autolog=True,
    framework="pytorch",
    autolog_kwargs={"log_every_n_epoch": 5},
)
@env.task
async def train_pytorch():
    ...
```

![Autolog](../../_static/images/integrations/mlflow/autolog.png)

## Run modes

The `run_mode` parameter controls how MLflow runs are created and shared across tasks:

| Mode               | Behavior                                                              |
| ------------------ | --------------------------------------------------------------------- |
| `"auto"` (default) | Reuse the parent's run if one exists, otherwise create a new run      |
| `"new"`            | Always create a new independent run                                   |
| `"nested"`         | Create a new run nested under the parent via `mlflow.parentRunId` tag |

### Sharing a run across tasks

With `run_mode="auto"` (the default), child tasks reuse the parent's MLflow run:

```python{hl_lines=[1, 5, 7]}
@mlflow_run
@env.task
async def parent_task():
    mlflow.log_param("stage", "parent")
    await child_task()  # Shares the same MLflow run

@mlflow_run
@env.task
async def child_task():
    mlflow.log_metric("child_metric", 1.0)  # Logged to the parent's run
```

### Creating independent runs

Use `run_mode="new"` when a task should always create its own top-level MLflow run, completely independent of any parent:

```python{hl_lines=[1]}
@mlflow_run(run_mode="new")
@env.task
async def standalone_experiment():
    mlflow.log_param("experiment_type", "baseline")
    mlflow.log_metric("accuracy", 0.95)
```

### Nested runs

Use `run_mode="nested"` to create a child run that appears under the parent in the MLflow UI. This works across processes and containers via the `mlflow.parentRunId` tag.

![Nested runs](../../_static/images/integrations/mlflow/mlflow_hpo.png)

This is the recommended pattern for hyperparameter optimization, where each trial should be tracked as a child of the parent study run:

```python{hl_lines=[1, 2, 15, "22-25"]}
from flyteplugins.mlflow import Mlflow

@mlflow_run(run_mode="nested")
@env.task(links=[Mlflow()])
async def run_trial(trial_number: int, n_estimators: int, max_depth: int) -> float:
    """Each trial creates a nested MLflow run under the parent."""
    mlflow.log_params({"n_estimators": n_estimators, "max_depth": max_depth})
    mlflow.log_param("trial_number", trial_number)

    model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    rmse = float(np.sqrt(mean_squared_error(y_val, model.predict(X_val))))
    mlflow.log_metric("rmse", rmse)
    return rmse

@mlflow_run
@env.task
async def hpo_search(n_trials: int = 30) -> str:
    """Parent run tracks the overall study."""
    run = get_mlflow_run()
    mlflow.log_param("n_trials", n_trials)

    # Run trials in parallel — each gets a nested MLflow run
    rmses = await asyncio.gather(
        *(run_trial(trial_number=i, **params) for i, params in enumerate(trial_params))
    )

    mlflow.log_metric("best_rmse", min(rmses))
    return run.info.run_id
```

![HPO](../../_static/images/integrations/mlflow/hpo.png)

## Workflow-level configuration

Use `mlflow_config()` with `flyte.with_runcontext()` to set MLflow configuration for an entire workflow. All `@mlflow_run`-decorated tasks in the workflow inherit these settings:

```python{hl_lines=[1, "4-8"]}
from flyteplugins.mlflow import mlflow_config

r = flyte.with_runcontext(
    custom_context=mlflow_config(
        tracking_uri="http://localhost:5000",
        experiment_id="846992856162999",
        tags={"team": "ml"},
    )
).run(train_model, learning_rate=0.001)
```

This eliminates the need to repeat `tracking_uri` and experiment settings on every `@mlflow_run` decorator.

### Per-task overrides

Use `mlflow_config()` as a context manager inside a task to override configuration for specific child tasks:

```python{hl_lines=[6]}
@mlflow_run
@env.task
async def parent_task():
    await shared_child()  # Inherits parent config

    with mlflow_config(run_mode="new", tags={"role": "independent"}):
        await independent_child()  # Gets its own run
```

### Configuration priority

Settings are resolved in priority order:

1. Explicit `@mlflow_run` decorator arguments
2. `mlflow_config()` context configuration
3. Environment variables (for `tracking_uri`)
4. MLflow defaults

## Distributed training

In distributed training, only rank 0 logs to MLflow by default. The plugin detects rank automatically from the `RANK` environment variable:

```python{hl_lines=[1, "4-6"]}
@mlflow_run
@env.task
async def distributed_train():
    # Only rank 0 creates an MLflow run and logs metrics.
    # Other ranks execute the task function directly without
    # creating an MLflow run or incurring any MLflow overhead.
    ...
```

On non-rank-0 workers, no MLflow run is created and `get_mlflow_run()` returns `None`. The task function still executes normally; only the MLflow instrumentation is skipped.

![Distributed training](../../_static/images/integrations/mlflow/distributed_training.png)

You can also set rank explicitly:

```python{hl_lines=[1]}
@mlflow_run(rank=0)
@env.task
async def train():
    ...
```

## MLflow UI links

The `Mlflow` link class displays links to the MLflow UI in the Flyte UI.

Since the MLflow run is created inside the task at execution time, the run URL cannot be determined before the task starts. Links are only shown when a run URL is already available from context, either because a parent task created the run, or because an explicit URL is provided.

The recommended pattern is for the parent task to create the MLflow run, and child tasks that inherit the run (via `run_mode="auto"`) display the link to that run. For nested runs (`run_mode="nested"`), children display a link to the parent run.

### Setup

Set `link_host` via `mlflow_config()` and attach `Mlflow()` links to child tasks:

```python{hl_lines=[4, 17]}
from flyteplugins.mlflow import Mlflow, mlflow_config

@mlflow_run
@env.task(links=[Mlflow()])
async def child_task():
    ...  # Link points to the parent's MLflow run

@mlflow_run
@env.task
async def parent_task():
    await child_task()

if __name__ == "__main__":
    r = flyte.with_runcontext(
        custom_context=mlflow_config(
            tracking_uri="http://localhost:5000",
            link_host="http://localhost:5000",
        )
    ).run(parent_task)
```

> [!NOTE]
> `Mlflow()` is instantiated without a `link` argument because the URL is auto-generated at runtime. When the parent task creates an MLflow run, the plugin builds the URL from `link_host` and the run's experiment/run IDs, then propagates it to child tasks via the Flyte context. Passing an explicit `link` would bypass this auto-generation.

### Custom URL templates

The default link format is:

```
{host}/#/experiments/{experiment_id}/runs/{run_id}
```

For platforms like Databricks that use a different URL structure, provide a custom template:

```python{hl_lines=[3]}
mlflow_config(
    link_host="https://dbc-xxx.cloud.databricks.com",
    link_template="{host}/ml/experiments/{experiment_id}/runs/{run_id}",
)
```

### Explicit links

If you know the run URL ahead of time, you can set it directly:

```python{hl_lines=[1]}
@env.task(links=[Mlflow(link="https://mlflow.example.com/#/experiments/1/runs/abc123")])
async def my_task():
    ...
```

### Link behavior by run mode

| Run mode   | Link behavior                                                                                  |
| ---------- | ---------------------------------------------------------------------------------------------- |
| `"auto"`   | Parent link propagates to child tasks sharing the run                                          |
| `"new"`    | Parent link is cleared; no link is shown until the task's own run is available to its children |
| `"nested"` | Parent link is kept and renamed to "MLflow (parent)"                                           |

## Automatic Flyte tags

When running inside Flyte, the plugin automatically tags MLflow runs with execution metadata:

| Tag                 | Description      |
| ------------------- | ---------------- |
| `flyte.action_name` | Task action name |
| `flyte.run_name`    | Flyte run name   |
| `flyte.project`     | Flyte project    |
| `flyte.domain`      | Flyte domain     |

These tags are merged with any user-provided tags.

## API reference

### `mlflow_run` and `mlflow_config`

`mlflow_run` is a decorator that manages MLflow runs for Flyte tasks. `mlflow_config` creates workflow-level configuration or per-task overrides. Both accept the same core parameters:

| Parameter         | Type             | Default  | Description                                                                   |
| ----------------- | ---------------- | -------- | ----------------------------------------------------------------------------- |
| `run_mode`        | `str`            | `"auto"` | `"auto"`, `"new"`, or `"nested"`                                              |
| `tracking_uri`    | `str`            | `None`   | MLflow tracking server URL                                                    |
| `experiment_name` | `str`            | `None`   | MLflow experiment name (raises `ValueError` if combined with `experiment_id`) |
| `experiment_id`   | `str`            | `None`   | MLflow experiment ID (raises `ValueError` if combined with `experiment_name`) |
| `run_name`        | `str`            | `None`   | Human-readable run name (raises `ValueError` if combined with `run_id`)       |
| `run_id`          | `str`            | `None`   | Explicit MLflow run ID (raises `ValueError` if combined with `run_name`)      |
| `tags`            | `dict[str, str]` | `None`   | Tags for the run                                                              |
| `autolog`         | `bool`           | `False`  | Enable MLflow autologging                                                     |
| `framework`       | `str`            | `None`   | Framework for autolog (e.g. `"sklearn"`, `"pytorch"`)                         |
| `log_models`      | `bool`           | `None`   | Log models automatically (requires `autolog`)                                 |
| `log_datasets`    | `bool`           | `None`   | Log datasets automatically (requires `autolog`)                               |
| `autolog_kwargs`  | `dict`           | `None`   | Extra parameters for `mlflow.autolog()`                                       |

Additional keyword arguments are passed to `mlflow.start_run()`.

`mlflow_run` also accepts:

| Parameter | Type  | Default | Description                                              |
| --------- | ----- | ------- | -------------------------------------------------------- |
| `rank`    | `int` | `None`  | Process rank for distributed training (only rank 0 logs) |

`mlflow_config` also accepts:

| Parameter       | Type  | Default | Description                                                                 |
| --------------- | ----- | ------- | --------------------------------------------------------------------------- |
| `link_host`     | `str` | `None`  | MLflow UI host for auto-generating links                                    |
| `link_template` | `str` | `None`  | Custom URL template (placeholders: `{host}`, `{experiment_id}`, `{run_id}`) |

### `get_mlflow_run`

Returns the current `mlflow.ActiveRun` if within a `@mlflow_run`-decorated task. Returns `None` otherwise.

```python
from flyteplugins.mlflow import get_mlflow_run

run = get_mlflow_run()
if run:
    print(run.info.run_id)
```

### `get_mlflow_context`

Returns the current `mlflow_config` settings from the Flyte context, or `None` if no MLflow configuration is set. Useful for inspecting the inherited configuration inside a task:

```python
from flyteplugins.mlflow import get_mlflow_context

@mlflow_run
@env.task
async def my_task():
    config = get_mlflow_context()
    if config:
        print(config.tracking_uri, config.experiment_id)
```

### `Mlflow`

Link class for displaying MLflow UI links in the Flyte console.

| Field  | Type  | Default    | Description                             |
| ------ | ----- | ---------- | --------------------------------------- |
| `name` | `str` | `"MLflow"` | Display name for the link               |
| `link` | `str` | `""`       | Explicit URL (bypasses auto-generation) |


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/omegaconf ===

# OmegaConf

[OmegaConf](https://omegaconf.readthedocs.io/) is a hierarchical configuration system used by many ML frameworks (and the foundation of [Hydra](../hydra/_index)). The `flyteplugins-omegaconf` plugin makes OmegaConf's `DictConfig` and `ListConfig` first-class types in Flyte tasks, so you can pass entire configs like plain dicts, YAML files or dataclass-backed structured configs between tasks without flattening them into individual scalar arguments.

The plugin enables:

- `DictConfig` and `ListConfig` as native task input and output types
- Round-tripping of structured configs (dataclass schemas) across task boundaries
- Preservation of OmegaConf-specific values: `MISSING` sentinels, `Enum`s, `pathlib.Path`s, `tuple`s, and `bytes`
- Resolved variable interpolations on the wire
- A YAML-rendered Flyte report tab for human-readable config inspection

## Installation

```bash
pip install flyteplugins-omegaconf
```

Installing the package automatically registers `DictConfig` and `ListConfig` with Flyte's `TypeEngine`. No manual setup is required.

If you are using the [Hydra plugin](../hydra/_index), `flyteplugins-omegaconf` is installed as a transitive dependency.

## Quick start

```python{hl_lines=[2, "8-9", "14-17"]}
import flyte
from omegaconf import DictConfig, OmegaConf

env = flyte.TaskEnvironment(name="training", image=...)

@env.task
async def train(cfg: DictConfig) -> float:
    return run_experiment(cfg.optimizer.lr, cfg.training.epochs)

@env.task
async def pipeline() -> float:
    cfg = OmegaConf.create(
        {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}}
    )
    return await train(cfg)
```

The config is serialized when `train` is invoked and reconstructed as a `DictConfig` inside the task. No type registration, manual encoding or schema declaration is required.

## When to use this plugin

Use `flyteplugins-omegaconf` when:

- You already use OmegaConf. For example, you have YAML configs, dataclass-based config trees or a Hydra app, and want to keep that representation intact across task boundaries.
- You want to pass a single composed config object instead of widening task signatures with dozens of scalar arguments.
- You want to enforce schema validation at the task entry point via dataclass-backed structured configs.
- You want resolved interpolations (`${other.value}`) to be materialized at submission time rather than at task runtime.

If you do not use OmegaConf elsewhere, prefer plain dataclasses, `pydantic.BaseModel` or `dict` for task inputs as they are supported by Flyte natively without an extra dependency.

## Building a DictConfig

Any of the standard OmegaConf construction methods produce a value the plugin can serialize.

### From a plain dict

```python{hl_lines=["1-3"]}
cfg = OmegaConf.create(
    {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}}
)
flyte.run(train, cfg=cfg)
```

### From a YAML file

```python{hl_lines=[1]}
cfg = OmegaConf.load("configs/training.yaml")
flyte.run(train, cfg=cfg)
```

The file is read locally on the submitter, not on the worker. If the YAML lives in your project tree and needs to be packaged into the task image, use `flyte.with_runcontext(copy_style="all").run(...)`.

### From a dataclass (structured config)

```python{hl_lines=["3-6", 8]}
from dataclasses import dataclass

@dataclass
class TrainConf:
    lr: float = 0.001
    epochs: int = 10

cfg = OmegaConf.structured(TrainConf())
flyte.run(train, cfg=cfg)
```

Structured configs are covered in detail in **OmegaConf > Structured configs** below.

### From a base config plus overrides

```python{hl_lines=["1-3"]}
base = OmegaConf.load("configs/training.yaml")
override = OmegaConf.create({"optimizer": {"lr": 0.01}})
cfg = OmegaConf.merge(base, override)
flyte.run(train, cfg=cfg)
```

This is the same pattern Hydra uses internally. See the [Hydra integration](../hydra/_index) for a full composition layer on top of this plugin.

## Variable interpolation

OmegaConf supports `${...}` interpolations that resolve relative to the config tree:

```python{hl_lines=[3, 4]}
cfg = OmegaConf.create(
    {
        "base_lr": 0.01,
        "optimizer": {"lr": "${base_lr}", "momentum": 0.9},
    }
)
flyte.run(train, cfg=cfg)
```

Interpolations are resolved at serialization time. By the time the task runs, `cfg.optimizer.lr` is the concrete float `0.01`, not the string `"${base_lr}"`. This means:

- The receiving task does not need any context that only existed in the submitter's environment.
- Resolved values appear in the Flyte I/O panel.
- A reference that fails to resolve at submission time fails fast, before any task runs.

If you need lazy resolution on the worker, resolve the reference yourself inside the task or pass the unresolved string through a normal `str` input.

## Nested and deeply structured configs

Nested configs are supported, including deeply structured OmegaConf objects.

```python{hl_lines=["1-13", 18]}
cfg = OmegaConf.create(
    {
        "experiment": {
            "model": {
                "encoder": {
                    "attention": {"num_heads": 8, "head_dim": 64},
                    "ffn": {"hidden_dim": 2048, "activation": "gelu"},
                },
                "decoder": {"num_layers": 6},
            }
        }
    }
)

@env.task
async def extract_leaf(cfg: DictConfig) -> int:
    return int(cfg.experiment.model.encoder.attention.num_heads)
```

## DictConfigs that contain lists

A `DictConfig` may hold list values; they are reconstructed as nested `ListConfig`s on the receiving side.

```python{hl_lines=[4, 5, 8, 9]}
cfg = OmegaConf.create(
    {
        "model": {
            "layer_sizes": [64, 128, 256, 512],
            "activations": ["relu", "relu", "relu", "sigmoid"],
        },
        "data": {
            "augmentations": ["random_flip", "random_crop", "color_jitter"],
            "input_size": [224, 224],
        },
    }
)

@env.task
async def double_layer_sizes(cfg: DictConfig) -> DictConfig:
    doubled = [size * 2 for size in cfg.model.layer_sizes]
    return OmegaConf.merge(cfg, {"model": {"layer_sizes": doubled}})
```

## ListConfig as input and output

`ListConfig` is symmetric with `DictConfig` and supports the same construction patterns.

### Lists of primitives

```python{hl_lines=[2]}
@env.task
async def scale_values(values: ListConfig, factor: float) -> ListConfig:
    return OmegaConf.create([v * factor for v in values])
```

### Building a schedule from another task

```python{hl_lines=[3, 7, 8]}
@env.task
async def build_lr_schedule(base_lr: float, num_stages: int) -> ListConfig:
    return OmegaConf.create([base_lr * (0.5 ** i) for i in range(num_stages)])

@env.task
async def train_with_schedule(cfg: DictConfig, lr_schedule: ListConfig) -> float:
    final_lr = float(lr_schedule[-1])
    ...
```

### Nested lists (list of lists)

```python{hl_lines=[1, 6]}
grid = OmegaConf.create([[0.001, 0.01, 0.1], [10, 20, 50]])

@env.task
async def flatten_grid(grid: ListConfig) -> ListConfig:
    flat = [item for sublist in OmegaConf.to_container(grid) for item in sublist]
    return OmegaConf.create(flat)
```

### Lists of DictConfigs

```python{hl_lines=["2-6"]}
configs = OmegaConf.create(
    [
        {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}},
        {"optimizer": {"lr": 0.01},  "training": {"epochs": 20}},
        {"optimizer": {"lr": 0.1},   "training": {"epochs": 5}},
    ]
)

@env.task
async def select_best_config(configs: ListConfig) -> DictConfig:
    best = max(OmegaConf.to_container(configs), key=lambda c: c["optimizer"]["lr"])
    return OmegaConf.create(best)
```

### Lists of dataclass instances

```python{hl_lines=["9-13"]}
@dataclass
class LayerConf:
    name: str
    width: int
    activation: str

layers = OmegaConf.create(
    [
        LayerConf(name="encoder", width=768, activation="gelu"),
        LayerConf(name="bottleneck", width=128, activation="relu"),
        LayerConf(name="decoder", width=768, activation="linear"),
    ]
)
```

Each element round-trips as a typed `DictConfig` backed by `LayerConf`, so the receiving task can call `OmegaConf.get_type(layers[0])` and access fields with attribute notation.

> **📝 Note**
>
> ListConfig is always plain. Even when its elements are dataclass-backed, the outer `ListConfig` does not carry a list-level schema as there is no structured (typed-element) `ListConfig` in OmegaConf. This affects only the outer container; nested elements retain their schemas.

## Structured configs

A structured config is a `DictConfig` that is bound to a Python dataclass. The dataclass acts as a schema: assigning a value of the wrong type raises `omegaconf.ValidationError`, and merging unknown keys raises an error instead of silently extending the config.

### Basic structured config

```python{hl_lines=["5-8", "11-14", 17, 20]}
from dataclasses import dataclass, field
from omegaconf import OmegaConf, DictConfig

@dataclass
class OptimizerConf:
    lr: float = 0.001
    weight_decay: float = 1e-4

@dataclass
class TrainConf:
    optimizer: OptimizerConf = field(default_factory=OptimizerConf)
    epochs: int = 10

cfg = OmegaConf.structured(TrainConf())
flyte.run(train, cfg=cfg)

# cfg.optimizer.lr = "oops"  # raises omegaconf.ValidationError
```

### Schema reconstruction in the receiving task

When a structured `DictConfig` is deserialized in a downstream task, the plugin operates in **Auto mode**: it reads the originating dataclass name from the wire payload and tries to import it. Two outcomes are possible:

- Dataclass importable in the receiving task: `cfg` is reconstructed as a `TrainConf`-backed `DictConfig`. `OmegaConf.get_type(cfg)` returns `TrainConf`, and type validation is enforced.
- Dataclass not importable: `cfg` falls back to a plain `DictConfig` carrying the raw values. `OmegaConf.get_type(cfg)` returns `dict`. The values are intact but the schema is lost.

To keep schemas across task hops, define dataclasses in modules that are importable from every task in the pipeline (for example, in a shared `configs.py` module bundled into the task image).

### Required (`MISSING`) fields

OmegaConf's `MISSING` sentinel marks a required field that has no default:

```python{hl_lines=[1, 5, "8-9", "12-13"]}
from omegaconf import MISSING

@dataclass
class TrainConf:
    data_path: str = MISSING
    epochs: int = 10

# Pass with MISSING still unset — serialization succeeds.
cfg = OmegaConf.structured(TrainConf())
flyte.run(train, cfg=cfg)

# Or fill it before passing.
cfg = OmegaConf.structured(TrainConf(data_path="/data/imagenet"))
flyte.run(train, cfg=cfg)
```

A config with an unset `MISSING` field serializes and deserializes successfully as the sentinel is preserved on the wire. Accessing the field on the receiving side raises `MissingMandatoryValue`.

> **📝 Note**
>
> Type annotations are preserved only in Auto mode. When the dataclass is importable on the receiving side, an unfilled `MISSING` field still carries its declared type (e.g. `StringNode` for `str`). When the plugin falls back to a plain `DictConfig` because the dataclass is not importable, the field becomes an `AnyNode` where the value is preserved, but the type annotation is not.

### Advanced field types

Beyond primitives and nested dataclasses, structured configs may declare fields of these types and they will round-trip with their schemas intact:

- `Enum` subclasses
- `pathlib.Path`
- `Optional[T]`
- `bytes`
- `dict[str, T]` where `T` is a dataclass
- `list[T]` where `T` is a dataclass

```python{hl_lines=["6-8", "20-35"]}
from enum import Enum
from pathlib import Path
from typing import Optional

class RunMode(Enum):
    TRAIN = "train"
    EVAL = "eval"

@dataclass
class CallbackConf:
    name: str = "early_stop"
    patience: int = 3
    monitor: str = MISSING

@dataclass
class AdvancedTrainConf:
    mode: RunMode = RunMode.TRAIN
    checkpoint_dir: Path = Path("/tmp/checkpoints")
    maybe_seed: Optional[int] = None
    payload: bytes = b"default-token"
    callbacks_by_name: dict[str, CallbackConf] = field(
        default_factory=lambda: {
            "early_stop": CallbackConf(name="early_stop", patience=3),
            "checkpoint": CallbackConf(name="checkpoint", monitor="val_loss"),
        }
    )
    callbacks: list[CallbackConf] = field(
        default_factory=lambda: [
            CallbackConf(name="lr_monitor", patience=2, monitor="lr"),
            CallbackConf(name="nan_guard", patience=1, monitor="loss"),
        ]
    )
```

Inside a downstream task:

```python
@env.task
async def inspect(cfg: DictConfig) -> str:
    assert OmegaConf.get_type(cfg) == AdvancedTrainConf
    assert OmegaConf.get_type(cfg.callbacks[0]) == CallbackConf
    assert isinstance(cfg.mode, RunMode)
    assert isinstance(cfg.checkpoint_dir, Path)
    assert isinstance(cfg.payload, bytes)
    return cfg.mode.value
```

### Merging overrides on top of a structured base

```python{hl_lines=[3, 11]}
@env.task
async def structured_merge_pipeline() -> str:
    base = OmegaConf.structured(TrainConf())
    overrides = OmegaConf.create(
        {
            "optimizer": {"lr": 0.05},
            "training": {"epochs": 100},
            "experiment_name": "sweep-run-1",
        }
    )
    cfg = OmegaConf.merge(base, overrides)
    return await validate_config(cfg)
```

Merging an unknown key against a structured config raises an error, so define every key the override layer might supply on the dataclass.

## Embedding rich Python values inside a plain DictConfig

A plain `DictConfig` (one not bound to a dataclass) can still hold Python values that OmegaConf does not natively model. The plugin preserves the following types end-to-end whether they appear in plain or structured configs:

- `pathlib.Path` and any subclass of `pathlib.PurePath`
- `enum.Enum` members
- `tuple` (round-trips as `tuple`, not `list`)
- `bytes`

```python{hl_lines=[1]}
cfg = OmegaConf.create({"model_path": Path("/opt/models/model.bin")})

@env.task
async def use_path(cfg: DictConfig) -> str:
    assert isinstance(cfg.model_path, Path)
    return f"model_path={cfg.model_path}"
```

If an `Enum`'s class cannot be imported in the receiving environment, the value is returned as the underlying primitive (`int`, `str`, ...) instead of the enum member.

## Reserved-looking keys

The plugin's wire format uses an internal payload marker (`__flyte_omegaconf__`), which means user-facing keys named `kind`, `values`, `name`, `value`, `type`, or `schema` round-trip unchanged:

```python{hl_lines=[1, 8]}
cfg = OmegaConf.create({"kind": "training-job", "values": {"lr": 0.001}})

@env.task
async def use_payload_shaped_config(cfg: DictConfig) -> str:
    # cfg.values resolves to DictConfig.values() — use bracket notation
    # to reach the user key named "values".
    return f"kind={cfg.kind} lr={cfg['values'].lr}"
```

The only practical consideration is Python's normal attribute-vs-method conflict: `cfg.values` is the `.values()` method, so reach for `cfg["values"]` when your config has a key with that name.

## YAML reports

The Flyte I/O panel displays the literal wire representation of a `DictConfig`.

![Wire Representation](../../_static/images/integrations/omegaconf/input.png)

For a YAML view, enable a Flyte report on the task and log the config with `log_yaml`:

```python{hl_lines=[1, 4, 6]}
from flyteplugins.omegaconf import log_yaml

@env.task(report=True)
async def train(cfg: DictConfig) -> DictConfig:
    await log_yaml.aio(cfg, title="Input config")
    ...
```

![YAML Report](../../_static/images/integrations/omegaconf/yaml_repr.png)

The plugin also exposes:

- `to_yaml(cfg)`: render an OmegaConf container as a YAML string.
- `to_html(cfg, title=...)`: wrap the YAML in escaped HTML for embedding in a custom report.
- `replace_yaml(cfg, ...)`: replace the contents of a report tab instead of appending.

```python
from flyteplugins.omegaconf.report import to_yaml, replace_yaml

text = to_yaml(cfg)
await replace_yaml.aio(cfg, tab="Final config")
```

`MISSING` fields appear as `???` in the YAML output, matching OmegaConf's own convention.

## Wire format

Both `DictConfig` and `ListConfig` are serialized as MessagePack blobs with the literal representation:

```
Literal(scalar=Scalar(binary=Binary(value=<msgpack bytes>, tag="msgpack")))
```

The msgpack payload uses an internal tagged structure to distinguish OmegaConf-specific concepts from raw values:

- A `DictConfig` payload includes the originating dataclass name (`builtins.dict` for plain configs) plus its values.
- `MISSING`, `Enum`, `Path`, and `tuple` values carry tagged shapes so they can be reconstructed faithfully.

You normally do not need to inspect this format. It is documented here because:

- The plugin serializes with `resolve=True`, so the wire representation always contains concrete values for `${...}` interpolations.
- Cache-key metadata is set via Flyte's `MESSAGEPACK` serialization format, so two tasks given equivalent configs hit the same cache entry.

## End-to-end example

The example below ties the pieces together: a structured `DictConfig` is created in a parent task, flows through several child tasks that read and modify it, and a `ListConfig` produced midway is consumed by a later stage. Each hop serializes and deserializes the config; the dataclass schema is recovered on the receiving side because `TrainConf` (and friends) are importable in every task in the pipeline.

```
from dataclasses import dataclass, field

import flyte
from omegaconf import DictConfig, ListConfig, OmegaConf

env = flyte.TaskEnvironment(
    name="omegaconf-pipeline-example",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-omegaconf"),
)

@dataclass
class OptimizerConf:
    lr: float = 0.001
    weight_decay: float = 1e-4

@dataclass
class DataConf:
    path: str = ""
    preprocessed: bool = False

@dataclass
class ResultsConf:
    val_loss: float = 0.0
    final_lr: float = 0.0
    num_lr_steps: int = 0

@dataclass
class TrainConf:
    optimizer: OptimizerConf = field(default_factory=OptimizerConf)
    data: DataConf = field(default_factory=DataConf)
    results: ResultsConf = field(default_factory=ResultsConf)
    epochs: int = 10
    batch_size: int = 32
    experiment: str = "baseline"

@env.task
async def preprocess(cfg: DictConfig, dataset: str) -> DictConfig:
    """First stage: fills in the data section of cfg."""
    return OmegaConf.merge(cfg, {"data": {"path": dataset, "preprocessed": True}})

@env.task
async def build_schedule(cfg: DictConfig) -> ListConfig:
    """Produces an LR schedule from cfg as a ListConfig."""
    lrs = [cfg.optimizer.lr * (0.5**i) for i in range(cfg.epochs)]
    return OmegaConf.create(lrs)

@env.task
async def train(cfg: DictConfig, lr_schedule: ListConfig) -> tuple[DictConfig, float]:
    """Simulates training. Returns the final cfg (with results filled in) and val loss."""
    final_lr = float(lr_schedule[-1])
    val_loss = final_lr * 10  # placeholder
    result_cfg = OmegaConf.merge(
        cfg,
        {
            "results": {
                "val_loss": val_loss,
                "final_lr": final_lr,
                "num_lr_steps": len(lr_schedule),
            }
        },
    )
    return result_cfg, val_loss

@env.task
async def evaluate(result_cfg: DictConfig, val_loss: float) -> str:
    """Final stage: formats a report from the result config."""
    return (
        f"experiment={result_cfg.experiment} "
        f"data={result_cfg.data.path} "
        f"val_loss={val_loss:.6f} "
        f"final_lr={result_cfg.results.final_lr:.6f} "
        f"lr_steps={result_cfg.results.num_lr_steps}"
    )

@env.task
async def training_pipeline(dataset: str) -> str:
    """Full pipeline: cfg flows preprocess, build_schedule, train and evaluate."""
    cfg = OmegaConf.structured(
        TrainConf(
            optimizer=OptimizerConf(lr=0.01, weight_decay=1e-5),
            epochs=5,
            batch_size=64,
            experiment="structured-cfg-pipeline",
        )
    )

    preprocessed_cfg = await preprocess(cfg, dataset=dataset)
    lr_schedule = await build_schedule(preprocessed_cfg)
    result_cfg, val_loss = await train(preprocessed_cfg, lr_schedule=lr_schedule)
    return await evaluate(result_cfg, val_loss=val_loss)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(training_pipeline, dataset="s3://my-bucket/imagenet")
    print(f"Run URL: {run.url}")
    print(f"Outputs: {run.outputs()}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/omegaconf/example.py*

For more focused examples such as plain `DictConfig` patterns, advanced `ListConfig` shapes, all `MISSING`/`Enum`/`Path`/`bytes` cases, see the [plugin repository](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/omegaconf/examples).


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/opentelemetry ===

# OpenTelemetry

`flyteplugins-otel` turns a Flyte run into an [OpenTelemetry](https://opentelemetry.io/) trace.

Every task becomes a span. Every **Tasks > Build tasks > Traces** becomes a child span inside it. Spans created by your own code or by any OpenTelemetry instrumentation library nest underneath without extra wiring. Export goes wherever OTLP goes: Grafana Tempo, Jaeger, Honeycomb, an OpenTelemetry Collector or several at once.

None of this is specific to agents or to LLM workloads. It is ordinary distributed tracing for ordinary Flyte tasks, plus two behaviors that exist because Flyte runs are durable and a stock OpenTelemetry setup has no way to model them:

- **A crashed and resumed run is one trace, not several:** Each attempt is a fresh process with a fresh OpenTelemetry SDK, so each would normally mint its own trace ID. The plugin derives the trace ID from the run instead, so every process converges on the same trace with no coordination.
- **Steps served from the durable log still appear.** A resumed run replays completed steps rather than re-executing them, so nothing instruments them and the trace would otherwise have holes exactly where durability did its job. The plugin records them as spans marked `flyte.replayed`.

**OpenTelemetry > Traces across crashes and resumes** covers both in detail.

## Installation

```bash
pip install flyteplugins-otel
```

OTLP over HTTP is included. gRPC ships separately:

```bash
pip install "flyteplugins-otel[grpc]"
```

## Quick start

Call `init()` once at module scope, then write tasks as you normally would:

```python{hl_lines=[2,6]}
import flyte
from flyteplugins.otel import init

# Module scope, not inside a task. The task span opens before the task body runs,
# so initializing from within the body means that task's own span is already missed.
init(service_name="my-service")

env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-otel"),
)

@flyte.trace
async def double(x: int) -> int:
    return x * 2

@env.task
async def main(n: int = 3) -> int:
    total = 0
    for i in range(n):
        total += await double(i)
    return total

if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(main, n=3).url)
```

That produces one trace shaped like this:

```text
main                          ← task span
├── double                    ← flyte.trace step span
├── double
└── double
```

![Flyte UI showing the run's action tree beside the console-exported JSON for one span](../../_static/images/integrations/opentelemetry/quick_start.png)

*The quick start run in the Flyte UI. On the left, the action tree shows `main` and its three `double` steps. On the right, the Logs tab holds `ConsoleSpanExporter` output for one `double` span, carrying the `flyte.*` attributes and the `parent_id` that nests it under the task span.*

With no arguments, `init()` reads the standard `OTEL_EXPORTER_OTLP_ENDPOINT` and `OTEL_EXPORTER_OTLP_HEADERS` variables, which is how most vendors document their setup. See **OpenTelemetry > Exporters and configuration** for pointing it at a real backend and for supplying credentials as a `flyte.Secret` instead of hardcoding them.

> [!WARNING] Call `init()` at module scope
> The task span opens before the task body runs, so calling `init()` from inside a task means
> that task's own span has already been missed. The symptom is a trace holding step spans with
> no task span to hang them from. The plugin logs a warning when it detects this.

## What becomes a span

| Flyte concept                                                    | Span                              | Parent                                             |
| ---------------------------------------------------------------- | --------------------------------- | -------------------------------------------------- |
| A task executing in its container                                | Task span, named after the task   | The inbound trace context if any; otherwise a root |
| A **Tasks > Build tasks > Traces** step | Step span                         | The task span that owns it                         |
| A step replayed from the durable log                             | Step span, `flyte.replayed=true`  | The task span of the attempt that replayed it      |
| A sub-action (a task calling another task)                       | Its own task span, in another pod | The calling task's span, via `custom_context`      |
| Anything an instrumentation library emits                        | Whatever that library emits       | The active span, which is the task or step span    |

Task lifecycle itself is not instrumented: there are no spans for scheduling, queueing or the control plane's decision to retry. A span starts when a container begins executing a task.

You will however, see HTTP client spans for Flyte's own calls to the control plane once tracing is on. Those come from Flyte's transport rather than from this plugin; **OpenTelemetry > Traces across crashes and resumes > Flyte's own control-plane spans** explains where they come from and how to switch them off.

## Span attributes

Every span the plugin emits carries the identifiers needed to get back to the run that produced it. These names are effectively public API: a Grafana data link queries on them to jump from a span into the Flyte UI and **OpenTelemetry > Exporters and configuration > Linking back from Grafana** is built on them.

| Attribute                                    | On         | Meaning                                           |
| -------------------------------------------- | ---------- | ------------------------------------------------- |
| `flyte.run_name`                             | All spans  | The run, and what the trace ID is derived from    |
| `flyte.action_name`                          | All spans  | The action that produced the span                 |
| `flyte.project`, `flyte.domain`, `flyte.org` | All spans  | Where the run lives                               |
| `flyte.task_name`                            | Task spans | The task being executed                           |
| `flyte.step_name`                            | Step spans | The traced function                               |
| `flyte.task_action_name`                     | Step spans | The task that owns the step                       |
| `flyte.replayed`                             | Step spans | Whether this step was served from the durable log |

## Your own spans

Spans you create with a plain OpenTelemetry tracer nest inside the task span automatically. Parenting in OpenTelemetry comes from the active context and the plugin keeps the task span active for the whole task body, so there is nothing to extract and no context to pass around:

```python
from opentelemetry import trace

tracer = trace.get_tracer("my.app")

@env.task
async def etl(rows: int = 100) -> int:
    with tracer.start_as_current_span("extract") as span:
        span.set_attribute("rows.requested", rows)
        extracted = rows

    with tracer.start_as_current_span("transform"):
        # Spans nest as deeply as you like; this one lands under transform.
        with tracer.start_as_current_span("validate"):
            transformed = extracted - 1

    with tracer.start_as_current_span("load") as span:
        span.set_attribute("rows.loaded", transformed)

    return transformed
```

The same is true of third-party auto-instrumentation. An HTTP client instrumentor, a database instrumentor or an LLM instrumentor needs no extra wiring:

```python
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor

init(service_name="my-service")
HTTPXClientInstrumentor().instrument()
```

Call `init()` before the other library so everything stays on one export pipeline. For libraries that want to be handed a tracer, `get_tracer()` returns the one the plugin built.

## What's next

- ****OpenTelemetry > Exporters and configuration****: point the plugin at a backend, adopt a tracer provider you already have, and link back from Grafana into the Flyte UI.
- ****OpenTelemetry > Traces across crashes and resumes****: trace context in and out of a run, run-derived trace IDs, and replayed steps.
- **[Grafana Agent Observability](../grafana-agent-observability/_index)**: add LLM generations, tool calls, token usage and cost on top of these traces.

> [!NOTE] Runnable examples
> The plugin ships [eight worked examples](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/otel/examples)
> covering console export, custom spans, nested tasks, adopting an existing provider, joining a
> caller's trace, HTTP auto-instrumentation, Grafana Cloud and a crash-and-resume trace. All but
> the last run either locally or on a cluster; the crash-and-resume one needs a cluster, because
> the replay it demonstrates comes from a platform retry.

## Subpages
- **OpenTelemetry > Exporters and configuration**
- **OpenTelemetry > Traces across crashes and resumes**


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/opentelemetry/configuration ===

# Exporters and configuration

`init()` builds an OTLP exporter by default because OTLP is what most backends document but nothing in the plugin requires it. Any `SpanExporter` works, several can run side by side, and a tracer provider you configured yourself is adopted whole.

## Environment variables

The lowest-friction setup is to configure nothing in code and let the standard OpenTelemetry variables do the work:

| Variable                             | Effect                                                                             |
| ------------------------------------ | ---------------------------------------------------------------------------------- |
| `OTEL_EXPORTER_OTLP_ENDPOINT`        | Where spans are sent. A base gateway URL is fine on HTTP; `/v1/traces` is appended |
| `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | Same, but traces-only. Takes precedence over the general endpoint                  |
| `OTEL_EXPORTER_OTLP_HEADERS`         | Export headers, in `k=v,k2=v2` form. This is where auth goes                       |
| `OTEL_EXPORTER_OTLP_PROTOCOL`        | `http/protobuf` (default) or `grpc`                                                |
| `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` | Same, but traces-only. Takes precedence                                            |
| `OTEL_SERVICE_NAME`                  | Value for `service.name` when `service_name` is not passed                         |

With those set, `init()` needs no arguments:

```python
from flyteplugins.otel import init

init()
```

Endpoints and credentials belong in a `flyte.Secret` rather than in your source. Attach them to the task environment as environment variables and the exporter picks them up:

```python{hl_lines=["5-6"]}
env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-otel"),
    secrets=[
        flyte.Secret(key="otlp_endpoint", as_env_var="OTEL_EXPORTER_OTLP_ENDPOINT"),
        flyte.Secret(key="otlp_headers", as_env_var="OTEL_EXPORTER_OTLP_HEADERS"),
    ],
)
```

See [Secrets](../../user-guide/tasks/task-configuration/secrets) for creating and managing them.

## Configuring in code

Everything the variables cover can also be passed directly:

```python
init(
    service_name="my-service",
    endpoint="https://otlp-gateway-prod-us-east-0.grafana.net/otlp",
    headers={"Authorization": "Basic <base64>"},
)
```

### `init()` parameters

| Parameter             | Default                             | What it does                                                                                                               |
| --------------------- | ----------------------------------- | -------------------------------------------------------------------------------------------------------------------------- |
| `service_name`        | `OTEL_SERVICE_NAME`, then `"flyte"` | Value for `service.name` on every span                                                                                     |
| `endpoint`            | The `OTEL_` variables               | OTLP endpoint. On HTTP a base gateway URL is fine and `/v1/traces` is appended; on gRPC the base endpoint is used as given |
| `headers`             | The `OTEL_` variables               | Export headers, as a mapping or the `k=v,k2=v2` string form                                                                |
| `protocol`            | `http/protobuf`                     | OTLP transport, `http/protobuf` or `grpc`. gRPC needs the `[grpc]` extra                                                   |
| `resource_attributes` | None                                | Extra resource attributes attached to every span                                                                           |
| `exporter`            | None                                | One `SpanExporter` or several, used instead of building an OTLP exporter                                                   |
| `tracer_provider`     | None                                | Adopt a provider you configured yourself. Cannot be combined with the arguments above                                      |
| `disable_batch`       | `False`                             | Export each span as it ends instead of batching                                                                            |
| `set_global`          | `True`                              | Install the provider as the global one, so other instrumentation shares it                                                 |

`init()` is idempotent: calling it a second time returns the observer registered by the first call and changes nothing.

### Choosing an exporter

Pass any `SpanExporter` or a list of them to fan out to several at once, which is useful for keeping a console exporter alongside a real backend while you develop:

```python
from opentelemetry.sdk.trace.export import ConsoleSpanExporter

init(exporter=[ConsoleSpanExporter(), JaegerExporter(...)])
```

Each exporter gets its own span processor, so they run independently.

### gRPC

The gRPC exporter ships as a separate distribution:

```bash
pip install "flyteplugins-otel[grpc]"
```

```python
init(endpoint="http://collector:4317", protocol="grpc")
```

Note the endpoint difference between transports: gRPC takes the base endpoint, HTTP wants the signal-specific path (which the plugin appends for you).

### Batching

Spans are batched by default, which is the right production setting. `disable_batch=True` exports each span as it ends: slower, but nothing is buffered when the process dies, which matters when the thing you are looking at is a crash.

```python
init(service_name="my-service", disable_batch=True)
```

## Adopting a tracer provider you already have

If your codebase already configures OpenTelemetry with its own resource, sampler and exporters, hand the provider over instead of letting the plugin build one:

```python{hl_lines=[11]}
from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter

provider = TracerProvider(resource=Resource.create({"service.name": "my-existing-service"}))
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)

# Sampler, resource and exporters are left exactly as configured above.
init(tracer_provider=provider)
```

Nothing about your setup changes. The plugin only wraps the provider's ID generator which is what lets trace IDs still be [derived from the run](./durable-traces#one-run-one-trace) so a crash and its resume share a trace.

Passing `tracer_provider` alongside `endpoint`, `headers`, `exporter`, `protocol` or `resource_attributes` raises a `ValueError` rather than silently overriding: those settings belong to the provider you configured.

## Grafana Cloud

Grafana Cloud is an ordinary OTLP backend so the same shape works for any other. Both values come from the OTLP section of the Grafana Cloud portal:

```python
env = flyte.TaskEnvironment(
    name="otel_grafana",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-otel"),
    secrets=[
        flyte.Secret(key="otlp_endpoint", as_env_var="OTEL_EXPORTER_OTLP_ENDPOINT"),
        flyte.Secret(key="otlp_headers", as_env_var="OTEL_EXPORTER_OTLP_HEADERS"),
    ],
)

# With the two variables set, init needs nothing else. Passing them explicitly:
#   init(
#       service_name="my-service",
#       endpoint="https://otlp-gateway-<zone>.grafana.net/otlp",
#       headers={"Authorization": "Basic <base64 instance_id:token>"},
#   )
init(service_name="my-service")
```

## Linking back from Grafana

`flyteplugins.otel.grafana` builds [links](../../user-guide/tasks/task-programming/links) from a Flyte action into Grafana, rendered on the action in the Flyte UI. They are plain URL builders with no Grafana dependency:

```python{hl_lines=[3]}
from flyteplugins.otel.grafana import GrafanaTrace

@env.task(links=(GrafanaTrace(host="https://myorg.grafana.net", datasource_uid="<tempo-uid>"),))
async def my_task() -> str:
    ...
```

![Flyte UI action summary with the Grafana trace link highlighted in its Links section](../../_static/images/integrations/opentelemetry/flyte_ui_link.png)

*`GrafanaTrace` renders in the action's **Links** section in the Flyte UI, on every run of the task.*

![Grafana Explore opened on the run's trace, with the TraceQL query already filled in](../../_static/images/integrations/opentelemetry/grafana_dashboard.png)

*Following it opens Grafana Explore with the query already scoped to this run, so you land on its spans instead of searching for the run name by hand.*

| Parameter        | Default           | What it does                                                    |
| ---------------- | ----------------- | --------------------------------------------------------------- |
| `host`           | Required          | Stack URL, for example `https://myorg.grafana.net`              |
| `datasource_uid` | Required          | UID of the Tempo datasource                                     |
| `name`           | `"Grafana trace"` | Label shown in the Flyte UI                                     |
| `lookback`       | `"now-7d"`        | Start of the Explore time range, in Grafana's relative syntax   |
| `action_scoped`  | `False`           | Narrow the query to the single action rather than the whole run |

The datasource UID is per-stack and not guessable. Find it under **Connections > Data sources** in Grafana; it is the last path segment of `/connections/datasources/edit/<uid>`.

Two design details worth knowing:

- The link runs a TraceQL query on `flyte.run_name` rather than addressing a trace by ID. That means it finds a run's spans whatever their trace IDs turn out to be, including runs whose trace context arrived from outside Flyte. Addressing by ID would depend on the derivation and break the moment something upstream propagated a context.
- It embeds a time range because Grafana Explore otherwise defaults to the last hour and a link to an older run would open on an empty pane.

For a link to a run's conversation in Grafana Agent Observability, see [`GrafanaAgentObservability`](../grafana-agent-observability/_index#linking-back-from-grafana). That one lives in `flyteplugins-agento11y` because it is that package's identity binding that makes a run addressable by conversation ID at all.

## Shutting down

`shutdown()` unregisters the observer and flushes pending spans. You rarely need it: the OpenTelemetry SDK registers its own exit hook, so a task that finishes normally flushes on the way out. Reach for it when you want to stop tracing inside a long-lived process or in tests.

```python
from flyteplugins.otel import shutdown

shutdown()
```

A provider you passed in with `tracer_provider=` is never shut down since it belongs to you.

## When nothing is configured

With no OTLP endpoint anywhere, no `OTEL_EXPORTER_OTLP_ENDPOINT`, no `endpoint=` and no explicit exporter, spans are recorded but not exported, and `init()` logs a warning once.

This is the normal state of the process that submits a run: it imports your module, and therefore runs `init()`, without having any reason to export. Nesting and context propagation behave exactly as they would otherwise; the spans are simply dropped instead of shipped.

If you run a local collector, point the variable at it explicitly (`http://localhost:4318`) rather than relying on the OTLP specification's default of the same address. Without the explicit setting the plugin assumes you meant nothing at all.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/opentelemetry/durable-traces ===

# Traces across crashes and resumes

A durable Flyte run is not one process. It crashes, resumes and retries, and each attempt starts a fresh OpenTelemetry SDK that knows nothing about the earlier ones.

Point a stock OpenTelemetry setup at a durable run and two things go wrong. Every attempt mints its own trace ID, so one run arrives at the backend as several unrelated traces. And every step the resumed run replayed out of its durable log is missing entirely because replayed steps never execute and so nothing instruments them. The trace ends up with holes in it exactly where durability did its job.

`flyteplugins-otel` fixes both and it does so without any coordination between the processes.

## Trace context, in and out

Flyte propagates a key-value `custom_context` through a run and into every sub-action. The plugin uses it as a [W3C trace context](https://www.w3.org/TR/trace-context/) carrier in both directions.

### Inbound: joining a trace that started outside Flyte

When a run is kicked off from inside an existing span, a web request, a scheduler, another service, you usually want the run to appear inside that trace rather than as a separate one. Inject a carrier into `custom_context` at submit time and the plugin picks it up: the task span starts under the caller's span instead of becoming a root.

```python{hl_lines=["24-25",27]}
import flyte
from opentelemetry import trace
from opentelemetry.propagate import inject

from flyteplugins.otel import init

init(service_name="my-service")
tracer = trace.get_tracer("my.caller")

env = flyte.TaskEnvironment(name="my_env")

# The task you submit. Flyte 2 has no separate workflow entrypoint: `run` takes the
# task itself, and any tasks it awaits become sub-actions of the same run.
@env.task
async def handle(url: str) -> str:
    return f"handled {url}"

if __name__ == "__main__":
    flyte.init_from_config()

    with tracer.start_as_current_span("incoming_request"):
        carrier: dict[str, str] = {}
        inject(carrier)

        run = flyte.with_runcontext(custom_context=carrier).run(handle, url="https://example.com")
        print(run.url)
```

The carrier is a plain `dict[str, str]`, which is exactly what `custom_context` expects.

`tracer` here is the caller's own tracer, not the plugin's. The plugin only needs `init()` to have run; it parents the task span off whatever `traceparent` arrives in the carrier, whoever produced it.

### Outbound: nested tasks, in other pods

Once the task span is open, the plugin publishes it back into `custom_context`. A child task running in a different pod nests under the task that spawned it with nothing passed by hand:

CODE0

CODE1

Because `custom_context` travels in the action's persisted inputs, this survives a resume as well.

Nothing in your task bodies has to call `extract` to get spans parented correctly; the plugin's own spans are already in the right place. Reach for `extract` only when you want to open your own spans under the incoming context.

## One run, one trace

When no trace context arrives from outside, the trace ID is derived from the run identity rather than generated randomly.

Every process computes the same 16 bytes from values it already has, the org, project, domain and run name, so spans recorded before a crash and spans recorded after the resume land in the same trace even though neither process ever spoke to the other.

Only the trace ID is derived. Span IDs stay random, which keeps each attempt a distinct subtree under the shared trace rather than a set of colliding IDs. A resumed run therefore reads as the attempt that crashed followed by the attempt that finished.

Deriving from the fully-qualified run identity rather than the run name alone means two runs that happen to share a name in different projects or domains stay distinct.

You can compute the same value yourself, which is useful for building your own links into a tracing backend:

CODE2

## Replayed steps

A resumed run serves already-completed [traced steps](../../user-guide/tasks/task-programming/traces) out of its durable log without re-executing them. The plugin records those as spans marked `flyte.replayed=true`.

They have no meaningful duration, because no work happened in this process, but they are present, so the trace is complete and you can see exactly which steps the resume skipped.

For an agent loop, this is also where the money is: a replayed step does not call the model again, so a resume does not pay for the generations the first attempt already bought.

## A worked crash and resume

The example below crashes partway through its first attempt. The retry resumes: steps 0 through 2 come back from the durable log as replayed spans, steps 3 and 4 actually execute, and both attempts share one trace.

CODE3

The resulting trace:

CODE4

![Grafana Tempo trace holding two attempts of one run, the second with microsecond replayed steps](../../_static/images/integrations/opentelemetry/replayed_trace.png)

*The same run in Grafana Tempo, found by a TraceQL query on `flyte.run_name` rather than by trace ID. Each attempt is its own subtree under a single trace. In the second, the microsecond `think` spans are replays served from the durable log and the 200 ms ones are the steps that actually executed. The `POST` spans are Flyte's own calls to the control plane.*

## Flyte's own control-plane spans

Once tracing is on, you will see `POST` client spans for Flyte's calls to the control plane, `Enqueue`, `CreateRun`, `UploadInputs`, alongside your own.

These do not come from this plugin. Flyte's HTTP transport takes an `enable_otel` flag that defaults to true and falls back to the global tracer provider when it is not given one. `init(set_global=True)`, the default, installs that provider, so the transport starts recording through it.

Mostly this is useful. Inside a task the spans nest under the task span, so you can see how much of a task's wall clock went on talking to the control plane, and the 401-then-200 pairs show the auth retry. Two things to be aware of:

- The volume scales with sub-action count, so a wide fan-out produces a lot of them.
- Calls made outside a task span, during submission, arrive as their own root traces rather than joining the run's trace.

There is no switch for this in the plugin, since the transport is Flyte's rather than the plugin's. `init(set_global=False)` keeps the provider out of the global slot, which stops the transport finding it, at the cost of other instrumentation not finding it either.

## Limitations

**Replayed spans have no duration:** The original timing is written to the control plane but does not come back over the channel a resumed run reads from. What you get is the step's presence, identity and outcome.

**Trace context rides in `custom_context`:** That is a flat string map which Flyte propagates wholesale, so the `traceparent` key is visible to task code and will be overwritten if something else writes that key.

**Nothing is emitted for control-plane lifecycle:** Scheduling, queueing and the retry decision itself are not instrumented. The spans you get start when a task's container begins executing it.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/pandera ===

# Pandera

The [Pandera](https://pandera.readthedocs.io/en/latest/) plugin validates dataframes at task boundaries using
[`DataFrameModel`](https://pandera.readthedocs.io/en/latest/dataframe_models.html) schemas. When a task receives or
returns a pandera-typed dataframe, the plugin automatically validates the data, raises or warns on schema violations,
and writes an HTML validation report to the Flyte deck.

Pandera supports multiple dataframe backends. The `flyteplugins-pandera` plugin handles:

| Pandera typing module | DataFrame library | Additional plugin |
|-|-|-|
| `pandera.typing.pandas` | pandas | - |
| `pandera.typing.polars` | Polars (eager and lazy) | `flyteplugins-polars` |
| `pandera.typing.pyspark_sql` | PySpark SQL | `flyteplugins-spark` |

## When to use this plugin

- You want compile-time-style guarantees that data flowing between tasks conforms to a declared schema
- You need column-level type, constraint, and statistical checks on task inputs and outputs
- You want automatic validation reports visible in the Flyte UI

## Installation

Install the plugin with the pandera extras for your dataframe backend:

### pandas

```bash
pip install flyteplugins-pandera 'pandera[pandas]'
```

### Polars

```bash
pip install flyteplugins-pandera flyteplugins-polars 'pandera[polars]'
```

### PySpark SQL

```bash
pip install flyteplugins-pandera flyteplugins-spark 'pandera[pyspark]'
```

## Defining schemas

Schemas are defined as Python classes that inherit from pandera's `DataFrameModel`. Each field declares a column name,
type, and optional constraints:

```python
import pandera.pandas as pa

class EmployeeSchema(pa.DataFrameModel):
    employee_id: int = pa.Field(ge=0)
    name: str

class EmployeeSchemaWithStatus(EmployeeSchema):
    status: str = pa.Field(isin=["active", "inactive"])
```

Schemas compose through inheritance: `EmployeeSchemaWithStatus` includes all columns from `EmployeeSchema` plus the
`status` column.

For full details on schema definition, including custom checks, regex column matching, and `Config` options, see the
[pandera DataFrameModel documentation](https://pandera.readthedocs.io/en/latest/dataframe_models.html).

## Using schemas in tasks

Annotate task inputs and outputs with pandera's generic `DataFrame` type. The plugin validates data on every
encode (output) and decode (input):

```python
import pandera.typing.pandas as pt

@env.task(report=True)
async def build_employees() -> pt.DataFrame[EmployeeSchema]:
    return pd.DataFrame({
        "employee_id": [1, 2, 3],
        "name": ["Ada", "Grace", "Barbara"],
    })

@env.task(report=True)
async def add_status(
    df: pt.DataFrame[EmployeeSchema],
) -> pt.DataFrame[EmployeeSchemaWithStatus]:
    return df.assign(status="active")
```

Setting `report=True` on the task makes validation reports visible as deck tabs in the Flyte UI.

## Error handling with `ValidationConfig`

By default, a validation failure raises an exception and fails the task. To downgrade failures to warnings instead,
annotate the parameter with `ValidationConfig(on_error="warn")`:

```python
from typing import Annotated
from flyteplugins.pandera import ValidationConfig

@env.task(report=True)
async def lenient_pass_through(
    df: Annotated[pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")],
) -> Annotated[pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")]:
    ...
```

| `on_error` value | Behavior |
|-|-|
| `"raise"` (default) | Validation failure raises `pandera.errors.SchemaError` and the task fails |
| `"warn"` | Validation failure logs a warning and writes the report, but the task continues |

You can mix `"raise"` and `"warn"` across inputs and outputs of the same task. For example, use `"warn"` on inputs
to accept best-effort data while still enforcing strict output contracts.

## Image configuration

Include the plugin in your task image. The exact setup depends on your dataframe backend:

### Pandas

```python
import flyte

img = flyte.Image.from_debian_base(
    python_version=(3, 12),
).with_pip_packages("flyteplugins-pandera")

env = flyte.TaskEnvironment(
    "pandera_pandas",
    image=img,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)
```

### Polars

```python
import flyte

img = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_pip_packages("flyteplugins-polars", "pandera[polars]")
)

env = flyte.TaskEnvironment(
    "pandera_polars",
    image=img,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)
```

### PySpark SQL

```python
import flyte
from flyteplugins.spark.task import Spark

image = (
    flyte.Image.from_base("apache/spark-py:v3.4.0")
    .clone(name="pandera-pyspark-sql", python_version=(3, 10), extendable=True)
    .with_pip_packages("flyteplugins-spark", "pandera[pyspark]")
)

spark_conf = Spark(
    spark_conf={
        "spark.driver.memory": "1000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
    },
)

env = flyte.TaskEnvironment(
    name="pandera_pyspark",
    plugin_config=spark_conf,
    image=image,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)
```

## Polars lazy frames

The Polars backend supports both `pt.DataFrame` (eager) and `pt.LazyFrame` (lazy). With lazy frames, pandera
validates the data when the frame is materialized at task I/O boundaries:

```python
import pandera.typing.polars as pt
import polars as pl

@env.task(report=True)
async def create_lazy() -> pt.LazyFrame[MetricsSchema]:
    return pl.LazyFrame({"item": ["x", "y"], "value": [3.0, 4.0]})

@env.task(report=True)
async def consume_lazy(
    lf: pt.LazyFrame[MetricsSchema],
) -> pt.DataFrame[MetricsSchema]:
    return lf.filter(pl.col("value") > 0.0).collect()
```

## Examples

### pandas

```python
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte",
#    "flyteplugins-pandera",
#    "pandera[pandas]",
# ]
# main = "main"
# ///

from __future__ import annotations

from typing import Annotated

import pandas as pd
import pandera.pandas as pa
import pandera.typing.pandas as pt
from flyteplugins.pandera import ValidationConfig

import flyte

img = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "flyteplugins-pandera", "pandera[pandas]"
)

env = flyte.TaskEnvironment(
    "pandera_pandas_schema",
    image=img,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

class EmployeeSchema(pa.DataFrameModel):
    employee_id: int = pa.Field(ge=0)
    name: str

class EmployeeSchemaWithStatus(EmployeeSchema):
    status: str = pa.Field(isin=["active", "inactive"])

# {{docs-fragment build_valid_employees}}
@env.task(report=True)
async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]:
    return pd.DataFrame(
        {
            "employee_id": [1, 2, 3],
            "name": ["Ada", "Grace", "Barbara"],
        }
    )
# {{/docs-fragment}}

# {{docs-fragment pass_through}}
@env.task(report=True)
async def pass_through(
    df: pt.DataFrame[EmployeeSchema],
) -> pt.DataFrame[EmployeeSchemaWithStatus]:
    return df.assign(status="active")
# {{/docs-fragment}}

# {{docs-fragment pass_through_with_error_warn}}
@env.task(report=True)
async def pass_through_with_error_warn(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")
]:
    del df["name"]
    return df
# {{/docs-fragment}}

# {{docs-fragment pass_through_with_error_raise}}
@env.task(report=True)
async def pass_through_with_error_raise(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")
]:
    del df["name"]
    return df
# {{/docs-fragment}}

@env.task(report=True)
async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]:
    df = await build_valid_employees()
    df2 = await pass_through(df)

    await pass_through_with_error_warn(df.drop(["employee_id"], axis="columns"))
    await pass_through_with_error_warn(df.assign(employee_id=-1))

    try:
        await pass_through_with_error_raise(df)
    except Exception as exc:
        print(exc)

    return df2

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()
    print("pandas pandera example OK:", run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pandas_schema.py*

### Polars

```python
# /// script
# requires-python = ">=3.12"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-pandera",
#    "flyteplugins-polars",
#    "pandera[polars]",
# ]
# main = "main"
# ///

from __future__ import annotations

from typing import Annotated

import pandera.polars as pa
import pandera.typing.polars as pt
import polars as pl
from flyteplugins.pandera import ValidationConfig

import flyte

img = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_pip_packages("flyteplugins-pandera", "flyteplugins-polars", "pandera[polars]")
)

env = flyte.TaskEnvironment(
    "pandera_polars_schema",
    image=img,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

class EmployeeSchema(pa.DataFrameModel):
    employee_id: int = pa.Field(ge=0)
    name: str

class EmployeeSchemaWithStatus(EmployeeSchema):
    status: str = pa.Field(isin=["active", "inactive"])

class MetricsSchema(pa.DataFrameModel):
    item: str
    value: float

# {{docs-fragment build_valid_employees}}
@env.task(report=True)
async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]:
    return pl.DataFrame(
        {
            "employee_id": [1, 2, 3],
            "name": ["Ada", "Grace", "Barbara"],
        }
    )
# {{/docs-fragment}}

# {{docs-fragment pass_through}}
@env.task(report=True)
async def pass_through(
    df: pt.DataFrame[EmployeeSchema],
) -> pt.DataFrame[EmployeeSchemaWithStatus]:
    return df.with_columns(pl.lit("active").alias("status"))
# {{/docs-fragment}}

@env.task(report=True)
async def pass_through_with_error_warn(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")
]:
    return df.drop("name")

@env.task(report=True)
async def pass_through_with_error_raise(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")
]:
    return df.drop("name")

# {{docs-fragment metrics_lazy}}
@env.task(report=True)
async def metrics_eager() -> pt.DataFrame[MetricsSchema]:
    return pl.DataFrame({"item": ["a", "b"], "value": [1.0, 2.0]})

@env.task(report=True)
async def metrics_lazy() -> pt.LazyFrame[MetricsSchema]:
    return pl.LazyFrame({"item": ["x", "y"], "value": [3.0, 4.0]})

@env.task(report=True)
async def filter_metrics(
    lf: pt.LazyFrame[MetricsSchema],
) -> pt.DataFrame[MetricsSchema]:
    return lf.filter(pl.col("value") > 0.0).collect()
# {{/docs-fragment}}

@env.task(report=True)
async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]:
    df = await build_valid_employees()
    df2 = await pass_through(df)

    await pass_through_with_error_warn(df.drop("employee_id"))
    await pass_through_with_error_warn(
        df.with_columns(pl.lit(-1).alias("employee_id"))
    )

    try:
        await pass_through_with_error_raise(df)
    except Exception as exc:
        print(exc)

    _ = await metrics_eager()
    lazy = await metrics_lazy()
    _ = await filter_metrics(lazy)

    return df2

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()
    print("polars pandera example OK:", run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/polars_schema.py*

### PySpark SQL

```python
# /// script
# requires-python = ">=3.10"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-pandera",
#    "flyteplugins-spark",
#    "pandera[pyspark]",
# ]
# main = "main"
# ///

from __future__ import annotations

from typing import Annotated, cast

import pandera.typing.pyspark_sql as pt
import pyspark.sql.types as T
from flyteplugins.pandera import ValidationConfig
from flyteplugins.spark.task import Spark
from pandera.pyspark import DataFrameModel, Field
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

import flyte

image = (
    flyte.Image.from_base("apache/spark-py:v3.4.0")
    .clone(name="pandera-pyspark-sql", python_version=(3, 10), extendable=True)
    .with_pip_packages(
        "flyteplugins-pandera",
        "flyteplugins-spark",
        "pandera[pyspark]",
    )
)

spark_conf = Spark(
    spark_conf={
        "spark.driver.memory": "1000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
        "spark.kubernetes.file.upload.path": "/opt/spark/work-dir",
        "spark.jars": (
            "https://storage.googleapis.com/hadoop-lib/gcs/"
            "gcs-connector-hadoop3-latest.jar,"
            "https://repo1.maven.org/maven2/org/apache/hadoop/"
            "hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,"
            "https://repo1.maven.org/maven2/com/amazonaws/"
            "aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar"
        ),
    },
)

env = flyte.TaskEnvironment(
    name="pandera_pyspark_sql_schema",
    plugin_config=spark_conf,
    image=image,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

# {{docs-fragment schemas}}
class EmployeeSchema(DataFrameModel):
    employee_id: int = Field(ge=0)
    name: str = Field()
    job_title: str = Field()

class EmployeeSchemaWithStatus(EmployeeSchema):
    status: str = Field(isin=["active", "inactive"])
# {{/docs-fragment}}

# {{docs-fragment build_valid_employees}}
@env.task(report=True)
async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]:
    spark = cast(SparkSession, flyte.ctx().data["spark_session"])
    data = [
        (1, "Ada", "Engineer"),
        (2, "Grace", "Mathematician"),
        (3, "Barbara", "Computer scientist"),
    ]
    schema = T.StructType(
        [
            T.StructField("employee_id", T.IntegerType(), False),
            T.StructField("name", T.StringType(), False),
            T.StructField("job_title", T.StringType(), False),
        ]
    )
    return spark.createDataFrame(data, schema=schema)
# {{/docs-fragment}}

# {{docs-fragment pass_through}}
@env.task(report=True)
async def pass_through(
    df: pt.DataFrame[EmployeeSchema],
) -> pt.DataFrame[EmployeeSchemaWithStatus]:
    return df.withColumn("status", F.lit("active"))
# {{/docs-fragment}}

@env.task(report=True)
async def pass_through_with_error_warn(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")
]:
    return df.drop("name")

@env.task(report=True)
async def pass_through_with_error_raise(
    df: Annotated[
        pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")
    ],
) -> Annotated[
    pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")
]:
    return df.drop("name")

@env.task(report=True)
async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]:
    df = await build_valid_employees()
    df2 = await pass_through(df)

    await pass_through_with_error_warn(df.drop("employee_id"))
    await pass_through_with_error_warn(df.withColumn("employee_id", F.lit(-1)))

    try:
        await pass_through_with_error_raise(df)
    except Exception as exc:
        print(exc)

    return df2

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait()
    print("pyspark_sql pandera example OK:", run.outputs()[0])
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pyspark_sql_schema.py*


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/papermill ===

# Papermill

The Papermill plugin lets you run Jupyter notebooks as Flyte tasks. It uses [papermill](https://papermill.readthedocs.io/) to parameterize and execute `.ipynb` files, capture their outputs as typed Flyte values, and render the executed notebook as an HTML report visible in the Flyte UI.

A `NotebookTask` behaves like any other Flyte task: it has typed inputs and outputs, participates in workflows, runs remotely, integrates with the Flyte type system (including `File`, `Dir`, and `DataFrame`), and can call other Flyte tasks from within the notebook.

## When to use this plugin

- Productionizing exploratory notebooks without rewriting them as Python modules
- Generating cell-by-cell HTML reports as task artifacts (charts, tables, narrative analysis)
- Letting data scientists iterate in notebooks while platform teams orchestrate them
- Running notebooks on Spark or with GPU/CPU resources configured on the task environment

## Installation

```bash
pip install flyteplugins-papermill
```

The plugin must also be installed in the task image. For example:

```python{hl_lines=["3-5"]}
import flyte

image = flyte.Image.from_debian_base(name="papermill-env").with_pip_packages(
    "flyteplugins-papermill"
)
env = flyte.TaskEnvironment(name="papermill_env", image=image)
```

## Quick start

```python{hl_lines=[1, 6, "9-15", 19]}
from flyteplugins.papermill import NotebookTask
import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_debian_base(name="my-env").with_pip_packages("flyteplugins-papermill"),
)

add_numbers = NotebookTask(
    name="add_numbers",
    notebook_path="notebooks/basic_math.ipynb",
    task_environment=env,
    inputs={"x": int, "y": float},
    outputs={"result": float},
)

@env.task
def workflow(x: int = 5, y: float = 3.14) -> float:
    return add_numbers(x=x, y=y)
```

`notebook_path` may be relative (resolved against the calling file's directory) or absolute.

## Notebook setup

Each notebook driven by a `NotebookTask` needs two specially tagged cells.

### `parameters` cell

Tag a cell with `parameters` and assign default values matching the names declared in `inputs={...}`. Papermill injects the actual values into a cell appended right after this one at execution time.

```python
# tagged: parameters
x = 0
y = 0.0
```

### `outputs` cell

Tag a cell with `outputs` and call `record_outputs(...)` as the last expression of the cell. The function returns a serialized representation of the values, which Jupyter captures as the cell's displayed output. `NotebookTask` then reads that captured output from the executed notebook to recover the typed values.

```python
# tagged: outputs
from flyteplugins.papermill import record_outputs

record_outputs(result=x + y)
```

`record_outputs` accepts any value that the Flyte type system supports such as primitives, `File`, `Dir`, `DataFrame`, dataclasses, etc. The output names and types must match the `outputs={...}` declaration on the `NotebookTask`.

> [!NOTE]
> Inputs and outputs have different type rules. Inputs are restricted to JSON-serializable primitives plus `File`/`Dir`/`DataFrame` because papermill's parameter mechanism is JSON-only. Outputs go through the full Flyte type engine inside the notebook via `record_outputs`, so dataclasses and any other Flyte-supported type work there.

If a notebook has no outputs, omit the `outputs` cell and don't pass `outputs` to `NotebookTask`. The notebook still runs and its HTML report is rendered, but no values are returned.

## Inputs and outputs

### Supported input types

Notebook parameters are passed through papermill, which only accepts JSON-serializable values. The plugin allows:

- Primitives: `int`, `float`, `str`, `bool`, `list`, `dict`, `None`
- Flyte I/O types: `flyte.io.File`, `flyte.io.Dir`, `flyte.io.DataFrame` (serialized to their path/URI strings)

Passing any other type raises `TypeError` at call time. Wrap unsupported values in a dataclass and serialize them to a primitive container, or write them to a `File`/`Dir` first.

### Complex types: File, Dir, DataFrame

`File`, `Dir` and `DataFrame` are passed to the notebook as plain path/URI strings. Reconstruct them inside the notebook with the provided helpers:

```python
from flyteplugins.papermill import load_file, load_dir, load_dataframe

# input_file, input_dir, input_df were injected as strings by papermill
f  = load_file(input_file)        # -> flyte.io.File
d  = load_dir(input_dir)          # -> flyte.io.Dir
df = load_dataframe(input_df)     # -> flyte.io.DataFrame (parquet by default)
```

`load_dataframe` accepts a `fmt` argument (default `"parquet"`) for non-parquet storage formats.

Jupyter supports top-level `await`, so use it directly for async I/O:

```python{hl_lines=[4, 5]}
import pandas as pd
from flyte.io import DataFrame

pdf = await df.open(pd.DataFrame).all()
output_df = await DataFrame.from_local(pdf)
```

To return a `DataFrame` from a notebook, materialize it as a `flyte.io.DataFrame` and pass it to `record_outputs`:

```python{hl_lines=[6, 7, 9]}
# tagged: outputs
import pandas as pd
from flyte.io import DataFrame
from flyteplugins.papermill import record_outputs

result_df = pd.DataFrame({"name": ["alice", "bob"], "score": [90, 75]})
output = await DataFrame.from_local(result_df)

record_outputs(filtered_df=output, row_count=len(result_df))
```

The same pattern applies to `File` (`await File.from_local(...)`) and `Dir` (`await Dir.from_local(...)`).

### Outputs: single, multiple, none

A `NotebookTask` returns:

- A single value when `outputs` has one entry
- A tuple in the order declared in `outputs` when there are multiple entries
- `None` when `outputs` is omitted

```python{hl_lines=[7, 12]}
# Multiple outputs
text_analysis = NotebookTask(
    name="text_analysis",
    notebook_path="notebooks/text.ipynb",
    task_environment=env,
    inputs={"text": str, "n": int},
    outputs={"repeated": str, "word_count": int, "char_count": int},
)

@env.task
def workflow(text: str, n: int) -> tuple[str, int, int]:
    repeated, word_count, char_count = text_analysis(text=text, n=n)
    return repeated, word_count, char_count
```

```python{hl_lines=[11]}
# No outputs — useful for side-effect-only notebooks (reports, exports)
printer = NotebookTask(
    name="printer",
    notebook_path="notebooks/print_report.ipynb",
    task_environment=env,
    inputs={"message": str},
)

@env.task
def report_workflow(message: str = "hello"):
    printer(message=message)
```

If a declared output is missing from `record_outputs(...)`, `NotebookTask` raises `TypeError` listing the missing names.

## Calling Flyte tasks from notebooks

You can call other Flyte tasks directly from inside a notebook. The plugin injects the parent task's runtime context into the notebook kernel at the start of execution, so task calls are routed through the Flyte controller automatically, so no manual setup required.

When running remotely, each task call is submitted to Flyte and appears as a separate node in the run graph. When running locally, the calls execute in-process as regular Python functions.

```python{hl_lines=[1, 4]}
# Inside a notebook cell
from my_tasks import expensive_task

result = await expensive_task(data=42)
```

Sync tasks can be called the same way:

```python{hl_lines=[3]}
from my_tasks import compute_total

total = compute_total(values=[1, 2, 3])
```

> [!NOTE]
> The setup cell that initializes the runtime context is injected automatically and stripped from the rendered HTML report and the uploaded `.ipynb` files, so it never shows up to users.

## Workflow patterns

### Chaining notebooks

Outputs from one `NotebookTask` can feed directly into another:

```python{hl_lines=[3, 4]}
@env.task
def chained_workflow(a: int, b: float, c: float) -> float:
    intermediate = step1_add(x=a, y=b)
    final = step2_add(x=int(intermediate), y=c)
    return final
```

### Mixing notebooks with regular tasks

`NotebookTask` composes with `@env.task` functions in either direction:

```python{hl_lines=["3-5"]}
@env.task
def mixed_workflow(n: int) -> float:
    doubled = double(n=n)              # regular task
    nb_result = notebook_add(x=doubled, y=100.0)  # notebook task
    return add(a=nb_result, b=0.5)     # regular task
```

### Inline definition

`NotebookTask` can be created inside a task function rather than at module scope. The resolver bakes the notebook path and type schemas into the task spec at registration time, so no module-level reference is required at execution.

```python{hl_lines=[3, 5]}
@env.task
def workflow(x: int = 3, y: float = 1.5) -> int:
    from flyteplugins.papermill import NotebookTask

    nb = NotebookTask(
        name="add_numbers",
        notebook_path="notebooks/basic_math.ipynb",
        task_environment=env,
        inputs={"x": int, "y": float},
        outputs={"result": float},
    )
    return nb(x=x, y=y)
```

### Calling from sync vs. async tasks

`NotebookTask` is internally synchronous. Papermill blocks while the notebook runs. Call it directly from a sync task or use `.aio()` from an async task:

```python{hl_lines=[2, 6, 7]}
@env.task
def sync_parent(x: int) -> float:
    return notebook(x=x)

@env.task
async def async_parent(x: int) -> float:
    return await notebook.aio(x=x)
```

### Running a NotebookTask directly as the entrypoint

A `NotebookTask` can be the workflow entrypoint without wrapping it in another task:

```python{hl_lines=[1, 11]}
nb = NotebookTask(
    name="add_numbers",
    notebook_path="notebooks/basic_math.ipynb",
    task_environment=env,
    inputs={"x": int, "y": float},
    outputs={"result": float},
)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.with_runcontext(mode="remote", copy_style="all").run(nb, x=3, y=1.5)
    print(run.url)
```

## Reports and notebook artifacts

### HTML report (default)

Every `NotebookTask` execution renders the executed notebook to HTML and logs it to the Flyte Report tab for that task. This happens whether the notebook succeeds or fails; see **Papermill > Reports and notebook artifacts > Failure reports** below. The report is on by default and requires no configuration.

![HTML Report](../../_static/images/integrations/papermill/default_report.png)

### Notebook artifacts

By default the executed notebook lives only inside the rendered HTML report. To get the source and executed `.ipynb` files as typed Flyte outputs (so downstream tasks can read them or so they show up as artifacts in the run UI), set `output_notebooks=True`:

```python{hl_lines=[7, 12]}
notebook = NotebookTask(
    name="analysis",
    notebook_path="notebooks/analysis.ipynb",
    task_environment=env,
    inputs={"x": int},
    outputs={"result": float},
    output_notebooks=True,
)

@env.task
def workflow(x: int = 5) -> tuple[float, File, File]:
    result, source_nb, executed_nb = notebook(x=x)
    return result, source_nb, executed_nb
```

When enabled, two outputs are appended to the task's interface automatically:

- `output_notebook`: The source `.ipynb` (no executed cell outputs)
- `output_notebook_executed`: The executed `.ipynb` (with cell outputs)

> [!WARNING]
> The names `output_notebook` and `output_notebook_executed` are reserved when `output_notebooks=True`. Don't use them as your own user output names.

### Clean reports

`report_mode=True` tells papermill to mark input cells with a `source_hidden` flag during execution. The plugin then strips those input cells from both the rendered HTML report and the uploaded `.ipynb` files, so only cell outputs (charts, tables, text) remain. This produces a clean stakeholder-facing report without exposing the underlying code.

```python{hl_lines=[3]}
notebook = NotebookTask(
    ...
    report_mode=True,
    output_notebooks=True,
)
```

![Clean Report](../../_static/images/integrations/papermill/clean_report.png)

### Failure reports

The HTML report is rendered even when the notebook fails. Papermill writes the output notebook cell-by-cell as it executes, so the partial notebook is on disk when an exception propagates out. The plugin renders this partial notebook to HTML and flushes it to the Flyte Report before re-raising the error, giving full visibility into which cell failed and what output the earlier cells produced.

This is especially useful for long-running notebooks: you can inspect partial results without re-running the whole pipeline.

![Failed Report](../../_static/images/integrations/papermill/failed_report.png)

## Spark notebooks

Pass `plugin_config=Spark(...)` to run a notebook inside a Spark driver pod managed by the Spark on Kubernetes Operator:

```python{hl_lines=["8-16"]}
from flyteplugins.papermill import NotebookTask
from flyteplugins.spark import Spark

spark_nb = NotebookTask(
    name="spark_analysis",
    notebook_path="notebooks/spark_analysis.ipynb",
    task_environment=env,
    plugin_config=Spark(
        spark_conf={
            "spark.executor.instances": "2",
            "spark.executor.memory": "2g",
            "spark.executor.cores": "1",
            "spark.driver.memory": "1g",
            "spark.driver.cores": "1",
        },
    ),
    inputs={"data": list},
    outputs={"total": int, "count": int},
)
```

Inside the notebook, build the `SparkSession` directly:

```python
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("FlyteSpark").getOrCreate()
```

> [!WARNING]
> `SparkContext.addPyFile()` is not called for notebook tasks. The notebook kernel runs in a subprocess that cannot share state with the parent task process, so dynamic code distribution via `addPyFile` is not supported. Executor pods use the same Docker image as the driver, so any package needed in UDFs must be installed in the image.

See the [Spark plugin](../spark/_index) page for the full `Spark` configuration reference.

## Local testing

Calling a `NotebookTask` as a regular Python function outside any Flyte runner executes the notebook synchronously through papermill and returns Python values:

```python
result = add_numbers(x=1, y=2.5)
```

In this mode:

- The notebook runs in-process (no remote submission)
- No HTML report is rendered (no task context)
- `File` and `Dir` outputs created inside the notebook resolve to local paths
- No plugin lifecycle hooks fire (so no Spark cluster is provisioned, etc.)

This makes iteration on notebook logic fast. You can run the task from a script, REPL or test without going through Flyte at all.

## Execution options

`NotebookTask` exposes the full set of papermill execution knobs. The snippet below shows example values. See **Papermill > `NotebookTask` reference** for defaults.

```python
NotebookTask(
    name="all_options",
    notebook_path="notebooks/basic_math.ipynb",
    task_environment=env,
    inputs={"x": int, "y": float},
    outputs={"result": float},
    kernel_name="python3",            # default None - use kernel from notebook metadata
    language=None,                    # rarely needed; overrides notebook language
    execution_timeout=300,            # default None - no per-cell timeout
    start_timeout=120,                # default 60 seconds to wait for kernel startup
    log_output=True,                  # default False; stream cell output to task log
    progress_bar=True,                # default True; tqdm-style progress in logs
    report_mode=False,                # default False; True hides input cells in report
    request_save_on_cell_execute=True,  # default True; save after every cell (nbclient)
    engine_name=None,                 # default None - nbclient
    engine_kwargs={"autosave_cell_every": 30},  # extra kwargs forwarded to engine
)
```

> [!NOTE]
> `request_save_on_cell_execute` is largely redundant in remote execution: the plugin always renders and uploads the partial notebook on failure, so crash diagnostics don't depend on it. Leave it on its default unless using a custom engine that requires it.

## `NotebookTask` reference

| Parameter                      | Default | Description                                                                               |
| ------------------------------ | ------- | ----------------------------------------------------------------------------------------- |
| `name`                         | -       | Task name                                                                                 |
| `notebook_path`                | -       | Path to the `.ipynb`, relative to the calling file or absolute                            |
| `task_environment`             | -       | `TaskEnvironment` for registration and remote execution                                   |
| `inputs`                       | `None`  | `{name: type}` dict of notebook inputs                                                    |
| `outputs`                      | `None`  | `{name: type}` dict of notebook outputs                                                   |
| `plugin_config`                | `None`  | Plugin config: currently only `Spark(...)` is supported. Sets the task type accordingly. |
| `kernel_name`                  | `None`  | Jupyter kernel name; `None` uses the kernel from notebook metadata                        |
| `engine_name`                  | `None`  | Papermill engine; `None` uses the default `nbclient` engine                               |
| `log_output`                   | `False` | Stream cell output to the task log                                                        |
| `start_timeout`                | `60`    | Seconds to wait for kernel startup                                                        |
| `execution_timeout`            | `None`  | Per-cell timeout in seconds; `None` means no timeout                                      |
| `report_mode`                  | `False` | Strip input cells from the report and uploaded `.ipynb`                                   |
| `request_save_on_cell_execute` | `True`  | Save notebook after every cell (nbclient engine only)                                     |
| `progress_bar`                 | `True`  | Show a tqdm-style progress bar during execution                                           |
| `language`                     | `None`  | Override notebook language (rarely needed)                                                |
| `engine_kwargs`                | `{}`    | Extra kwargs forwarded to the papermill engine                                            |
| `output_notebooks`             | `False` | Upload source and executed `.ipynb` as `File` task outputs                                |

## Helper functions

These are imported from `flyteplugins.papermill` and called from inside the notebook.

| Function                             | Purpose                                                                                                             |
| ------------------------------------ | ------------------------------------------------------------------------------------------------------------------- |
| `record_outputs(**kwargs)`           | Records outputs from the `outputs`-tagged cell. Must be the cell's last expression. Accepts any Flyte-typed values. |
| `load_file(path)`                    | Reconstructs a `flyte.io.File` from the path string injected by papermill.                                          |
| `load_dir(path)`                     | Reconstructs a `flyte.io.Dir` from the path string injected by papermill.                                           |
| `load_dataframe(uri, fmt="parquet")` | Reconstructs a `flyte.io.DataFrame` from the URI string injected by papermill.                                      |


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/polars ===

# Polars

The Polars plugin adds native support for [Polars](https://pola.rs/) `pl.DataFrame` (eager) and `pl.LazyFrame` (lazy) values as task inputs and outputs. Frames are serialized to and from [Parquet](https://parquet.apache.org/) automatically, so you can pass Polars data between tasks with no manual conversion. Just annotate your task signatures with the Polars types.

Installing the plugin registers encode/decode handlers with Flyte's `flyte.io.DataFrame` transformer engine. That also means a `pl.DataFrame` can be exchanged with the generic `flyte.io.DataFrame` type and with other dataframe backends (pandas, PySpark) through the same Parquet interchange.

## When to use this plugin

- High-performance dataframe processing with Polars' query engine
- Passing large tabular datasets between tasks efficiently via Parquet
- Deferred, optimized computation with `pl.LazyFrame`
- Interoperating with `flyte.io.DataFrame` or other dataframe libraries in the same workflow

## Installation

```bash
pip install flyteplugins-polars
```

Add the plugin to your task image. Installing it registers the Polars type handlers automatically; no explicit registration call is needed:

```
import flyte

image = flyte.Image.from_debian_base(name="polars").with_pip_packages("flyteplugins-polars")

env = flyte.TaskEnvironment(
    name="polars_env",
    image=image,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/polars/polars_example.py*

## Using Polars DataFrames

Annotate task inputs and outputs with `pl.DataFrame`. The plugin encodes returned frames to Parquet and decodes them back on the receiving task:

```
import polars as pl

@env.task
def make_dataframe() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "name": ["Alice", "Bob", "Charlie"],
            "category": ["A", "B", "A"],
            "salary": [55000.0, 75000.0, 72000.0],
            "active": [True, False, True],
        }
    )

@env.task
def summarize(df: pl.DataFrame) -> pl.DataFrame:
    return (
        df.filter(pl.col("active"))
        .group_by("category")
        .agg(pl.col("salary").mean().alias("avg_salary"), pl.len().alias("count"))
        .sort("category")
    )

@env.task
def main() -> pl.DataFrame:
    return summarize(make_dataframe())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/polars/polars_example.py*

Run it with:

```
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/polars/polars_example.py*

## Using LazyFrames

`pl.LazyFrame` is supported the same way and lets Polars defer and optimize the query until the frame is materialized:

```
@env.task
def lazy_summary(lf: pl.LazyFrame) -> pl.LazyFrame:
    return (
        lf.filter(pl.col("active"))
        .group_by("category")
        .agg(pl.col("salary").mean().alias("avg_salary"))
        .sort("category")
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/polars/polars_example.py*

> [!NOTE]
> When a task returns a `pl.LazyFrame` and you want the caller to receive it as a `pl.LazyFrame` (rather than an eagerly collected frame), run with `preserve_original_types=True`:
>
> ```python
> run = flyte.with_runcontext(preserve_original_types=True).run(lazy_summary, lf=my_lazyframe)
> ```

## Interoperating with `flyte.io.DataFrame`

Because the Polars handlers register against the shared dataframe transformer engine, a task can accept the generic `flyte.io.DataFrame` and return a Polars frame, or vice versa. Convert an in-memory Polars frame to a `flyte.io.DataFrame` with `flyte.io.DataFrame.wrap_df()` (preferred over deprecated `from_df()`):

```
import flyte.io

@env.task
def to_flyte_df(df: pl.DataFrame) -> flyte.io.DataFrame:
    return flyte.io.DataFrame.wrap_df(df)

@env.task
def from_flyte_df(df: flyte.io.DataFrame) -> pl.DataFrame:
    return df  # returned to the caller as a Polars DataFrame
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/polars/polars_example.py*

This makes it straightforward to mix Polars with pandas or PySpark tasks in the same workflow: each side declares the dataframe type it wants, and Flyte handles the Parquet interchange.

## Common use cases

- **ETL and feature engineering**: filter, join, and aggregate large tables with Polars' fast query engine across task boundaries.
- **Deferred pipelines**: build up a `pl.LazyFrame` query plan and let Polars optimize it before materialization.
- **Mixed-backend workflows**: bridge Polars and pandas/PySpark tasks through `flyte.io.DataFrame`.

## API reference

See the [Polars API reference](../../api-reference/integrations/polars/_index) for the full list of encode/decode handlers.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/pytorch ===

# PyTorch

The PyTorch plugin lets you run distributed [PyTorch](https://pytorch.org/) training jobs natively on Kubernetes. It uses the [Kubeflow Training Operator](https://github.com/kubeflow/training-operator) to manage multi-node training with PyTorch's elastic launch (`torchrun`).

## When to use this plugin

- Single-node or multi-node distributed training with `DistributedDataParallel` (DDP)
- Elastic training that can scale up and down during execution
- Any workload that uses `torch.distributed` for data-parallel or model-parallel training

## Installation

```bash
pip install flyteplugins-pytorch
```

## Configuration

Create an `Elastic` configuration and pass it as `plugin_config` to a `TaskEnvironment`:

```python
from flyteplugins.pytorch import Elastic

torch_env = flyte.TaskEnvironment(
    name="torch_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")),
    plugin_config=Elastic(
        nnodes=2,
        nproc_per_node=1,
    ),
    image=image,
)
```

### `Elastic` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `nnodes` | `int` or `str` | **Required.** Number of nodes. Use an int for a fixed count or a range string (e.g., `"2:4"`) for elastic training |
| `nproc_per_node` | `int` | **Required.** Number of processes (workers) per node |
| `rdzv_backend` | `str` | Rendezvous backend: `"c10d"` (default), `"etcd"`, or `"etcd-v2"` |
| `max_restarts` | `int` | Maximum worker group restarts (default: `3`) |
| `monitor_interval` | `int` | Agent health check interval in seconds (default: `3`) |
| `run_policy` | `RunPolicy` | Job run policy (cleanup, TTL, deadlines, retries) |

### `RunPolicy` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `clean_pod_policy` | `str` | Pod cleanup policy: `"None"`, `"all"`, or `"Running"` |
| `ttl_seconds_after_finished` | `int` | Seconds to keep pods after job completion |
| `active_deadline_seconds` | `int` | Maximum time the job can run (seconds) |
| `backoff_limit` | `int` | Number of retries before marking the job as failed |

### NCCL tuning parameters

The plugin includes built-in NCCL timeout tuning to reduce failure-detection latency (PyTorch defaults to 1800 seconds):

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `nccl_heartbeat_timeout_sec` | `int` | `300` | NCCL heartbeat timeout (seconds) |
| `nccl_async_error_handling` | `bool` | `False` | Enable async NCCL error handling |
| `nccl_collective_timeout_sec` | `int` | `None` | Timeout for NCCL collective operations |
| `nccl_enable_monitoring` | `bool` | `True` | Enable NCCL monitoring |

### Writing a distributed training task

Tasks using this plugin do not need to be `async`. Initialize the process group and use `DistributedDataParallel` as you normally would with `torchrun`:

```python
import torch
import torch.distributed
from torch.nn.parallel import DistributedDataParallel as DDP

@torch_env.task
def train(epochs: int) -> float:
    torch.distributed.init_process_group("gloo")
    model = DDP(MyModel())
    # ... training loop ...
    return final_loss
```

> [!NOTE]
> When `nnodes=1`, the task runs as a regular Python task (no Kubernetes training job is created). Set `nnodes >= 2` for multi-node distributed training.

## Example

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-pytorch",
#    "torch"
# ]
# main = "torch_distributed_train"
# params = "3"
# ///

import typing

import torch
import torch.distributed
import torch.nn as nn
import torch.optim as optim
from flyteplugins.pytorch.task import Elastic
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset

import flyte

image = flyte.Image.from_debian_base(name="torch").with_pip_packages("flyteplugins-pytorch", pre=True)

torch_env = flyte.TaskEnvironment(
    name="torch_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")),
    plugin_config=Elastic(
        nproc_per_node=1,
        # if you want to do local testing set nnodes=1
        nnodes=2,
    ),
    image=image,
)

class LinearRegressionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

def prepare_dataloader(rank: int, world_size: int, batch_size: int = 2) -> DataLoader:
    """
    Prepare a DataLoader with a DistributedSampler so each rank
    gets a shard of the dataset.
    """
    # Dummy dataset
    x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
    y_train = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
    dataset = TensorDataset(x_train, y_train)

    # Distributed-aware sampler
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)

    return DataLoader(dataset, batch_size=batch_size, sampler=sampler)

def train_loop(epochs: int = 3) -> float:
    """
    A simple training loop for linear regression.
    """
    torch.distributed.init_process_group("gloo")
    model = DDP(LinearRegressionModel())

    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()

    dataloader = prepare_dataloader(
        rank=rank,
        world_size=world_size,
        batch_size=64,
    )

    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    final_loss = 0.0

    for _ in range(epochs):
        for x, y in dataloader:
            outputs = model(x)
            loss = criterion(outputs, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            final_loss = loss.item()
        if torch.distributed.get_rank() == 0:
            print(f"Loss: {final_loss}")

    return final_loss

@torch_env.task
def torch_distributed_train(epochs: int) -> typing.Optional[float]:
    """
    A nested task that sets up a simple distributed training job using PyTorch's
    """
    print("starting launcher")
    loss = train_loop(epochs=epochs)
    print("Training complete")
    return loss

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(torch_distributed_train, epochs=3)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pytorch/pytorch_example.py*

## API reference

See the [PyTorch API reference](../../api-reference/integrations/pytorch/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/ray ===

# Ray

The Ray plugin lets you run [Ray](https://www.ray.io/) jobs natively on Kubernetes. Flyte provisions a transient Ray cluster for each task execution using [KubeRay](https://github.com/ray-project/kuberay) and tears it down on completion.

## When to use this plugin

- Distributed Python workloads (parallel computation, data processing)
- ML training with Ray Train or hyperparameter tuning with Ray Tune
- Ray Serve inference workloads
- Any workload that benefits from Ray's actor model or task parallelism

## Installation

```bash
pip install flyteplugins-ray
```

Your task image must also include a compatible version of Ray:

```python
image = (
    flyte.Image.from_debian_base(name="ray")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray")
)
```

## Configuration

Create a `RayJobConfig` and pass it as `plugin_config` to a `TaskEnvironment`:

```python
from flyteplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
)
```

### `RayJobConfig` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `worker_node_config` | `List[WorkerNodeConfig]` | **Required.** List of worker group configurations |
| `head_node_config` | `HeadNodeConfig` | Head node configuration (optional) |
| `enable_autoscaling` | `bool` | Enable Ray autoscaler (default: `False`) |
| `runtime_env` | `dict` | Ray runtime environment (pip packages, env vars, etc.) |
| `address` | `str` | Connect to an existing Ray cluster instead of provisioning one |
| `shutdown_after_job_finishes` | `bool` | Shut down the cluster after the job completes (default: `False`) |
| `ttl_seconds_after_finished` | `int` | Seconds to keep the cluster after completion before cleanup |

### `WorkerNodeConfig` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `group_name` | `str` | **Required.** Name of this worker group |
| `replicas` | `int` | **Required.** Number of worker replicas |
| `min_replicas` | `int` | Minimum replicas (for autoscaling) |
| `max_replicas` | `int` | Maximum replicas (for autoscaling) |
| `ray_start_params` | `Dict[str, str]` | Ray start parameters for workers |
| `requests` | `Resources` | Resource requests per worker |
| `limits` | `Resources` | Resource limits per worker |
| `pod_template` | `PodTemplate` | Full pod template (mutually exclusive with `requests`/`limits`) |

### `HeadNodeConfig` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `ray_start_params` | `Dict[str, str]` | Ray start parameters for the head node |
| `requests` | `Resources` | Resource requests for the head node |
| `limits` | `Resources` | Resource limits for the head node |
| `pod_template` | `PodTemplate` | Full pod template (mutually exclusive with `requests`/`limits`) |

### Connecting to an existing cluster

To connect to an existing Ray cluster instead of provisioning a new one, set the `address` parameter:

```python
ray_config = RayJobConfig(
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    address="ray://existing-cluster:10001",
)
```

## Examples

The following example shows how to configure Ray in a `TaskEnvironment`. Flyte automatically provisions a Ray cluster for each task using this configuration:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-ray",
#    "ray[default]==2.46.0"
# ]
# main = "hello_ray_nested"
# params = "3"
# ///

import asyncio
import typing

import ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

import flyte.remote
import flyte.storage

@ray.remote
def f(x):
    return x * x

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

image = (
    flyte.Image.from_debian_base(name="ray")
    .with_apt_packages("wget")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray", "pip", "mypy")
)

task_env = flyte.TaskEnvironment(
    name="hello_ray", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")),
    depends_on=[task_env],
)

@task_env.task()
async def hello_ray():
    await asyncio.sleep(20)
    print("Hello from the Ray task!")

@ray_env.task
async def hello_ray_nested(n: int = 3) -> typing.List[int]:
    print("running ray task")
    t = asyncio.create_task(hello_ray())
    futures = [f.remote(i) for i in range(n)]
    res = ray.get(futures)
    await t
    return res

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(hello_ray_nested)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray_example.py*

The next example demonstrates how Flyte can create ephemeral Ray clusters and run a subtask that connects to an existing Ray cluster:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-ray",
#    "ray[default]==2.46.0"
# ]
# main = "create_ray_cluster"
# params = ""
# ///

import os
import typing

import ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

import flyte.storage

@ray.remote
def f(x):
    return x * x

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=3600,
)

image = (
    flyte.Image.from_debian_base(name="ray")
    .with_apt_packages("wget")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray")
)

task_env = flyte.TaskEnvironment(
    name="ray_client", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
ray_env = flyte.TaskEnvironment(
    name="ray_cluster",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(2, 4), memory=("2000Mi", "4000Mi")),
    depends_on=[task_env],
)

@task_env.task()
async def hello_ray(cluster_ip: str) -> typing.List[int]:
    """
    Run a simple Ray task that connects to an existing Ray cluster.
    """
    ray.init(address=f"ray://{cluster_ip}:10001")
    futures = [f.remote(i) for i in range(5)]
    res = ray.get(futures)
    return res

@ray_env.task
async def create_ray_cluster() -> str:
    """
    Create a Ray cluster and return the head node IP address.
    """
    print("creating ray cluster")
    cluster_ip = os.getenv("MY_POD_IP")
    if cluster_ip is None:
        raise ValueError("MY_POD_IP environment variable is not set")
    return f"{cluster_ip}"

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(create_ray_cluster)
    run.wait()
    print("run url:", run.url)
    print("cluster created, running ray task")
    print("ray address:", run.outputs()[0])
    run = flyte.run(hello_ray, cluster_ip=run.outputs()[0])
    print("run url:", run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray_existing_example.py*

## API reference

See the [Ray API reference](../../api-reference/integrations/ray/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/snowflake ===

# Snowflake

The Snowflake connector lets you run SQL queries against [Snowflake](https://www.snowflake.com/) directly from Flyte tasks. Queries are submitted asynchronously and polled for completion, so they don't block a worker while waiting for results.

The connector supports:

- Parameterized SQL queries with typed inputs
- Key-pair and password-based authentication
- Returns query results as DataFrames
- Automatic links to the Snowflake query dashboard in the Flyte UI
- Query cancellation on task abort

## Installation

```bash
pip install flyteplugins-snowflake
```

This installs the Snowflake Python connector and the `cryptography` library for key-pair authentication.

## Quick start

Here's a minimal example that runs a SQL query on Snowflake:

```python {hl_lines=[2, 4, 12]}
from flyte.io import DataFrame
from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig

config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

count_users = Snowflake(
    name="count_users",
    query_template="SELECT COUNT(*) FROM users",
    plugin_config=config,
    output_dataframe_type=DataFrame,
)
```

This defines a task called `count_users` that runs `SELECT COUNT(*) FROM users` on the configured Snowflake instance. When executed, the connector:

1. Connects to Snowflake using the provided configuration
2. Submits the query asynchronously
3. Polls until the query completes or fails
4. Provides a link to the query in the Snowflake dashboard

![Snowflake Link](../../_static/images/integrations/snowflake/ui.png)

To run the task, create a `TaskEnvironment` from it and execute it locally or remotely:

```python {hl_lines=3}
import flyte

snowflake_env = flyte.TaskEnvironment.from_task("snowflake_env", count_users)

if __name__ == "__main__":
    flyte.init_from_config()

    # Run locally (connector runs in-process, requires credentials and packages locally)
    run = flyte.with_runcontext(mode="local").run(count_users)

    # Run remotely (connector runs as a service in your data plane)
    run = flyte.with_runcontext(mode="remote").run(count_users)

    print(run.url)
```

> [!NOTE]
> The `TaskEnvironment` created by `from_task` does not need an image or pip packages. Snowflake tasks are connector tasks, which means the query executes on the connector service, not in your task container. In `local` mode, the connector runs in-process and requires `flyteplugins-snowflake` and credentials to be available on your machine. In `remote` mode, the connector runs as a service in your data plane.

## Configuration

The `SnowflakeConfig` dataclass defines the connection settings for your Snowflake instance.

### Required fields

| Field       | Type  | Description                                             |
| ----------- | ----- | ------------------------------------------------------- |
| `account`   | `str` | Snowflake account identifier (e.g. `"myorg-myaccount"`) |
| `database`  | `str` | Target database name                                    |
| `schema`    | `str` | Target schema name (e.g. `"PUBLIC"`)                    |
| `warehouse` | `str` | Compute warehouse to use for query execution            |
| `user`      | `str` | Snowflake username                                      |

### Additional connection parameters

Use `connection_kwargs` to pass any additional parameters supported by the [Snowflake Python connector](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api). This is a dictionary that gets forwarded directly to `snowflake.connector.connect()`.

Common options include:

| Parameter       | Type  | Description                                                                |
| --------------- | ----- | -------------------------------------------------------------------------- |
| `role`          | `str` | Snowflake role to use for the session                                      |
| `authenticator` | `str` | Authentication method (e.g. `"snowflake"`, `"externalbrowser"`, `"oauth"`) |
| `token`         | `str` | OAuth token when using `authenticator="oauth"`                             |
| `login_timeout` | `int` | Timeout in seconds for the login request                                   |

Example with a role:

```python {hl_lines=8}
config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    connection_kwargs={
        "role": "DATA_ANALYST",
    },
)
```

## Authentication

The connector supports two authentication approaches: key-pair authentication, and password-based or other authentication methods provided through `connection_kwargs`.

### Key-pair authentication

Key-pair authentication is the recommended approach for automated workloads. Pass the names of the Flyte secrets containing the private key and optional passphrase:

```python {hl_lines=[5, 6]}
query = Snowflake(
    name="secure_query",
    query_template="SELECT * FROM sensitive_data",
    plugin_config=config,
    snowflake_private_key="my-snowflake-private-key",
    snowflake_private_key_passphrase="my-snowflake-pk-passphrase",
)
```

The `snowflake_private_key` parameter is the name of the secret (or secret key) that contains your PEM-encoded private key. The `snowflake_private_key_passphrase` parameter is the name of the secret (or secret key) that contains the passphrase, if your key is encrypted. If your key is not encrypted, omit the passphrase parameter.

The connector decodes the PEM key and converts it to DER format for Snowflake authentication.

> [!NOTE]
> If your credentials are stored in a secret group, you can pass `secret_group` to the `Snowflake` task. The plugin expects `snowflake_private_key` and
> `snowflake_private_key_passphrase` to be keys within the same secret group.

### Password authentication

Send the password via `connection_kwargs`:

```python {hl_lines=8}
config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    connection_kwargs={
        "password": "my-password",
    },
)
```

### OAuth authentication

For OAuth-based authentication, specify the authenticator and token:

```python {hl_lines=["8-9"]}
config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    connection_kwargs={
        "authenticator": "oauth",
        "token": "<oauth-token>",
    },
)
```

## Query templating

Use the `inputs` parameter to define typed inputs for your query. Input values are bound using the `%(param)s` syntax supported by the [Snowflake Python connector](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api), which handles type conversion and escaping automatically.

### Supported input types

The `inputs` dictionary maps parameter names to Python values. Supported scalar types include `str`, `int`, `float`, and `bool`.

To insert multiple rows in a single query, you can also provide lists as input values. When using list inputs, be sure to set `batch=True` on the `Snowflake` task. This enables automatic batching, where the inputs are expanded and sent as a single multi-row query instead of you having to write multiple individual statements.

### Batched `INSERT` with list inputs

When `batch=True` is enabled, a parameterized `INSERT` query with list inputs is automatically expanded into a multi-row `VALUES` statement.

Example:

```python
query = "INSERT INTO t (a, b) VALUES (%(a)s, %(b)s)"
inputs = {"a": [1, 2], "b": ["x", "y"]}
```

This is expanded into:

```sql
INSERT INTO t (a, b)
VALUES (%(a_0)s, %(b_0)s), (%(a_1)s, %(b_1)s)
```

with the following flattened parameters:

```python
flat_params = {
    "a_0": 1,
    "b_0": "x",
    "a_1": 2,
    "b_1": "y",
}
```

#### Constraints

- The query must contain exactly one `VALUES (...)` clause.
- All list inputs must have the same non-zero length.

### Parameterized `SELECT`

```python {hl_lines=[5, 7]}
from flyte.io import DataFrame

events_by_date = Snowflake(
    name="events_by_date",
    query_template="SELECT * FROM events WHERE event_date = %(event_date)s",
    plugin_config=config,
    inputs={"event_date": str},
    output_dataframe_type=DataFrame,
)
```

You can call the task with the required inputs:

```python {hl_lines=3}
@env.task
async def fetch_events() -> DataFrame:
    return await events_by_date(event_date="2024-01-15")
```

### Multiple inputs

You can define multiple input parameters of different types:

```python {hl_lines=["4-8", "12-15"]}
filtered_events = Snowflake(
    name="filtered_events",
    query_template="""
        SELECT * FROM events
        WHERE event_date >= %(start_date)s
          AND event_date <= %(end_date)s
          AND region = %(region)s
          AND score > %(min_score)s
    """,
    plugin_config=config,
    inputs={
        "start_date": str,
        "end_date": str,
        "region": str,
        "min_score": float,
    },
    output_dataframe_type=DataFrame,
)
```

> [!NOTE]
> The query template is normalized before execution: newlines and tabs are replaced with spaces, and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution.

## Retrieving query results

If your query produces output, set `output_dataframe_type` to capture the results. `output_dataframe_type` accepts `DataFrame` from `flyte.io`. This is a meta-wrapper type that represents tabular results and can be materialized into a concrete DataFrame implementation using `open()` where you specify the target type and `all()`.

```python {hl_lines=13}
from flyte.io import DataFrame

top_customers = Snowflake(
    name="top_customers",
    query_template="""
        SELECT customer_id, SUM(amount) AS total_spend
        FROM orders
        GROUP BY customer_id
        ORDER BY total_spend DESC
        LIMIT 100
    """,
    plugin_config=config,
    output_dataframe_type=DataFrame,
)
```

At present, only `pandas.DataFrame` is supported. The results are returned directly when you call the task:

```python {hl_lines=6}
import pandas as pd

@env.task
async def analyze_top_customers() -> dict:
    df = await top_customers()
    pandas_df = await df.open(pd.DataFrame).all()
    total_spend = pandas_df["total_spend"].sum()
    return {"total_spend": float(total_spend)}
```

If you specify `pandas.DataFrame` as the `output_dataframe_type`, you do not need to call `open()` and `all()` to materialize the results.

```python {hl_lines=[1, 13, "18-19"]}
import pandas as pd

top_customers = Snowflake(
    name="top_customers",
    query_template="""
        SELECT customer_id, SUM(amount) AS total_spend
        FROM orders
        GROUP BY customer_id
        ORDER BY total_spend DESC
        LIMIT 100
    """,
    plugin_config=config,
    output_dataframe_type=pd.DataFrame,
)

@env.task
async def analyze_top_customers() -> dict:
    df = await top_customers()
    total_spend = df["total_spend"].sum()
    return {"total_spend": float(total_spend)}
```

> [!NOTE]
> Be sure to inject the `SNOWFLAKE_PRIVATE_KEY` and `SNOWFLAKE_PRIVATE_KEY_PASSPHRASE` environment variables as secrets into your downstream tasks, as they must have access to Snowflake credentials in order to retrieve the DataFrame results. More on how you can **Tasks > Configure tasks > Secrets**.

If you don't need query results (for example, `DDL` statements or `INSERT` queries), omit `output_dataframe_type`.

## End-to-end example

Here's a complete workflow that uses the Snowflake connector as part of a data pipeline. The workflow creates a staging table, inserts records, queries aggregated results and processes them in a downstream task.

```
import flyte
from flyte.io import DataFrame
from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig

config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    connection_kwargs={
        "role": "ETL_ROLE",
    },
)

# Step 1: Create the staging table if it doesn't exist
create_staging = Snowflake(
    name="create_staging",
    query_template="""
        CREATE TABLE IF NOT EXISTS staging.daily_events (
            event_id STRING,
            event_date DATE,
            user_id STRING,
            event_type STRING,
            payload VARIANT
        )
    """,
    plugin_config=config,
    snowflake_private_key="snowflake",
    snowflake_private_key_passphrase="snowflake_passphrase",
)

# Step 2: Insert a record into the staging table
insert_events = Snowflake(
    name="insert_event",
    query_template="""
        INSERT INTO staging.daily_events (event_id, event_date, user_id, event_type)
        VALUES (%(event_id)s, %(event_date)s, %(user_id)s, %(event_type)s)
    """,
    plugin_config=config,
    inputs={
        "event_id": list[str],
        "event_date": list[str],
        "user_id": list[str],
        "event_type": list[str],
    },
    snowflake_private_key="snowflake",
    snowflake_private_key_passphrase="snowflake_passphrase",
    batch=True,
)

# Step 3: Query aggregated results for a given date
daily_summary = Snowflake(
    name="daily_summary",
    query_template="""
        SELECT
            event_type,
            COUNT(*) AS event_count,
            COUNT(DISTINCT user_id) AS unique_users
        FROM staging.daily_events
        WHERE event_date = %(report_date)s
        GROUP BY event_type
        ORDER BY event_count DESC
    """,
    plugin_config=config,
    inputs={"report_date": str},
    output_dataframe_type=DataFrame,
    snowflake_private_key="snowflake",
    snowflake_private_key_passphrase="snowflake_passphrase",
)

# Create environments for all Snowflake tasks
snowflake_env = flyte.TaskEnvironment.from_task(
    "snowflake_env", create_staging, insert_events, daily_summary
)

# Main pipeline environment that depends on the Snowflake task environments
env = flyte.TaskEnvironment(
    name="analytics_env",
    resources=flyte.Resources(memory="512Mi"),
    image=flyte.Image.from_debian_base(name="analytics").with_pip_packages(
        "flyteplugins-snowflake", pre=True
    ),
    secrets=[
        flyte.Secret(key="snowflake", as_env_var="SNOWFLAKE_PRIVATE_KEY"),
        flyte.Secret(
            key="snowflake_passphrase", as_env_var="SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"
        ),
    ],
    depends_on=[snowflake_env],
)

# Step 4: Process the results in Python
@env.task
async def generate_report(summary: DataFrame) -> dict:
    import pandas as pd

    df = await summary.open(pd.DataFrame).all()
    total_events = df["event_count"].sum()
    top_event = df.iloc[0]["event_type"]
    return {
        "total_events": int(total_events),
        "top_event_type": top_event,
        "event_types_count": len(df),
    }

# Compose the pipeline
@env.task
async def run_daily_pipeline(
    event_ids: list[str],
    event_dates: list[str],
    user_ids: list[str],
    event_types: list[str],
) -> dict:
    await create_staging()
    await insert_events(
        event_id=event_ids,
        event_date=event_dates,
        user_id=user_ids,
        event_type=event_types,
    )
    summary = await daily_summary(report_date=event_dates[0])
    return await generate_report(summary=summary)

if __name__ == "__main__":
    flyte.init_from_config()

    # Run locally
    run = flyte.with_runcontext(mode="local").run(
        run_daily_pipeline,
        event_ids=["event-1", "event-2"],
        event_dates=["2023-01-01", "2023-01-02"],
        user_ids=["user-1", "user-2"],
        event_types=["click", "view"],
    )

    # Or run remotely
    run = flyte.with_runcontext(mode="remote").run(
        run_daily_pipeline,
        event_ids=["event-1", "event-2"],
        event_dates=["2023-01-01", "2023-01-02"],
        user_ids=["user-1", "user-2"],
        event_types=["click", "view"],
    )

    print(run.url)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/snowflake/example.py*


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/spark ===

# Spark

The Spark plugin lets you run [Apache Spark](https://spark.apache.org/) jobs natively on Kubernetes. Flyte manages the full cluster lifecycle: provisioning a transient Spark cluster for each task execution, running the job, and tearing the cluster down on completion.

Under the hood, the plugin uses the [Spark on Kubernetes Operator](https://github.com/GoogleCloudPlatform/spark-on-k8s-operator) to create and manage Spark applications. No external Spark service or long-running cluster is required.

## When to use this plugin

- Large-scale data processing and ETL pipelines
- Jobs that benefit from Spark's distributed execution engine (Spark SQL, PySpark, Spark MLlib)
- Workloads that need Hadoop-compatible storage access (S3, GCS, HDFS)

## Installation

```bash
pip install flyteplugins-spark
```

## Configuration

Create a `Spark` configuration and pass it as `plugin_config` to a `TaskEnvironment`:

```python
from flyteplugins.spark import Spark

spark_config = Spark(
    spark_conf={
        "spark.driver.memory": "3000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
    },
)

spark_env = flyte.TaskEnvironment(
    name="spark_env",
    plugin_config=spark_config,
    image=image,
)
```

### `Spark` parameters

| Parameter | Type | Description |
|-----------|------|-------------|
| `spark_conf` | `Dict[str, str]` | Spark configuration key-value pairs (e.g., executor memory, cores, instances) |
| `hadoop_conf` | `Dict[str, str]` | Hadoop configuration key-value pairs (e.g., S3/GCS access settings) |
| `executor_path` | `str` | Path to the Python binary for PySpark executors |
| `applications_path` | `str` | Path to the main Spark application file |
| `driver_pod` | `PodTemplate` | Pod template for the Spark driver pod |
| `executor_pod` | `PodTemplate` | Pod template for the Spark executor pods |

### Accessing the Spark session

Inside a Spark task, the `SparkSession` is available through the task context:

```python
from flyte._context import internal_ctx

@spark_env.task
async def my_spark_task() -> float:
    ctx = internal_ctx()
    spark = ctx.data.task_context.data["spark_session"]
    # Use spark as a normal SparkSession
    df = spark.read.parquet("s3://my-bucket/data.parquet")
    return df.count()
```

### Overriding configuration at runtime

You can override Spark configuration for individual task calls using `.override()`:

```python
from copy import deepcopy

updated_config = deepcopy(spark_config)
updated_config.spark_conf["spark.executor.instances"] = "4"

result = await my_spark_task.override(plugin_config=updated_config)()
```

## Example

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-spark"
# ]
# main = "hello_spark_nested"
# params = "3"
# ///

import random
from copy import deepcopy
from operator import add

from flyteplugins.spark.task import Spark

import flyte.remote
from flyte._context import internal_ctx

image = (
    flyte.Image.from_base("apache/spark-py:v3.4.0")
    .clone(name="spark", python_version=(3, 10), registry="ghcr.io/flyteorg")
    .with_pip_packages("flyteplugins-spark", pre=True)
)

task_env = flyte.TaskEnvironment(
    name="get_pi", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)

spark_conf = Spark(
    spark_conf={
        "spark.driver.memory": "3000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
        "spark.kubernetes.file.upload.path": "/opt/spark/work-dir",
        "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar",
    },
)

spark_env = flyte.TaskEnvironment(
    name="spark_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")),
    plugin_config=spark_conf,
    image=image,
    depends_on=[task_env],
)

def f(_):
    x = random.random() * 2 - 1
    y = random.random() * 2 - 1
    return 1 if x**2 + y**2 <= 1 else 0

@task_env.task
async def get_pi(count: int, partitions: int) -> float:
    return 4.0 * count / partitions

@spark_env.task
async def hello_spark_nested(partitions: int = 3) -> float:
    n = 1 * partitions
    ctx = internal_ctx()
    spark = ctx.data.task_context.data["spark_session"]
    count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add)

    return await get_pi(count, partitions)

@task_env.task
async def spark_overrider(executor_instances: int = 3, partitions: int = 4) -> float:
    updated_spark_conf = deepcopy(spark_conf)
    updated_spark_conf.spark_conf["spark.executor.instances"] = str(executor_instances)
    return await hello_spark_nested.override(plugin_config=updated_spark_conf)(partitions=partitions)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(hello_spark_nested)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/spark/spark_example.py*

## API reference

See the [Spark API reference](../../api-reference/integrations/spark/_index) for full details.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb ===

# Weights & Biases

[Weights & Biases](https://wandb.ai) (W&B) is a platform for tracking machine learning experiments, visualizing metrics and optimizing hyperparameters. This plugin integrates W&B with Flyte, enabling you to:

- Automatically initialize W&B runs in your tasks without boilerplate
- Link directly from the Flyte UI to your W&B runs and sweeps
- Share W&B runs across parent and child tasks
- Track distributed training jobs across multiple GPUs and nodes
- Run hyperparameter sweeps with parallel agents

## Installation

```bash
pip install flyteplugins-wandb
```

You also need a W&B API key. Store it as a Flyte secret so your tasks can authenticate with W&B.

## Quick start

Here's a minimal example that logs metrics to W&B from a Flyte task:

```
import flyte

from flyteplugins.wandb import get_wandb_run, wandb_config, wandb_init

env = flyte.TaskEnvironment(
    name="wandb-example",
    image=flyte.Image.from_debian_base(name="wandb-example").with_pip_packages(
        "flyteplugins-wandb"
    ),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

@wandb_init
@env.task
async def train_model() -> str:
    wandb_run = get_wandb_run()

    # Your training code here
    for epoch in range(10):
        loss = 1.0 / (epoch + 1)
        wandb_run.log({"epoch": epoch, "loss": loss})

    return "Training complete"

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.with_runcontext(
        custom_context=wandb_config(
            project="my-project",
            entity="my-team",
        ),
    ).run(train_model)

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/quick_start.py*

This example demonstrates the core pattern:

1. **Define a task environment** with the plugin installed and your W&B API key as a secret
2. **Decorate your task** with `@wandb_init` (must be the outermost decorator, above `@env.task`)
3. **Access the run** with `get_wandb_run()` to log metrics
4. **Provide configuration** via `wandb_config()` when running the task

The plugin handles calling `wandb.init()` and `wandb.finish()` for you, and automatically adds a link to the W&B run in the Flyte UI.

![UI](../../_static/images/integrations/wandb/ui.png)

## What's next

This integration guide is split into focused sections, depending on how you want to use Weights & Biases with Flyte:

- ****Weights & Biases > Experiments****: Create and manage W&B runs from Flyte tasks.
- ****Weights & Biases > Distributed training****: Track experiments across multi-GPU and multi-node training jobs.
- ****Weights & Biases > Sweeps****: Run hyperparameter searches and manage sweep execution from Flyte tasks.
- ****Weights & Biases > Downloading logs****: Download logs and execution metadata from Weights & Biases.
- ****Weights & Biases > Constraints and best practices****: Learn about limitations, edge cases and recommended patterns.
- ****Weights & Biases > Manual integration****: Use Weights & Biases directly in Flyte tasks without decorators or helpers.

> **📝 Note**
>
> We've included [additional examples](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/wandb/examples) developed while testing edge cases of the plugin.

## Subpages
- **Weights & Biases > Experiments**
- **Weights & Biases > Distributed training**
- **Weights & Biases > Sweeps**
- **Weights & Biases > Downloading logs**
- **Weights & Biases > Constraints and best practices**
- **Weights & Biases > Manual integration**


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/experiments ===

# Experiments

The `@wandb_init` decorator automatically initializes a W&B run when your task executes and finishes it when the task completes. This section covers the different ways to use it.

## Basic usage

Apply `@wandb_init` as the outermost decorator on your task:

```python {hl_lines=1}
@wandb_init
@env.task
async def my_task() -> str:
    run = get_wandb_run()
    run.log({"metric": 42})
    return "done"
```

The decorator:

- Calls `wandb.init()` before your task code runs
- Calls `wandb.finish()` after your task completes (or fails)
- Adds a link to the W&B run in the Flyte UI

You can also use it on synchronous tasks:

```python {hl_lines=[1, 3]}
@wandb_init
@env.task
def my_sync_task() -> str:
    run = get_wandb_run()
    run.log({"metric": 42})
    return "done"
```

## Accessing the run object

Use `get_wandb_run()` to access the current W&B run object:

```python {hl_lines=6}
from flyteplugins.wandb import get_wandb_run

@wandb_init
@env.task
async def train() -> str:
    run = get_wandb_run()

    # Log metrics
    run.log({"loss": 0.5, "accuracy": 0.9})

    # Access run properties
    print(f"Run ID: {run.id}")
    print(f"Run URL: {run.url}")
    print(f"Project: {run.project}")

    # Log configuration
    run.config.update({"learning_rate": 0.001, "batch_size": 32})

    return run.id
```

## Parent-child task relationships

When a parent task calls child tasks, the plugin can share the same W&B run across all of them. This is useful for tracking an entire workflow in a single run.

```python {hl_lines=[1, 9, 16]}
@wandb_init
@env.task
async def child_task(x: int) -> int:
    run = get_wandb_run()
    run.log({"child_metric": x * 2})
    return x * 2

@wandb_init
@env.task
async def parent_task() -> int:
    run = get_wandb_run()
    run.log({"parent_metric": 100})

    # Child task shares the parent's run by default
    result = await child_task(5)

    return result
```

By default (`run_mode="auto"`), child tasks reuse their parent's W&B run. All metrics logged by the parent and children appear in the same run in the W&B UI.

## Run modes

The `run_mode` parameter controls how tasks create or reuse W&B runs. There are three modes:

| Mode             | Behavior                                                                   |
| ---------------- | -------------------------------------------------------------------------- |
| `auto` (default) | Create a new run if no parent run exists, otherwise reuse the parent's run |
| `new`            | Always create a new run, even if a parent run exists                       |
| `shared`         | Always reuse the parent's run (fails if no parent run exists)              |

### Using `run_mode="new"` for independent runs

```python {hl_lines=1}
@wandb_init(run_mode="new")
@env.task
async def independent_child(x: int) -> int:
    run = get_wandb_run()
    # This task gets its own separate run
    run.log({"independent_metric": x})
    return x

@wandb_init
@env.task
async def parent_task() -> str:
    run = get_wandb_run()
    parent_run_id = run.id

    # This child creates its own run
    await independent_child(5)

    # Parent's run is unchanged
    assert run.id == parent_run_id
    return parent_run_id
```

### Using `run_mode="shared"` for explicit sharing

```python {hl_lines=1}
@wandb_init(run_mode="shared")
@env.task
async def must_share_run(x: int) -> int:
    # This task requires a parent run to exist
    # It will fail if called as a top-level task
    run = get_wandb_run()
    run.log({"shared_metric": x})
    return x
```

## Configuration with `wandb_config`

Use `wandb_config()` to configure W&B runs. You can set it at the workflow level or override it for specific tasks, allowing you to provide configuration values at runtime.

### Workflow-level configuration

```python {hl_lines=["5-9"]}
if __name__ == "__main__":
    flyte.init_from_config()

    flyte.with_runcontext(
        custom_context=wandb_config(
            project="my-project",
            entity="my-team",
            tags=["experiment-1", "production"],
            config={"model": "resnet50", "dataset": "imagenet"},
        ),
    ).run(train_task)
```

### Overriding configuration for child tasks

Use `wandb_config()` as a context manager to override settings for specific child task calls:

```python {hl_lines=[8, 12]}
@wandb_init
@env.task
async def parent_task() -> str:
    run = get_wandb_run()
    run.log({"parent_metric": 100})

    # Override tags and config for this child call
    with wandb_config(tags=["special-run"], config={"learning_rate": 0.01}):
        await child_task(10)

    # Override run_mode for this child call
    with wandb_config(run_mode="new"):
        await child_task(20)  # Gets its own run

    return "done"
```

## Using traces with W&B runs

Flyte traces can access the parent task's W&B run without needing the `@wandb_init` decorator. This is useful for helper functions that should log to the same run:

```python {hl_lines=[1, 3]}
@flyte.trace
async def log_validation_metrics(accuracy: float, f1: float):
    run = get_wandb_run()
    if run:
        run.log({"val_accuracy": accuracy, "val_f1": f1})

@wandb_init
@env.task
async def train_and_validate() -> str:
    run = get_wandb_run()

    # Training loop
    for epoch in range(10):
        run.log({"train_loss": 1.0 / (epoch + 1)})

    # Trace logs to the same run
    await log_validation_metrics(accuracy=0.95, f1=0.92)

    return "done"
```

> **📝 Note**
>
> Do not apply `@wandb_init` to traces. Traces automatically access the parent task's run via `get_wandb_run()`.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/distributed_training ===

# Distributed training

When running distributed training jobs, multiple processes run simultaneously across GPUs. The `@wandb_init` decorator automatically detects distributed training environments and coordinates W&B logging across processes.

The plugin:

- Auto-detects distributed context from environment variables (set by launchers like `torchrun`)
- Controls which processes initialize W&B runs based on the `run_mode` and `rank_scope` parameters
- Generates unique run IDs that distinguish between workers and ranks
- Adds links to W&B runs in the Flyte UI

## Quick start

Here's a minimal single-node example that logs metrics from a distributed training task. By default (`run_mode="auto"`, `rank_scope="global"`), only rank 0 logs to W&B:

```
import flyte
import torch
import torch.distributed
from flyteplugins.pytorch.task import Elastic
from flyteplugins.wandb import get_wandb_run, wandb_config, wandb_init

image = flyte.Image.from_debian_base(name="torch-wandb").with_pip_packages(
    "flyteplugins-wandb", "flyteplugins-pytorch"
)

env = flyte.TaskEnvironment(
    name="distributed_env",
    image=image,
    resources=flyte.Resources(gpu="A100:2"),
    plugin_config=Elastic(nproc_per_node=2, nnodes=1),
    secrets=flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY"),
)

@wandb_init
@env.task
def train() -> float:
    torch.distributed.init_process_group("nccl")

    # Only rank 0 gets a W&B run object; others get None
    run = get_wandb_run()

    # Simulate training
    for step in range(100):
        loss = 1.0 / (step + 1)

        # Safe to call on all ranks - only rank 0 actually logs
        if run:
            run.log({"loss": loss, "step": step})

    torch.distributed.destroy_process_group()
    return loss

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.with_runcontext(
        custom_context=wandb_config(project="my-project", entity="my-team")
    ).run(train)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/distributed_training_quick_start.py*

A few things to note:

1. Use the `Elastic` plugin to configure distributed training (number of processes, nodes)
2. Apply `@wandb_init` as the outermost decorator
3. Check if `run` is not None before logging - only the primary rank has a run object in `auto` mode

> **📝 Note**
>
> The `if run:` check is always safe regardless of run mode. In `shared` and `new` modes all ranks get a run object, but the check doesn't hurt and keeps your code portable across modes.

![Single-node auto](../../_static/images/integrations/wandb/single_node_auto_flyte.png)

## Run modes in distributed training

The `run_mode` parameter controls how W&B runs are created across distributed processes. The behavior differs between single-node (one machine, multiple GPUs) and multi-node (multiple machines) setups.

### Single-node behavior

| Mode             | Which ranks log       | Result                                 |
| ---------------- | --------------------- | -------------------------------------- |
| `auto` (default) | Only rank 0           | 1 W&B run                              |
| `shared`         | All ranks to same run | 1 W&B run with metrics labeled by rank |
| `new`            | Each rank separately  | N W&B runs (grouped in UI)             |

### Multi-node behavior

For multi-node training, the `rank_scope` parameter controls the granularity of W&B runs:

- **`global`** (default): Treat all workers as one unit
- **`worker`**: Treat each worker/node independently

The combination of `run_mode` and `rank_scope` determines logging behavior:

| `run_mode` | `rank_scope` | Who initializes W&B    | W&B Runs | Grouping |
| ---------- | ------------ | ---------------------- | -------- | -------- |
| `auto`     | `global`     | Global rank 0 only     | 1        | -        |
| `auto`     | `worker`     | Local rank 0 per worker | N        | -        |
| `shared`   | `global`     | All ranks (shared globally) | 1        | -        |
| `shared`   | `worker`     | All ranks (shared per worker) | N        | -        |
| `new`      | `global`     | All ranks              | N × M    | 1 group  |
| `new`      | `worker`     | All ranks              | N × M    | N groups |

Where `N` = number of workers/nodes, `M` = processes per worker.

### Choosing run mode and rank scope

- **`auto`** (recommended): Use when you want clean dashboards with minimal runs. Most metrics (loss, accuracy) are the same across ranks after gradient synchronization, so logging from one rank is sufficient.
- **`shared`**: Use when you need to compare metrics across ranks in a single view. Each rank's metrics are labeled with an `x_label` identifier. Useful for debugging load imbalance or per-GPU throughput.
- **`new`**: Use when you need completely separate runs per GPU, for example to track GPU-specific metrics or compare training dynamics across devices.

For multi-node training:

- Use **`rank_scope="global"`** (default) for most cases. A single consolidated run across all nodes is sufficient since metrics like loss and accuracy converge after gradient synchronization.
- Use **`rank_scope="worker"`** for debugging and per-node analysis. This is useful when you need to inspect data distribution across nodes, compare predictions from different workers, or track metrics on individual batches outside the main node.

## Single-node multi-GPU

For single-node distributed training, configure the `Elastic` plugin with `nnodes=1` and set `nproc_per_node` to your GPU count.

### Basic example with `auto` mode

```python {hl_lines=["6-7", 13, 18, 30]}
import os

import torch
import torch.distributed
import flyte
from flyteplugins.pytorch.task import Elastic
from flyteplugins.wandb import wandb_init, get_wandb_run

env = flyte.TaskEnvironment(
    name="single_node_env",
    image=image,
    resources=flyte.Resources(gpu="A100:4"),
    plugin_config=Elastic(nproc_per_node=4, nnodes=1),
    secrets=flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY"),
)

@wandb_init # run_mode="auto" (default)
@env.task
def train_single_node() -> float:
    torch.distributed.init_process_group("nccl")
    rank = torch.distributed.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))

    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    run = get_wandb_run()

    # Training loop - only rank 0 logs
    for epoch in range(10):
        loss = train_epoch(model, dataloader, device)

        if run:
            run.log({"epoch": epoch, "loss": loss})

    torch.distributed.destroy_process_group()
    return loss
```

### Using `shared` mode for per-rank metrics

When you need to see metrics from all GPUs in a single run, use `run_mode="shared"`:

```python {hl_lines=[3, 13, 19]}
import os

@wandb_init(run_mode="shared")
@env.task
def train_with_per_gpu_metrics() -> float:
    torch.distributed.init_process_group("nccl")
    rank = torch.distributed.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))

    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # In shared mode, all ranks get a run object
    run = get_wandb_run()

    for step in range(1000):
        loss, throughput = train_step(model, batch, device)

        # Each rank logs with automatic x_label identification
        if run:
            run.log({
                "loss": loss,
                "throughput_samples_per_sec": throughput,
                "gpu_memory_used": torch.cuda.memory_allocated(device),
            })

    torch.distributed.destroy_process_group()
    return loss
```

![Single-node shared](../../_static/images/integrations/wandb/single_node_shared_flyte.png)

In the W&B UI, metrics from each rank appear with distinct labels, allowing you to compare GPU utilization and throughput across devices.

![Single-node shared W&B UI](../../_static/images/integrations/wandb/single_node_shared_wandb.png)

### Using `new` mode for per-rank runs

When you need completely separate W&B runs for each GPU, use `run_mode="new"`. Each rank gets its own run, and runs are grouped together in the W&B UI:

```python {hl_lines=[1, "11-12"]}
@wandb_init(run_mode="new")  # Each rank gets its own run
@env.task
def train_per_rank() -> float:
    torch.distributed.init_process_group("nccl")
    rank = torch.distributed.get_rank()
    # ...

    # Each rank has its own W&B run
    run = get_wandb_run()

    # Run IDs: {base}-rank-{rank}
    # All runs are grouped under {base} in W&B UI
    run.log({"train/loss": loss.item(), "rank": rank})
    # ...
```

With `run_mode="new"`:

- Each rank creates its own W&B run
- Run IDs follow the pattern `{run_name}-{action_name}-rank-{rank}`
- All runs are grouped together in the W&B UI for comparison

## Multi-node training with `Elastic`

For multi-node distributed training, set `nnodes` to your node count. The `rank_scope` parameter controls whether you get a single W&B run across all nodes (`global`) or one run per node (`worker`).

### Global scope (default): Single run across all nodes

With `run_mode="auto"` and `rank_scope="global"` (both defaults), only global rank 0 initializes W&B, resulting in a single run for the entire distributed job:

```python {hl_lines=["11-12", "27-30", "35", "59-60", "95-98"]}
import os

import torch
import torch.distributed
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

import flyte
from flyteplugins.pytorch.task import Elastic
from flyteplugins.wandb import wandb_init, wandb_config, get_wandb_run

image = flyte.Image.from_debian_base(name="torch-wandb").with_pip_packages(
    "flyteplugins-wandb", "flyteplugins-pytorch", pre=True
)

multi_node_env = flyte.TaskEnvironment(
    name="multi_node_env",
    image=image,
    resources=flyte.Resources(
        cpu=(1, 2),
        memory=("1Gi", "10Gi"),
        gpu="A100:4",
        shm="auto",
    ),
    plugin_config=Elastic(
        nproc_per_node=4,  # GPUs per node
        nnodes=2,          # Number of nodes
    ),
    secrets=flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY"),
)

@wandb_init  # rank_scope="global" by default → 1 run total
@multi_node_env.task
def train_multi_node(epochs: int, batch_size: int) -> float:
    torch.distributed.init_process_group("nccl")

    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))

    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Model with DDP
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])

    # Distributed data loading
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)

    optimizer = optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # Only global rank 0 gets a W&B run
    run = get_wandb_run()

    for epoch in range(epochs):
        sampler.set_epoch(epoch)
        model.train()

        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            if run and batch_idx % 100 == 0:
                run.log({
                    "train/loss": loss.item(),
                    "train/epoch": epoch,
                    "train/batch": batch_idx,
                })

        if run:
            run.log({"train/epoch_complete": epoch})

    # Barrier ensures all ranks finish before cleanup
    torch.distributed.barrier()
    torch.distributed.destroy_process_group()

    return loss.item()

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.with_runcontext(
        custom_context=wandb_config(
            project="multi-node-training",
            tags=["distributed", "multi-node"],
        )
    ).run(train_multi_node, epochs=10, batch_size=32)
```

With this configuration:

- Two nodes run the task, each with 4 GPUs (8 total processes)
- Only global rank 0 creates a W&B run
- Run ID follows the pattern `{run_name}-{action_name}`
- The Flyte UI shows a single link to the W&B run

### Worker scope: One run per node

Use `rank_scope="worker"` when you want each node to have its own W&B run for per-node analysis:

```python {hl_lines=[1, 8]}
@wandb_init(rank_scope="worker")  # 1 run per worker/node
@multi_node_env.task
def train_per_worker(epochs: int, batch_size: int) -> float:
    torch.distributed.init_process_group("nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    # ...

    # Local rank 0 of each worker gets a W&B run
    run = get_wandb_run()

    if run:
        # Each worker logs to its own run
        run.log({"train/loss": loss.item()})
    # ...
```

With `run_mode="auto"`, `rank_scope="worker"`:

- Each node's local rank 0 creates a W&B run
- Run IDs follow the pattern `{run_name}-{action_name}-worker-{worker_index}`
- The Flyte UI shows links to each worker's W&B run

![Multi-node](../../_static/images/integrations/wandb/multi_node.png)

### Shared mode: All ranks log to the same run

Use `run_mode="shared"` when you need metrics from all ranks in a single view. Each rank's metrics are labeled with an `x_label` identifier.

#### Shared + global scope (1 run total)

```python {hl_lines=[1, 7]}
@wandb_init(run_mode="shared")  # All ranks log to 1 shared run
@multi_node_env.task
def train_shared_global() -> float:
    torch.distributed.init_process_group("nccl")
    # ...

    # All ranks get a run object, all log to the same run
    run = get_wandb_run()

    # Each rank logs with automatic x_label identification
    run.log({"train/loss": loss.item(), "rank": rank})
    # ...
```

#### Shared + worker scope (N runs, 1 per node)

```python {hl_lines=[1, 7, 10]}
@wandb_init(run_mode="shared", rank_scope="worker")  # 1 shared run per worker
@multi_node_env.task
def train_shared_worker() -> float:
    torch.distributed.init_process_group("nccl")
    # ...

    # All ranks get a run object, grouped by worker
    run = get_wandb_run()

    # Ranks on the same worker share a run
    run.log({"train/loss": loss.item(), "local_rank": local_rank})
    # ...
```

### New mode: Separate run per rank

Use `run_mode="new"` when you need completely separate runs per GPU. Runs are grouped in the W&B UI for easy comparison.

#### New + global scope (N×M runs, 1 group)

```python {hl_lines=[1, 7, 10]}
@wandb_init(run_mode="new")  # Each rank gets its own run, all in 1 group
@multi_node_env.task
def train_new_global() -> float:
    torch.distributed.init_process_group("nccl")
    # ...

    # Each rank has its own run
    run = get_wandb_run()

    # Run IDs: {base}-rank-{global_rank}
    run.log({"train/loss": loss.item()})
    # ...
```

#### New + worker scope (N×M runs, N groups)

```python {hl_lines=[1, 7, 10]}
@wandb_init(run_mode="new", rank_scope="worker")  # Each rank gets own run, grouped per worker
@multi_node_env.task
def train_new_worker() -> float:
    torch.distributed.init_process_group("nccl")
    # ...

    # Each rank has its own run, grouped by worker
    run = get_wandb_run()

    # Run IDs: {base}-worker-{idx}-rank-{local_rank}
    run.log({"train/loss": loss.item()})
    # ...
```

## How it works

The plugin automatically detects distributed training by checking environment variables set by distributed launchers like `torchrun`:

| Environment variable | Description                                              |
| -------------------- | -------------------------------------------------------- |
| `RANK`               | Global rank across all processes                         |
| `WORLD_SIZE`         | Total number of processes                                |
| `LOCAL_RANK`         | Rank within the current node                             |
| `LOCAL_WORLD_SIZE`   | Number of processes on the current node                  |
| `GROUP_RANK`         | Node/worker index (0 for first node, 1 for second, etc.) |

When these variables are present, the plugin:

1. **Determines which ranks should initialize W&B** based on `run_mode` and `rank_scope`
2. **Generates unique run IDs** that include worker and rank information
4. **Creates UI links** for each W&B run (single link with `rank_scope="global"`, one per worker with `rank_scope="worker"`)

The plugin automatically adapts to your training setup, eliminating the need for manual distributed configuration.

### Run ID patterns

| Scenario                     | Run ID Pattern                                | Group                    |
| ---------------------------- | --------------------------------------------- | ------------------------ |
| Single-node auto/shared      | `{base}`                                      | -                        |
| Single-node new              | `{base}-rank-{rank}`                          | `{base}`                 |
| Multi-node auto/shared (global) | `{base}`                                   | -                        |
| Multi-node auto/shared (worker) | `{base}-worker-{idx}`                      | -                        |
| Multi-node new (global)      | `{base}-rank-{global_rank}`                   | `{base}`                 |
| Multi-node new (worker)      | `{base}-worker-{idx}-rank-{local_rank}`       | `{base}-worker-{idx}`    |

Where `{base}` = `{run_name}-{action_name}`


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/sweeps ===

# Sweeps

W&B sweeps automate hyperparameter optimization by running multiple trials with different parameter combinations. The `@wandb_sweep` decorator creates a sweep and makes it easy to run trials in parallel using Flyte's distributed execution.

## Creating a sweep

Use `@wandb_sweep` to create a W&B sweep when the task executes:

```
import flyte
import wandb
from flyteplugins.wandb import (
    get_wandb_sweep_id,
    wandb_config,
    wandb_init,
    wandb_sweep,
    wandb_sweep_config,
)

env = flyte.TaskEnvironment(
    name="wandb-example",
    image=flyte.Image.from_debian_base(name="wandb-example").with_pip_packages(
        "flyteplugins-wandb"
    ),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

@wandb_init
def objective():
    """Objective function that W&B calls for each trial."""
    wandb_run = wandb.run
    config = wandb_run.config

    # Simulate training with hyperparameters from the sweep
    for epoch in range(config.epochs):
        loss = 1.0 / (config.learning_rate * config.batch_size) + epoch * 0.1
        wandb_run.log({"epoch": epoch, "loss": loss})

@wandb_sweep
@env.task
async def run_sweep() -> str:
    sweep_id = get_wandb_sweep_id()

    # Run 10 trials
    wandb.agent(sweep_id, function=objective, count=10)

    return sweep_id

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.with_runcontext(
        custom_context={
            **wandb_config(project="my-project", entity="my-team"),
            **wandb_sweep_config(
                method="random",
                metric={"name": "loss", "goal": "minimize"},
                parameters={
                    "learning_rate": {"min": 0.0001, "max": 0.1},
                    "batch_size": {"values": [16, 32, 64, 128]},
                    "epochs": {"values": [5, 10, 20]},
                },
            ),
        },
    ).run(run_sweep)

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/sweep.py*

The `@wandb_sweep` decorator:

- Creates a W&B sweep when the task starts
- Makes the sweep ID available via `get_wandb_sweep_id()`
- Adds a link to the main sweeps page in the Flyte UI

Use `wandb_sweep_config()` to define the sweep parameters. This is passed to W&B's sweep API.

> **📝 Note**
>
> Random and Bayesian searches run indefinitely, and the sweep remains in the `Running` state until you stop it.
> You can stop a running sweep from the Weights & Biases UI or from the command line.

## Running parallel agents

Flyte's distributed execution makes it easy to run multiple sweep agents in parallel, each on its own compute resources:

```
import asyncio
from datetime import timedelta

import flyte
import wandb
from flyteplugins.wandb import (
    get_wandb_sweep_id,
    wandb_config,
    wandb_init,
    wandb_sweep,
    wandb_sweep_config,
    get_wandb_context,
)

env = flyte.TaskEnvironment(
    name="wandb-parallel-sweep-example",
    image=flyte.Image.from_debian_base(
        name="wandb-parallel-sweep-example"
    ).with_pip_packages("flyteplugins-wandb"),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

@wandb_init
def objective():
    wandb_run = wandb.run
    config = wandb_run.config

    for epoch in range(config.epochs):
        loss = 1.0 / (config.learning_rate * config.batch_size) + epoch * 0.1
        wandb_run.log({"epoch": epoch, "loss": loss})

@wandb_sweep
@env.task
async def sweep_agent(agent_id: int, sweep_id: str, count: int = 5) -> int:
    """Single agent that runs a subset of trials."""
    wandb.agent(
        sweep_id, function=objective, count=count, project=get_wandb_context().project
    )
    return agent_id

@wandb_sweep
@env.task
async def run_parallel_sweep(total_trials: int = 20, trials_per_agent: int = 5) -> str:
    """Orchestrate multiple agents running in parallel."""
    sweep_id = get_wandb_sweep_id()

    num_agents = (total_trials + trials_per_agent - 1) // trials_per_agent

    # Launch agents in parallel, each with its own resources
    agent_tasks = [
        sweep_agent.override(
            resources=flyte.Resources(cpu="2", memory="4Gi"),
            retries=3,
            timeout=timedelta(minutes=30),
        )(agent_id=i, sweep_id=sweep_id, count=trials_per_agent)
        for i in range(num_agents)
    ]

    await asyncio.gather(*agent_tasks)
    return sweep_id

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.with_runcontext(
        custom_context={
            **wandb_config(project="my-project", entity="my-team"),
            **wandb_sweep_config(
                method="random",
                metric={"name": "loss", "goal": "minimize"},
                parameters={
                    "learning_rate": {"min": 0.0001, "max": 0.1},
                    "batch_size": {"values": [16, 32, 64]},
                    "epochs": {"values": [5, 10, 20]},
                },
            ),
        },
    ).run(
        run_parallel_sweep,
        total_trials=20,
        trials_per_agent=5,
    )

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/parallel_sweep.py*

This pattern provides:

- **Distributed execution**: Each agent runs on separate compute nodes
- **Resource allocation**: Specify CPU, memory, and GPU per agent
- **Fault tolerance**: Failed agents can retry without affecting others
- **Timeout protection**: Prevent runaway trials

> **📝 Note**
>
> `run_parallel_sweep` links to the main Weights & Biases sweeps page and `sweep_agent` links to the specific sweep URL because we cannot determine the sweep ID at link rendering time.

![Sweep](../../_static/images/integrations/wandb/sweep.png)

## Writing objective functions

The objective function is called by `wandb.agent()` for each trial. It must be a regular Python function decorated with `@wandb_init`:

```python {hl_lines=["1-2", "5-6"]}
@wandb_init
def objective():
    """Objective function for sweep trials."""
    # Access hyperparameters from wandb.run.config
    run = wandb.run
    config = run.config

    # Your training code
    model = create_model(
        learning_rate=config.learning_rate,
        hidden_size=config.hidden_size,
    )

    for epoch in range(config.epochs):
        train_loss = train_epoch(model)
        val_loss = validate(model)

        # Log metrics - W&B tracks these for the sweep
        run.log({
            "epoch": epoch,
            "train_loss": train_loss,
            "val_loss": val_loss,
        })

    # The final val_loss is used by the sweep to rank trials
```

Key points:

- Use `@wandb_init` on the objective function (not `@env.task`)
- Access hyperparameters via `wandb.run.config` (not `get_wandb_run()` since this is outside Flyte context)
- Log the metric specified in `wandb_sweep_config(metric=...)` so the sweep can optimize it
- The function is called multiple times by `wandb.agent()`, once per trial


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/downloading_logs ===

# Downloading logs

This integration enables downloading Weights & Biases run data, including metrics history, summary data, and synced files.

## Automatic download

Set `download_logs=True` to automatically download run data after your task completes:

```python {hl_lines=1}
@wandb_init(download_logs=True)
@env.task
async def train_with_download():
    run = get_wandb_run()

    for epoch in range(10):
        run.log({"loss": 1.0 / (epoch + 1)})

    return run.id
```

The downloaded data is traced by Flyte and appears as a `Dir` output in the Flyte UI. Downloaded files include:

- `summary.json`: Final summary metrics
- `metrics_history.json`: Step-by-step metrics history
- Any files synced by W&B (`requirements.txt`, `wandb_metadata.json`, etc.)

You can also set `download_logs=True` in `wandb_config()`:

```python {hl_lines=5}
flyte.with_runcontext(
    custom_context=wandb_config(
        project="my-project",
        entity="my-team",
        download_logs=True,
    ),
).run(train_task)
```

![Logs](../../_static/images/integrations/wandb/logs.png)

For sweeps, set `download_logs=True` on `@wandb_sweep` or `wandb_sweep_config()` to download all trial data:

```python {hl_lines=1}
@wandb_sweep(download_logs=True)
@env.task
async def run_sweep():
    sweep_id = get_wandb_sweep_id()
    wandb.agent(sweep_id, function=objective, count=10)
    return sweep_id
```

![Sweep Logs](../../_static/images/integrations/wandb/sweep_logs.png)

## Accessing run directories during execution

Use `get_wandb_run_dir()` to access the local W&B run directory during task execution. This is useful for writing custom files that get synced to W&B:

```python {hl_lines=[1, 7, "18-19"]}
from flyteplugins.wandb import get_wandb_run_dir

@wandb_init
@env.task
def train_with_artifacts():
    run = get_wandb_run()
    local_dir = get_wandb_run_dir()

    # Train your model
    for epoch in range(10):
        run.log({"loss": 1.0 / (epoch + 1)})

    # Save model checkpoint to the run directory
    model_path = f"{local_dir}/model_checkpoint.pt"
    torch.save(model.state_dict(), model_path)

    # Save custom metrics file
    with open(f"{local_dir}/custom_metrics.json", "w") as f:
        json.dump({"final_accuracy": 0.95}, f)

    return run.id
```

Files written to the run directory are automatically synced to W&B and can be accessed later via the W&B UI or by setting `download_logs=True`.

> **📝 Note**
>
> `get_wandb_run_dir()` accesses the local directory without making network calls. Files written here may have a brief delay before appearing in the W&B cloud.


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/constraints_and_best_practices ===

# Constraints and best practices

## Decorator ordering

`@wandb_init` and `@wandb_sweep` must be the **outermost decorators**, applied after `@env.task`:

```python
# Correct
@wandb_init
@env.task
async def my_task():
    ...

# Incorrect - will not work
@env.task
@wandb_init
async def my_task():
    ...
```

## Traces cannot use decorators

Do not apply `@wandb_init` to traces. Traces automatically access the parent task's run via `get_wandb_run()`:

```python
# Correct
@flyte.trace
async def my_trace():
    run = get_wandb_run()
    if run:
        run.log({"metric": 42})

# Incorrect - don't decorate traces
@wandb_init
@flyte.trace
async def my_trace():
    ...
```

## Maximum sweep agents

[W&B limits sweeps to a maximum of 20 concurrent agents](https://docs.wandb.ai/models/sweeps/existing-project#3-launch-agents).

## Configuration priority

Configuration is merged with the following priority (highest to lowest):

1. Decorator parameters (`@wandb_init(project="...")`)
2. Context manager (`with wandb_config(...)`)
3. Workflow-level context (`flyte.with_runcontext(custom_context=wandb_config(...))`)
4. Auto-generated values (run ID from Flyte context)

## Run ID generation

When no explicit `id` is provided, the plugin generates run IDs using the pattern:

```
{run_name}-{action_name}
```

This ensures unique, predictable IDs that can be matched between the `Wandb` link class and manual `wandb.init()` calls.

## Sync delay for local files

Files written to the run directory (via `get_wandb_run_dir()`) are synced to W&B asynchronously. There may be a brief delay before they appear in the W&B cloud or can be downloaded via `download_wandb_run_dir()`.

## Shared run mode requirements

When using `run_mode="shared"`, the task requires a parent task to have already created a W&B run. Calling a task with `run_mode="shared"` as a top-level task will fail.

## Objective functions for sweeps

Objective functions passed to `wandb.agent()` should:

- Be regular Python functions (not Flyte tasks)
- Be decorated with `@wandb_init`
- Access hyperparameters via `wandb.run.config` (not `get_wandb_run()`)
- Log the metric specified in `wandb_sweep_config(metric=...)` so the sweep can optimize it

## Error handling

The plugin raises standard exceptions:

- `RuntimeError`: When `download_wandb_run_dir()` is called without a run ID and no active run exists
- `wandb.errors.AuthenticationError`: When `WANDB_API_KEY` is not set or invalid
- `wandb.errors.CommError`: When a run cannot be found in the W&B cloud


=== PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/manual ===

# Manual integration

If you need more control over W&B initialization, you can use the `Wandb` and `WandbSweep` link classes directly instead of the decorators. This lets you call `wandb.init()` and `wandb.finish()` yourself while still getting automatic links in the Flyte UI.

## Using the Wandb link class

Add a `Wandb` link to your task to generate a link to the W&B run in the Flyte UI:

```
import flyte
import wandb
from flyteplugins.wandb import Wandb

env = flyte.TaskEnvironment(
    name="wandb-manual-init-example",
    image=flyte.Image.from_debian_base(
        name="wandb-manual-init-example"
    ).with_pip_packages("flyteplugins-wandb"),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

@env.task(
    links=(
        Wandb(
            project="my-project",
            entity="my-team",
            run_mode="new",
            # No id parameter - link will auto-generate from run_name-action_name
        ),
    )
)
async def train_model(learning_rate: float) -> str:
    ctx = flyte.ctx()

    # Generate run ID matching the link's auto-generated ID
    run_id = f"{ctx.action.run_name}-{ctx.action.name}"

    # Manually initialize W&B
    wandb_run = wandb.init(
        project="my-project",
        entity="my-team",
        id=run_id,
        config={"learning_rate": learning_rate},
    )

    # Your training code
    for epoch in range(10):
        loss = 1.0 / (learning_rate * (epoch + 1))
        wandb_run.log({"epoch": epoch, "loss": loss})

    # Manually finish the run
    wandb_run.finish()

    return wandb_run.id

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.with_runcontext().run(
        train_model,
        learning_rate=0.01,
    )

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/init_manual.py*

### With a custom run ID

If you want to use your own run ID, specify it in both the link and the `wandb.init()` call:

```python {hl_lines=[6, 14]}
@env.task(
    links=(
        Wandb(
            project="my-project",
            entity="my-team",
            id="my-custom-run-id",
        ),
    )
)
async def train_with_custom_id() -> str:
    run = wandb.init(
        project="my-project",
        entity="my-team",
        id="my-custom-run-id",  # Must match the link's ID
        resume="allow",
    )

    # Training code...
    run.finish()
    return run.id
```

### Adding links at runtime with override

You can also add links when calling a task using `.override()`:

```python {hl_lines=9}
@env.task
async def train_model(learning_rate: float) -> str:
    # ... training code with manual wandb.init() ...
    return run.id

# Add link when running the task
result = await train_model.override(
    links=(Wandb(project="my-project", entity="my-team", run_mode="new"),)
)(learning_rate=0.01)
```

## Using the `WandbSweep` link class

Use `WandbSweep` to add a link to a W&B sweep:

```
import flyte
import wandb
from flyteplugins.wandb import WandbSweep

env = flyte.TaskEnvironment(
    name="wandb-manual-sweep-example",
    image=flyte.Image.from_debian_base(
        name="wandb-manual-sweep-example"
    ).with_pip_packages("flyteplugins-wandb"),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

def objective():
    with wandb.init(project="my-project", entity="my-team") as wandb_run:
        config = wandb_run.config

        for epoch in range(config.epochs):
            loss = 1.0 / (config.learning_rate * config.batch_size) + epoch * 0.1
            wandb_run.log({"epoch": epoch, "loss": loss})

@env.task(
    links=(
        WandbSweep(
            project="my-project",
            entity="my-team",
        ),
    )
)
async def manual_sweep() -> str:
    # Manually create the sweep
    sweep_config = {
        "method": "random",
        "metric": {"name": "loss", "goal": "minimize"},
        "parameters": {
            "learning_rate": {"min": 0.0001, "max": 0.1},
            "batch_size": {"values": [16, 32, 64]},
            "epochs": {"value": 10},
        },
    }

    sweep_id = wandb.sweep(sweep_config, project="my-project", entity="my-team")

    # Run the sweep
    wandb.agent(sweep_id, function=objective, count=10, project="my-project")

    return sweep_id

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.with_runcontext().run(manual_sweep)

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/sweep_manual.py*

The link will point to the project's sweeps page. If you have the sweep ID, you can specify it in the link:

```python {hl_lines=6}
@env.task(
    links=(
        WandbSweep(
            project="my-project",
            entity="my-team",
            id="known-sweep-id",
        ),
    )
)
async def resume_sweep() -> str:
    # Resume an existing sweep
    wandb.agent("known-sweep-id", function=objective, count=10)
    return "known-sweep-id"
```


=== PAGE: https://www.union.ai/docs/v2/flyte/api-reference ===

# Reference

This section provides the reference material for the Flyte SDK and CLI.

To get started, add `flyte` to your project

```shell
$ uv pip install --no-cache --upgrade flyte
```

This will install both the Flyte SDK and CLI.

### **Flyte SDK**

The Flyte SDK provides the core Python API for building workflows and apps on your Union instance.

### **Flyte CLI**

The Flyte CLI is the command-line interface for interacting with your Union instance.

### **Flyte agent plugins**

A portable agent-harness plugin — skills and MCP servers for authoring, migrating, and deploying Flyte with an AI agent.

Migrating from Flyte 1? See **Migration > From Flyte 1 to 2 > Migration overview** in the User Guide.

## Subpages
- **LLM-optimized documentation** - LLM-optimized documentation for Union.ai and Flyte, provided at four levels of granularity and following the llms.txt convention so AI coding agents and search engines can consume the docs.
- **Flyte CLI**
- **Flyte SDK**
- **Flyte agent plugins**
- **Integrations**


=== PAGE: https://www.union.ai/docs/v2/flyte/api-reference/flyte-context ===

LLM-optimized documentation for Union.ai and Flyte, provided at four levels of granularity and following the llms.txt convention so AI coding agents and search engines can consume the docs.

# LLM-optimized documentation

This site provides LLM-optimized documentation at four levels of granularity,
designed for use by AI coding agents such as
[Claude Code](https://docs.anthropic.com/en/docs/claude-code),
[Cursor](https://www.cursor.com/),
[Windsurf](https://windsurf.com/),
and similar tools.
These files also follow the [`llms.txt` convention](https://llmstxt.org/),
making them discoverable by AI search engines.

These files are not linked from the pages they cover. They are addressed by convention:
append `/page.md` to any page URL, or `/section.md` to a section URL. Start from the
`llms.txt` index below, which lists every page and every available bundle.

All links within LLM-optimized files use absolute URLs (`https://www.union.ai/docs/...`),
so files work correctly when copied locally and used outside the docs site.

## Per-page markdown (`page.md`)

Every page on this site has a parallel LLM-optimized version in clean Markdown,
accessible at the same URL path with `/page.md` appended.
For example, this page is at:

* `https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/`

and its LLM-optimized version is at:

* **LLM-optimized documentation**

Section landing pages include a `## Subpages` table listing child pages with their H2/H3 headings,
making it easy to identify the right page to fetch.

## Section bundles (`section.md`)

For key documentation sections, a curated bundle file concatenates all pages in the section
into a single `section.md` file.

These are accessible at the same URL path as the top page of the section, with `/section.md` appended.

These `section.md` files are sized to fit within modern LLM context windows
and are ideal for pasting into a prompt or adding to project context.

Available bundle files:

{{< llm-readable-list >}}

## Page index (`llms.txt`)

The `llms.txt` file is a compact index of all LLM-optimized pages, organized by section.
Each page entry includes the H2/H3 headings found on that page, so an agent can identify
the right page to fetch without downloading it first.

Sections that have a `section.md` bundle are marked in the index.

Download it and append its contents to the `AGENTS.md`, `CLAUDE.md` or similar file in your project root.
Make sure you append the index into a file that is **loaded into context by default** by your coding tool.
Adding it as a skill or tool is less effective because the agent must decide to load it
rather than having the information always available.

* [`llms.txt`](https://www.union.ai/docs/v2/flyte/llms.txt) (~50K tokens)

> [!NOTE]
> You are viewing the **Flyte OSS** docs.
> To get the `llms.txt` for a different product variant, use the variant selector at the top of the page.

## Full documentation (`llms-full.txt`)

The `llms-full.txt` file contains the entire Flyte version 2.0 documentation as a single Markdown file.
This file is very large and is not suitable for direct inclusion in an LLM context window,
but it may be useful for RAG-based tools.

* [`llms-full.txt`](https://www.union.ai/docs/v2/flyte/llms-full.txt) (~2M tokens)

> [!NOTE]
> You are viewing the **Flyte OSS** docs.
> To get the `llms-full.txt` for a different product variant, use the variant selector at the top of the page.


=== PAGE: https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli ===

# Flyte CLI

This is the command line interface for Flyte.

| Object | Action |
| ------ | -- |
| `action` | **Flyte CLI > flyte > flyte abort > flyte abort action**, **Flyte CLI > flyte > flyte get > flyte get action**  |
| `run` | **Flyte CLI > flyte > flyte abort > flyte abort run**, **Flyte CLI > flyte > flyte get > flyte get run**  |
| `config` | **Flyte CLI > flyte > flyte create > flyte create config**, **Flyte CLI > flyte > flyte get > flyte get config**  |
| `project` | **Flyte CLI > flyte > flyte create > flyte create project**, **Flyte CLI > flyte > flyte get > flyte get project**, **Flyte CLI > flyte update > flyte update project**  |
| `secret` | **Flyte CLI > flyte > flyte create > flyte create secret**, **Flyte CLI > flyte > flyte delete > flyte delete secret**, **Flyte CLI > flyte > flyte get > flyte get secret**  |
| `trigger` | **Flyte CLI > flyte > flyte create > flyte create trigger**, **Flyte CLI > flyte > flyte delete > flyte delete trigger**, **Flyte CLI > flyte get trigger**, **Flyte CLI > flyte update > flyte update trigger**  |
| `app` | **Flyte CLI > flyte > flyte delete > flyte delete app**, **Flyte CLI > flyte > flyte get > flyte get app**, **Flyte CLI > flyte update > flyte update app**  |
| `devbox` | **Flyte CLI > flyte > flyte delete > flyte delete devbox**, **Flyte CLI > flyte start > flyte start devbox**, **Flyte CLI > flyte stop > flyte stop devbox**  |
| `local-cache` | **Flyte CLI > flyte > flyte delete > flyte delete local-cache**  |
| `settings` | **Flyte CLI > flyte > flyte edit > flyte edit settings**, **Flyte CLI > flyte > flyte get > flyte get settings**  |
| `docs` | **Flyte CLI > flyte > flyte gen > flyte gen docs**  |
| `condition` | **Flyte CLI > flyte > flyte get > flyte get condition**, **Flyte CLI > flyte signal > flyte signal condition**  |
| `io` | **Flyte CLI > flyte > flyte get > flyte get io**  |
| `logs` | **Flyte CLI > flyte > flyte get > flyte get logs**  |
| `task` | **Flyte CLI > flyte get task**  |
| `hf-model` | **Flyte CLI > flyte prefetch > flyte prefetch hf-model**  |
| `deployed-task` | **Flyte CLI > flyte run > flyte run deployed-task**  |
| `tui` | **Flyte CLI > flyte start > flyte start tui**  |

| Action | On |
| ------ | -- |
| `abort` | **Flyte CLI > flyte > flyte abort > flyte abort action**, **Flyte CLI > flyte > flyte abort > flyte abort run**  |
| **Flyte CLI > flyte > flyte build** | - |
| `create` | **Flyte CLI > flyte > flyte create > flyte create config**, **Flyte CLI > flyte > flyte create > flyte create project**, **Flyte CLI > flyte > flyte create > flyte create secret**, **Flyte CLI > flyte > flyte create > flyte create trigger**  |
| `delete` | **Flyte CLI > flyte > flyte delete > flyte delete app**, **Flyte CLI > flyte > flyte delete > flyte delete devbox**, **Flyte CLI > flyte > flyte delete > flyte delete local-cache**, **Flyte CLI > flyte > flyte delete > flyte delete secret**, **Flyte CLI > flyte > flyte delete > flyte delete trigger**  |
| **Flyte CLI > flyte > flyte deploy** | - |
| `edit` | **Flyte CLI > flyte > flyte edit > flyte edit settings**  |
| `gen` | **Flyte CLI > flyte > flyte gen > flyte gen docs**  |
| `get` | **Flyte CLI > flyte > flyte get > flyte get action**, **Flyte CLI > flyte > flyte get > flyte get app**, **Flyte CLI > flyte > flyte get > flyte get condition**, **Flyte CLI > flyte > flyte get > flyte get config**, **Flyte CLI > flyte > flyte get > flyte get io**, **Flyte CLI > flyte > flyte get > flyte get logs**, **Flyte CLI > flyte > flyte get > flyte get project**, **Flyte CLI > flyte > flyte get > flyte get run**, **Flyte CLI > flyte > flyte get > flyte get secret**, **Flyte CLI > flyte > flyte get > flyte get settings**, **Flyte CLI > flyte get task**, **Flyte CLI > flyte get trigger**  |
| `prefetch` | **Flyte CLI > flyte prefetch > flyte prefetch hf-model**  |
| **Flyte CLI > flyte rerun** | - |
| `run` | **Flyte CLI > flyte run > flyte run deployed-task**  |
| **Flyte CLI > flyte serve** | - |
| `signal` | **Flyte CLI > flyte signal > flyte signal condition**  |
| `start` | **Flyte CLI > flyte start > flyte start devbox**, **Flyte CLI > flyte start > flyte start tui**  |
| `stop` | **Flyte CLI > flyte stop > flyte stop devbox**  |
| `update` | **Flyte CLI > flyte update > flyte update app**, **Flyte CLI > flyte update > flyte update project**, **Flyte CLI > flyte update > flyte update trigger**  |
| **Flyte CLI > flyte whoami** | - |

## flyte

**`flyte [OPTIONS] COMMAND [ARGS]...`**

The Flyte CLI is the command line interface for working with the Flyte SDK and backend.

It follows a simple verb/noun structure,
where the top-level commands are verbs that describe the action to be taken,
and the subcommands are nouns that describe the object of the action.

The root command can be used to configure the CLI for persistent settings,
such as the endpoint, organization, and verbosity level.

Set endpoint and organization:

```bash
$ flyte --endpoint <endpoint> --org <org> get project <project_name>
```

Increase verbosity level (This is useful for debugging,
this will show more logs and exception traces):

```bash
$ flyte -vvv get logs <run-name>
```

Override the default config file:

```bash
$ flyte --config /path/to/config.yaml run ...
```

* [Documentation](https://www.union.ai/docs/flyte/user-guide/)
* [GitHub](https://github.com/flyteorg/flyte): Please leave a star if you like Flyte!
* [Slack](https://slack.flyte.org): Join the community and ask questions.
* [Issues](https://github.com/flyteorg/flyte/issues)

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--version` | `boolean` | `False` | Show the version and exit. |
| `--endpoint` | `text` | `Sentinel.UNSET` | The endpoint to connect to. This will override any configuration file and simply use `pkce` to connect. |
| `--insecure` | `boolean` |  | Use an insecure connection to the endpoint. If not specified, the CLI will use TLS. |
| `--image-builder` `--builder` | `choice` |  | Image builder to use for building images. Overrides the config file setting. If not specified, the builder from the config file (image.builder) is used, falling back to 'local'. |
| `--auth-type` | `choice` |  | Authentication type to use for the Flyte backend. Defaults to 'pkce'. |
| `-v` `--verbose` | `integer` | `0` | Show verbose messages and exception traces. Repeating multiple times increases the verbosity (e.g., -vvv). |
| `--org` | `text` | `Sentinel.UNSET` | The organization to which the command applies. |
| `-c` `--config` | `file` | `Sentinel.UNSET` | Path to the configuration file to use. If not specified, the default configuration file is used. |
| `--output-format` `-of` | `choice` | `table` | Output format for commands that support it. Defaults to 'table'. |
| `--log-format` | `choice` | `console` | Formatting for logs, defaults to 'console' which is meant to be human readable. 'json' is meant for machine parsing. |
| `--user-log-level` | `choice` | `info` | Log level for user task logs. Independent of the internal Flyte log level (-v). |
| `--reset-root-logger` | `boolean` | `False` | If set, the root logger will be reset to use Flyte logging style |
| `--no-progress` | `boolean` | `False` | Disable the animated progress spinner — useful in CI / non-interactive logs. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte abort

**`flyte abort COMMAND [ARGS]...`**

Abort an ongoing process.

#### flyte abort action

**`flyte abort action [OPTIONS] RUN_NAME ACTION_NAME`**

Abort an action associated with a run.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--reason` | `text` | `Manually aborted from the CLI` | The reason to abort the run. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte abort run

**`flyte abort run [OPTIONS] RUN_NAME`**

Abort a run.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--reason` | `text` | `Manually aborted from the CLI` | The reason to abort the run. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte build

**`flyte build [OPTIONS] COMMAND [ARGS]...`**

Build the environments defined in a python file or directory. This will build the images associated with the
environments.

To build the image for a single named environment:

```bash
flyte build hello.py my_env
```

To build the images for all environments in a file (without naming one), use the `--all` flag:

```bash
flyte build --all hello.py
```

To recursively build all environments in a directory and its subdirectories, use the `--recursive` flag:

```bash
flyte build --all --recursive ./src
```

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--copy-style` | `choice` | `loaded_modules` | Copy style of the eventual deploy. Must match the deploy's --copy-style so the image content hash — and therefore the registry tag — lines up. |
| `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. |
| `--recursive` `-r` | `boolean` | `False` | Recursively build all environments in the current directory and its subdirectories. |
| `--all` | `boolean` | `False` | Build the images for all environments in the file or directory, ignoring the file name. |
| `--ignore-load-errors` `-i` | `boolean` | `False` | Ignore errors when loading environments, especially when using --recursive or --all. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte create

**`flyte create COMMAND [ARGS]...`**

Create resources in a Flyte deployment.

#### flyte create config

**`flyte create config [OPTIONS]`**

Creates a configuration file for Flyte CLI.
If the `--output` option is not specified, it will create a file named `config.yaml` in the current directory.
If the file already exists, it will raise an error unless the `--force` option is used.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--endpoint` | `text` | `Sentinel.UNSET` | Endpoint of the Flyte backend. |
| `--insecure` | `boolean` | `False` | Use an insecure connection to the Flyte backend. |
| `--org` | `text` | `Sentinel.UNSET` | Organization to use. This will override the organization in the configuration file. |
| `-o` `--output` | `path` | `.flyte/config.yaml` | Path to the output directory where the configuration will be saved. Defaults to current directory. |
| `--force` | `boolean` | `False` | Force overwrite of the configuration file if it already exists. |
| `--image-builder` `--builder` | `choice` | `local` | Image builder to use for building images. Defaults to 'local'. |
| `--registry` | `text` |  | Container registry to use as the base registry when building images (e.g. 'ghcr.io/my-org'). When set, this overrides the built-in default base registry. Equivalent to the 'image.registry' config entry or the FLYTE_IMAGE_REGISTRY environment variable. |
| `--auth-type` | `choice` |  | Authentication type to use for the Flyte backend. Defaults to 'pkce'. |
| `--local-persistence` | `boolean` | `False` | Enable SQLite persistence for local run metadata, allowing past runs to be browsed via 'flyte start tui'. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte create project

**`flyte create project [OPTIONS]`**

Create a new project.


Example usage:

```bash
flyte create project --id my_project_id --name "My Project"
flyte create project --id my_project_id --name "My Project" --description "My project" -l team=ml -l env=prod
```

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--id` | `text` | `Sentinel.UNSET` | Unique identifier for the project (immutable). |
| `--name` | `text` | `Sentinel.UNSET` | Display name for the project. |
| `--description` | `text` | `` | Description for the project. |
| `--label` `-l` | `text` | `Sentinel.UNSET` | Labels as key=value pairs. Can be specified multiple times. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte create secret

**`flyte create secret [OPTIONS] NAME`**

Create a new secret. The name of the secret is required. For example:

CODE7

If you don't provide a `--value` flag, you will be prompted to enter the
secret value in the terminal.

CODE8

If `--from-file` is specified, the value will be read from the file instead of being provided directly:

CODE9

The `--type` option can be used to create specific types of secrets.
Either `regular` or `image_pull` can be specified.
Secrets intended to access container images should be specified as `image_pull`.
Other secrets should be specified as `regular`.
If no type is specified, `regular` is assumed.

For image pull secrets, you have several options:

1. Interactive mode (prompts for registry, username, password):
CODE10

2. With explicit credentials:
CODE11

3. Lastly, you can create a secret from your existing Docker installation (i.e., you've run `docker login` in
the past) and you just want to pull from those credentials. Since you may have logged in to multiple registries,
you can specify which registries to include. If no registries are specified, all registries are added.
CODE12

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--value` | `text` | `Sentinel.UNSET` | Secret value Mutually exclusive with from_file, from_docker_config, registry. |
| `--from-file` | `path` | `Sentinel.UNSET` | Path to the file with the binary secret. Mutually exclusive with value, from_docker_config, registry. |
| `--type` | `choice` | `regular` | Type of the secret. |
| `--from-docker-config` | `boolean` | `False` | Create image pull secret from Docker config file (only for --type image_pull). Mutually exclusive with value, from_file, registry, username, password. |
| `--docker-config-path` | `path` | `Sentinel.UNSET` | Path to Docker config file (defaults to ~/.docker/config.json or $DOCKER_CONFIG). Requires from_docker_config. |
| `--registries` | `text` | `Sentinel.UNSET` | Comma-separated list of registries to include (only with --from-docker-config). |
| `--registry` | `text` | `Sentinel.UNSET` | Registry hostname (e.g., ghcr.io, docker.io) for explicit credentials (only for --type image_pull). Mutually exclusive with value, from_file, from_docker_config. |
| `--username` | `text` | `Sentinel.UNSET` | Username for the registry (only with --registry). |
| `--password` | `text` | `Sentinel.UNSET` | Password for the registry (only with --registry). If not provided, will prompt. |
| `--cluster-pool` | `text` |  | Scope the secret to a cluster pool. Mutually exclusive with --project and --domain. Mutually exclusive with project, domain. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte create trigger

**`flyte create trigger [OPTIONS] TASK_NAME NAME`**

Create a new trigger for a task. The task name and trigger name are required.

Example:

CODE13

This will create a trigger that runs every day at midnight.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--schedule` | `text` | `Sentinel.UNSET` | Cron schedule for the trigger. Defaults to every minute. |
| `--description` | `text` | `` | Description of the trigger. |
| `--auto-activate` | `boolean` | `True` | Whether the trigger should not be automatically activated. Defaults to True. |
| `--trigger-time-var` | `text` | `trigger_time` | Variable name for the trigger time in the task inputs. Defaults to 'trigger_time'. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte delete

**`flyte delete COMMAND [ARGS]...`**

Remove resources from a Flyte deployment.

#### flyte delete app

**`flyte delete app [OPTIONS] NAME`**

Delete apps from a Flyte deployment.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte delete devbox

**`flyte delete devbox [OPTIONS]`**

Stop and remove the local Flyte devbox cluster container.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--volume` | `boolean` | `False` | Also delete the Docker volume used for persistent storage. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte delete local-cache

**`flyte delete local-cache`**

Delete the entire local cache directory (~/.flyte/local-cache).

This removes the local SQLite cache used for image lookups, bundle uploads,
run history, and task caching.

#### flyte delete secret

**`flyte delete secret [OPTIONS] NAME`**

Delete a secret. The name of the secret is required.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--cluster-pool` | `text` |  | Scope the secret to a cluster pool. Mutually exclusive with --project and --domain. Mutually exclusive with project, domain. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte delete trigger

**`flyte delete trigger [OPTIONS] NAME TASK_NAME`**

Delete a trigger. The name of the trigger is required.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte deploy

**`flyte deploy [OPTIONS] COMMAND [ARGS]...`**

Deploy one or more environments from a python file.

This command will create or update environments in the Flyte system, registering
all tasks and their dependencies.

Example usage:

CODE14

Arguments to the deploy command are provided right after the `deploy` command and before the file name.

To deploy all environments in a file, use the `--all` flag:

CODE15

To recursively deploy all environments in a directory and its subdirectories, use the `--recursive` flag:

CODE16

You can combine `--all` and `--recursive` to deploy everything:

CODE17

You can provide image mappings with `--image` flag. This allows you to specify
the image URI for the task environment during CLI execution without changing
the code. Any images defined with `Image.from_ref_name("name")` will resolve to the
corresponding URIs you specify here.

CODE18

If the image name is not provided, it is regarded as a default image and will
be used when no image is specified in TaskEnvironment:

CODE19

You can specify multiple image arguments:

CODE20

To deploy a specific version, use the `--version` flag:

CODE21

To preview what would be deployed without actually deploying, use the `--dry-run` flag:

CODE22

You can specify the `--config` flag to point to a specific Flyte cluster:

CODE23

You can override the default configured project and domain:

CODE24

If loading some files fails during recursive deployment, you can use the `--ignore-load-errors` flag
to continue deploying the environments that loaded successfully:

CODE25

Other arguments to the deploy command are listed below.

To see the environments available in a file, use `--help` after the file name:

CODE26

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--version` | `text` | `Sentinel.UNSET` | Version of the environment to deploy |
| `--dry-run` `--dryrun` | `boolean` | `False` | Dry run. Do not actually call the backend service. |
| `--copy-style` | `choice` | `loaded_modules` | Copy style to use when running the task |
| `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. |
| `--recursive` `-r` | `boolean` | `False` | Recursively deploy all environments in the current directory |
| `--all` | `boolean` | `False` | Deploy all environments in the current directory, ignoring the file name |
| `--ignore-load-errors` `-i` | `boolean` | `False` | Ignore errors when loading environments especially when using --recursive or --all. |
| `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. |
| `--image` | `text` | `Sentinel.UNSET` | Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte edit

**`flyte edit COMMAND [ARGS]...`**

#### flyte edit settings

**`flyte edit settings [OPTIONS]`**

Edit hierarchical settings interactively — or apply a YAML file directly.

**Interactive mode** (default). Opens settings in your ``$EDITOR``. Three
comment tiers appear:

- ``###`` section headers and the scope line
- ``##`` per-field descriptions and inline metadata
- ``#`` inactive settings (uncomment the single ``#`` to activate)

If the edited YAML fails to parse, the editor reopens with an error
header so you can fix the syntax without losing your edits. If you
decline to reopen — or if the server rejects the update — your buffer
is saved under ``~/.flyte/settings-edit-<timestamp>.yaml``.

**Non-interactive mode**: pass ``--from-file <path>`` to skip the editor
entirely. The file's contents are parsed, the diff is printed, and the
overrides are applied without a confirmation prompt. Ideal for
CI/automation.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--from-file` `-f` | `file` |  | Apply overrides from a YAML file and skip the editor. The file can be produced by `flyte get settings` (comment markers are honoured) or be a plain YAML mapping of flat dot-notation keys to values. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte gen

**`flyte gen COMMAND [ARGS]...`**

Generate documentation.

#### flyte gen docs

**`flyte gen docs [OPTIONS]`**

Generate documentation.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--type` | `text` | `Sentinel.UNSET` | Type of documentation (valid: markdown) |
| `--plugin-variants` | `text` |  | Hugo variant names for plugin commands (e.g., 'union'). When set, plugin command sections and index entries are wrapped in {{&lt; variant >}} shortcodes. Core commands appear unconditionally. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte get

**`flyte get COMMAND [ARGS]...`**

Retrieve resources from a Flyte deployment.

You can get information about projects, runs, tasks, actions, secrets, logs and input/output values.

Each command supports optional parameters to filter or specify the resource you want to retrieve.

Using a `get` subcommand without any arguments will retrieve a list of available resources to get.
For example:

* `get project` (without specifying a project), will list all projects.
* `get project my_project` will return the details of the project named `my_project`.

In some cases, a partially specified command will act as a filter and return available further parameters.
For example:

* `get action my_run` will return all actions for the run named `my_run`.
* `get action my_run my_action` will return the details of the action named `my_action` for the run `my_run`.

#### flyte get action

**`flyte get action [OPTIONS] RUN_NAME [ACTION_NAME]`**

Get all actions for a run or details for a specific action.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--in-phase` | `choice` | `Sentinel.UNSET` | Filter actions by their phase. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get app

**`flyte get app [OPTIONS] [NAME]`**

Get a list of all apps, or details of a specific app by name.

Apps are long-running services deployed on the Flyte platform.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--limit` | `integer` | `100` | Limit the number of apps to fetch when listing. |
| `--only-mine` | `boolean` | `False` | Show only apps created by the current user (you). |
| `--status` | `choice` |  | Filter apps by deployment status. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get condition

**`flyte get condition [OPTIONS] RUN_NAME [ACTION_NAME]`**

List conditions (paused condition actions) for a run, optionally filtered to a
specific parent action.

Each condition corresponds to a condition action registered via
``flyte.new_condition(...)`` from a workflow. Use ``flyte signal condition`` to
resolve one.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get config

**`flyte get config`**

Shows the automatically detected configuration to connect with the remote backend.

The configuration will include the endpoint, organization, and other settings that are used by the CLI.

#### flyte get io

**`flyte get io [OPTIONS] RUN_NAME [ACTION_NAME]`**

Get the inputs and outputs of a run or action.
If only the run name is provided, it will show the inputs and outputs of the root action of that run.
If an action name is provided, it will show the inputs and outputs for that action.
If `--inputs-only` or `--outputs-only` is specified, it will only show the inputs or outputs respectively.

Examples:

CODE27

CODE28

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--inputs-only` `-i` | `boolean` | `False` | Show only inputs |
| `--outputs-only` `-o` | `boolean` | `False` | Show only outputs |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get logs

**`flyte get logs [OPTIONS] RUN_NAME [ACTION_NAME]`**

Stream logs for the provided run or action.
If only the run is provided, only the logs for the parent action will be streamed:

CODE29

If you want to see the logs for a specific action, you can provide the action name as well:

CODE30

By default, logs will be shown in the raw format and will scroll the terminal.
If automatic scrolling and only tailing `--lines` number of lines is desired, use the `--pretty` flag:

CODE31

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--lines` `-l` | `integer` | `30` | Number of lines to show, only useful for --pretty |
| `--show-ts` | `boolean` | `False` | Show timestamps |
| `--pretty` | `boolean` | `False` | Show logs in an auto-scrolling box, where number of lines is limited to `--lines` |
| `--attempt` `-a` | `integer` |  | Attempt number to show logs for, defaults to the latest attempt. |
| `--filter-system` | `boolean` | `False` | Filter all system logs from the output. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get project

**`flyte get project [OPTIONS] [NAME]`**

Get a list of all projects, or details of a specific project by name.

By default, only active (unarchived) projects are shown. Use `--archived` to
show archived projects instead.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--archived` | `boolean` | `False` | Show archived projects instead of active ones. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get run

**`flyte get run [OPTIONS] [NAME]`**

Get a list of all runs, or details of a specific run by name.

The run details will include information about the run, its status, but only the root action will be shown.

If you want to see the actions for a run, use `get action <run_name>`.

You can filter runs by task name and optionally task version:

CODE32

You can filter runs by their user-defined labels:

CODE33

You can show only runs that have a paused action (waiting on a human in the loop):

CODE34

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--limit` | `integer` | `100` | Limit the number of runs to fetch when listing. |
| `--in-phase` | `choice` | `Sentinel.UNSET` | Filter runs by their status. |
| `--only-mine` | `boolean` | `False` | Show only runs created by the current user (you). |
| `--paused-only` | `boolean` | `False` | Show only runs that have a paused action (waiting on a human in the loop). |
| `--task-name` | `text` |  | Filter runs by task name. |
| `--task-version` | `text` |  | Filter runs by task version. |
| `--created-after` | `datetime` |  | Show runs created at or after this datetime (UTC). Accepts ISO dates, 'now', 'today', or 'now - 1 day'. |
| `--created-before` | `datetime` |  | Show runs created before this datetime (UTC). |
| `--updated-after` | `datetime` |  | Show runs updated at or after this datetime (UTC). Accepts ISO dates, 'now', 'today', or 'now - 1 day'. |
| `--updated-before` | `datetime` |  | Show runs updated before this datetime (UTC). |
| `--with-label` | `text` | `()` | Filter runs that have this label key=value. Can be specified multiple times (AND semantics). |
| `--with-label-key` | `text` | `()` | Filter runs that have this label key present (existence check). Can be specified multiple times. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get secret

**`flyte get secret [OPTIONS] [NAME]`**

Get a list of all secrets, or details of a specific secret by name.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--cluster-pool` | `text` |  | Scope the secret to a cluster pool. Mutually exclusive with --project and --domain. Mutually exclusive with project, domain. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get settings

**`flyte get settings [OPTIONS]`**

Get settings for a scope as editable YAML.

Renders three sections:


* Local overrides — uncommented, applied at this scope.
* Inherited settings — commented, with the scope they come from.
* Available settings — commented placeholders for every key that
  isn't set anywhere yet, so you can see what can be configured.


Examples:

CODE35

Use `flyte edit settings` to interactively modify these values.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--to-file` `-o` | `file` |  | Write the scope's YAML to this file instead of printing it. The file round-trips through `flyte edit settings --from-file`. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get task

**`flyte get task [OPTIONS] [NAME] [VERSION]`**

Retrieve a list of all tasks, or details of a specific task by name and version.

Currently, both `name` and `version` are required to get a specific task.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--limit` | `integer` | `100` | Limit the number of tasks to fetch. |
| `--entrypoint` | `boolean` | `False` | Show only entrypoint tasks. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte get trigger

**`flyte get trigger [OPTIONS] [TASK_NAME] [NAME]`**

Get a list of all triggers, or details of a specific trigger by name.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--limit` | `integer` | `100` | Limit the number of triggers to fetch. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte prefetch

**`flyte prefetch COMMAND [ARGS]...`**

Prefetch artifacts from remote registries.

These commands help you download and prefetch artifacts like HuggingFace models
to your Flyte storage for faster access during task execution.

#### flyte prefetch hf-model

**`flyte prefetch hf-model [OPTIONS] REPO`**

Prefetch a HuggingFace model to Flyte storage.

Downloads a model from the HuggingFace Hub and prefetches it to your configured
Flyte storage backend. This is useful for:

- Pre-fetching large models before running inference tasks
- Sharding models for tensor-parallel inference
- Avoiding repeated downloads during development

**Basic Usage:**

CODE36

**With Sharding:**

Create a shard config file (shard_config.yaml):

CODE37

Then run:

CODE38bash
--shard-config shard_config.yaml \
--accelerator A100:8 \
--hf-token-key HF_TOKEN
CODE39

**Wait for Completion:**

CODE40

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--raw-data-path` | `text` |  | Object store path to store the model. If not provided, the model will be stored using the default path generated by Flyte storage layer. |
| `--artifact-name` | `text` |  | Artifact name to use for the stored model. Must only contain alphanumeric characters, underscores, and hyphens. If not provided, the repo name will be used (replacing '.' with '-'). |
| `--architecture` | `text` | `Sentinel.UNSET` | Model architecture, as given in HuggingFace config.json. |
| `--task` | `text` | `auto` | Model task, e.g., 'generate', 'classify', 'embed', 'score', etc. Refer to vLLM docs. 'auto' will try to discover this automatically. |
| `--modality` | `text` | `('text',)` | Modalities supported by the model, e.g., 'text', 'image', 'audio', 'video'. Can be specified multiple times. |
| `--format` | `text` | `Sentinel.UNSET` | Model serialization format, e.g., safetensors, onnx, torchscript, joblib, etc. |
| `--model-type` | `text` | `Sentinel.UNSET` | Model type, e.g., 'transformer', 'xgboost', 'custom', etc. For HuggingFace models, this is auto-determined from config.json['model_type']. |
| `--short-description` | `text` | `Sentinel.UNSET` | Short description of the model. |
| `--force` | `integer` | `0` | Force store of the model. Increment value (--force=1, --force=2, ...) to force a new store. |
| `--wait` | `boolean` | `False` | Wait for the model to be stored before returning. |
| `--hf-token-key` | `text` | `HF_TOKEN` | Name of the Flyte secret containing your HuggingFace token. Note: This is not the HuggingFace token itself, but the name of the secret in the Flyte secret store. |
| `--cpu` | `text` | `2` | CPU request for the prefetch task (e.g., '2', '4', '2,4' for 2-4 CPUs). |
| `--mem` | `text` | `8Gi` | Memory request for the prefetch task (e.g., '16Gi', '64Gi', '16Gi,64Gi' for 16-64GB). |
| `--gpu` | `choice` |  | The gpu to use for downloading and (optionally) sharding the model. Format: '{type}:{quantity}' (e.g., 'A100:8', 'L4:1'). |
| `--disk` | `text` | `50Gi` | Disk storage request for the prefetch task (e.g., '100Gi', '500Gi'). |
| `--shm` | `text` |  | Shared memory request for the prefetch task (e.g., '100Gi', 'auto'). |
| `--shard-config` | `path` | `Sentinel.UNSET` | Path to a YAML file containing sharding configuration. The file should have 'engine' (currently only 'vllm') and 'args' keys. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte rerun

**`flyte rerun [OPTIONS] RUN_NAME`**

Re-run an existing run RUN_NAME with its original code and inputs.

Fetches the prior run's task + inputs from the platform (no local code needed) and launches a
new run that returns the same way ``flyte run`` does. ``--recover`` reuses the prior run's
succeeded actions (re-running only what failed or changed); ``--force-rerun-action`` forces
named actions to re-execute anyway. To re-run with *new* local code (reusing the prior run's
inputs), use ``flyte run <file> <task> --rerun-from <run>``.

Examples:

CODE41

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project for the new run (defaults to config). |
| `-d` `--domain` | `text` |  | Domain for the new run (defaults to config). |
| `--name` | `text` |  | Name for the new run (a random name is generated if unset). |
| `-e` `--env` | `text` | `Sentinel.UNSET` | Env var KEY=VALUE for the new run. Repeatable. |
| `--label` | `text` | `Sentinel.UNSET` | Label KEY=VALUE for the new run. Repeatable. |
| `--follow` `-f` | `boolean` | `False` | Stream the parent action logs after launch. |
| `--recover` | `boolean` | `False` | Recover from this run: reuse its succeeded actions, re-run only what failed or changed. |
| `--force-rerun-action` | `text` | `Sentinel.UNSET` | With --recover: name of an action to re-execute even though it succeeded in the source run. Repeatable. A listed parent re-enqueues its children (list them too to force the whole subtree); unknown names are ignored. |
| `--allow-missing-outputs` | `boolean` | `False` | Proceed when the source run's outputs were cleaned up from storage, using its inputs URI directly. The inputs cannot be verified from the client — if they were deleted too, the new run fails at runtime. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte run

**`flyte run [OPTIONS] COMMAND [ARGS]...`**

Run a task from a python file or deployed task.

Example usage:

CODE42

Arguments to the run command are provided right after the `run` command and before the file name.
Arguments for the task itself are provided after the task name.

To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote
Flyte environment:

CODE43

You can provide image mappings with `--image` flag. This allows you to specify
the image URI for the task environment during CLI execution without changing
the code. Any images defined with `Image.from_ref_name("name")` will resolve to the
corresponding URIs you specify here.

CODE44

If the image name is not provided, it is regarded as a default image and will
be used when no image is specified in TaskEnvironment:

CODE45

You can specify multiple image arguments:

CODE46

To run tasks that you've already deployed to Flyte, use the deployed-task command:

CODE47

To run a specific version of a deployed task, use the `env.task:version` syntax:

CODE48

You can specify the `--config` flag to point to a specific Flyte cluster:

CODE49

You can override the default configured project and domain:

CODE50

You can discover what deployed tasks are available by running:

CODE51

To run an arbitrary Python script on a remote cluster (without defining a task), use `python-script`:

CODE52

You can also install extra packages and wait for completion:

CODE53

Other arguments to the run command are listed below.

Arguments for the task itself are provided after the task name and can be retrieved using `--help`. For example:

CODE54

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--local` | `boolean` | `False` | Run the task locally |
| `--copy-style` | `choice` | `loaded_modules` | Copy style to use when running the task |
| `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. |
| `--raw-data-path` | `text` | `Sentinel.UNSET` | Override the output prefix used to store offloaded data types. e.g. s3://bucket/ |
| `--service-account` | `text` | `Sentinel.UNSET` | Kubernetes service account. If not provided, the configured default will be used |
| `--name` | `text` | `Sentinel.UNSET` | Name of the run. If not provided, a random name will be generated. |
| `--follow` `-f` | `boolean` | `False` | Wait and watch logs for the parent action. If not provided, the CLI will exit after successfully launching a remote execution with a link to the UI. |
| `--tui` | `boolean` | `False` | Show interactive TUI for local execution (requires flyte[tui]). |
| `--image` | `text` | `Sentinel.UNSET` | Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times. |
| `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. |
| `--run-project` | `text` |  | Run the remote task in this project, only applicable when using `deployed-task` subcommand. |
| `--run-domain` | `text` |  | Run the remote task in this domain, only applicable when using `deployed-task` subcommand. |
| `--debug` | `boolean` | `False` | Run the task as a VSCode debug task. Starts a code-server in the container so you can connect via the UI to interactively debug/run the task. |
| `--env` `-e` | `text` | `Sentinel.UNSET` | Environment variable to set on the run context. Format: KEY=VALUE. Can be specified multiple times, e.g. `-e LOG_LEVEL=debug -e FOO=bar`. |
| `--max-action-concurrency` | `integer range` |  | Maximum number of actions that can run concurrently within the run. If not provided, the platform default (run.max_action_concurrency setting) applies. |
| `--label` | `text` | `Sentinel.UNSET` | User-defined label to attach to the run. Format: KEY=VALUE. Can be specified multiple times, e.g. `--label team=ml --label env=prod`. |
| `--recover-from` | `text` |  | Recover a fresh run from a prior run: reuse its succeeded actions and re-run only what failed or changed. Remote-only. |
| `--force-rerun-action` | `text` | `Sentinel.UNSET` | With --recover-from: name of an action to re-execute even though it succeeded in the prior run. Repeatable. A listed parent re-enqueues its children (list them too to force the whole subtree); unknown names are ignored. |
| `--rerun-from` | `text` |  | Re-run an existing run with THIS local code, reusing that run's inputs (no per-task input flags are needed). Remote-only. |
| `--queue` | `text` |  | Queue (cluster) to send the run to. Overrides any queue set on the task. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte run deployed-task

**`flyte run deployed-task [OPTIONS] COMMAND [ARGS]...`**

Run remote task from the Flyte backend

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte serve

**`flyte serve [OPTIONS] COMMAND [ARGS]...`**

Serve an app from a Python file using flyte.serve().

This command allows you to serve apps defined with `flyte.app.AppEnvironment`
in your Python files. The serve command will deploy the app to the Flyte backend
and start it, making it accessible via a URL.

Example usage:

CODE55

**Local serving:** Use the `--local` flag to serve the app on localhost without
deploying to the Flyte backend. This is useful for local development and testing:

CODE56

Arguments to the serve command are provided right after the `serve` command and before the file name.

To follow the logs of the served app, use the `--follow` flag:

CODE57

Note: Log streaming is not yet fully implemented and will be added in a future release.

You can provide image mappings with `--image` flag. This allows you to specify
the image URI for the app environment during CLI execution without changing
the code. Any images defined with `Image.from_ref_name("name")` will resolve to the
corresponding URIs you specify here.

CODE58

If the image name is not provided, it is regarded as a default image and will
be used when no image is specified in AppEnvironment:

CODE59

You can specify multiple image arguments:

CODE60

You can specify the `--config` flag to point to a specific Flyte cluster:

CODE61

You can override the default configured project and domain:

CODE62

Other arguments to the serve command are listed below.

Note: This pattern is primarily useful for serving apps defined in tasks.
Serving deployed apps is not currently supported through this CLI command.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--copy-style` | `choice` | `loaded_modules` | Copy style to use when serving the app |
| `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. |
| `--service-account` | `text` | `Sentinel.UNSET` | Kubernetes service account. If not provided, the configured default will be used |
| `--name` | `text` | `Sentinel.UNSET` | Name of the app deployment. If not provided, the app environment name will be used. |
| `--follow` `-f` | `boolean` | `False` | Wait and watch logs for the app. If not provided, the CLI will exit after successfully deploying the app with a link to the UI. |
| `--image` | `text` | `Sentinel.UNSET` | Image to be used in the serve. Format: imagename=imageuri. Can be specified multiple times. |
| `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. |
| `--env-var` `-e` | `text` | `Sentinel.UNSET` | Environment variable to set in the app. Format: KEY=VALUE. Can be specified multiple times. Example: --env-var LOG_LEVEL=DEBUG --env-var DATABASE_URL=postgresql://... |
| `--local` | `boolean` | `False` | Serve the app locally on localhost instead of deploying to the Flyte backend. The app will be served on the port defined in the AppEnvironment. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte signal

**`flyte signal COMMAND [ARGS]...`**

Signal a paused condition action.

#### flyte signal condition

**`flyte signal condition [OPTIONS] RUN_NAME ACTION_NAME [VALUE]`**

Signal a paused condition action.

The condition's declared payload type and prompt are read from the
backend. If VALUE is omitted the condition's prompt is displayed and a
typed interactive prompt is shown to collect the payload. When VALUE is
provided it's coerced to the expected type (``true``/``false`` for bool,
integer literals for int, decimal literals for float, any string for str).

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte start

**`flyte start COMMAND [ARGS]...`**

Start various Flyte services.

#### flyte start devbox

**`flyte start devbox [OPTIONS]`**

Start a local Flyte devbox cluster.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--image` | `text` |  | Docker image to use for the devbox cluster. |
| `--dev` | `boolean` | `False` | Enable dev mode inside the devbox cluster (sets FLYTE_DEV=True). |
| `--gpu` | `boolean` | `False` | Pass host GPUs into the devbox container (adds --gpus all to docker run). Requires an NVIDIA-enabled host. Defaults --image to a GPU-capable image if --image is not explicitly set. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte start tui

**`flyte start tui [OPTIONS]`**

Launch the Flyte TUI. Install with ``pip install flyte[tui]``.

The mode is chosen from the resolved config:

* Remote (config has an endpoint, or FLYTE_API_KEY is set): browse a remote
  Flyte v2 cluster — projects, runs, actions, logs, tasks, apps, and triggers.
  ``flyte start tui --config remote.yaml``
* Local (no endpoint): explore past local runs recorded with persistence.
  ``flyte start tui --config local.yaml``

Local persistence can be enabled in 2 ways:

1. In the config, to record every local run:
   ``flyte create config --endpoint ... --local-persistence``
2. Via ``flyte.init(local_persistence=True)``, recording ``flyte.run`` runs
   that are local and within the active ``flyte.init``.

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `-c` `--config` | `file` |  | Path to the Flyte configuration file. Defaults to ~/.flyte/config.yaml. |
| `--poll-interval` | `float` | `2.0` | Seconds between run detail refreshes while browsing a remote run. Remote mode only. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte stop

**`flyte stop COMMAND [ARGS]...`**

Stop various Flyte services.

#### flyte stop devbox

**`flyte stop devbox`**

Pause the local Flyte devbox cluster without removing it.

### flyte update

**`flyte update COMMAND [ARGS]...`**

Update various flyte entities.

#### flyte update app

**`flyte update app [OPTIONS] NAME`**

Update an app by starting or stopping it.


Example usage:

CODE63

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--activate` `--deactivate` | `boolean` |  | Activate or deactivate app. |
| `--wait` | `boolean` | `False` | Wait for the app to reach the desired state. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte update project

**`flyte update project [OPTIONS] ID`**

Update a project's name, description, labels, or archive state.


Example usage:

CODE64

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--name` | `text` |  | Update the project display name. |
| `--description` | `text` |  | Update the project description. |
| `--label` `-l` | `text` | `Sentinel.UNSET` | Set labels as key=value pairs. Can be specified multiple times. Replaces all existing labels. |
| `--archive` `--unarchive` | `boolean` |  | Archive or unarchive the project. |
| `--help` | `boolean` | `False` | Show this message and exit. |

#### flyte update trigger

**`flyte update trigger [OPTIONS] NAME TASK_NAME`**

Update a trigger.


Example usage:

CODE65

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--activate` `--deactivate` | `boolean` | `Sentinel.UNSET` | Activate or deactivate the trigger. |
| `-p` `--project` | `text` |  | Project to which this command applies. |
| `-d` `--domain` | `text` |  | Domain to which this command applies. |
| `--help` | `boolean` | `False` | Show this message and exit. |

### flyte whoami

**`flyte whoami`**

Display the current user information.


=== PAGE: https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk ===

# Flyte SDK

These are the docs for Flyte SDK version 2.0

Flyte is the core Python SDK for the Union and Flyte platforms.

## Directory

### Classes

| Class | Description |
|-|-|
| **Flyte SDK > flyte > AsyncFunctionTaskTemplate** | A task template that wraps an asynchronous functions. |
| **Flyte SDK > flyte > Backoff** | Exponential backoff policy applied between user retries. |
| **Flyte SDK > flyte > BaseCheckpoint** | Base type for task checkpoint helpers. |
| **Flyte SDK > flyte > Cache** | Cache configuration for a task. |
| **Flyte SDK > flyte > Checkpoint** | Checkpoint helper using `flyte.io.File` for all checkpoint blob I/O (load/save, async and sync). |
| **Flyte SDK > flyte > ConditionWebhook** | Webhook configuration for a condition notification. |
| **Flyte SDK > flyte > Cron** | Cron-based automation schedule for use with `Trigger`. |
| **Flyte SDK > flyte > Device** | Represents a device type, its quantity and partition if applicable. |
| **Flyte SDK > flyte > Environment** | Base class for execution environments, shared by `TaskEnvironment` and. |
| **Flyte SDK > flyte > FixedRate** | Fixed-rate (interval-based) automation schedule for use with `Trigger`. |
| **Flyte SDK > flyte > Image** | Container image specification built using a fluent, two-step pattern:. |
| **Flyte SDK > flyte > ImageBuild** | Result of an image build operation. |
| **Flyte SDK > flyte > PodTemplate** | Custom PodTemplate specification for a Task. |
| **Flyte SDK > flyte > Resources** | Resources such as CPU, Memory, and GPU that can be allocated to a task. |
| **Flyte SDK > flyte > RetryStrategy** | Retry strategy for a task. |
| **Flyte SDK > flyte > ReusePolicy** | Configure a task environment for container reuse across multiple task invocations. |
| **Flyte SDK > flyte > Secret** | Secrets are used to inject sensitive information into tasks or image build context. |
| **Flyte SDK > flyte > TaskEnvironment** | Define an execution environment for a set of tasks. |
| **Flyte SDK > flyte > TaskTemplate** | Task template is a template for a task that can be executed. |
| **Flyte SDK > flyte > Timeout** | Timeout bounds for a task. |
| **Flyte SDK > flyte > Trigger** | Specification for a scheduled trigger that can be associated with any Flyte task. |
| **Flyte SDK > flyte.ai.agents > AccessDenied** | Raised when a write targets a read-only or reserved prefix. |
| **Flyte SDK > flyte.ai.agents > Agent** | A flyte-native tool-use agent harness. |
| **Flyte SDK > flyte.ai.agents > AgentEvent** | Lightweight event emitted by the agent loop. |
| **Flyte SDK > flyte.ai.agents > AgentResult** | Outcome of a single agent invocation. |
| **Flyte SDK > flyte.ai.agents > AgentTool** | A normalized tool descriptor used by `Agent`. |
| **Flyte SDK > flyte.ai.agents > ConcurrencyError** | Raised when an ``expected_sha`` precondition does not match the current state. |
| [`flyte.ai.agents.LLMMessage`](flyte.ai.agents/llmmessage/page.md) | Provider-agnostic shape returned by `LLMCallable`. |
| [`flyte.ai.agents.MCPServerSpec`](flyte.ai.agents/mcpserverspec/page.md) | Declarative spec for a remote MCP server that exposes tools. |
| [`flyte.ai.agents.MemoryMeta`](flyte.ai.agents/memorymeta/page.md) | Per-file metadata sidecar (sha256, actor, timestamp, …) for a memory entry. |
| [`flyte.ai.agents.MemoryStore`](flyte.ai.agents/memorystore/page.md) | Conversation transcript + path-addressed artifact memory backed by `flyte.io.Dir`. |
| [`flyte.ai.agents.MemoryStoreError`](flyte.ai.agents/memorystoreerror/page.md) | Base class for `MemoryStore` errors. |
| [`flyte.ai.agents.ToolFn`](flyte.ai.agents/toolfn/page.md) | The tool under invocation, handed to a `ToolCallHandler`. |
| [`flyte.ai.agents.agent.Agent`](flyte.ai.agents.agent/agent/page.md) | A flyte-native tool-use agent harness. |
| [`flyte.ai.agents.agent.AgentEvent`](flyte.ai.agents.agent/agentevent/page.md) | Lightweight event emitted by the agent loop. |
| [`flyte.ai.agents.memory.AccessDenied`](flyte.ai.agents.memory/accessdenied/page.md) | Raised when a write targets a read-only or reserved prefix. |
| [`flyte.ai.agents.memory.ConcurrencyError`](flyte.ai.agents.memory/concurrencyerror/page.md) | Raised when an ``expected_sha`` precondition does not match the current state. |
| [`flyte.ai.agents.memory.MemoryMeta`](flyte.ai.agents.memory/memorymeta/page.md) | Per-file metadata sidecar (sha256, actor, timestamp, …) for a memory entry. |
| [`flyte.ai.agents.memory.MemoryStore`](flyte.ai.agents.memory/memorystore/page.md) | Conversation transcript + path-addressed artifact memory backed by `flyte.io.Dir`. |
| [`flyte.ai.agents.memory.MemoryStoreError`](flyte.ai.agents.memory/memorystoreerror/page.md) | Base class for `MemoryStore` errors. |
| [`flyte.ai.agents.protocol.AgentResult`](flyte.ai.agents.protocol/agentresult/page.md) | Outcome of a single agent invocation. |
| [`flyte.ai.chat.AgentChatAppEnvironment`](flyte.ai.chat/agentchatappenvironment/page.md) | An `AppEnvironment` that spins up a FastAPI chat. |
| [`flyte.ai.chat.CustomTheme`](flyte.ai.chat/customtheme/page.md) | Declarative color theme for the Agent Chat UI. |
| [`flyte.ai.chat.app.AgentChatAppEnvironment`](flyte.ai.chat.app/agentchatappenvironment/page.md) | An `AppEnvironment` that spins up a FastAPI chat. |
| [`flyte.ai.chat.app.CustomTheme`](flyte.ai.chat.app/customtheme/page.md) | Declarative color theme for the Agent Chat UI. |
| [`flyte.ai.mcp.FlyteMCPAppEnvironment`](flyte.ai.mcp/flytemcpappenvironment/page.md) | Serve a Flyte-facing MCP server over HTTP (FastMCP + Starlette + Uvicorn). |
| [`flyte.ai.mcp.MCPAppEnvironment`](flyte.ai.mcp/mcpappenvironment/page.md) | Serve a FastMCP server over HTTP (Starlette + Uvicorn) or over stdio. |
| [`flyte.app.AppEndpoint`](flyte.app/appendpoint/page.md) | Embed an upstream app's endpoint as an app parameter. |
| [`flyte.app.AppEnvironment`](flyte.app/appenvironment/page.md) | Configure a long-running app environment for APIs, dashboards, or model servers. |
| [`flyte.app.ConnectorEnvironment`](flyte.app/connectorenvironment/page.md) | Configure a connector environment for custom Flyte connectors. |
| [`flyte.app.DeployedAppEnvironment`](flyte.app/deployedappenvironment/page.md) |  |
| [`flyte.app.Domain`](flyte.app/domain/page.md) | Subdomain to use for the domain. |
| [`flyte.app.Link`](flyte.app/link/page.md) | Custom links to add to the app. |
| [`flyte.app.Parameter`](flyte.app/parameter/page.md) | Parameter for application. |
| [`flyte.app.Port`](flyte.app/port/page.md) |  |
| [`flyte.app.RunOutput`](flyte.app/runoutput/page.md) | Use a run's output for app parameters. |
| [`flyte.app.Scaling`](flyte.app/scaling/page.md) | Controls replica count and autoscaling behavior for app environments. |
| [`flyte.app.Timeouts`](flyte.app/timeouts/page.md) | Timeout configuration for the application. |
| [`flyte.app.extras.FastAPIAppEnvironment`](flyte.app.extras/fastapiappenvironment/page.md) |  |
| [`flyte.app.extras.FastAPIPassthroughAuthMiddleware`](flyte.app.extras/fastapipassthroughauthmiddleware/page.md) | FastAPI middleware that automatically sets Flyte auth metadata from request headers. |
| [`flyte.app.extras.FlyteWebhookAppEnvironment`](flyte.app.extras/flytewebhookappenvironment/page.md) | A pre-built FastAPI app environment for common Flyte webhook operations. |
| [`flyte.clustered.ClusterFailurePolicy`](flyte.clustered/clusterfailurepolicy/page.md) | Failure and restart policy for the JobSet as a whole. |
| [`flyte.clustered.ClusteredTaskEnvironment`](flyte.clustered/clusteredtaskenvironment/page.md) | A TaskEnvironment that emits a Kubernetes JobSet for distributed multi-node training. |
| [`flyte.clustered.ClusteredTaskTemplate`](flyte.clustered/clusteredtasktemplate/page.md) | Task template for ``ClusteredTaskEnvironment``. |
| [`flyte.clustered.TorchRun`](flyte.clustered/torchrun/page.md) | TorchRun launcher configuration for a ClusteredTaskEnvironment. |
| [`flyte.config.Config`](flyte.config/config/page.md) | This the parent configuration object and holds all the underlying configuration object types. |
| [`flyte.connectors.AsyncConnector`](flyte.connectors/asyncconnector/page.md) | This is the base class for all async connectors, and it defines the interface that all connectors must implement. |
| [`flyte.connectors.AsyncConnectorExecutorMixin`](flyte.connectors/asyncconnectorexecutormixin/page.md) | This mixin class is used to run the connector task locally, and it's only used for local execution. |
| [`flyte.connectors.ConnectorRegistry`](flyte.connectors/connectorregistry/page.md) | This is the registry for all connectors. |
| [`flyte.connectors.ConnectorService`](flyte.connectors/connectorservice/page.md) |  |
| [`flyte.connectors.Resource`](flyte.connectors/resource/page.md) | This is the output resource of the job. |
| [`flyte.connectors.ResourceMeta`](flyte.connectors/resourcemeta/page.md) | This is the metadata for the job. |
| [`flyte.errors.ActionAbortedError`](flyte.errors/actionabortederror/page.md) | This error is raised when an action was aborted, externally. |
| [`flyte.errors.ActionNotFoundError`](flyte.errors/actionnotfounderror/page.md) | This error is raised when the user tries to access an action that does not exist. |
| [`flyte.errors.BaseRuntimeError`](flyte.errors/baseruntimeerror/page.md) | Base class for all Union runtime errors. |
| [`flyte.errors.CodeBundleError`](flyte.errors/codebundleerror/page.md) | This error is raised when the code bundle cannot be created, for example when no files are found to bundle. |
| [`flyte.errors.ConditionAlreadyExistsError`](flyte.errors/conditionalreadyexistserror/page.md) | This error is raised when the user tries to create a condition that already exists within the action. |
| [`flyte.errors.ConditionFailedError`](flyte.errors/conditionfailederror/page.md) | This error is raised when a condition fails during execution. |
| [`flyte.errors.ConditionNotFoundError`](flyte.errors/conditionnotfounderror/page.md) | This error is raised when the user tries to access a condition that does not exist. |
| [`flyte.errors.ConditionTimedoutError`](flyte.errors/conditiontimedouterror/page.md) | This error is raised when a condition is not signaled within its specified timeout. |
| [`flyte.errors.CustomError`](flyte.errors/customerror/page.md) | This error is raised when the user raises a custom error. |
| [`flyte.errors.DeploymentError`](flyte.errors/deploymenterror/page.md) | This error is raised when the deployment of a task fails, or some preconditions for deployment are not met. |
| [`flyte.errors.ImageBuildError`](flyte.errors/imagebuilderror/page.md) | This error is raised when the image build fails. |
| [`flyte.errors.ImagePullBackOffError`](flyte.errors/imagepullbackofferror/page.md) | This error is raised when the image cannot be pulled. |
| [`flyte.errors.InitializationError`](flyte.errors/initializationerror/page.md) | This error is raised when the Union system is tried to access without being initialized. |
| [`flyte.errors.InlineIOMaxBytesBreached`](flyte.errors/inlineiomaxbytesbreached/page.md) | This error is raised when the inline IO max bytes limit is breached. |
| [`flyte.errors.InvalidImageNameError`](flyte.errors/invalidimagenameerror/page.md) | This error is raised when the image name is invalid. |
| [`flyte.errors.InvalidPackageError`](flyte.errors/invalidpackageerror/page.md) | Raised when an invalid system package is detected during image build. |
| [`flyte.errors.LogsNotYetAvailableError`](flyte.errors/logsnotyetavailableerror/page.md) | This error is raised when the logs are not yet available for a task. |
| [`flyte.errors.ModuleLoadError`](flyte.errors/moduleloaderror/page.md) | This error is raised when the module cannot be loaded, either because it does not exist or because of a. |
| [`flyte.errors.NonRecoverableError`](flyte.errors/nonrecoverableerror/page.md) | Raised when an error is encountered that is not recoverable. |
| [`flyte.errors.NotInTaskCont