diff --git a/src/perplexity/lib/__init__.py b/src/perplexity/lib/__init__.py new file mode 100644 index 0000000..6d1932b --- /dev/null +++ b/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", +] diff --git a/src/perplexity/lib/model_capabilities.py b/src/perplexity/lib/model_capabilities.py new file mode 100644 index 0000000..066009a --- /dev/null +++ b/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." + ) diff --git a/src/perplexity/resources/chat/completions.py b/src/perplexity/resources/chat/completions.py index 5931b13..1e7cc44 100644 --- a/src/perplexity/resources/chat/completions.py +++ b/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( diff --git a/src/perplexity/types/stream_chunk.py b/src/perplexity/types/stream_chunk.py index e2a9309..ffa557a 100644 --- a/src/perplexity/types/stream_chunk.py +++ b/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 diff --git a/tests/test_model_capabilities.py b/tests/test_model_capabilities.py new file mode 100644 index 0000000..6153901 --- /dev/null +++ b/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"]