Skip to content

Commit a6d4c33

Browse files
committed
Add max turns parameter & more oai tool support
- Also adds corresponding unit tests
1 parent 38f882b commit a6d4c33

File tree

7 files changed

+1080
-174
lines changed

7 files changed

+1080
-174
lines changed

src/agentex/lib/adk/providers/_modules/openai.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ async def run_agent(
8888
mcp_timeout_seconds: int | None = None,
8989
input_guardrails: list[InputGuardrail] | None = None,
9090
output_guardrails: list[OutputGuardrail] | None = None,
91+
max_turns: int | None = None,
9192
) -> SerializableRunResult | RunResult:
9293
"""
9394
Run an agent without streaming or TaskMessage creation.
@@ -114,6 +115,7 @@ async def run_agent(
114115
mcp_timeout_seconds: Optional param to set the timeout threshold for the MCP servers. Defaults to 5 seconds.
115116
input_guardrails: Optional list of input guardrails to run on initial user input.
116117
output_guardrails: Optional list of output guardrails to run on final agent output.
118+
max_turns: Maximum number of turns the agent can take. Uses Runner's default if None.
117119
118120
Returns:
119121
Union[SerializableRunResult, RunResult]: SerializableRunResult when in Temporal, RunResult otherwise.
@@ -136,6 +138,7 @@ async def run_agent(
136138
mcp_timeout_seconds=mcp_timeout_seconds,
137139
input_guardrails=input_guardrails,
138140
output_guardrails=output_guardrails,
141+
max_turns=max_turns,
139142
)
140143
return await ActivityHelpers.execute_activity(
141144
activity_name=OpenAIActivityName.RUN_AGENT,
@@ -163,6 +166,7 @@ async def run_agent(
163166
mcp_timeout_seconds=mcp_timeout_seconds,
164167
input_guardrails=input_guardrails,
165168
output_guardrails=output_guardrails,
169+
max_turns=max_turns,
166170
)
167171

168172
async def run_agent_auto_send(
@@ -191,6 +195,7 @@ async def run_agent_auto_send(
191195
mcp_timeout_seconds: int | None = None,
192196
input_guardrails: list[InputGuardrail] | None = None,
193197
output_guardrails: list[OutputGuardrail] | None = None,
198+
max_turns: int | None = None,
194199
) -> SerializableRunResult | RunResult:
195200
"""
196201
Run an agent with automatic TaskMessage creation.
@@ -216,6 +221,7 @@ async def run_agent_auto_send(
216221
mcp_timeout_seconds: Optional param to set the timeout threshold for the MCP servers. Defaults to 5 seconds.
217222
input_guardrails: Optional list of input guardrails to run on initial user input.
218223
output_guardrails: Optional list of output guardrails to run on final agent output.
224+
max_turns: Maximum number of turns the agent can take. Uses Runner's default if None.
219225
220226
Returns:
221227
Union[SerializableRunResult, RunResult]: SerializableRunResult when in Temporal, RunResult otherwise.
@@ -239,6 +245,7 @@ async def run_agent_auto_send(
239245
mcp_timeout_seconds=mcp_timeout_seconds,
240246
input_guardrails=input_guardrails,
241247
output_guardrails=output_guardrails,
248+
max_turns=max_turns,
242249
)
243250
return await ActivityHelpers.execute_activity(
244251
activity_name=OpenAIActivityName.RUN_AGENT_AUTO_SEND,
@@ -267,6 +274,7 @@ async def run_agent_auto_send(
267274
mcp_timeout_seconds=mcp_timeout_seconds,
268275
input_guardrails=input_guardrails,
269276
output_guardrails=output_guardrails,
277+
max_turns=max_turns,
270278
)
271279

272280
async def run_agent_streamed(
@@ -291,6 +299,7 @@ async def run_agent_streamed(
291299
mcp_timeout_seconds: int | None = None,
292300
input_guardrails: list[InputGuardrail] | None = None,
293301
output_guardrails: list[OutputGuardrail] | None = None,
302+
max_turns: int | None = None,
294303
) -> RunResultStreaming:
295304
"""
296305
Run an agent with streaming enabled but no TaskMessage creation.
@@ -320,6 +329,7 @@ async def run_agent_streamed(
320329
mcp_timeout_seconds: Optional param to set the timeout threshold for the MCP servers. Defaults to 5 seconds.
321330
input_guardrails: Optional list of input guardrails to run on initial user input.
322331
output_guardrails: Optional list of output guardrails to run on final agent output.
332+
max_turns: Maximum number of turns the agent can take. Uses Runner's default if None.
323333
324334
Returns:
325335
RunResultStreaming: The result of the agent run with streaming.
@@ -352,6 +362,7 @@ async def run_agent_streamed(
352362
mcp_timeout_seconds=mcp_timeout_seconds,
353363
input_guardrails=input_guardrails,
354364
output_guardrails=output_guardrails,
365+
max_turns=max_turns,
355366
)
356367

357368
async def run_agent_streamed_auto_send(
@@ -380,6 +391,7 @@ async def run_agent_streamed_auto_send(
380391
mcp_timeout_seconds: int | None = None,
381392
input_guardrails: list[InputGuardrail] | None = None,
382393
output_guardrails: list[OutputGuardrail] | None = None,
394+
max_turns: int | None = None,
383395
) -> SerializableRunResultStreaming | RunResultStreaming:
384396
"""
385397
Run an agent with streaming enabled and automatic TaskMessage creation.
@@ -405,6 +417,7 @@ async def run_agent_streamed_auto_send(
405417
output_type: Optional output type.
406418
tool_use_behavior: Optional tool use behavior.
407419
mcp_timeout_seconds: Optional param to set the timeout threshold for the MCP servers. Defaults to 5 seconds.
420+
max_turns: Maximum number of turns the agent can take. Uses Runner's default if None.
408421
409422
Returns:
410423
Union[SerializableRunResultStreaming, RunResultStreaming]: SerializableRunResultStreaming when in Temporal, RunResultStreaming otherwise.
@@ -428,6 +441,7 @@ async def run_agent_streamed_auto_send(
428441
mcp_timeout_seconds=mcp_timeout_seconds,
429442
input_guardrails=input_guardrails,
430443
output_guardrails=output_guardrails,
444+
max_turns=max_turns
431445
)
432446
return await ActivityHelpers.execute_activity(
433447
activity_name=OpenAIActivityName.RUN_AGENT_STREAMED_AUTO_SEND,
@@ -456,4 +470,5 @@ async def run_agent_streamed_auto_send(
456470
mcp_timeout_seconds=mcp_timeout_seconds,
457471
input_guardrails=input_guardrails,
458472
output_guardrails=output_guardrails,
473+
max_turns=max_turns,
459474
)

0 commit comments

Comments
 (0)