diff --git a/src/minisweagent/exceptions.py b/src/minisweagent/exceptions.py index b11e404c9..14801add0 100644 --- a/src/minisweagent/exceptions.py +++ b/src/minisweagent/exceptions.py @@ -18,6 +18,10 @@ class TimeExceeded(LimitsExceeded): """Raised when the agent has exceeded its wall-clock time limit.""" +class ProviderTimeout(InterruptAgentFlow): + """Raised when the model provider does not return within its timeout.""" + + class UserInterruption(InterruptAgentFlow): """Raised when the user interrupts the agent.""" diff --git a/src/minisweagent/models/litellm_model.py b/src/minisweagent/models/litellm_model.py index 32bdbc791..bec578cc0 100644 --- a/src/minisweagent/models/litellm_model.py +++ b/src/minisweagent/models/litellm_model.py @@ -7,9 +7,13 @@ from typing import Any, Literal import litellm +import openai +import requests +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse from pydantic import BaseModel -from minisweagent.exceptions import FormatError +from minisweagent.exceptions import FormatError, ProviderTimeout from minisweagent.models import GLOBAL_MODEL_STATS from minisweagent.models.utils.actions_toolcall import ( BASH_TOOL, @@ -24,11 +28,27 @@ logger = logging.getLogger("litellm_model") +def _is_timeout_exception(e: Exception) -> bool: + if isinstance( + e, + ( + TimeoutError, + litellm.exceptions.Timeout, + openai.APITimeoutError, + requests.exceptions.Timeout, + ), + ): + return True + return any("timeout" in cls.__name__.lower() for cls in type(e).mro()) or "timed out" in str(e).lower() + + class LitellmModelConfig(BaseModel): model_name: str """Model name. Highly recommended to include the provider in the model name, e.g., `anthropic/claude-sonnet-4-5-20250929`.""" model_kwargs: dict[str, Any] = {} """Additional arguments passed to the API.""" + provider_timeout: float | None = 5.0 + """Default read timeout in seconds for streaming provider requests. Set to null to use the provider or LiteLLM default.""" litellm_model_registry: Path | str | None = os.getenv("LITELLM_MODEL_REGISTRY_PATH") """Model registry for cost tracking and model metadata. See the local model guide (https://mini-swe-agent.com/latest/models/local_models/) for more details.""" set_cache_control: Literal["default_end"] | None = None @@ -53,6 +73,7 @@ class LitellmModel: litellm.exceptions.PermissionDeniedError, litellm.exceptions.ContextWindowExceededError, litellm.exceptions.AuthenticationError, + ProviderTimeout, KeyboardInterrupt, ] @@ -63,15 +84,97 @@ def __init__(self, *, config_class: Callable = LitellmModelConfig, **kwargs): def _query(self, messages: list[dict[str, str]], **kwargs): try: - return litellm.completion( + response = litellm.completion( model=self.config.model_name, messages=messages, tools=[BASH_TOOL], - **(self.config.model_kwargs | kwargs), + **self._model_kwargs(**kwargs), ) + return self._consume_stream_response(response) if isinstance(response, CustomStreamWrapper) else response except litellm.exceptions.AuthenticationError as e: e.message += " You can permanently set your API key with `mini-extra config set KEY VALUE`." raise e + except Exception as e: + self._raise_provider_timeout(e) + + def _model_kwargs(self, **kwargs) -> dict[str, Any]: + model_kwargs = self.config.model_kwargs | kwargs + if ( + self.config.provider_timeout is not None + and model_kwargs.get("stream") is True + and not {"timeout", "request_timeout"} & model_kwargs.keys() + ): + model_kwargs["timeout"] = openai.Timeout(10.0, read=self.config.provider_timeout) + return model_kwargs + + def _raise_provider_timeout(self, e: Exception) -> None: + if not _is_timeout_exception(e): + raise e + raise ProviderTimeout( + self.format_message( + role="exit", + content="ProviderTimeout: model provider did not respond within the configured timeout.", + extra={"exit_status": "ProviderTimeout", "submission": "", "exception_str": str(e)}, + ) + ) from e + + def _consume_stream_response(self, stream: CustomStreamWrapper) -> ModelResponse: + content: list[str] = [] + tool_calls_by_index: dict[int, dict[str, str]] = {} + finish_reason = None + response_id = None + model = self.config.model_name + created = None + usage = None + for chunk in stream: + response_id = getattr(chunk, "id", response_id) + model = getattr(chunk, "model", model) + created = getattr(chunk, "created", created) + usage = getattr(chunk, "usage", None) or usage + if not chunk.choices: + continue + choice = chunk.choices[0] + finish_reason = choice.finish_reason or finish_reason + delta = choice.delta + if delta.content: + content.append(delta.content) + for tool_call in delta.tool_calls or []: + index = int(tool_call.index or 0) + state = tool_calls_by_index.setdefault( + index, {"id": "", "type": "function", "name": "", "arguments": ""} + ) + if tool_call.id: + state["id"] = tool_call.id + if tool_call.type: + state["type"] = tool_call.type + if tool_call.function: + if tool_call.function.name: + state["name"] = tool_call.function.name + if tool_call.function.arguments: + state["arguments"] += tool_call.function.arguments + tool_calls = [ + ChatCompletionMessageToolCall( + id=state["id"], + type=state["type"], + function=Function(name=state["name"], arguments=state["arguments"]), + ) + for _, state in sorted(tool_calls_by_index.items()) + ] + response = ModelResponse( + id=response_id, + created=created, + model=model, + usage=usage, + choices=[ + Choices( + finish_reason=finish_reason, + message=Message(content="".join(content) or None, tool_calls=tool_calls or None), + ) + ], + ) + if usage is None: + response._mswea_streamed_without_usage = True + return response def _prepare_messages_for_api(self, messages: list[dict]) -> list[dict]: prepared = [{k: v for k, v in msg.items() if k != "extra"} for msg in messages] @@ -105,6 +208,8 @@ def query(self, messages: list[dict[str, str]], **kwargs) -> dict: return message def _calculate_cost(self, response) -> dict[str, float]: + if getattr(response, "_mswea_streamed_without_usage", False) is True: + return {"cost": 0.0} try: cost = litellm.cost_calculator.completion_cost(response, model=self.config.model_name) if cost <= 0.0: diff --git a/src/minisweagent/models/litellm_response_model.py b/src/minisweagent/models/litellm_response_model.py index e3d959361..b185ddac1 100644 --- a/src/minisweagent/models/litellm_response_model.py +++ b/src/minisweagent/models/litellm_response_model.py @@ -43,11 +43,13 @@ def _query(self, messages: list[dict[str, str]], **kwargs): model=self.config.model_name, input=messages, tools=[BASH_TOOL_RESPONSE_API], - **(self.config.model_kwargs | kwargs), + **self._model_kwargs(**kwargs), ) except litellm.exceptions.AuthenticationError as e: e.message += " You can permanently set your API key with `mini-extra config set KEY VALUE`." raise e + except Exception as e: + self._raise_provider_timeout(e) def query(self, messages: list[dict[str, str]], **kwargs) -> dict: for attempt in retry(logger=logger, abort_exceptions=self.abort_exceptions): diff --git a/src/minisweagent/models/litellm_textbased_model.py b/src/minisweagent/models/litellm_textbased_model.py index db7e042a9..6e9b46f90 100644 --- a/src/minisweagent/models/litellm_textbased_model.py +++ b/src/minisweagent/models/litellm_textbased_model.py @@ -19,12 +19,12 @@ def __init__(self, **kwargs): def _query(self, messages: list[dict[str, str]], **kwargs): try: - return litellm.completion( - model=self.config.model_name, messages=messages, **(self.config.model_kwargs | kwargs) - ) + return litellm.completion(model=self.config.model_name, messages=messages, **self._model_kwargs(**kwargs)) except litellm.exceptions.AuthenticationError as e: e.message += " You can permanently set your API key with `mini-extra config set KEY VALUE`." raise e + except Exception as e: + self._raise_provider_timeout(e) def _parse_actions(self, response: dict) -> list[dict]: """Parse actions from the model response. Raises FormatError if not exactly one action.""" diff --git a/tests/agents/test_default.py b/tests/agents/test_default.py index 48aa5b0be..ff670377a 100644 --- a/tests/agents/test_default.py +++ b/tests/agents/test_default.py @@ -5,7 +5,7 @@ from minisweagent.agents.default import DefaultAgent from minisweagent.environments.local import LocalEnvironment -from minisweagent.exceptions import FormatError +from minisweagent.exceptions import FormatError, ProviderTimeout from minisweagent.models.test_models import ( DeterministicModel, DeterministicResponseAPIToolcallModel, @@ -477,6 +477,17 @@ def query(self, messages, **kwargs): return output +class _ProviderTimeoutModel(DeterministicToolcallModel): + def query(self, messages, **kwargs): + raise ProviderTimeout( + { + "role": "exit", + "content": "ProviderTimeout", + "extra": {"exit_status": "ProviderTimeout", "submission": ""}, + } + ) + + def test_repeated_format_errors_terminate_cleanly(toolcall_config): """With max_consecutive_format_errors set, a run that keeps producing no-tool-call / truncation turns stops cleanly with exit_status=RepeatedFormatError instead of looping until the budget is @@ -492,6 +503,13 @@ def test_repeated_format_errors_terminate_cleanly(toolcall_config): assert agent.n_calls == 2 # stopped at the 2nd consecutive error, didn't burn all 5 +def test_provider_timeout_terminates_cleanly(toolcall_config): + agent = DefaultAgent(model=_ProviderTimeoutModel(outputs=[]), env=LocalEnvironment(), **toolcall_config) + info = agent.run("Test provider timeout") + assert info["exit_status"] == "ProviderTimeout" + assert agent.n_calls == 1 + + def test_format_error_counter_resets_on_success(toolcall_config): """A successful tool call between format errors resets the consecutive counter, so isolated errors don't accumulate to the termination threshold.""" diff --git a/tests/models/test_litellm_model.py b/tests/models/test_litellm_model.py index d4426a780..bc83ca915 100644 --- a/tests/models/test_litellm_model.py +++ b/tests/models/test_litellm_model.py @@ -1,8 +1,14 @@ +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import MagicMock, patch +import litellm +import openai import pytest -from minisweagent.exceptions import FormatError +from minisweagent.exceptions import FormatError, ProviderTimeout from minisweagent.models.litellm_model import LitellmModel, LitellmModelConfig from minisweagent.models.utils.actions_toolcall import BASH_TOOL @@ -11,6 +17,9 @@ class TestLitellmModelConfig: def test_default_format_error_template(self): assert LitellmModelConfig(model_name="test").format_error_template == "{{ error }}" + def test_default_provider_timeout(self): + assert LitellmModelConfig(model_name="test").provider_timeout == 5.0 + def _mock_litellm_response(tool_calls): mock_response = MagicMock() @@ -37,6 +46,256 @@ def test_query_includes_bash_tool(self, mock_cost, mock_completion): mock_completion.assert_called_once() assert mock_completion.call_args.kwargs["tools"] == [BASH_TOOL] + assert "timeout" not in mock_completion.call_args.kwargs + + @patch("minisweagent.models.litellm_model.litellm.completion") + @patch("minisweagent.models.litellm_model.litellm.cost_calculator.completion_cost") + def test_query_includes_stream_provider_timeout(self, mock_cost, mock_completion): + tool_call = MagicMock() + tool_call.function.name = "bash" + tool_call.function.arguments = '{"command": "echo test"}' + tool_call.id = "call_1" + mock_completion.return_value = _mock_litellm_response([tool_call]) + mock_cost.return_value = 0.001 + + model = LitellmModel(model_name="gpt-4", model_kwargs={"stream": True}) + model.query([{"role": "user", "content": "test"}]) + + timeout = mock_completion.call_args.kwargs["timeout"] + assert isinstance(timeout, openai.Timeout) + assert timeout.connect == 10.0 + assert timeout.read == 5.0 + + @patch("minisweagent.models.litellm_model.litellm.completion") + @patch("minisweagent.models.litellm_model.litellm.cost_calculator.completion_cost") + def test_query_preserves_explicit_timeout(self, mock_cost, mock_completion): + tool_call = MagicMock() + tool_call.function.name = "bash" + tool_call.function.arguments = '{"command": "echo test"}' + tool_call.id = "call_1" + mock_completion.return_value = _mock_litellm_response([tool_call]) + mock_cost.return_value = 0.001 + + model = LitellmModel(model_name="gpt-4", model_kwargs={"timeout": 30}) + model.query([{"role": "user", "content": "test"}]) + + assert mock_completion.call_args.kwargs["timeout"] == 30 + + @patch("minisweagent.models.litellm_model.litellm.completion") + @patch("minisweagent.models.litellm_model.litellm.cost_calculator.completion_cost") + def test_query_preserves_explicit_request_timeout(self, mock_cost, mock_completion): + tool_call = MagicMock() + tool_call.function.name = "bash" + tool_call.function.arguments = '{"command": "echo test"}' + tool_call.id = "call_1" + mock_completion.return_value = _mock_litellm_response([tool_call]) + mock_cost.return_value = 0.001 + + model = LitellmModel(model_name="gpt-4", model_kwargs={"request_timeout": 30}) + model.query([{"role": "user", "content": "test"}]) + + assert "timeout" not in mock_completion.call_args.kwargs + assert mock_completion.call_args.kwargs["request_timeout"] == 30 + + @patch("minisweagent.models.litellm_model.litellm.completion") + def test_query_timeout_raises_provider_timeout(self, mock_completion): + mock_completion.side_effect = TimeoutError("request timed out") + + model = LitellmModel(model_name="gpt-4") + with pytest.raises(ProviderTimeout) as exc: + model.query([{"role": "user", "content": "test"}]) + assert exc.value.messages[0]["role"] == "exit" + assert exc.value.messages[0]["extra"]["exit_status"] == "ProviderTimeout" + + @pytest.mark.parametrize( + "timeout_error", + [ + TimeoutError("request timed out"), + litellm.exceptions.Timeout("request timed out", "gpt-4", "openai"), + ], + ) + @patch("minisweagent.models.litellm_model.litellm.completion") + def test_query_recognizes_provider_timeout_types(self, mock_completion, timeout_error): + mock_completion.side_effect = timeout_error + + model = LitellmModel(model_name="gpt-4") + with pytest.raises(ProviderTimeout): + model.query([{"role": "user", "content": "test"}]) + + def test_query_times_out_when_stream_stalls_mid_tool_call(self): + class StalledToolCallStream(BaseHTTPRequestHandler): + def do_POST(self): + self.server.seen_request = True + self.server.request_body = json.loads( + self.rfile.read(int(self.headers.get("content-length", "0"))).decode() + ) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.end_headers() + for event in [ + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": ""}, + } + ] + }, + "finish_reason": None, + } + ], + }, + ]: + self.wfile.write(f"data: {json.dumps(event)}\n\n".encode()) + self.wfile.flush() + time.sleep(5) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), StalledToolCallStream) + server.seen_request = False + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + model = LitellmModel( + model_name="openai/gpt-4o", + provider_timeout=0.2, + model_kwargs={ + "api_key": "fake-key", + "api_base": f"http://127.0.0.1:{server.server_port}/v1", + "stream": True, + }, + cost_tracking="ignore_errors", + ) + start = time.monotonic() + with pytest.raises(ProviderTimeout): + model.query([{"role": "user", "content": "test"}]) + assert server.seen_request + assert server.request_body["stream"] is True + assert time.monotonic() - start < 4 + finally: + server.shutdown() + server.server_close() + + def test_query_accumulates_streamed_tool_call_arguments(self): + class ToolCallStream(BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers.get("content-length", "0"))) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.end_headers() + for event in [ + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": ""}, + } + ] + }, + "finish_reason": None, + } + ], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"arguments": '{"command":"'}, + } + ] + }, + "finish_reason": None, + } + ], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"arguments": 'echo ok"}'}, + } + ] + }, + "finish_reason": None, + } + ], + }, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}], + }, + ]: + self.wfile.write(f"data: {json.dumps(event)}\n\n".encode()) + self.wfile.flush() + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), ToolCallStream) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + model = LitellmModel( + model_name="openai/gpt-4o", + model_kwargs={ + "api_key": "fake-key", + "api_base": f"http://127.0.0.1:{server.server_port}/v1", + "stream": True, + }, + ) + result = model.query([{"role": "user", "content": "test"}]) + assert result["extra"]["actions"] == [{"command": "echo ok", "tool_call_id": "call_1"}] + assert result["extra"]["cost"] == 0.0 + finally: + server.shutdown() + server.server_close() @patch("minisweagent.models.litellm_model.litellm.completion") @patch("minisweagent.models.litellm_model.litellm.cost_calculator.completion_cost") diff --git a/tests/models/test_litellm_response_model.py b/tests/models/test_litellm_response_model.py new file mode 100644 index 000000000..0d9c069df --- /dev/null +++ b/tests/models/test_litellm_response_model.py @@ -0,0 +1,28 @@ +from unittest.mock import patch + +import openai +import pytest + +from minisweagent.exceptions import ProviderTimeout +from minisweagent.models.litellm_response_model import LitellmResponseModel + + +def test_query_includes_default_provider_timeout(): + model = LitellmResponseModel(model_name="gpt-4", model_kwargs={"stream": True}) + + with patch("minisweagent.models.litellm_response_model.litellm.responses") as mock_responses: + model._query([{"role": "user", "content": "test"}]) + assert isinstance(mock_responses.call_args.kwargs["timeout"], openai.Timeout) + assert mock_responses.call_args.kwargs["timeout"].read == 5.0 + + +def test_query_timeout_raises_provider_timeout(): + model = LitellmResponseModel(model_name="gpt-4") + + with patch( + "minisweagent.models.litellm_response_model.litellm.responses", + side_effect=TimeoutError("request timed out"), + ): + with pytest.raises(ProviderTimeout) as exc_info: + model._query([{"role": "user", "content": "test"}]) + assert exc_info.value.messages[0]["extra"]["exit_status"] == "ProviderTimeout" diff --git a/tests/models/test_litellm_textbased_model.py b/tests/models/test_litellm_textbased_model.py index 475d3a6b9..04b5a0e40 100644 --- a/tests/models/test_litellm_textbased_model.py +++ b/tests/models/test_litellm_textbased_model.py @@ -4,8 +4,10 @@ from unittest.mock import Mock, patch import litellm +import openai import pytest +from minisweagent.exceptions import ProviderTimeout from minisweagent.models import GLOBAL_MODEL_STATS from minisweagent.models.litellm_textbased_model import LitellmTextbasedModel @@ -32,6 +34,24 @@ def side_effect(*args, **kwargs): assert "You can permanently set your API key with `mini-extra config set KEY VALUE`." in str(exc_info.value) +def test_query_includes_default_provider_timeout(): + model = LitellmTextbasedModel(model_name="gpt-4", model_kwargs={"stream": True}) + + with patch("litellm.completion") as mock_completion: + model._query([{"role": "user", "content": "test"}]) + assert isinstance(mock_completion.call_args.kwargs["timeout"], openai.Timeout) + assert mock_completion.call_args.kwargs["timeout"].read == 5.0 + + +def test_query_timeout_raises_provider_timeout(): + model = LitellmTextbasedModel(model_name="gpt-4") + + with patch("litellm.completion", side_effect=TimeoutError("request timed out")): + with pytest.raises(ProviderTimeout) as exc_info: + model._query([{"role": "user", "content": "test"}]) + assert exc_info.value.messages[0]["extra"]["exit_status"] == "ProviderTimeout" + + def test_model_registry_loading(): """Test that custom model registry is loaded and registered when provided.""" model_costs = {