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..2fc908e --- /dev/null +++ b/src/perplexity/lib/model_capabilities.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Final, FrozenSet, Iterable + +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", + } +) + +MODELS_SUPPORTING_REASONING_EFFORT: Final[FrozenSet[str]] = frozenset( + { + "sonar-deep-research", + "sonar-reasoning-pro", + } +) + +MODELS_SUPPORTING_SEARCH_CONTEXT_SIZE: Final[FrozenSet[str]] = frozenset( + { + "sonar-pro", + "sonar-reasoning-pro", + } +) + +MAX_SEARCH_DOMAIN_FILTERS: Final[int] = 3 + + +class UnsupportedParameterError(ValueError): + """Raised when a request parameter is not supported for the selected model.""" + + +class ReturnImagesUnsupportedError(UnsupportedParameterError): + """Raised when return_images=True is requested for a model that does not return images.""" + + +def _normalize_model(model: str) -> str: + return model.strip().lower() + + +def _is_omitted(value: object) -> bool: + return value is omit or value is None or value is False + + +def model_supports_return_images(model: str) -> bool: + normalized = _normalize_model(model) + if normalized in MODELS_WITHOUT_RETURN_IMAGES: + return False + if normalized in MODELS_SUPPORTING_RETURN_IMAGES: + return True + return True + + +def model_supports_reasoning_effort(model: str) -> bool: + return _normalize_model(model) in MODELS_SUPPORTING_REASONING_EFFORT + + +def model_supports_search_context_size(model: str) -> bool: + return _normalize_model(model) in MODELS_SUPPORTING_SEARCH_CONTEXT_SIZE + + +def validate_return_images(*, model: str, return_images: object) -> None: + if _is_omitted(return_images): + 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." + ) + + +def validate_reasoning_effort(*, model: str, reasoning_effort: object) -> None: + if _is_omitted(reasoning_effort): + return + if not model_supports_reasoning_effort(model): + raise UnsupportedParameterError( + f"Model {model!r} does not support reasoning_effort. " + f"Use one of {sorted(MODELS_SUPPORTING_REASONING_EFFORT)} instead." + ) + + +def validate_search_context_size(*, model: str, search_context_size: object) -> None: + if _is_omitted(search_context_size): + return + if not model_supports_search_context_size(model): + raise UnsupportedParameterError( + f"Model {model!r} does not support search_context_size. " + f"Use one of {sorted(MODELS_SUPPORTING_SEARCH_CONTEXT_SIZE)} instead." + ) + + +def validate_search_domain_filter(*, model: str, search_domain_filter: object) -> None: + if _is_omitted(search_domain_filter): + return + if not isinstance(search_domain_filter, Iterable) or isinstance(search_domain_filter, (str, bytes)): + raise UnsupportedParameterError("search_domain_filter must be a sequence of domain strings") + domains = list(search_domain_filter) + if len(domains) > MAX_SEARCH_DOMAIN_FILTERS: + raise UnsupportedParameterError( + f"search_domain_filter accepts at most {MAX_SEARCH_DOMAIN_FILTERS} domains; got {len(domains)}" + ) + + +def validate_chat_completion_params( + *, + model: str, + return_images: object = omit, + reasoning_effort: object = omit, + search_context_size: object = omit, + search_domain_filter: object = omit, +) -> None: + validate_return_images(model=model, return_images=return_images) + validate_reasoning_effort(model=model, reasoning_effort=reasoning_effort) + validate_search_context_size(model=model, search_context_size=search_context_size) + validate_search_domain_filter(model=model, search_domain_filter=search_domain_filter) diff --git a/src/perplexity/resources/chat/completions.py b/src/perplexity/resources/chat/completions.py index 5931b13..216e009 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_chat_completion_params 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,19 @@ def create( extra_body: Body | None = None, timeout: float | httpx.Timeout | None | NotGiven = not_given, ) -> StreamChunk | Stream[StreamChunk]: + search_context_size = omit + if web_search_options is not omit and web_search_options is not None: + if isinstance(web_search_options, dict): + search_context_size = web_search_options.get("search_context_size", omit) + else: + search_context_size = getattr(web_search_options, "search_context_size", omit) + validate_chat_completion_params( + model=model, + return_images=return_images, + reasoning_effort=reasoning_effort, + search_context_size=search_context_size, + search_domain_filter=search_domain_filter, + ) return self._post( "/chat/completions", body=maybe_transform( @@ -788,6 +802,19 @@ async def create( extra_body: Body | None = None, timeout: float | httpx.Timeout | None | NotGiven = not_given, ) -> StreamChunk | AsyncStream[StreamChunk]: + search_context_size = omit + if web_search_options is not omit and web_search_options is not None: + if isinstance(web_search_options, dict): + search_context_size = web_search_options.get("search_context_size", omit) + else: + search_context_size = getattr(web_search_options, "search_context_size", omit) + validate_chat_completion_params( + model=model, + return_images=return_images, + reasoning_effort=reasoning_effort, + search_context_size=search_context_size, + search_domain_filter=search_domain_filter, + ) 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..3176d9e --- /dev/null +++ b/tests/test_model_capabilities.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import pytest + +from perplexity.lib.model_capabilities import ( + ReturnImagesUnsupportedError, + UnsupportedParameterError, + model_supports_reasoning_effort, + model_supports_search_context_size, + validate_chat_completion_params, + validate_return_images, + validate_search_domain_filter, +) + + +def test_sonar_rejects_reasoning_effort() -> None: + with pytest.raises(UnsupportedParameterError, match="reasoning_effort"): + validate_chat_completion_params(model="sonar", reasoning_effort="high") + + +def test_sonar_deep_research_allows_reasoning_effort() -> None: + validate_chat_completion_params(model="sonar-deep-research", reasoning_effort="high") + + +def test_sonar_rejects_search_context_size() -> None: + with pytest.raises(UnsupportedParameterError, match="search_context_size"): + validate_chat_completion_params(model="sonar", search_context_size="high") + + +def test_sonar_pro_allows_search_context_size() -> None: + validate_chat_completion_params(model="sonar-pro", search_context_size="medium") + + +def test_search_domain_filter_rejects_too_many_domains() -> None: + with pytest.raises(UnsupportedParameterError, match="at most 3"): + validate_search_domain_filter( + model="sonar-pro", + search_domain_filter=["a.com", "b.com", "c.com", "d.com"], + ) + + +def test_model_supports_flags() -> None: + assert model_supports_reasoning_effort("sonar-reasoning-pro") + assert not model_supports_reasoning_effort("sonar") + assert model_supports_search_context_size("sonar-pro") + assert not model_supports_search_context_size("sonar") + + +def test_return_images_still_raises_for_sonar() -> None: + with pytest.raises(ReturnImagesUnsupportedError): + validate_return_images(model="sonar", return_images=True) + + +def test_sonar_pro_supports_return_images() -> None: + from perplexity.lib.model_capabilities import model_supports_return_images + + assert model_supports_return_images("sonar-pro") + + +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"]