fix: validate return_images support per model (#61)

src/perplexity/lib/__init__.py

@@ -0,0 +1,17 @@
+"""Hand-maintained helpers that are not modified by the Stainless generator."""
+
+from .model_capabilities import (
+    MODELS_SUPPORTING_RETURN_IMAGES,
+    MODELS_WITHOUT_RETURN_IMAGES,
+    ReturnImagesUnsupportedError,
+    model_supports_return_images,
+    validate_return_images,
+)
+
+__all__ = [
+    "MODELS_SUPPORTING_RETURN_IMAGES",
+    "MODELS_WITHOUT_RETURN_IMAGES",
+    "ReturnImagesUnsupportedError",
+    "model_supports_return_images",
+    "validate_return_images",
+]

src/perplexity/lib/model_capabilities.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Final, FrozenSet
+
+from perplexity._types import Omit, omit
+
+MODELS_SUPPORTING_RETURN_IMAGES: Final[FrozenSet[str]] = frozenset(
+    {
+        "sonar-pro",
+        "sonar-reasoning-pro",
+        "sonar-deep-research",
+    }
+)
+
+MODELS_WITHOUT_RETURN_IMAGES: Final[FrozenSet[str]] = frozenset(
+    {
+        "sonar",
+    }
+)
+
+
+class ReturnImagesUnsupportedError(ValueError):
+    """Raised when return_images=True is requested for a model that does not return images."""
+
+
+def model_supports_return_images(model: str) -> bool:
+    normalized = model.strip().lower()
+    if normalized in MODELS_WITHOUT_RETURN_IMAGES:
+        return False
+    if normalized in MODELS_SUPPORTING_RETURN_IMAGES:
+        return True
+    # Unknown/new models: allow the request but callers should verify response.images.
+    return True
+
+
+def validate_return_images(*, model: str, return_images: object) -> None:
+    if return_images is omit or return_images is None or return_images is False:
+        return
+    if return_images is True and not model_supports_return_images(model):
+        raise ReturnImagesUnsupportedError(
+            f"Model {model!r} does not return images when return_images=True. "
+            f"Use one of {sorted(MODELS_SUPPORTING_RETURN_IMAGES)} instead."
+        )

src/perplexity/resources/chat/completions.py

@@ -19,6 +19,7 @@
 )
 from ..._streaming import Stream, AsyncStream
 from ...types.chat import completion_create_params
+from ...lib.model_capabilities import validate_return_images
 from ..._base_client import make_request_options
 from ...types.stream_chunk import StreamChunk
 from ...types.shared_params.chat_message_input import ChatMessageInput
@@ -369,6 +370,7 @@ def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> StreamChunk | Stream[StreamChunk]:
+        validate_return_images(model=model, return_images=return_images)
         return self._post(
             "/chat/completions",
             body=maybe_transform(
@@ -788,6 +790,7 @@ async def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> StreamChunk | AsyncStream[StreamChunk]:
+        validate_return_images(model=model, return_images=return_images)
         return await self._post(
             "/chat/completions",
             body=await async_maybe_transform(

src/perplexity/types/stream_chunk.py

@@ -22,6 +22,8 @@ class StreamChunk(BaseModel):
 
     citations: Optional[List[str]] = None
 
+    images: Optional[List[str]] = None
+
     object: Optional[str] = None
 
     search_results: Optional[List[APIPublicSearchResult]] = None

tests/test_model_capabilities.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import pytest
+
+from perplexity.lib.model_capabilities import (
+    ReturnImagesUnsupportedError,
+    model_supports_return_images,
+    validate_return_images,
+)
+
+
+def test_sonar_does_not_support_return_images() -> None:
+    assert not model_supports_return_images("sonar")
+
+
+def test_sonar_pro_supports_return_images() -> None:
+    assert model_supports_return_images("sonar-pro")
+
+
+def test_validate_return_images_raises_for_sonar() -> None:
+    with pytest.raises(ReturnImagesUnsupportedError, match="sonar"):
+        validate_return_images(model="sonar", return_images=True)
+
+
+def test_validate_return_images_allows_sonar_pro() -> None:
+    validate_return_images(model="sonar-pro", return_images=True)
+
+
+def test_validate_return_images_ignores_false_and_omit() -> None:
+    validate_return_images(model="sonar", return_images=False)
+    validate_return_images(model="sonar", return_images=None)
+
+
+def test_stream_chunk_parses_images_field() -> None:
+    from perplexity.types import StreamChunk
+
+    chunk = StreamChunk.model_validate(
+        {
+            "id": "id",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": ""},
+                    "message": {"role": "assistant", "content": "hi"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "created": 1,
+            "model": "sonar-pro",
+            "images": ["https://example.com/a.png"],
+        }
+    )
+    assert chunk.images == ["https://example.com/a.png"]