532 lines
20 KiB
Python
532 lines
20 KiB
Python
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from actions.utils.openai_utils import (
|
|
MODEL_COSTS,
|
|
OPENAI_MODEL_DEFAULT,
|
|
PR_REVIEW_MODEL_DEFAULT,
|
|
_is_anthropic_model,
|
|
_openai_usage_cost,
|
|
_response_tool_calls,
|
|
get_agent_response,
|
|
get_response,
|
|
get_review_model,
|
|
remove_outer_codeblocks,
|
|
)
|
|
|
|
|
|
def test_default_models():
|
|
"""Test canonical default models are priced so max_cost budgets stay enforceable."""
|
|
assert OPENAI_MODEL_DEFAULT == "gpt-5.6-luna"
|
|
assert PR_REVIEW_MODEL_DEFAULT == "gpt-5.6-terra"
|
|
assert OPENAI_MODEL_DEFAULT in MODEL_COSTS # unpriced models disable max_cost budgets
|
|
assert PR_REVIEW_MODEL_DEFAULT in MODEL_COSTS
|
|
assert MODEL_COSTS["gpt-5.6-sol"] == (5.00, 30.00)
|
|
assert MODEL_COSTS["gpt-5.6-terra"] == (2.50, 15.00)
|
|
assert MODEL_COSTS["gpt-5.6-luna"] == (1.00, 6.00)
|
|
|
|
|
|
def test_gpt_56_cost_includes_cache_write_and_long_context_rates():
|
|
"""GPT-5.6 cache writes bill at 125%, with long requests at 2x input and 1.5x output."""
|
|
usage = {
|
|
"input_tokens": 1000,
|
|
"input_tokens_details": {"cached_tokens": 200, "cache_write_tokens": 300},
|
|
"output_tokens": 100,
|
|
}
|
|
expected = ((1000 - 200 * 0.9 + 300 * 0.25) * 5.00 + 100 * 30.00) / 1e6
|
|
assert _openai_usage_cost(usage, "gpt-5.6-sol") == expected
|
|
usage["input_tokens"] = 272001
|
|
expected = ((272001 - 200 * 0.9 + 300 * 0.25) * 5.00 * 2 + 100 * 30.00 * 1.5) / 1e6
|
|
assert _openai_usage_cost(usage, "gpt-5.6-sol") == expected
|
|
turns = [{"input_tokens": 150000, "output_tokens": 0}] * 2
|
|
assert sum(_openai_usage_cost(turn, "gpt-5.6-luna") for turn in turns) == 0.3
|
|
assert _openai_usage_cost({"input_tokens": 300000, "output_tokens": 0}, "gpt-5.6-luna") == 0.6
|
|
old_model_expected = ((1000 - 200 * 0.9) * 5.00 + 100 * 30.00) / 1e6
|
|
assert _openai_usage_cost({**usage, "input_tokens": 1000}, "gpt-5.5") == old_model_expected
|
|
|
|
|
|
def test_is_anthropic_model():
|
|
"""Test model provider detection."""
|
|
assert _is_anthropic_model("claude-sonnet-4-6") is True
|
|
assert _is_anthropic_model("claude-haiku-4-5-20251001") is True
|
|
assert _is_anthropic_model("claude-opus-4-7") is True
|
|
assert _is_anthropic_model("gpt-5.6-terra") is False
|
|
assert _is_anthropic_model("gpt-5-mini-2025-08-07") is False
|
|
|
|
|
|
def test_response_tool_calls():
|
|
"""Test Responses API tool-call item naming, including hosted tools without a name field."""
|
|
output_items = [
|
|
{"type": "function_call", "name": "lookup_value"},
|
|
{"type": "web_search_call"},
|
|
{"type": "message"},
|
|
{"type": "function_call_output"},
|
|
]
|
|
assert _response_tool_calls(output_items) == ["lookup_value", "web_search"]
|
|
|
|
|
|
def test_remove_outer_codeblocks():
|
|
"""Test removing outer code block markers from strings."""
|
|
# Test with python code block
|
|
input_str = "```python\ndef test():\n return True\n```"
|
|
expected = "def test():\n return True"
|
|
assert remove_outer_codeblocks(input_str) == expected
|
|
|
|
# Test with no language specified
|
|
input_str = "```\ndef test():\n return True\n```"
|
|
expected = "def test():\n return True"
|
|
assert remove_outer_codeblocks(input_str) == expected
|
|
|
|
# Test with no code blocks
|
|
input_str = "def test():\n return True"
|
|
assert remove_outer_codeblocks(input_str) == input_str
|
|
|
|
|
|
def test_get_review_model_override():
|
|
"""Test review model override logic."""
|
|
with patch("actions.utils.openai_utils.REVIEW_MODEL", "claude-opus-4-7"), patch(
|
|
"actions.utils.openai_utils.MODEL", "gpt-5.6-terra"
|
|
):
|
|
assert get_review_model() == "claude-opus-4-7"
|
|
|
|
|
|
def test_get_review_model_fallback():
|
|
"""Test review model fallback to default model."""
|
|
with patch("actions.utils.openai_utils.REVIEW_MODEL", None):
|
|
assert get_review_model() == PR_REVIEW_MODEL_DEFAULT
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_response(mock_post):
|
|
"""Test OpenAI Responses API completion function with mocked response."""
|
|
# Setup mock response with Responses API structure
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.elapsed.total_seconds.return_value = 1.5
|
|
mock_response.json.return_value = {
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"content": [{"type": "output_text", "text": "Test response from OpenAI"}],
|
|
}
|
|
]
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
# Test with basic messages
|
|
messages = [{"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "Hello"}]
|
|
|
|
# Use a context manager for the environment variable
|
|
with patch.dict("os.environ", {"OPENAI_API_KEY": "test-key"}, clear=False), patch(
|
|
"actions.utils.openai_utils.OPENAI_API_KEY", "test-key"
|
|
):
|
|
result = get_response(messages, check_links=False)
|
|
|
|
assert result == "Test response from OpenAI"
|
|
mock_post.assert_called_once()
|
|
|
|
|
|
@patch("time.sleep")
|
|
@patch("requests.post")
|
|
def test_get_response_read_timeout_propagates(mock_post, mock_sleep):
|
|
"""Test a read timeout is NOT retried: the request may have completed server-side and re-POSTing double-bills."""
|
|
mock_post.side_effect = requests.exceptions.ReadTimeout()
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"):
|
|
try:
|
|
get_response([{"role": "user", "content": "Hello"}], check_links=False, retries=2)
|
|
raise AssertionError("ReadTimeout should propagate")
|
|
except requests.exceptions.ReadTimeout:
|
|
pass
|
|
assert mock_post.call_count == 1 # no re-POST of a possibly-billed request
|
|
|
|
|
|
@patch("time.sleep")
|
|
@patch("requests.post")
|
|
def test_get_response_retries_rate_limits(mock_post, mock_sleep):
|
|
"""Test a 429 response is retried with a longer backoff than server errors."""
|
|
limited = MagicMock()
|
|
limited.status_code = 429
|
|
limited.elapsed.total_seconds.return_value = 0.1
|
|
ok = MagicMock()
|
|
ok.status_code = 200
|
|
ok.elapsed.total_seconds.return_value = 1.0
|
|
ok.json.return_value = {"output": [{"type": "message", "content": [{"type": "output_text", "text": "recovered"}]}]}
|
|
mock_post.side_effect = [limited, ok]
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"):
|
|
result = get_response([{"role": "user", "content": "Hello"}], check_links=False, retries=1)
|
|
|
|
assert result == "recovered"
|
|
assert mock_post.call_count == 2
|
|
mock_sleep.assert_called_once_with(10) # 10 * 2**0, not the 2**0 server-error backoff
|
|
|
|
|
|
@patch("requests.post")
|
|
@patch("actions.utils.openai_utils.check_links_in_string")
|
|
def test_get_response_with_link_check(mock_check_links, mock_post):
|
|
"""Test get_response with link checking."""
|
|
# Setup mocks with Responses API structure
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.elapsed.total_seconds.return_value = 2.0
|
|
mock_response.json.return_value = {
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"content": [{"type": "output_text", "text": "Response with https://example.com link"}],
|
|
}
|
|
]
|
|
}
|
|
mock_post.return_value = mock_response
|
|
mock_check_links.return_value = (True, [])
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
|
|
# Use a context manager for the environment variable
|
|
with patch.dict("os.environ", {"OPENAI_API_KEY": "test-key"}, clear=False), patch(
|
|
"actions.utils.openai_utils.OPENAI_API_KEY", "test-key"
|
|
):
|
|
result = get_response(messages)
|
|
|
|
assert result == "Response with https://example.com link"
|
|
mock_check_links.assert_called_once()
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_agent_response_calls_function_tools(mock_post):
|
|
"""Test iterative Responses API agent calls local tools and returns structured output."""
|
|
first_response = MagicMock()
|
|
first_response.status_code = 200
|
|
first_response.elapsed.total_seconds.return_value = 1.0
|
|
first_response.json.return_value = {
|
|
"id": "resp_first",
|
|
"output": [
|
|
{
|
|
"type": "function_call",
|
|
"call_id": "call_123",
|
|
"name": "lookup_value",
|
|
"arguments": '{"value": "abc"}',
|
|
},
|
|
{"type": "web_search_call"},
|
|
],
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"input_tokens_details": {"cached_tokens": 4, "cache_write_tokens": 3},
|
|
"output_tokens": 5,
|
|
},
|
|
}
|
|
second_response = MagicMock()
|
|
second_response.status_code = 200
|
|
second_response.elapsed.total_seconds.return_value = 1.0
|
|
second_response.json.return_value = {
|
|
"id": "resp_second",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"content": [{"type": "output_text", "text": '{"comments": [], "summary": "done"}'}],
|
|
}
|
|
],
|
|
"usage": {
|
|
"input_tokens": 20,
|
|
"input_tokens_details": {"cached_tokens": 8, "cache_write_tokens": 6},
|
|
"output_tokens": 7,
|
|
},
|
|
}
|
|
mock_post.side_effect = [first_response, second_response]
|
|
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {"comments": {"type": "array"}, "summary": {"type": "string"}},
|
|
"required": ["comments", "summary"],
|
|
"additionalProperties": False,
|
|
}
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"name": "lookup_value",
|
|
"description": "Lookup a value.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"value": {"type": "string"}},
|
|
"required": ["value"],
|
|
"additionalProperties": False,
|
|
},
|
|
"strict": True,
|
|
}
|
|
]
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), patch("builtins.print") as mock_print:
|
|
result = get_agent_response(
|
|
[{"role": "user", "content": "review"}],
|
|
tools=tools,
|
|
tool_handlers={"lookup_value": lambda value: {"found": value}},
|
|
text_format={"format": {"type": "json_schema", "name": "review", "strict": True, "schema": schema}},
|
|
retries=0,
|
|
)
|
|
|
|
assert result == {"comments": [], "summary": "done"}
|
|
assert mock_post.call_count == 2
|
|
first_payload = mock_post.call_args_list[0].kwargs["json"]
|
|
assert first_payload["store"] is True
|
|
assert first_payload["service_tier"] == "default"
|
|
assert first_payload["reasoning"] == {"effort": "low"}
|
|
assert first_payload["context_management"] == [{"type": "compaction", "compact_threshold": 200_000}]
|
|
assert first_payload["prompt_cache_key"].startswith("agent-run:")
|
|
assert "include" not in first_payload
|
|
assert "previous_response_id" not in first_payload
|
|
assert first_payload["input"] == [{"role": "user", "content": "review"}]
|
|
second_payload = mock_post.call_args_list[1].kwargs["json"]
|
|
assert second_payload["previous_response_id"] == "resp_first"
|
|
assert second_payload["prompt_cache_key"] == first_payload["prompt_cache_key"]
|
|
second_input = second_payload["input"]
|
|
assert second_input == [
|
|
{
|
|
"type": "function_call_output",
|
|
"call_id": "call_123",
|
|
"output": '{"found": "abc"}',
|
|
}
|
|
]
|
|
printed = "\n".join(str(c.args[0]) for c in mock_print.call_args_list if c.args)
|
|
assert "turn 1/6, 2 tools (lookup_value, web_search)" in printed
|
|
assert "turn 2/6, 0 tools" in printed
|
|
assert "30→12 tokens (40% cached), $0.01" in printed
|
|
assert "agent total, 2 turns, 2 tools (lookup_value, web_search)" in printed
|
|
assert "Agent tool turn" not in printed # tool names live in the per-turn usage line now
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_agent_response_summarizes_after_max_turns(mock_post):
|
|
"""Test max-turn exhaustion makes a final no-tool synthesis call."""
|
|
tool_response = MagicMock()
|
|
tool_response.status_code = 200
|
|
tool_response.elapsed.total_seconds.return_value = 1.0
|
|
tool_response.json.return_value = {
|
|
"id": "resp_tool",
|
|
"output": [
|
|
{
|
|
"type": "function_call",
|
|
"call_id": "call_123",
|
|
"name": "lookup_value",
|
|
"arguments": '{"value": "abc"}',
|
|
}
|
|
],
|
|
"usage": {"input_tokens": 10, "output_tokens": 5},
|
|
}
|
|
final_response = MagicMock()
|
|
final_response.status_code = 200
|
|
final_response.elapsed.total_seconds.return_value = 1.0
|
|
final_response.json.return_value = {
|
|
"id": "resp_final",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"content": [{"type": "output_text", "text": '{"comments": [], "summary": "synthesized"}'}],
|
|
}
|
|
],
|
|
"usage": {"input_tokens": 20, "output_tokens": 7},
|
|
}
|
|
mock_post.side_effect = [tool_response, requests.exceptions.ConnectTimeout(), final_response]
|
|
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {"comments": {"type": "array"}, "summary": {"type": "string"}},
|
|
"required": ["comments", "summary"],
|
|
"additionalProperties": False,
|
|
}
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"name": "lookup_value",
|
|
"description": "Lookup a value.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"value": {"type": "string"}},
|
|
"required": ["value"],
|
|
"additionalProperties": False,
|
|
},
|
|
"strict": True,
|
|
}
|
|
]
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), patch(
|
|
"actions.utils.openai_utils.time.sleep"
|
|
) as mock_sleep:
|
|
result = get_agent_response(
|
|
[{"role": "user", "content": "review"}],
|
|
tools=tools,
|
|
tool_handlers={"lookup_value": lambda value: f"raw tool output for {value}"},
|
|
text_format={"format": {"type": "json_schema", "name": "review", "strict": True, "schema": schema}},
|
|
max_turns=1,
|
|
retries=0,
|
|
)
|
|
|
|
assert result == {"comments": [], "summary": "synthesized"}
|
|
assert mock_post.call_count == 3
|
|
mock_sleep.assert_called_once_with(1)
|
|
final_payload = mock_post.call_args_list[2].kwargs["json"]
|
|
assert final_payload["tools"] == tools
|
|
assert final_payload["tool_choice"] == "none"
|
|
assert final_payload["previous_response_id"] == "resp_tool"
|
|
assert final_payload["input"][0] == {
|
|
"type": "function_call_output",
|
|
"call_id": "call_123",
|
|
"output": "raw tool output for abc",
|
|
}
|
|
assert "Synthesize the gathered tool results" in final_payload["input"][-1]["content"]
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_response_anthropic(mock_post):
|
|
"""Test Anthropic Messages API completion function with mocked response."""
|
|
# Setup mock response with Anthropic Messages API structure
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.elapsed.total_seconds.return_value = 1.5
|
|
mock_response.json.return_value = {
|
|
"content": [{"type": "text", "text": "Test response from Claude"}],
|
|
"usage": {
|
|
"input_tokens": 50,
|
|
"cache_read_input_tokens": 900,
|
|
"cache_creation_input_tokens": 50,
|
|
"output_tokens": 20,
|
|
},
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
messages = [{"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "Hello"}]
|
|
|
|
with patch.dict("os.environ", {"ANTHROPIC_API_KEY": "test-key"}, clear=False), patch(
|
|
"actions.utils.openai_utils.ANTHROPIC_API_KEY", "test-key"
|
|
), patch("builtins.print") as mock_print:
|
|
result = get_response(messages, check_links=False, model="claude-sonnet-4-6", background=True)
|
|
|
|
assert result == "Test response from Claude"
|
|
printed = "\n".join(str(c.args[0]) for c in mock_print.call_args_list if c.args)
|
|
# Cache reads/writes fold into input and reads count as cached, matching ultralytics/assistant normalization
|
|
assert "1000→20 tokens (90% cached), $0.00087" in printed
|
|
mock_post.assert_called_once()
|
|
# Verify Anthropic endpoint was called
|
|
call_args = mock_post.call_args
|
|
assert call_args[0][0] == "https://api.anthropic.com/v1/messages"
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_agent_response_stops_at_cost_budget(mock_post):
|
|
"""Test the cost budget aborts rather than synthesizing from incomplete tool evidence."""
|
|
tool_response = MagicMock()
|
|
tool_response.status_code = 200
|
|
tool_response.elapsed.total_seconds.return_value = 1.0
|
|
tool_response.json.return_value = {
|
|
"id": "resp_tool",
|
|
"output": [
|
|
{
|
|
"type": "function_call",
|
|
"call_id": "call_123",
|
|
"name": "lookup_value",
|
|
"arguments": '{"value": "abc"}',
|
|
}
|
|
],
|
|
"usage": {"input_tokens": 1_000_000, "output_tokens": 0}, # $5.00 for gpt-5.6-sol, over any small budget
|
|
}
|
|
mock_post.return_value = tool_response
|
|
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {"comments": {"type": "array"}, "summary": {"type": "string"}},
|
|
"required": ["comments", "summary"],
|
|
"additionalProperties": False,
|
|
}
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"name": "lookup_value",
|
|
"description": "Lookup a value.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"value": {"type": "string"}},
|
|
"required": ["value"],
|
|
"additionalProperties": False,
|
|
},
|
|
"strict": True,
|
|
}
|
|
]
|
|
|
|
def forbidden_handler(value):
|
|
raise AssertionError("tool handlers must not run once the cost budget is reached")
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), pytest.raises(
|
|
RuntimeError, match=r"cost budget.*reached"
|
|
):
|
|
get_agent_response(
|
|
[{"role": "user", "content": "review"}],
|
|
tools=tools,
|
|
tool_handlers={"lookup_value": forbidden_handler},
|
|
text_format={"format": {"type": "json_schema", "name": "review", "strict": True, "schema": schema}},
|
|
model="gpt-5.6-sol",
|
|
max_turns=8,
|
|
max_cost=1.00,
|
|
retries=0,
|
|
)
|
|
|
|
mock_post.assert_called_once()
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_agent_response_rejects_incomplete_response(mock_post):
|
|
"""Test terminal incomplete responses cannot become review evidence."""
|
|
response = MagicMock(status_code=200)
|
|
response.elapsed.total_seconds.return_value = 1.0
|
|
response.json.return_value = {"id": "resp_incomplete", "status": "incomplete", "incomplete_details": "limit"}
|
|
mock_post.return_value = response
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), pytest.raises(
|
|
RuntimeError, match="ended with limit"
|
|
):
|
|
get_agent_response([{"role": "user", "content": "review"}], tools=[], tool_handlers={}, retries=0)
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_response_rejects_incomplete_response(mock_post):
|
|
"""Test synchronous Responses completions reject terminal incomplete output."""
|
|
response = MagicMock(status_code=200)
|
|
response.elapsed.total_seconds.return_value = 1.0
|
|
response.json.return_value = {"id": "resp_incomplete", "status": "incomplete", "incomplete_details": "limit"}
|
|
mock_post.return_value = response
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), pytest.raises(
|
|
RuntimeError, match="ended with limit"
|
|
):
|
|
get_response([{"role": "user", "content": "summary"}], check_links=False, retries=0)
|
|
|
|
|
|
@patch("requests.post")
|
|
def test_get_agent_response_rejects_failed_tool(mock_post):
|
|
"""Test failed evidence tools abort the agent run."""
|
|
response = MagicMock(status_code=200)
|
|
response.elapsed.total_seconds.return_value = 1.0
|
|
response.json.return_value = {
|
|
"id": "resp_tool",
|
|
"status": "completed",
|
|
"output": [{"type": "function_call", "call_id": "call_123", "name": "read_file", "arguments": "{}"}],
|
|
"usage": {"input_tokens": 10, "output_tokens": 5},
|
|
}
|
|
mock_post.return_value = response
|
|
|
|
def failed_read():
|
|
raise RuntimeError("read failed")
|
|
|
|
with patch("actions.utils.openai_utils.OPENAI_API_KEY", "test-key"), pytest.raises(
|
|
RuntimeError, match="read failed"
|
|
):
|
|
get_agent_response(
|
|
[{"role": "user", "content": "review"}],
|
|
tools=[],
|
|
tool_handlers={"read_file": failed_read},
|
|
retries=0,
|
|
)
|