fix: expand chat completion capability validation (#64 follow-up)

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,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)

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(

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,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"]