feat: add image upstream deadline

apps/ai/providers/openai_compatible.py

@@ -20,6 +20,7 @@ from .utils import (
     extract_image_from_response,
     extract_text_from_response,
     extract_titles_from_response,
+    image_request_timeout,
     image_bytes_to_data_url,
     normalize_api_url,
     request_timeout,
@@ -79,6 +80,16 @@ class BaseHttpProvider:
     def _read_timeout(self, model: ResolvedModel, resolution: str) -> int:
         return self._timeout(model, resolution)[1]
 
+    def _image_timeout(self, model: ResolvedModel, resolution: str) -> tuple[int, int]:
+        return image_request_timeout(
+            model.connect_timeout_seconds,
+            model.timeout_seconds,
+            resolution,
+        )
+
+    def _image_read_timeout(self, model: ResolvedModel, resolution: str) -> int:
+        return self._image_timeout(model, resolution)[1]
+
 
 class ChatCompletionsProvider(BaseHttpProvider):
     def capabilities(self) -> set[str]:
@@ -142,14 +153,14 @@ class ChatCompletionsProvider(BaseHttpProvider):
             url,
             headers=self._headers(model, json=True),
             json=payload,
-            timeout=self._timeout(model, resolution),
+            timeout=self._image_timeout(model, resolution),
         )
         response.raise_for_status()
         raw = response.json()
         image_bytes = extract_image_from_response(
             raw,
             session=self.session,
-            timeout=self._read_timeout(model, resolution),
+            timeout=self._image_read_timeout(model, resolution),
         )
         if not image_bytes:
             raise AiResponseParseError("AI response did not contain an image")
@@ -217,14 +228,14 @@ class GeminiProvider(ChatCompletionsProvider):
             url,
             headers=self._headers(model, json=True),
             json=payload,
-            timeout=self._timeout(model, resolution),
+            timeout=self._image_timeout(model, resolution),
         )
         response.raise_for_status()
         raw = response.json()
         image_bytes = extract_image_from_response(
             raw,
             session=self.session,
-            timeout=self._read_timeout(model, resolution),
+            timeout=self._image_read_timeout(model, resolution),
         )
         if not image_bytes:
             raise AiResponseParseError("AI response did not contain an image")
@@ -266,14 +277,14 @@ class ImagesGenerationProvider(BaseHttpProvider):
             url,
             headers=self._headers(model, json=True),
             json=payload,
-            timeout=self._timeout(model, resolution),
+            timeout=self._image_timeout(model, resolution),
         )
         response.raise_for_status()
         raw = response.json()
         image_bytes = extract_image_from_response(
             raw,
             session=self.session,
-            timeout=self._read_timeout(model, resolution),
+            timeout=self._image_read_timeout(model, resolution),
         )
         if not image_bytes:
             raise AiResponseParseError("AI response did not contain an image")
@@ -317,14 +328,14 @@ class ImagesEditsProvider(BaseHttpProvider):
             headers=self._headers(model),
             data=data,
             files=files,
-            timeout=self._timeout(model, resolution),
+            timeout=self._image_timeout(model, resolution),
         )
         response.raise_for_status()
         raw = response.json()
         image_bytes = extract_image_from_response(
             raw,
             session=self.session,
-            timeout=self._read_timeout(model, resolution),
+            timeout=self._image_read_timeout(model, resolution),
         )
         if not image_bytes:
             raise AiResponseParseError("AI response did not contain an image")

apps/ai/providers/utils.py

@@ -8,6 +8,8 @@ from typing import Any, Iterable
 from urllib.parse import urljoin, urlparse
 
 import requests
+from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
 
 from .base import AiProviderConfigError, AiResponseParseError
 
@@ -26,6 +28,7 @@ SUPPORTED_API_TYPES = {
 }
 
 RESOLUTION_TIMEOUTS = {"512": 180, "1K": 240, "2K": 360, "4K": 600}
+DEFAULT_IMAGE_UPSTREAM_DEADLINE_SECONDS = 180
 BASE64_KEYS = {"image_base64", "base64", "b64_json", "data"}
 IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
 
@@ -44,6 +47,34 @@ def request_timeout(connect_timeout: int, read_timeout: int, resolution: str) ->
     return connect_timeout, resolved_read_timeout
 
 
+def image_request_timeout(connect_timeout: int, read_timeout: int, resolution: str) -> tuple[int, int]:
+    resolved_connect_timeout, resolved_read_timeout = request_timeout(
+        connect_timeout,
+        read_timeout,
+        resolution,
+    )
+    return resolved_connect_timeout, cap_image_read_timeout(resolved_read_timeout)
+
+
+def cap_image_read_timeout(read_timeout: int) -> int:
+    deadline = image_upstream_deadline_seconds()
+    if deadline <= 0:
+        return read_timeout
+    return min(read_timeout, deadline)
+
+
+def image_upstream_deadline_seconds() -> int:
+    try:
+        value = getattr(
+            settings,
+            "AI_IMAGE_UPSTREAM_DEADLINE_SECONDS",
+            DEFAULT_IMAGE_UPSTREAM_DEADLINE_SECONDS,
+        )
+    except ImproperlyConfigured:
+        value = DEFAULT_IMAGE_UPSTREAM_DEADLINE_SECONDS
+    return int(value)
+
+
 def detect_api_type(url: str, api_type: str = API_AUTO) -> str:
     if api_type and api_type != API_AUTO:
         if api_type not in SUPPORTED_API_TYPES: