Files
Eduard van Valkenburg 1acd242550 Python: Add AgentLoopMiddleware for re-running agents in a loop (#6174)
* Python: Add AgentLoopMiddleware for re-running agents in a loop

Add `AgentLoopMiddleware`, an `AgentMiddleware` that re-runs the wrapped
agent in a loop. A single configurable class covers three common patterns,
each with a convenience classmethod factory:

- Ralph loop (`.ralph(...)`): no exit criteria, with feedback tracking
  (`record_feedback`/`progress`), progress injection (`inject_progress`),
  optional fresh context per iteration (`fresh_context`), and an early-stop
  completion signal (`is_complete`).
- Predicate (`.with_predicate(...)`): loop while a `should_continue` callable
  returns True (e.g. paired with `todos_remaining`/`background_tasks_running`).
- Judge (`.with_judge(...)`): a second chat client decides whether the original
  request was answered, using a `JudgeVerdict` structured-output response.

The loop also auto-resolves pending function-approval / user-input requests via
an `on_approval_request` callable (bounded by `max_approval_rounds`), and the
next iteration's input is controlled by `next_message`. Supports both streaming
and non-streaming runs.

Exports `AgentLoopMiddleware`, `JudgeVerdict`, `todos_remaining`, and
`background_tasks_running`. Adds tests, a sample, and docs.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Refine AgentLoopMiddleware API and sample

- with_judge: add criteria list with {{criteria}} templating into judge
  instructions plus an agent-side instruction; add fresh_context, additional
  judge feedback relay; default judge max_iterations.
- should_continue is now required and positional; supports (bool, str|None)
  feedback tuples surfaced to next_message/record_feedback via feedback kwarg.
- Judge forwards full multi-modal request and response messages.
- Default max_iterations=10 (explicit None = unbounded); removed is_complete and
  Ralph terminology; ShouldContinueResult is a real TypeAlias.
- Sample: stream all loops, print iteration counts via injected user-block
  boundaries (robust to function calling), <role>: content formatting, per-method
  expected output, and a looping todo sample.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Fix CI checks for AgentLoopMiddleware

- Resolve pyright errors in _loop.py: drop the always-true final_result None
  check (the while loop always assigns it) and cast finish_reason to the
  AgentResponse constructor's expected type.
- Apply pyupgrade --py310-plus: import TypeAlias from typing.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Resolve mypy/pyright disagreement on finish_reason

pyright infers AgentResponse.finish_reason as including str and rejects the
direct assignment, while mypy considers a cast redundant. Drop the cast and
suppress only pyright with a targeted reportArgumentType ignore, satisfying
both type checkers.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Add todo+judge AgentLoopMiddleware sample

Add a second AgentLoopMiddleware sample that composes two criteria in one
should_continue predicate: a TodoProvider check (evaluated first) and a
report-style judge chat client (evaluated once todos are complete) that grades
the assembled report against shared requirements. Register it in the middleware
samples README.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Compose todo+judge loops as two middleware

Rework the todo+judge sample to compose two AgentLoopMiddleware on the agent
itself (middleware=[judge_loop, todo_loop]) instead of a single hand-written
predicate. The inner todos_remaining loop drafts the report todo-by-todo and the
outer with_judge loop re-runs it until an editor chat client judges the report
publication-ready, reusing the built-in helpers.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Reset session for fresh_context loops via snapshot/restore

AgentLoopMiddleware.fresh_context previously only reset context.messages,
so with an attached session each iteration still reloaded the local
transcript or re-threaded the service-side conversation id and the model
saw the accumulated history. Snapshot the session once before the loop
(via to_dict) and restore it (from_dict + field copy) between iterations,
so every pass starts from the pre-loop baseline. The final iteration's
pass is persisted (no restore after the terminating iteration), so a
subsequent agent.run continues from there.

Removed the obsolete warning, updated docstrings and core AGENTS.md, and
added tests: a snapshot/restore round-trip, a session-reset
streaming x fresh_context x inject_progress x store matrix across multiple
runs and loop iterations, and response_format parsing across the loop.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Updated samples and docstrings

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-12 14:35:54 +00:00

130 lines
5.1 KiB
Python

# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, AgentLoopMiddleware, AgentSession, TodoProvider, todos_remaining
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
Agent Loop Middleware: todo loop (should_continue via a provider helper)
This sample demonstrates ``AgentLoopMiddleware`` driven by a ``should_continue`` predicate built from
a ``TodoProvider``. The ``todos_remaining`` helper keeps the agent running while it still has open
todo items, so the agent plans work on its first turn and completes one item per turn afterwards.
``max_iterations`` bounds the loop as a safety cap, and a single session keeps the todo state across
iterations. After the run the sample prints the todos the agent created.
The loop is run with streaming, so the injected messages between iterations show up as ``user``
updates; the stream is printed as ``<role>: <content>`` lines.
Environment variables:
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint URL
FOUNDRY_MODEL — Model deployment name
Authentication:
Run ``az login`` before running this sample.
"""
async def todo_loop(client: FoundryChatClient) -> None:
"""Loop while a TodoProvider still has open items."""
print("\n=== Callable criterion (loop while todos remain) ===")
# 1. A TodoProvider gives the agent tools to plan and track work as todo items.
todo_provider = TodoProvider()
# 2. ``todos_remaining`` builds a ``should_continue`` predicate that returns True while any todo
# item is still open. ``max_iterations`` guarantees the loop stops even if the agent stalls.
loop = AgentLoopMiddleware(
should_continue=todos_remaining(todo_provider),
max_iterations=6,
)
agent = Agent(
client=client,
name="planner",
instructions=(
"You are a writing assistant working through a todo list. "
"On your FIRST turn, break the task into todo items using your todo tools and stop "
"(do not start writing yet). On EACH SUBSEQUENT turn, complete exactly ONE remaining "
"todo item, write its content, and mark it done using your tools — never complete more "
"than one item per turn. When every item is done, give a brief final summary."
),
context_providers=[todo_provider],
middleware=[loop],
)
# 3. Reuse a single session so todo state persists across loop iterations. Each contiguous
# ``user`` block marks the boundary into the next iteration, so we count loop iterations by
# those boundaries — robust to the function calling this loop relies on (the todo tools issue
# several model calls per iteration, but tool calls/results are never ``user`` updates).
session = AgentSession()
prompt = "Plan and write a short 3-section blog post about Rayleigh scattering."
iterations = 1
in_user_block = False
assistant_open = False
async for update in agent.run(prompt, session=session, stream=True):
if update.role == "user":
if not in_user_block:
iterations += 1
in_user_block = True
assistant_open = False
print(f"\nuser: {update.text}", flush=True)
continue
in_user_block = False
if update.text:
if not assistant_open:
print("\nassistant: ", end="", flush=True)
assistant_open = True
print(update.text, end="", flush=True)
print(f"\n\nCompleted in {iterations} iteration(s).")
# 4. Inspect the todos the agent created, loaded from the same store the loop predicate uses.
items = await todo_provider.store.load_items(session, source_id=todo_provider.source_id)
print("\nTodos after the run:")
for item in items:
mark = "x" if item.is_complete else " "
print(f" [{mark}] {item.id}. {item.title}")
async def main() -> None:
async with AzureCliCredential() as credential:
client = FoundryChatClient(credential=credential)
await todo_loop(client)
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output (abridged; exact text varies by model):
=== Callable criterion (loop while todos remain) ===
assistant: Here is my plan. I'll create todos for each section.
user: Progress so far:
- Here is my plan. I'll create todos for each section.
user: Continue working on the task. If it is complete, say so.
assistant: Section 1 drafted. Marking it done.
user: Progress so far:
- Section 1 drafted. Marking it done.
user: Continue working on the task. If it is complete, say so.
assistant: Section 2 drafted. Marking it done.
user: Progress so far:
- Section 2 drafted. Marking it done.
user: Continue working on the task. If it is complete, say so.
assistant: Section 3 drafted. Marking it done.
Completed in 4 iteration(s).
Todos after the run:
[x] 1. Draft "What light is" section
[x] 2. Draft "How Rayleigh scattering works" section
[x] 3. Draft "Why the sky is blue" section
"""