Augment an AsyncFlow application

This recipe adds agent decisions to an application that already runs on AsyncFlow, without rewriting it as an agent graph. The application keeps its scientific state, its tasks and its hard policies; an agent is called at a fixed decision point, as one more AsyncFlow task. flowgentic.agent.bind_agent makes that call.

The recipe comes from the agent-augmented demo, examples/ai-hpc-coupling/augmented_application.py, which adds a planner, an analyst and a supervisor to the RADICAL baseline. Its run is tested in tests/integration/test_ai_hpc_augmented_demo.py, and bind_agent in tests/unit/test_agent_adapter.py.

Under ADR

An application controlled by ADR takes the same decision in ADR's @decide instead. ADR is not released, and Flowgentic 0.3.0 does not support ADR mode.

1. Keep the application's tasks on AsyncFlow

The application's service and tasks stay ordinary AsyncFlow code, each on a named backend:

# 1. Existing RADICAL services and tasks remain direct AsyncFlow code.
@flow.function_task(service=True, backend="ai")
async def load_surrogate() -> SurrogateService:
    return await SurrogateService.load()

service_future = load_surrogate()
service = await service_future

@flow.function_task(backend="ai")
async def propose_candidates(
    center: float,
    radius: float,
    cycle: int,
) -> dict[str, Any]:
    return await service.propose(center, radius, settings.batch_size, cycle)

@flow.function_task(backend="compute")
async def simulate(
    x: float,
    cycle: int,
    candidate_index: int,
    fail_once: bool = False,
) -> dict[str, Any]:
    return await evaluate_candidate(x, cycle, candidate_index, fail_once)

@flow.function_task(backend="ai")
async def update_surrogate(
    results: list[dict[str, Any]],
    radius: float,
    cycle: int,
) -> dict[str, Any]:
    return await service.assimilate(results, radius, cycle)

2. Write the agent component

An agent component is any object with an async ainvoke(context) method that returns a mapping or a Pydantic model. Inside it, the component can use any framework. In the demo each role is a private, one-node LangGraph graph:

@dataclass(frozen=True)
class LangGraphAgentComponent:
    """A role-specific LangGraph hidden behind Flowgentic's agent interface."""

    identity: str
    graph: CompiledStateGraph

    async def ainvoke(self, context: Mapping[str, Any]) -> dict[str, Any]:
        """Run the component's private graph and return its structured decision."""
        result = await self.graph.ainvoke({"context": dict(context)})
        return result["decision"]

def _build_component(
    identity: str,
    decide: Callable[[dict[str, Any]], Awaitable[BaseModel]],
) -> LangGraphAgentComponent:
    graph = StateGraph(ComponentState)

    async def reason(state: ComponentState) -> dict[str, Any]:
        decision = await decide(state["context"])
        return {"decision": decision.model_dump()}

    graph.add_node("reason", reason)
    graph.add_edge(START, "reason")
    graph.add_edge("reason", END)
    return LangGraphAgentComponent(identity, graph.compile())

decide is the role's decision model: scripted in rehearsal mode, an LLM with structured output in live mode.

3. Bind each component to an AsyncFlow task

bind_agent(flow, component, identity=..., framework=..., backend=...) registers the component as an AsyncFlow function task on the named backend ("ai" by default) and returns an AgentTask:

# 2. Only agent components cross the Flowgentic/framework boundary.
planner = bind_agent(
    flow,
    components.planner,
    identity="planner",
    framework=components.framework,
    backend="ai",
)
analyst = bind_agent(
    flow,
    components.analyst,
    identity="analyst",
    framework=components.framework,
    backend="ai",
)
supervisor = bind_agent(
    flow,
    components.supervisor,
    identity="supervisor",
    framework=components.framework,
    backend="ai",
)

Each call to AgentTask.ainvoke(context) submits one AsyncFlow task and returns a dictionary:

Key	Content
agent	The identity given to bind_agent
framework	The framework label given to bind_agent
decision	The component's output as a plain mapping (a Pydantic model is dumped with model_dump())
event	name (agent.<identity>), backend, cycle (the context's cycle, or -1), started, finished and duration

A component that returns anything other than a mapping or a Pydantic model raises TypeError.

With OpenTelemetry configured, each call is also an invoke_agent <identity> span. The span carries the AsyncFlow task's uid and backend, and the spans of the component's own model calls nest under it; see Telemetry.

4. Call the agent at the decision point

The application builds the agent's context, calls it, and then checks its decision against its own hard policy before acting on it:

proposal = await propose_candidates(
    state["center"], state["radius"], state["cycle"]
)
plan = await planner.ainvoke(
    {
        "cycle": state["cycle"],
        "strategy": state["strategy"],
        "center": state["center"],
        "radius": state["radius"],
        "best_x": state["best_x"],
        "best_value": state["best_value"],
        "uncertainty": state["uncertainty"],
        "expected_count": expected_count,
        "lower_bound": lower_bound,
        "upper_bound": upper_bound,
        "surrogate_candidates": proposal["candidates"],
        "previous_results": state["results"],
    }
)
planner_decision = PlannerDecision.model_validate(plan["decision"])
candidates, candidate_policy = enforce_candidate_policy(
    planner_decision,
    proposal["candidates"],
    expected_count=expected_count,
    lower_bound=lower_bound,
    upper_bound=upper_bound,
)

enforce_candidate_policy returns the accepted candidates: the application, not the agent, decides what runs. The analyst and supervisor are called the same way later in the cycle.

Placement on a named backend

The task runs on the backend named in bind_agent. This test binds a component to ai and checks that the task ran there and not on compute:

async def test_bind_agent_runs_on_the_named_asyncflow_backend() -> None:
    compute = await LocalExecutionBackend(
        ThreadPoolExecutor(max_workers=1), name="compute"
    )
    ai = await LocalExecutionBackend(ThreadPoolExecutor(max_workers=1), name="ai")
    flow = await WorkflowEngine.create(backend=[compute, ai])
    try:
        agent = bind_agent(
            flow,
            FakeAgentComponent(),
            identity="lead-planner",
            framework="test-framework",
            backend="ai",
        )

        result = await agent.ainvoke({"cycle": 2, "candidate": 3.0})

        assert result["decision"] == {"choice": 3.0, "rationale": "bounded"}
        assert [task["name"] for task in ai.tasks.values()] == ["agent_lead_planner"]
        assert not compute.tasks
    finally:
        await flow.shutdown()

The task's name is agent_<identity>, with hyphens replaced by underscores.

Run the demo

uv run --frozen --extra examples examples/ai-hpc-coupling/run_augmented_demo.sh application

Agent-augmented RADICAL AI-HPC coupling demo
application/ADR -> AsyncFlow <-> Flowgentic <-> agent components
controller=application  agent_mode=rehearsal  model=rehearsal-agents
target=2.75  budget=24

persistent AI service started: surrogate-service-01 (model loads=1)
cycle= 1  controller=application  simulations=4  retries=0  best_x=4.0000  uncertainty=2.2000  budget=4/24
         agents: next=explore  recommendation=continue planner=Use the balanced strategy with a surrogate-informed batch while respect…
cycle= 2  controller=application  simulations=4  retries=1  best_x=3.2667  uncertainty=1.2100  budget=8/24
         agents: next=explore  recommendation=continue planner=Use the explore strategy with a surrogate-informed batch while respecti…
cycle= 3  controller=application  simulations=4  retries=0  best_x=2.8633  uncertainty=0.6655  budget=12/24
         agents: next=balanced recommendation=continue planner=Use the explore strategy with a surrogate-informed batch while respecti…
cycle= 4  controller=application  simulations=4  retries=0  best_x=2.6415  uncertainty=0.3660  budget=16/24
         agents: next=exploit  recommendation=continue planner=Use the balanced strategy with a surrogate-informed batch while respect…
cycle= 5  controller=application  simulations=4  retries=0  best_x=2.7696  uncertainty=0.2013  budget=20/24
         agents: next=exploit  recommendation=stop     planner=Use the exploit strategy with a surrogate-informed batch while respecti…

Proof points
  application and scientific cycle remain direct AsyncFlow
  LangGraph is confined to the three agent components
  Flowgentic binds only agent components to the ai backend
  observed compute parallelism: 4
  guarded agent decisions recorded: 15
  stopped because: scientific uncertainty goal reached

The summary is written to examples/ai-hpc-coupling/augmented_demo_results/campaign_summary.json. Compare it with the baseline, examples/ai-hpc-coupling/run_radical_demo.sh application: both reach the same uncertainty after 20 simulations.