uranusjr commented on code in PR #76015:
URL: https://github.com/apache/airflow/pull/76015#discussion_r4439154957
##########
providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py:
##########
@@ -139,6 +153,97 @@ async def function(arguments: dict[str, Any]) ->
ToolResult:
)
+class _ParallelModel(Model):
+ """
+ Asks for one ``lookup`` per key in a single response, then answers with
their results.
+
+ The calls get the ids ``call-<key>``. ``fail_on_answer`` makes the request
that would
+ answer raise, as a crashed worker would leave it.
+ """
+
+ def __init__(self, keys: tuple[str, ...] = ("a", "b"), *, fail_on_answer:
bool = False) -> None:
+ self.keys = keys
+ self.fail_on_answer = fail_on_answer
+ self.calls = 0
+
+ async def get_response(self, system_instructions, input, *args: Any,
**kwargs: Any) -> ModelResponse:
+ self.calls += 1
+ outputs = [
+ item["output"]
+ for item in (input if isinstance(input, list) else [])
+ if isinstance(item, dict) and item.get("type") ==
"function_call_output"
+ ]
+ output: list[ResponseFunctionToolCall | ResponseOutputMessage]
+ if not outputs:
+ output = [
+ ResponseFunctionToolCall(
+ id=f"fc-{key}",
+ call_id=f"call-{key}",
+ name="lookup",
+ arguments=json.dumps({"key": key}),
+ type="function_call",
+ )
+ for key in self.keys
+ ]
+ else:
+ if self.fail_on_answer:
+ raise RuntimeError("worker lost")
+ text = ResponseOutputText(text=f"answer: {sorted(outputs)}",
type="output_text", annotations=[])
+ output = [
+ ResponseOutputMessage(
+ id="msg", role="assistant", status="completed",
type="message", content=[text]
+ )
+ ]
+ return ModelResponse(output=list(output), usage=Usage(),
response_id=f"resp-{self.calls}")
+
+ def stream_response(self, *args: Any, **kwargs: Any):
+ raise NotImplementedError
+
+
+class _SleepingTools:
+ """A ``lookup`` tool that sleeps ``body_delays[key]`` seconds, fails for
the keys in ``failing``, and counts the calls per key."""
Review Comment:
This is too long for a single-line docstring.
--
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]
