uranusjr commented on code in PR #76015:
URL: https://github.com/apache/airflow/pull/76015#discussion_r4450678749
##########
providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py:
##########
@@ -519,3 +629,353 @@ def
test_outside_a_task_it_runs_the_agent_as_runner_run_does(self):
assert result.final_output == "answer: k0=42"
assert model.calls == 2
+
+
+class _SlowRunHooks(RunHooks):
+ """Makes every tool call wait ``delays[call id]`` seconds before it is
invoked."""
+
+ def __init__(self, delays: dict[str, float]) -> None:
+ self.delays = delays
+
+ async def on_tool_start(self, context, agent, tool) -> None:
+ await asyncio.sleep(self.delays[context.tool_call_id])
+
+
+class _SlowAgentHooks(AgentHooks):
+ """Makes every tool call of the agent wait ``delays[call id]`` seconds
before it is invoked."""
+
+ def __init__(self, delays: dict[str, float]) -> None:
+ self.delays = delays
+
+ async def on_tool_start(self, context, agent, tool) -> None:
+ await asyncio.sleep(self.delays[context.tool_call_id])
+
+
+def _attempt_parallel(
+ storage: MemoryStorage,
+ model: Model,
+ tools: _SleepingTools,
+ *,
+ agent_hooks: AgentHooks | None = None,
+ guardrail_delays: dict[str, float] | None = None,
+ rejected_calls: set[str] | None = None,
+ **kwargs: Any,
+):
+ """Run one attempt of an agent whose model asks for several tools at
once."""
+ function_tools = tools.tools()
+ if guardrail_delays is not None or rejected_calls is not None:
+
+ @tool_input_guardrail
+ async def wait(data):
+ await asyncio.sleep((guardrail_delays or
{}).get(data.context.tool_call_id, 0))
+ if rejected_calls and data.context.tool_call_id in rejected_calls:
+ return ToolGuardrailFunctionOutput.reject_content("stopped by
a guardrail")
+ return ToolGuardrailFunctionOutput.allow()
+
+ for function_tool_ in function_tools:
+ function_tool_.tool_input_guardrails = [wait]
+ agent = Agent(name="analyst", model=model, tools=function_tools,
hooks=agent_hooks)
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ return run_durably_sync(agent, "go", run_config=NO_TRACING, **kwargs),
journal
+
+
+class TestParallelToolCalls:
+ def
test_parallel_tool_calls_replay_when_their_bodies_finish_in_a_different_order(self):
+ """
+ Characterizes the documented claim that parallel calls keep the order
they started in.
+
+ The model's response reserves a position for each tool call, so how
long a tool body takes does not matter.
+ """
+ storage = MemoryStorage()
+ tools = _SleepingTools({"a": 0.05, "b": 0.01})
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt_parallel(storage, _ParallelModel(fail_on_answer=True),
tools)
+ assert tools.calls == {"a": 1, "b": 1}
+
+ tools.body_delays = {"a": 0.01, "b": 0.05}
+ retry_model = _ParallelModel()
+ _, journal = _attempt_parallel(storage, retry_model, tools)
+
+ assert journal.stats.replayed == {"model": 1, "tool": 2}
+ assert tools.calls == {"a": 1, "b": 1}
+ assert retry_model.calls == 1
+
+ def test_only_the_parallel_tool_call_that_failed_runs_again_on_retry(self):
+ storage = MemoryStorage()
+ tools = _SleepingTools(failing={"b"})
+ with pytest.raises(Exception, match="lookup b failed"):
+ _attempt_parallel(storage, _ParallelModel(("a", "b", "c")), tools)
+ assert tools.calls == {"a": 1, "b": 1, "c": 1}
+
+ tools.failing = set()
+ retry_model = _ParallelModel(("a", "b", "c"))
+ _, journal = _attempt_parallel(storage, retry_model, tools)
+
+ assert tools.calls == {"a": 1, "b": 2, "c": 1}
+ assert journal.stats.replayed == {"model": 1, "tool": 2}
+ assert retry_model.calls == 1
+
+ def
test_a_replayed_tool_call_a_guardrail_stops_is_not_counted_as_replayed(self):
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+ rejected: set[str] = set()
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt_parallel(storage, _ParallelModel(fail_on_answer=True),
tools, rejected_calls=rejected)
+ assert tools.calls == {"a": 1, "b": 1}
+
+ rejected.add("call-b")
+ retry_model = _ParallelModel()
+ _, journal = _attempt_parallel(storage, retry_model, tools,
rejected_calls=rejected)
+
+ assert journal.stats.replayed == {"model": 1, "tool": 1}
+ assert tools.calls == {"a": 1, "b": 1}
+ assert retry_model.calls == 1
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("call_id", ["call-shared", ""],
ids=["duplicate-id", "empty-id"])
+ async def
test_tool_calls_sharing_a_call_id_each_replay_their_own_result(self, call_id):
+ """The SDK rejects a repeated id itself, but a model that emitted one
must not make the calls trade steps."""
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+ response = ModelResponse(
+ output=[
+ ResponseFunctionToolCall(
+ id=f"fc-{key}",
+ call_id=call_id,
+ name="lookup",
+ arguments=json.dumps({"key": key}),
+ type="function_call",
+ )
+ for key in ("a", "b")
+ ],
+ usage=Usage(),
+ response_id="resp",
+ )
+
+ async def attempt(order: tuple[str, ...]) -> tuple[dict[str, Any],
DurableJournal]:
+ journal = DurableJournal(storage)
+ durable_run = journal.start_run()
+ durable_tools = [_get_durable_tool(tool, durable_run) for tool in
tools.tools()]
+ _DurableModel(mock.create_autospec(Model, instance=True),
durable_run)._reserve_tool_steps(
+ response, durable_tools
+ )
+ outputs = {}
+ for key in order:
+ context = ToolContext(
+ context=None,
+ usage=Usage(),
+ tool_name="lookup",
+ tool_call_id=call_id,
+ tool_arguments=json.dumps({"key": key}),
+ )
+ outputs[key] = await durable_tools[0].on_invoke_tool(context,
json.dumps({"key": key}))
+ _RESERVED_TOOL_STEPS.pop(durable_run, None)
+ return outputs, journal
+
+ await attempt(("a", "b"))
+ assert tools.calls == {"a": 1, "b": 1}
+
+ outputs, journal = await attempt(("b", "a"))
+
+ assert tools.calls == {"a": 1, "b": 1}
+ assert outputs == {"a": "a=42", "b": "b=42"}
+ assert journal.stats.replayed["tool"] == 2
+
+ def test_only_wrapped_tools_get_a_reserved_step(self):
+ """
+ A tool the model is handed but the durable run does not wrap never
takes a reserved step.
+
+ MCP tools and a handoff agent's tools reach the model this way, so
reserving for them would leave
+ positions nobody takes.
+ """
+
+ @function_tool
+ def plain_lookup(key: str) -> str:
+ """Look a key up."""
+ return key
+
+ journal = DurableJournal(MemoryStorage())
+ durable_run = journal.start_run()
+ durable_tool = _get_durable_tool(_SleepingTools().tools()[0],
durable_run)
+ response = ModelResponse(
+ output=[
+ ResponseFunctionToolCall(
+ id=f"fc-{call_id}",
+ call_id=call_id,
+ name=name,
+ arguments=json.dumps({"key": "a"}),
+ type="function_call",
+ )
+ for call_id, name in (("call-durable", "lookup"),
("call-plain", "plain_lookup"))
+ ],
+ usage=Usage(),
+ response_id="resp",
+ )
+
+ _DurableModel(mock.create_autospec(Model, instance=True),
durable_run)._reserve_tool_steps(
+ response, [durable_tool, plain_lookup]
+ )
+
+ assert list(_RESERVED_TOOL_STEPS[durable_run]) == ["call-durable"]
+ assert durable_run.position == 1
+ _RESERVED_TOOL_STEPS.pop(durable_run, None)
+
+ def test_a_finished_run_does_not_keep_its_reserved_tool_steps(self):
+ @function_tool(needs_approval=True)
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ return key
+
+ journal = DurableJournal(MemoryStorage(), clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(
+ Agent(name="analyst", model=_ScriptedModel(), tools=[lookup]),
"go", run_config=NO_TRACING
+ )
+
+ # The tool never ran, so its reservation was never taken by the tool.
Review Comment:
```suggestion
```
--
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]
