aarmens702-hub opened a new pull request, #77073:
URL: https://github.com/apache/airflow/pull/77073
related: #76788
## Human Summary
Follow-up to #76788, where I proposed an eval operator and kaxil suggested a
smaller change: a public way to get the agent a task builds.
Today, evaluating an agent task means rebuilding the agent by hand, so the
eval runs against a copy that can drift from what the Dag actually runs.
This PR adds `build_agent()` to:
- `AgentOperator` (covers `@task.agent`)
- `LLMOperator` (covers `@task.llm`)
The branch, file analysis, SQL, and schema compare operators each construct
their agent differently, so each gets its own `build_agent()` override. Without
it, the inherited method would silently return the wrong agent.
The eval example in the local development guide now uses `build_agent()`.
## AI Summary
<details><summary>Click here</summary>
Follow-up to the discussion in #76788. To eval an agent task today you
rebuild the agent by hand (the local development guide's notebook agent), so
the eval checks a copy that can drift from what the task actually runs. This
adds a public `build_agent()` that returns the agent the task's `execute()`
runs.
- **`AgentOperator`** (and so `@task.agent`): `_build_agent()` becomes
`build_agent()`, with masking, durability, prompt caching and tool logging
applied as before. Nothing outside the provider used the private name.
- **`LLMOperator`** (and so `@task.llm`): new `build_agent()`, which
`execute()` now calls.
- **The four `LLMOperator` subclasses build their agent differently**, so
each overrides `build_agent()` and its `execute()` calls it at the same point
as before:
- `LLMBranchOperator`: the output type is one of the downstream task IDs.
Raises the existing "no downstream tasks" error if there are none.
- `LLMFileAnalysisOperator`: the file analysis system prompt.
- `LLMSQLQueryOperator` and `LLMSchemaCompareOperator`: read the table
schemas for the system prompt, so they need their database connections, as
`execute()` does.
Example (from the updated docs):
```python
agent = orders_report().get_task("orders_question").build_agent()
report = dataset.evaluate_sync(lambda question:
agent.run_sync(question).output)
```
On the two questions from the issue:
1. **Unrendered templates:** `build_agent()` doesn't raise. Templates are
only rendered in a task run, so a templated `system_prompt` is passed as
written; the docstring and docs say so. Raising would block evals of any task
with a Jinja prompt.
2. **Docs:** the "Check the answers themselves" section of
`local_development.rst` now uses `build_agent()` with a pydantic-evals example
on the guide's own `orders_report` Dag.
Tests: 11 new tests in `test_agent.py` and `test_llm.py`. They check that
`build_agent()` returns the agent `execute()` builds, that it works on a task
taken from a Dag (operator and decorator forms) with a templated system prompt
passed as written, and that each subclass builds its own output type and system
prompt. All `common.ai` operator and decorator tests pass (853). Across the
full `common.ai` unit suite, the only failures are 50 sandbox tests that fail
the same way on main in my local environment. ruff and mypy are clean on the
changed files.
</details>
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: Claude Code following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
Drafted-by: Claude Code (Opus 5.5); reviewed by @aarmens702-hub before
posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]