feat: add image upstream deadline

apps/ai/tests.py

@@ -13,7 +13,7 @@ from apps.ai.importers import import_ai_models_config
 from apps.ai.models import AiConfigAuditLog, AiModel, ModelAlias
 from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type
 from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider
-from apps.ai.providers.utils import resolution_to_size
+from apps.ai.providers.utils import image_request_timeout, resolution_to_size
 
 
 TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("ascii")
@@ -60,6 +60,11 @@ class ProviderUtilsTests(SimpleTestCase):
         self.assertEqual(resolution_to_size("1k"), "1024x1024")
         self.assertEqual(resolution_to_size("512px"), "512x512")
 
+    @override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
+    def test_image_request_timeout_caps_read_timeout_to_deadline(self):
+        self.assertEqual(image_request_timeout(30, 0, "4K"), (30, 180))
+        self.assertEqual(image_request_timeout(30, 120, "4K"), (30, 120))
+
 
 class ChatCompletionsProviderTests(SimpleTestCase):
     def test_generate_text_builds_chat_payload_and_cleans_titles(self):
@@ -192,6 +197,64 @@ class ChatCompletionsProviderTests(SimpleTestCase):
         self.assertEqual(content[0], {"type": "text", "text": "Generate product image"})
         self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
 
+    @override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
+    def test_generate_image_caps_post_and_download_timeouts(self):
+        session = FakeSession(
+            FakeResponse(
+                {
+                    "choices": [
+                        {
+                            "message": {
+                                "content": [
+                                    {
+                                        "type": "image_url",
+                                        "image_url": {"url": "https://cdn.example.com/out.png"},
+                                    }
+                                ]
+                            }
+                        }
+                    ]
+                }
+            ),
+            FakeResponse({}, content=b"generated-image"),
+        )
+        provider = ChatCompletionsProvider(session=session)
+        model = ResolvedModel(
+            name="Slow image",
+            url="https://api.vectorengine.ai/v1/chat/completions",
+            model="image-model",
+            api_key="test-key",
+            api_type="chat",
+            timeout_seconds=0,
+            connect_timeout_seconds=30,
+        )
+
+        result = provider.generate_image("Generate product image", model, resolution="4K")
+
+        self.assertEqual(result.image, b"generated-image")
+        self.assertEqual(session.posts[0]["timeout"], (30, 180))
+        self.assertEqual(session.gets[0]["timeout"], 180)
+
+    @override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
+    def test_generate_text_does_not_use_image_deadline(self):
+        session = FakeSession(
+            FakeResponse({"choices": [{"message": {"content": "1. Red Dress"}}]})
+        )
+        provider = ChatCompletionsProvider(session=session)
+        model = ResolvedModel(
+            name="Slow text",
+            url="https://api.vectorengine.ai/v1/chat/completions",
+            model="text-model",
+            api_key="test-key",
+            api_type="chat",
+            timeout_seconds=0,
+            connect_timeout_seconds=30,
+        )
+
+        provider.generate_text("Generate titles", model, resolution="4K")
+
+        self.assertEqual(session.posts[0]["timeout"], (30, 600))
+
 
 class ImagesEditsProviderTests(SimpleTestCase):
     def test_generate_image_builds_multipart_request_and_parses_base64(self):