fix: optimize tool_choice reset logic and fix lint errors

- Refactor tool_choice reset to target only problematic edge cases
- Replace manual ModelSettings recreation with dataclasses.replace
- Fix line length and error handling lint issues in tests
This commit is contained in:
xianghuijin 2025-03-22 14:10:09 +08:00
parent d169d79288
commit bbcda753df
2 changed files with 137 additions and 147 deletions

View file

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import dataclasses
import inspect import inspect
from collections.abc import Awaitable from collections.abc import Awaitable
from dataclasses import dataclass from dataclasses import dataclass
@ -51,7 +52,7 @@ from .model_settings import ModelSettings
from .models.interface import ModelTracing from .models.interface import ModelTracing
from .run_context import RunContextWrapper, TContext from .run_context import RunContextWrapper, TContext
from .stream_events import RunItemStreamEvent, StreamEvent from .stream_events import RunItemStreamEvent, StreamEvent
from .tool import ComputerTool, FunctionTool, FunctionToolResult from .tool import ComputerTool, FunctionTool, FunctionToolResult, Tool
from .tracing import ( from .tracing import (
SpanError, SpanError,
Trace, Trace,
@ -208,34 +209,22 @@ class RunImpl:
new_step_items.extend(computer_results) new_step_items.extend(computer_results)
# Reset tool_choice to "auto" after tool execution to prevent infinite loops # Reset tool_choice to "auto" after tool execution to prevent infinite loops
if (processed_response.functions or processed_response.computer_actions): if processed_response.functions or processed_response.computer_actions:
# Reset agent's model_settings tools = agent.tools
if agent.model_settings.tool_choice == "required" or isinstance(agent.model_settings.tool_choice, str): # Only reset in the problematic scenarios where loops are likely unintentional
# Create a new model_settings to avoid modifying the original shared instance if cls._should_reset_tool_choice(agent.model_settings, tools):
agent.model_settings = ModelSettings( agent.model_settings = dataclasses.replace(
temperature=agent.model_settings.temperature, agent.model_settings,
top_p=agent.model_settings.top_p, tool_choice="auto"
frequency_penalty=agent.model_settings.frequency_penalty,
presence_penalty=agent.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=agent.model_settings.parallel_tool_calls,
truncation=agent.model_settings.truncation,
max_tokens=agent.model_settings.max_tokens,
) )
# Also reset run_config's model_settings if it exists if (
if run_config.model_settings and (run_config.model_settings.tool_choice == "required" or run_config.model_settings and
isinstance(run_config.model_settings.tool_choice, str)): cls._should_reset_tool_choice(run_config.model_settings, tools)
# Create a new model_settings for run_config ):
run_config.model_settings = ModelSettings( run_config.model_settings = dataclasses.replace(
temperature=run_config.model_settings.temperature, run_config.model_settings,
top_p=run_config.model_settings.top_p, tool_choice="auto"
frequency_penalty=run_config.model_settings.frequency_penalty,
presence_penalty=run_config.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=run_config.model_settings.parallel_tool_calls,
truncation=run_config.model_settings.truncation,
max_tokens=run_config.model_settings.max_tokens,
) )
# Second, check if there are any handoffs # Second, check if there are any handoffs
@ -328,6 +317,24 @@ class RunImpl:
next_step=NextStepRunAgain(), next_step=NextStepRunAgain(),
) )
@classmethod
def _should_reset_tool_choice(cls, model_settings: ModelSettings, tools: list[Tool]) -> bool:
if model_settings is None or model_settings.tool_choice is None:
return False
# for specific tool choices
if (
isinstance(model_settings.tool_choice, str) and
model_settings.tool_choice not in ["auto", "required", "none"]
):
return True
# for one tool and required tool choice
if model_settings.tool_choice == "required":
return len(tools) == 1
return False
@classmethod @classmethod
def process_model_response( def process_model_response(
cls, cls,

View file

@ -1,13 +1,15 @@
from unittest import mock
import asyncio import asyncio
import dataclasses
import json import json
from typing import List from unittest import mock
from agents import Agent, ModelSettings, RunConfig, function_tool, Runner
from agents.models.interface import ModelResponse
from agents.items import Usage
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from agents import Agent, ModelSettings, RunConfig, Runner, function_tool
from agents.items import Usage
from agents.models.interface import ModelResponse
from agents.tool import Tool
@function_tool @function_tool
def echo(text: str) -> str: def echo(text: str) -> str:
@ -15,6 +17,23 @@ def echo(text: str) -> str:
return text return text
def should_reset_tool_choice(model_settings: ModelSettings, tools: list[Tool]) -> bool:
if model_settings is None or model_settings.tool_choice is None:
return False
# for specific tool choices
if (
isinstance(model_settings.tool_choice, str) and
model_settings.tool_choice not in ["auto", "required", "none"]
):
return True
# for one tool and required tool choice
if model_settings.tool_choice == "required":
return len(tools) == 1
return False
# Mock model implementation that always calls tools when tool_choice is set # Mock model implementation that always calls tools when tool_choice is set
class MockModel: class MockModel:
def __init__(self, tool_call_counter): def __init__(self, tool_call_counter):
@ -60,7 +79,7 @@ class TestToolChoiceReset:
# Create an agent with tool_choice="required" # Create an agent with tool_choice="required"
agent = Agent( agent = Agent(
name="Test agent", name="Test agent",
tools=[echo], tools=[echo], # Only one tool
model_settings=ModelSettings(tool_choice="required"), model_settings=ModelSettings(tool_choice="required"),
) )
@ -77,31 +96,22 @@ class TestToolChoiceReset:
# Execute our code under test # Execute our code under test
if processed_response.functions: if processed_response.functions:
# Reset agent's model_settings # Apply the targeted reset logic
if agent.model_settings.tool_choice == "required" or isinstance(agent.model_settings.tool_choice, str): tools = agent.tools
agent.model_settings = ModelSettings( if should_reset_tool_choice(agent.model_settings, tools):
temperature=agent.model_settings.temperature, agent.model_settings = dataclasses.replace(
top_p=agent.model_settings.top_p, agent.model_settings,
frequency_penalty=agent.model_settings.frequency_penalty, tool_choice="auto" # Reset to auto
presence_penalty=agent.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=agent.model_settings.parallel_tool_calls,
truncation=agent.model_settings.truncation,
max_tokens=agent.model_settings.max_tokens,
) )
# Also reset run_config's model_settings if it exists # Also reset run_config's model_settings if it exists
if run_config.model_settings and (run_config.model_settings.tool_choice == "required" or if (
isinstance(run_config.model_settings.tool_choice, str)): run_config.model_settings and
run_config.model_settings = ModelSettings( should_reset_tool_choice(run_config.model_settings, tools)
temperature=run_config.model_settings.temperature, ):
top_p=run_config.model_settings.top_p, run_config.model_settings = dataclasses.replace(
frequency_penalty=run_config.model_settings.frequency_penalty, run_config.model_settings,
presence_penalty=run_config.model_settings.presence_penalty, tool_choice="auto" # Reset to auto
tool_choice="auto", # Reset to auto
parallel_tool_calls=run_config.model_settings.parallel_tool_calls,
truncation=run_config.model_settings.truncation,
max_tokens=run_config.model_settings.max_tokens,
) )
# Check that tool_choice was reset to "auto" # Check that tool_choice was reset to "auto"
@ -115,7 +125,7 @@ class TestToolChoiceReset:
instructions="You are a test agent", instructions="You are a test agent",
tools=[echo], tools=[echo],
model="gpt-4-0125-preview", model="gpt-4-0125-preview",
model_settings=ModelSettings(tool_choice="echo"), model_settings=ModelSettings(tool_choice="echo"), # Specific function name
) )
# Execute our code under test # Execute our code under test
@ -129,31 +139,22 @@ class TestToolChoiceReset:
# Execute our code under test # Execute our code under test
if processed_response.functions: if processed_response.functions:
# Reset agent's model_settings # Apply the targeted reset logic
if agent.model_settings.tool_choice == "required" or isinstance(agent.model_settings.tool_choice, str): tools = agent.tools
agent.model_settings = ModelSettings( if should_reset_tool_choice(agent.model_settings, tools):
temperature=agent.model_settings.temperature, agent.model_settings = dataclasses.replace(
top_p=agent.model_settings.top_p, agent.model_settings,
frequency_penalty=agent.model_settings.frequency_penalty, tool_choice="auto" # Reset to auto
presence_penalty=agent.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=agent.model_settings.parallel_tool_calls,
truncation=agent.model_settings.truncation,
max_tokens=agent.model_settings.max_tokens,
) )
# Also reset run_config's model_settings if it exists # Also reset run_config's model_settings if it exists
if run_config.model_settings and (run_config.model_settings.tool_choice == "required" or if (
isinstance(run_config.model_settings.tool_choice, str)): run_config.model_settings and
run_config.model_settings = ModelSettings( should_reset_tool_choice(run_config.model_settings, tools)
temperature=run_config.model_settings.temperature, ):
top_p=run_config.model_settings.top_p, run_config.model_settings = dataclasses.replace(
frequency_penalty=run_config.model_settings.frequency_penalty, run_config.model_settings,
presence_penalty=run_config.model_settings.presence_penalty, tool_choice="auto" # Reset to auto
tool_choice="auto", # Reset to auto
parallel_tool_calls=run_config.model_settings.parallel_tool_calls,
truncation=run_config.model_settings.truncation,
max_tokens=run_config.model_settings.max_tokens,
) )
# Check that tool_choice was reset to "auto" # Check that tool_choice was reset to "auto"
@ -179,31 +180,22 @@ class TestToolChoiceReset:
# Execute our code under test # Execute our code under test
if processed_response.functions: if processed_response.functions:
# Reset agent's model_settings # Apply the targeted reset logic
if agent.model_settings.tool_choice == "required" or isinstance(agent.model_settings.tool_choice, str): tools = agent.tools
agent.model_settings = ModelSettings( if should_reset_tool_choice(agent.model_settings, tools):
temperature=agent.model_settings.temperature, agent.model_settings = dataclasses.replace(
top_p=agent.model_settings.top_p, agent.model_settings,
frequency_penalty=agent.model_settings.frequency_penalty, tool_choice="auto" # Reset to auto
presence_penalty=agent.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=agent.model_settings.parallel_tool_calls,
truncation=agent.model_settings.truncation,
max_tokens=agent.model_settings.max_tokens,
) )
# Also reset run_config's model_settings if it exists # Also reset run_config's model_settings if it exists
if run_config.model_settings and (run_config.model_settings.tool_choice == "required" or if (
isinstance(run_config.model_settings.tool_choice, str)): run_config.model_settings and
run_config.model_settings = ModelSettings( should_reset_tool_choice(run_config.model_settings, tools)
temperature=run_config.model_settings.temperature, ):
top_p=run_config.model_settings.top_p, run_config.model_settings = dataclasses.replace(
frequency_penalty=run_config.model_settings.frequency_penalty, run_config.model_settings,
presence_penalty=run_config.model_settings.presence_penalty, tool_choice="auto" # Reset to auto
tool_choice="auto", # Reset to auto
parallel_tool_calls=run_config.model_settings.parallel_tool_calls,
truncation=run_config.model_settings.truncation,
max_tokens=run_config.model_settings.max_tokens,
) )
# Check that tool_choice remains "auto" # Check that tool_choice remains "auto"
@ -214,7 +206,7 @@ class TestToolChoiceReset:
# Create an agent with default model_settings # Create an agent with default model_settings
agent = Agent( agent = Agent(
name="Test agent", name="Test agent",
tools=[echo], tools=[echo], # Only one tool
model_settings=ModelSettings(tool_choice=None), model_settings=ModelSettings(tool_choice=None),
) )
@ -229,31 +221,22 @@ class TestToolChoiceReset:
# Execute our code under test # Execute our code under test
if processed_response.functions: if processed_response.functions:
# Reset agent's model_settings # Apply the targeted reset logic
if agent.model_settings.tool_choice == "required" or isinstance(agent.model_settings.tool_choice, str): tools = agent.tools
agent.model_settings = ModelSettings( if should_reset_tool_choice(agent.model_settings, tools):
temperature=agent.model_settings.temperature, agent.model_settings = dataclasses.replace(
top_p=agent.model_settings.top_p, agent.model_settings,
frequency_penalty=agent.model_settings.frequency_penalty, tool_choice="auto" # Reset to auto
presence_penalty=agent.model_settings.presence_penalty,
tool_choice="auto", # Reset to auto
parallel_tool_calls=agent.model_settings.parallel_tool_calls,
truncation=agent.model_settings.truncation,
max_tokens=agent.model_settings.max_tokens,
) )
# Also reset run_config's model_settings if it exists # Also reset run_config's model_settings if it exists
if run_config.model_settings and (run_config.model_settings.tool_choice == "required" or if (
isinstance(run_config.model_settings.tool_choice, str)): run_config.model_settings and
run_config.model_settings = ModelSettings( should_reset_tool_choice(run_config.model_settings, tools)
temperature=run_config.model_settings.temperature, ):
top_p=run_config.model_settings.top_p, run_config.model_settings = dataclasses.replace(
frequency_penalty=run_config.model_settings.frequency_penalty, run_config.model_settings,
presence_penalty=run_config.model_settings.presence_penalty, tool_choice="auto" # Reset to auto
tool_choice="auto", # Reset to auto
parallel_tool_calls=run_config.model_settings.parallel_tool_calls,
truncation=run_config.model_settings.truncation,
max_tokens=run_config.model_settings.max_tokens,
) )
# Check that run_config's tool_choice was reset to "auto" # Check that run_config's tool_choice was reset to "auto"
@ -298,6 +281,6 @@ class TestToolChoiceReset:
assert tool_call_counter["count"] < 5 assert tool_call_counter["count"] < 5
assert not tool_call_counter["potential_infinite_loop"] assert not tool_call_counter["potential_infinite_loop"]
except asyncio.TimeoutError: except asyncio.TimeoutError as err:
# If we hit a timeout, the test failed - we likely have an infinite loop # If we hit a timeout, the test failed - we likely have an infinite loop
assert False, "Timeout occurred, potential infinite loop detected" raise AssertionError("Timeout occurred, potential infinite loop detected") from err