inspect-ai 0.3.72__py3-none-any.whl → 0.3.73__py3-none-any.whl

inspect_ai/model/_providers/google.py

@@ -1,4 +1,3 @@
-import asyncio
 import functools
 import hashlib
 import json
@@ -9,6 +8,7 @@ from logging import getLogger
 from typing import Any
 
 # SDK Docs: https://googleapis.github.io/python-genai/
+import anyio
 from google.genai import Client  # type: ignore
 from google.genai.errors import APIError, ClientError  # type: ignore
 from google.genai.types import (  # type: ignore
@@ -26,6 +26,7 @@ from google.genai.types import (  # type: ignore
     GenerationConfig,
     HarmBlockThreshold,
     HarmCategory,
+    HttpOptions,
     Part,
     SafetySetting,
     SafetySettingDict,
@@ -49,6 +50,7 @@ from inspect_ai._util.content import (
     ContentVideo,
 )
 from inspect_ai._util.error import PrerequisiteError
+from inspect_ai._util.http import is_retryable_http_status
 from inspect_ai._util.images import file_as_data
 from inspect_ai._util.kvstore import inspect_kvstore
 from inspect_ai._util.trace import trace_message
@@ -69,6 +71,7 @@ from inspect_ai.model import (
 )
 from inspect_ai.model._model_call import ModelCall
 from inspect_ai.model._providers.util import model_base_url
+from inspect_ai.model._providers.util.hooks import HttpHooks, urllib3_hooks
 from inspect_ai.tool import (
     ToolCall,
     ToolChoice,
@@ -199,11 +202,15 @@ class GoogleGenAIAPI(ModelAPI):
         tool_choice: ToolChoice,
         config: GenerateConfig,
     ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
+        # generate request_id
+        request_id = urllib3_hooks().start_request()
+
         # Create google-genai types.
         gemini_contents = await as_chat_messages(self.client, input)
         gemini_tools = chat_tools(tools) if len(tools) > 0 else None
         gemini_tool_config = chat_tool_config(tool_choice) if len(tools) > 0 else None
         parameters = GenerateContentConfig(
+            http_options=HttpOptions(headers={HttpHooks.REQUEST_ID_HEADER: request_id}),
             temperature=config.temperature,
             top_p=config.top_p,
             top_k=config.top_k,
@@ -219,6 +226,11 @@ class GoogleGenAIAPI(ModelAPI):
                 self.client, input
             ),
         )
+        if config.response_schema is not None:
+            parameters.response_mime_type = "application/json"
+            parameters.response_schema = schema_from_param(
+                config.response_schema.json_schema, nullable=None
+            )
 
         response: GenerateContentResponse | None = None
 
@@ -230,10 +242,9 @@ class GoogleGenAIAPI(ModelAPI):
                 tools=gemini_tools,
                 tool_config=gemini_tool_config,
                 response=response,
+                time=urllib3_hooks().end_request(request_id),
             )
 
-        # TODO: would need to monkey patch AuthorizedSession.request
-
         try:
             response = await self.client.aio.models.generate_content(
                 model=self.model_name,
@@ -252,11 +263,25 @@ class GoogleGenAIAPI(ModelAPI):
         return output, model_call()
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
-        # see https://cloud.google.com/storage/docs/retry-strategy
-        return isinstance(ex, APIError) and (
-            ex.code in (408, 429, 429) or ex.code >= 500
-        )
+    def should_retry(self, ex: Exception) -> bool:
+        import requests  # type: ignore
+
+        # standard http errors
+        if isinstance(ex, APIError):
+            return is_retryable_http_status(ex.status)
+
+        # low-level requests exceptions
+        elif isinstance(ex, requests.exceptions.RequestException):
+            return isinstance(
+                ex,
+                (
+                    requests.exceptions.ConnectionError
+                    | requests.exceptions.ConnectTimeout
+                    | requests.exceptions.ChunkedEncodingError
+                ),
+            )
+        else:
+            return False
 
     @override
     def connection_key(self) -> str:
@@ -296,6 +321,7 @@ def build_model_call(
     tools: list[Tool] | None,
     tool_config: ToolConfig | None,
     response: GenerateContentResponse | None,
+    time: float | None,
 ) -> ModelCall:
     return ModelCall.create(
         request=dict(
@@ -307,6 +333,7 @@ def build_model_call(
         ),
         response=response if response is not None else {},
         filter=model_call_filter,
+        time=time,
     )
 
 
@@ -464,7 +491,9 @@ def chat_tools(tools: list[ToolInfo]) -> list[Tool]:
 
 
 # https://ai.google.dev/gemini-api/tutorials/extract_structured_data#define_the_schema
-def schema_from_param(param: ToolParam | ToolParams, nullable: bool = False) -> Schema:
+def schema_from_param(
+    param: ToolParam | ToolParams, nullable: bool | None = False
+) -> Schema:
     if isinstance(param, ToolParams):
         param = ToolParam(
             type=param.type, properties=param.properties, required=param.required
@@ -529,10 +558,13 @@ def chat_tool_config(tool_choice: ToolChoice) -> ToolConfig:
 
 
 def completion_choice_from_candidate(candidate: Candidate) -> ChatCompletionChoice:
-    # check for completion text
-    content = ""
     # content can be None when the finish_reason is SAFETY
-    if candidate.content is not None:
+    if candidate.content is None:
+        content = ""
+    # content.parts can be None when the finish_reason is MALFORMED_FUNCTION_CALL
+    elif candidate.content.parts is None:
+        content = ""
+    else:
         content = " ".join(
             [
                 part.text
@@ -680,6 +712,8 @@ def finish_reason_to_stop_reason(finish_reason: FinishReason) -> StopReason:
         ):
             return "content_filter"
         case _:
+            # Note: to avoid adding another option to StopReason,
+            # this includes FinishReason.MALFORMED_FUNCTION_CALL
             return "unknown"
 
 
@@ -775,7 +809,7 @@ async def file_for_content(
             file=BytesIO(content_bytes), config=dict(mime_type=mime_type)
         )
         while upload.state.name == "PROCESSING":
-            await asyncio.sleep(3)
+            await anyio.sleep(3)
             upload = client.files.get(name=upload.name)
         if upload.state.name == "FAILED":
             trace(f"Failed to upload file '{upload.name}: {upload.error}")

inspect_ai/model/_providers/groq.py

@@ -5,8 +5,9 @@ from typing import Any, Dict, Iterable, List, Optional
 
 import httpx
 from groq import (
+    APIStatusError,
+    APITimeoutError,
     AsyncGroq,
-    RateLimitError,
 )
 from groq.types.chat import (
     ChatCompletion,
@@ -25,10 +26,10 @@ from typing_extensions import override
 
 from inspect_ai._util.constants import (
     BASE_64_DATA_REMOVED,
-    DEFAULT_MAX_RETRIES,
     DEFAULT_MAX_TOKENS,
 )
 from inspect_ai._util.content import Content, ContentReasoning, ContentText
+from inspect_ai._util.http import is_retryable_http_status
 from inspect_ai._util.images import file_as_data_uri
 from inspect_ai._util.url import is_http_url
 from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
@@ -54,7 +55,7 @@ from .util import (
     environment_prerequisite_error,
     model_base_url,
 )
-from .util.tracker import HttpxTimeTracker
+from .util.hooks import HttpxHooks
 
 GROQ_API_KEY = "GROQ_API_KEY"
 
@@ -84,18 +85,12 @@ class GroqAPI(ModelAPI):
         self.client = AsyncGroq(
             api_key=self.api_key,
             base_url=model_base_url(base_url, "GROQ_BASE_URL"),
-            max_retries=(
-                config.max_retries
-                if config.max_retries is not None
-                else DEFAULT_MAX_RETRIES
-            ),
-            timeout=config.timeout if config.timeout is not None else 60.0,
             **model_args,
             http_client=httpx.AsyncClient(limits=httpx.Limits(max_connections=None)),
         )
 
         # create time tracker
-        self._time_tracker = HttpxTimeTracker(self.client._client)
+        self._http_hooks = HttpxHooks(self.client._client)
 
     @override
     async def close(self) -> None:
@@ -109,7 +104,7 @@ class GroqAPI(ModelAPI):
         config: GenerateConfig,
     ) -> tuple[ModelOutput, ModelCall]:
         # allocate request_id (so we can see it from ModelCall)
-        request_id = self._time_tracker.start_request()
+        request_id = self._http_hooks.start_request()
 
         # setup request and response for ModelCall
         request: dict[str, Any] = {}
@@ -120,7 +115,7 @@ class GroqAPI(ModelAPI):
                 request=request,
                 response=response,
                 filter=model_call_filter,
-                time=self._time_tracker.end_request(request_id),
+                time=self._http_hooks.end_request(request_id),
             )
 
         messages = await as_groq_chat_messages(input)
@@ -137,7 +132,7 @@ class GroqAPI(ModelAPI):
         request = dict(
             messages=messages,
             model=self.model_name,
-            extra_headers={HttpxTimeTracker.REQUEST_ID_HEADER: request_id},
+            extra_headers={HttpxHooks.REQUEST_ID_HEADER: request_id},
             **params,
         )
 
@@ -215,8 +210,13 @@ class GroqAPI(ModelAPI):
         ]
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
-        return isinstance(ex, RateLimitError)
+    def should_retry(self, ex: Exception) -> bool:
+        if isinstance(ex, APIStatusError):
+            return is_retryable_http_status(ex.status_code)
+        elif isinstance(ex, APITimeoutError):
+            return True
+        else:
+            return False
 
     @override
     def connection_key(self) -> str:

inspect_ai/model/_providers/hf.py

@@ -1,15 +1,19 @@
-import asyncio
+import concurrent
+import concurrent.futures
 import copy
 import functools
 import gc
 import json
 import os
 import time
+from concurrent.futures import Future
 from dataclasses import dataclass
+from logging import getLogger
 from queue import Empty, Queue
 from threading import Thread
 from typing import Any, Literal, Protocol, cast
 
+import anyio
 import numpy as np
 import torch  # type: ignore
 from torch import Tensor  # type: ignore
@@ -23,6 +27,7 @@ from typing_extensions import override
 
 from inspect_ai._util.constants import DEFAULT_MAX_TOKENS
 from inspect_ai._util.content import ContentText
+from inspect_ai._util.trace import trace_action
 from inspect_ai.tool import ToolChoice, ToolInfo
 
 from .._chat_message import ChatMessage, ChatMessageAssistant
@@ -38,6 +43,9 @@ from .._model_output import (
 )
 from .util import ChatAPIHandler, HFHandler
 
+logger = getLogger(__name__)
+
+
 HF_TOKEN = "HF_TOKEN"
 
 
@@ -385,8 +393,7 @@ class GenerateOutput:
 @dataclass
 class _QueueItem:
     input: GenerateInput
-    future: asyncio.Future[GenerateOutput]
-    loop: asyncio.AbstractEventLoop
+    future: Future[GenerateOutput]
 
 
 batch_thread: Thread | None = None
@@ -402,25 +409,26 @@ async def batched_generate(input: GenerateInput) -> GenerateOutput:
         batch_thread.start()
 
     # enqueue the job
-    loop = asyncio.get_event_loop()
-    future: asyncio.Future[GenerateOutput] = loop.create_future()
-    batch_queue.put(_QueueItem(input=input, future=future, loop=loop))
-
-    # await the job
-    await future
+    future = Future[GenerateOutput]()
+    batch_queue.put(_QueueItem(input=input, future=future))
 
-    # return it
-    return future.result()
+    # await the future
+    with trace_action(logger, "HF Batched Generate", "HF Batched Generate"):
+        while True:
+            try:
+                return future.result(timeout=0.01)
+            except concurrent.futures.TimeoutError:
+                pass
+            await anyio.sleep(1)
 
 
 def process_batches() -> None:
     while True:
         # drain the queue (wait until no new messages have shown up for 2 seconds)
-        inputs: list[tuple[GenerateInput, asyncio.Future[GenerateOutput]]] = []
+        inputs: list[tuple[GenerateInput, Future[GenerateOutput]]] = []
         while True:
             try:
                 input = batch_queue.get(timeout=2)
-                loop = input.loop
                 inputs.append((input.input, input.future))
                 if len(inputs) == input.input.batch_size:
                     # max batch size reached
@@ -480,8 +488,7 @@ def process_batches() -> None:
                 # asyncio futures are not thread safe, so we need to pass the event loop
                 # down to this point, so we can mark the future as done in a thread safe manner.
                 # see: https://docs.python.org/3/library/asyncio-dev.html#concurrency-and-multithreading
-                loop.call_soon_threadsafe(
-                    future.set_result,
+                future.set_result(
                     GenerateOutput(
                         output=output,
                         input_tokens=input_tokens,
@@ -489,13 +496,13 @@ def process_batches() -> None:
                         total_tokens=input_tokens + output_tokens,
                         logprobs=logprobs[i] if logprobs is not None else None,
                         time=total_time,
-                    ),
+                    )
                 )
 
         except Exception as ex:
             for inp in inputs:
                 future = inp[1]
-                loop.call_soon_threadsafe(future.set_exception, ex)
+                future.set_exception(ex)
 
 
 def extract_logprobs(

inspect_ai/model/_providers/mistral.py

@@ -7,6 +7,7 @@ from httpcore import ReadTimeout
 from httpx import ReadTimeout as AsyncReadTimeout
 from mistralai import (
     ContentChunk,
+    DocumentURLChunk,
     FunctionCall,
     FunctionName,
     ImageURL,
@@ -22,6 +23,12 @@ from mistralai.models import (
     ChatCompletionChoice as MistralChatCompletionChoice,
 )
 from mistralai.models import Function as MistralFunction
+from mistralai.models import (
+    JSONSchema as MistralJSONSchema,
+)
+from mistralai.models import (
+    ResponseFormat as MistralResponseFormat,
+)
 from mistralai.models import SDKError
 from mistralai.models import SystemMessage as MistralSystemMessage
 from mistralai.models import Tool as MistralTool
@@ -38,11 +45,9 @@ from typing_extensions import override
 
 # TODO: Migration guide:
 # https://github.com/mistralai/client-python/blob/main/MIGRATION.md
-from inspect_ai._util.constants import (
-    DEFAULT_TIMEOUT,
-    NO_CONTENT,
-)
+from inspect_ai._util.constants import NO_CONTENT
 from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.http import is_retryable_http_status
 from inspect_ai._util.images import file_as_data_uri
 from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
 
@@ -61,7 +66,7 @@ from .._model_output import (
     StopReason,
 )
 from .util import environment_prerequisite_error, model_base_url
-from .util.tracker import HttpxTimeTracker
+from .util.hooks import HttpxHooks
 
 AZURE_MISTRAL_API_KEY = "AZURE_MISTRAL_API_KEY"
 AZUREAI_MISTRAL_API_KEY = "AZUREAI_MISTRAL_API_KEY"
@@ -127,16 +132,12 @@ class MistralAPI(ModelAPI):
         config: GenerateConfig,
     ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
         # create client
-        with Mistral(
-            api_key=self.api_key,
-            timeout_ms=(config.timeout if config.timeout else DEFAULT_TIMEOUT) * 1000,
-            **self.model_args,
-        ) as client:
+        with Mistral(api_key=self.api_key, **self.model_args) as client:
             # create time tracker
-            time_tracker = HttpxTimeTracker(client.sdk_configuration.async_client)
+            http_hooks = HttpxHooks(client.sdk_configuration.async_client)
 
             # build request
-            request_id = time_tracker.start_request()
+            request_id = http_hooks.start_request()
             request: dict[str, Any] = dict(
                 model=self.model_name,
                 messages=await mistral_chat_messages(input),
@@ -144,7 +145,7 @@ class MistralAPI(ModelAPI):
                 tool_choice=(
                     mistral_chat_tool_choice(tool_choice) if len(tools) > 0 else None
                 ),
-                http_headers={HttpxTimeTracker.REQUEST_ID_HEADER: request_id},
+                http_headers={HttpxHooks.REQUEST_ID_HEADER: request_id},
             )
             if config.temperature is not None:
                 request["temperature"] = config.temperature
@@ -154,6 +155,18 @@ class MistralAPI(ModelAPI):
                 request["max_tokens"] = config.max_tokens
             if config.seed is not None:
                 request["random_seed"] = config.seed
+            if config.response_schema is not None:
+                request["response_format"] = MistralResponseFormat(
+                    type="json_schema",
+                    json_schema=MistralJSONSchema(
+                        name=config.response_schema.name,
+                        description=config.response_schema.description,
+                        schema_definition=config.response_schema.json_schema.model_dump(
+                            exclude_none=True
+                        ),
+                        strict=config.response_schema.strict,
+                    ),
+                )
 
             # prepare response for inclusion in model call
             response: dict[str, Any] = {}
@@ -169,7 +182,7 @@ class MistralAPI(ModelAPI):
                 return ModelCall.create(
                     request=req,
                     response=response,
-                    time=time_tracker.end_request(request_id),
+                    time=http_hooks.end_request(request_id),
                 )
 
             # send request
@@ -205,12 +218,13 @@ class MistralAPI(ModelAPI):
             ), model_call()
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
-        return (
-            isinstance(ex, SDKError)
-            and ex.status_code == 429
-            or isinstance(ex, ReadTimeout | AsyncReadTimeout)
-        )
+    def should_retry(self, ex: Exception) -> bool:
+        if isinstance(ex, SDKError):
+            return is_retryable_http_status(ex.status_code)
+        elif isinstance(ex, ReadTimeout | AsyncReadTimeout):
+            return True
+        else:
+            return False
 
     @override
     def connection_key(self) -> str:
@@ -462,6 +476,8 @@ def completion_content_chunk(content: ContentChunk) -> Content:
         raise TypeError("ReferenceChunk content is not supported by Inspect.")
     elif isinstance(content, TextChunk):
         return ContentText(text=content.text)
+    elif isinstance(content, DocumentURLChunk):
+        return ContentText(text=content.document_url)
     else:
         if isinstance(content.image_url, str):
             return ContentImage(image=content.image_url)

inspect_ai/model/_providers/openai.py

@@ -7,25 +7,22 @@ import httpx
 from openai import (
     DEFAULT_CONNECTION_LIMITS,
     DEFAULT_TIMEOUT,
-    APIConnectionError,
+    APIStatusError,
     APITimeoutError,
     AsyncAzureOpenAI,
     AsyncOpenAI,
     BadRequestError,
-    InternalServerError,
     RateLimitError,
 )
 from openai._types import NOT_GIVEN
-from openai.types.chat import (
-    ChatCompletion,
-)
+from openai.types.chat import ChatCompletion
 from typing_extensions import override
 
-from inspect_ai._util.constants import DEFAULT_MAX_RETRIES
 from inspect_ai._util.error import PrerequisiteError
+from inspect_ai._util.http import is_retryable_http_status
 from inspect_ai._util.logger import warn_once
 from inspect_ai.model._openai import chat_choices_from_openai
-from inspect_ai.model._providers.util.tracker import HttpxTimeTracker
+from inspect_ai.model._providers.util.hooks import HttpxHooks
 from inspect_ai.tool import ToolChoice, ToolInfo
 
 from .._chat_message import ChatMessage
@@ -130,9 +127,6 @@ class OpenAIAPI(ModelAPI):
                 api_key=self.api_key,
                 azure_endpoint=base_url,
                 azure_deployment=model_name,
-                max_retries=(
-                    config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
-                ),
                 http_client=http_client,
                 **model_args,
             )
@@ -140,15 +134,12 @@ class OpenAIAPI(ModelAPI):
             self.client = AsyncOpenAI(
                 api_key=self.api_key,
                 base_url=model_base_url(base_url, "OPENAI_BASE_URL"),
-                max_retries=(
-                    config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
-                ),
                 http_client=http_client,
                 **model_args,
             )
 
         # create time tracker
-        self._time_tracker = HttpxTimeTracker(self.client._client)
+        self._http_hooks = HttpxHooks(self.client._client)
 
     def is_azure(self) -> bool:
         return self.service == "azure"
@@ -186,7 +177,7 @@ class OpenAIAPI(ModelAPI):
             )
 
         # allocate request_id (so we can see it from ModelCall)
-        request_id = self._time_tracker.start_request()
+        request_id = self._http_hooks.start_request()
 
         # setup request and response for ModelCall
         request: dict[str, Any] = {}
@@ -197,7 +188,7 @@ class OpenAIAPI(ModelAPI):
                 request=request,
                 response=response,
                 filter=image_url_filter,
-                time=self._time_tracker.end_request(request_id),
+                time=self._http_hooks.end_request(request_id),
             )
 
         # unlike text models, vision models require a max_tokens (and set it to a very low
@@ -216,7 +207,7 @@ class OpenAIAPI(ModelAPI):
             tool_choice=openai_chat_tool_choice(tool_choice)
             if len(tools) > 0
             else NOT_GIVEN,
-            extra_headers={HttpxTimeTracker.REQUEST_ID_HEADER: request_id},
+            extra_headers={HttpxHooks.REQUEST_ID_HEADER: request_id},
             **self.completion_params(config, len(tools) > 0),
         )
 
@@ -266,17 +257,21 @@ class OpenAIAPI(ModelAPI):
         return chat_choices_from_openai(response, tools)
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
+    def should_retry(self, ex: Exception) -> bool:
         if isinstance(ex, RateLimitError):
             # Do not retry on these rate limit errors
             # The quota exceeded one is related to monthly account quotas.
-            if "You exceeded your current quota" not in ex.message:
+            if "You exceeded your current quota" in ex.message:
+                warn_once(logger, f"OpenAI quota exceeded, not retrying: {ex.message}")
+                return False
+            else:
                 return True
-        elif isinstance(
-            ex, (APIConnectionError | APITimeoutError | InternalServerError)
-        ):
+        elif isinstance(ex, APIStatusError):
+            return is_retryable_http_status(ex.status_code)
+        elif isinstance(ex, APITimeoutError):
             return True
-        return False
+        else:
+            return False
 
     @override
     def connection_key(self) -> str:
@@ -315,8 +310,6 @@ class OpenAIAPI(ModelAPI):
             params["temperature"] = 1
         if config.top_p is not None:
             params["top_p"] = config.top_p
-        if config.timeout is not None:
-            params["timeout"] = float(config.timeout)
         if config.num_choices is not None:
             params["n"] = config.num_choices
         if config.logprobs is not None:
@@ -331,6 +324,18 @@ class OpenAIAPI(ModelAPI):
             and not self.is_o1_mini()
         ):
             params["reasoning_effort"] = config.reasoning_effort
+        if config.response_schema is not None:
+            params["response_format"] = dict(
+                type="json_schema",
+                json_schema=dict(
+                    name=config.response_schema.name,
+                    schema=config.response_schema.json_schema.model_dump(
+                        exclude_none=True
+                    ),
+                    description=config.response_schema.description,
+                    strict=config.response_schema.strict,
+                ),
+            )
 
         return params
 

inspect_ai/model/_providers/openai_o1.py

@@ -107,7 +107,7 @@ def chat_messages(
 ) -> list[ChatCompletionMessageParam]:
     # o1 does not allow system messages so convert system -> user
     messages: list[ChatMessage] = [
-        ChatMessageUser(content=message.content)
+        ChatMessageUser(id=message.id, content=message.content)
         if message.role == "system"
         else message
         for message in input

inspect_ai/model/_providers/providers.py

@@ -148,7 +148,7 @@ def cf() -> type[ModelAPI]:
 def mistral() -> type[ModelAPI]:
     FEATURE = "Mistral API"
     PACKAGE = "mistralai"
-    MIN_VERSION = "1.5.0"
+    MIN_VERSION = "1.5.1"
 
     # verify we have the package
     try:

inspect_ai/model/_providers/together.py

@@ -34,8 +34,8 @@ from .util import (
     chat_api_input,
     chat_api_request,
     environment_prerequisite_error,
-    is_chat_api_rate_limit,
     model_base_url,
+    should_retry_chat_api_error,
 )
 
 
@@ -186,7 +186,6 @@ class TogetherRESTAPI(ModelAPI):
             url=f"{chat_url}",
             headers={"Authorization": f"Bearer {self.api_key}"},
             json=json,
-            config=config,
         )
 
         if "error" in response:
@@ -215,8 +214,8 @@ class TogetherRESTAPI(ModelAPI):
             return ModelOutput(model=model, choices=choices, usage=usage)
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
-        return is_chat_api_rate_limit(ex)
+    def should_retry(self, ex: Exception) -> bool:
+        return should_retry_chat_api_error(ex)
 
     # cloudflare enforces rate limits by model for each account
     @override