data-designer-engine 0.9.4__tar.gz → 0.9.5__tar.gz

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: data-designer-engine
-Version: 0.9.4
+Version: 0.9.5
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -15,14 +15,11 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Requires-Dist: anyascii<1,>=0.3.3
-Requires-Dist: anyio<5,>=4.14.2
 Requires-Dist: asciichartpy<2,>=1.5.25
-Requires-Dist: chardet<6,>=3.0.2
-Requires-Dist: cryptography<51,>=50.0.0
-Requires-Dist: data-designer-config==0.9.4
+Requires-Dist: data-designer-config==0.9.5
 Requires-Dist: duckdb<2,>=1.5.0
 Requires-Dist: faker<21,>=20.1.0
-Requires-Dist: fsspec<2026,>=2025.3.0
+Requires-Dist: fsspec<2027,>=2026.6.0
 Requires-Dist: httpx-retries<1,>=0.4.2
 Requires-Dist: httpx<1,>=0.27.2
 Requires-Dist: huggingface-hub<2,>=1.0.1
@@ -30,7 +27,7 @@ Requires-Dist: jinja2<4,>=3.1.6
 Requires-Dist: json-repair<1,>=0.48.0
 Requires-Dist: jsonpath-rust-bindings<2,>=1.0
 Requires-Dist: jsonschema<5,>=4.0.0
-Requires-Dist: lxml<7,>=6.1.0
+Requires-Dist: lxml<7,>=6.1.1
 Requires-Dist: marko<3,>=2.1.2
 Requires-Dist: mcp<2,>=1.29.0
 Requires-Dist: networkx<4,>=3.0
@@ -38,13 +35,10 @@ Requires-Dist: numpy<3,>=1.23.5
 Requires-Dist: pandas<4,>=2.3.3
 Requires-Dist: pyarrow<25,>=24
 Requires-Dist: pydantic<3,>=2.9.2
-Requires-Dist: pyjwt[crypto]<3,>=2.15.0
-Requires-Dist: python-multipart<1,>=0.0.32
 Requires-Dist: rich<15,>=13.7.1
 Requires-Dist: ruff<1,>=0.14.10
 Requires-Dist: scipy<2,>=1.11.0
 Requires-Dist: sqlfluff<5,>=4.1.0
-Requires-Dist: starlette<2,>=1.2.0
 Requires-Dist: tiktoken<1,>=0.8.0
 Requires-Dist: wcwidth<1,>=0.2.13
 Description-Content-Type: text/markdown

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/pyproject.toml

@@ -34,14 +34,11 @@ bump = true
 [tool.hatch.metadata.hooks.uv-dynamic-versioning]
 dependencies = [
     "anyascii>=0.3.3,<1",
-    "anyio>=4.14.2,<5", # 4.14.2 fixes AnyIO security advisories pulled in by HTTP clients
     "asciichartpy>=1.5.25,<2",
-    "chardet>=3.0.2,<6",  # Pulled in by sqlfluff
-    "cryptography>=50.0.0,<51", # 50.0.0 fixes certificate and PKCS#7 security advisories
     "data-designer-config=={{ version }}",
     "duckdb>=1.5.0,<2",
     "faker>=20.1.0,<21",
-    "fsspec>=2025.3.0,<2026",
+    "fsspec>=2026.6.0,<2027", # 2026.6.0 fixes CVE-2026-104851
     "httpx>=0.27.2,<1",
     "httpx-retries>=0.4.2,<1",
     "huggingface-hub>=1.0.1,<2",
@@ -49,7 +46,7 @@ dependencies = [
     "json-repair>=0.48.0,<1",
     "jsonpath-rust-bindings>=1.0,<2",
     "jsonschema>=4.0.0,<5",
-    "lxml>=6.1.0,<7",
+    "lxml>=6.1.1,<7", # 6.1.1 fixes CVE-2026-49825
     "marko>=2.1.2,<3",
     "mcp>=1.29.0,<2", # 1.29.0 fixes MCP security advisories
     "networkx>=3.0,<4",
@@ -57,13 +54,10 @@ dependencies = [
     "pandas>=2.3.3,<4",
     "pyarrow>=24,<25", # 24.0.0 fixes PyArrow security advisories
     "pydantic>=2.9.2,<3",
-    "pyjwt[crypto]>=2.15.0,<3", # 2.15.0 fixes PyJWT security advisories pulled in by mcp
-    "python-multipart>=0.0.32,<1", # 0.0.32 addresses multipart security advisories pulled in by mcp
     "rich>=13.7.1,<15",
     "ruff>=0.14.10,<1",
     "scipy>=1.11.0,<2",
     "sqlfluff>=4.1.0,<5",
-    "starlette>=1.2.0,<2", # 1.2.0 fixes security advisory pulled in by mcp
     "tiktoken>=0.8.0,<1",
     "wcwidth>=0.2.13,<1",
 ]

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/src/data_designer/engine/models/clients/adapters/anthropic.py

@@ -45,6 +45,7 @@ class AnthropicClient(HttpModelClient):
             "stop",
             "max_tokens",
             "tools",
+            "tool_choice",
             "n",
             "response_format",
             "frequency_penalty",

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/src/data_designer/engine/models/clients/adapters/anthropic_translation.py

@@ -33,6 +33,12 @@ _UNSUPPORTED_MEDIA_BLOCK_MODALITIES: dict[str, str] = {
     "video_url": "video",
     "input_video": "video",
 }
+_OPENAI_TOOL_CHOICE_TO_ANTHROPIC_TYPE: dict[str, str] = {
+    "auto": "auto",
+    "required": "any",
+    "none": "none",
+}
+_ANTHROPIC_TOOL_CHOICE_TYPES = frozenset({"auto", "any", "none", "tool"})
 
 
 class UnsupportedAnthropicMediaBlockError(ValueError):
@@ -77,6 +83,9 @@ def build_anthropic_payload(request: ChatCompletionRequest) -> dict[str, Any]:
     if request.tools:
         payload["tools"] = [translate_tool_definition(tool) for tool in request.tools]
 
+    if request.tool_choice is not None:
+        payload["tool_choice"] = translate_tool_choice(request.tool_choice)
+
     if request.stop is not None:
         if isinstance(request.stop, str):
             payload["stop_sequences"] = [request.stop]
@@ -276,6 +285,43 @@ def translate_tool_definition(tool: dict[str, Any]) -> dict[str, Any]:
     return translated_tool
 
 
+def translate_tool_choice(tool_choice: Any) -> dict[str, Any]:
+    """Translate an OpenAI-style ``tool_choice`` into Anthropic's format.
+
+    Anthropic-native values (``{"type": "auto" | "any" | "none"}`` or
+    ``{"type": "tool", "name": ...}``) pass through unchanged.
+
+    Raises:
+        ValueError: If the value has no Anthropic equivalent.
+    """
+    if isinstance(tool_choice, str):
+        choice_type = _OPENAI_TOOL_CHOICE_TO_ANTHROPIC_TYPE.get(tool_choice)
+        if choice_type is None:
+            raise ValueError(
+                f"Anthropic tool_choice must be 'auto', 'required', 'none', or a named function, got: {tool_choice!r}"
+            )
+        return {"type": choice_type}
+
+    if not isinstance(tool_choice, dict):
+        raise ValueError(f"Anthropic tool_choice must be a string or an object, got: {tool_choice!r}")
+
+    choice_type = tool_choice.get("type")
+    if choice_type == "function":
+        function = tool_choice.get("function")
+        name = function.get("name") if isinstance(function, dict) else None
+        if not isinstance(name, str) or not name:
+            raise ValueError(f"Anthropic tool_choice is missing a function name, got: {tool_choice!r}")
+        return {"type": "tool", "name": name}
+
+    if not isinstance(choice_type, str) or choice_type not in _ANTHROPIC_TOOL_CHOICE_TYPES:
+        raise ValueError(f"Anthropic tool_choice has an unsupported type, got: {tool_choice!r}")
+    if choice_type == "tool":
+        name = tool_choice.get("name")
+        if not isinstance(name, str) or not name:
+            raise ValueError(f"Anthropic tool_choice is missing a tool name, got: {tool_choice!r}")
+    return dict(tool_choice)
+
+
 def translate_tool_calls(tool_calls: Any) -> list[dict[str, Any]]:
     if not isinstance(tool_calls, list):
         raise ValueError(f"Anthropic assistant tool calls must be a list, got: {tool_calls!r}")

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/src/data_designer/engine/models/clients/errors.py

@@ -8,6 +8,7 @@ import email.utils
 import time
 from enum import Enum
 
+import data_designer.lazy_heavy_imports as lazy
 from data_designer.engine.models.clients.types import HttpResponse
 
 
@@ -204,11 +205,18 @@ def _parse_http_date_as_delay(value: str) -> float | None:
 
 
 def infer_error_kind_from_exception(exc: Exception) -> ProviderErrorKind:
-    """Infer a ``ProviderErrorKind`` from an exception's type name.
+    """Infer a ``ProviderErrorKind`` from HTTPX types or an exception's type name.
 
     Used by adapters to classify transport-level exceptions (timeouts,
     connection failures, etc.) that don't carry an HTTP status code.
     """
+    # Include stale keep-alive failures so the request executor can retry them.
+    # These can also occur after a POST was processed. We accept possible duplicate
+    # generation or billing to recover from transient failures, matching the
+    # existing transport retry policy.
+    # Local protocol errors are client mistakes and should not be retried.
+    if isinstance(exc, (lazy.httpx.NetworkError, lazy.httpx.RemoteProtocolError)):
+        return ProviderErrorKind.API_CONNECTION
     type_name = type(exc).__name__.lower()
     if "timeout" in type_name:
         return ProviderErrorKind.TIMEOUT

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/src/data_designer/engine/models/clients/types.py

@@ -61,6 +61,7 @@ class ChatCompletionRequest:
     model: str
     messages: list[dict[str, Any]]
     tools: list[dict[str, Any]] | None = None
+    tool_choice: str | dict[str, Any] | None = None
     n: int | None = None
     temperature: float | None = None
     top_p: float | None = None

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/src/data_designer/engine/models/facade.py

@@ -91,10 +91,11 @@ def _build_generation_validation_error(
 
 # Known keyword arguments extracted into request fields for each modality.
 # Note: `extra_body` and `extra_headers` appear in every set but receive special
-# treatment in `consolidate_kwargs` (merged with provider-level overrides) and in
-# `TransportKwargs` (extra_body is either flattened into the request body or
-# preserved as a nested dict depending on the adapter; extra_headers are
-# forwarded as HTTP headers).  They are NOT regular model parameters.
+# treatment in `consolidate_kwargs` (`extra_body` is merged key by key: model
+# config, then the call, then the provider) and in `TransportKwargs`
+# (extra_body keys are merged into the request body and take precedence over
+# same-named request fields; extra_headers are forwarded as HTTP headers).
+# They are NOT regular model parameters.
 _COMPLETION_REQUEST_FIELDS = frozenset(
     {
         "temperature",
@@ -108,6 +109,7 @@ _COMPLETION_REQUEST_FIELDS = frozenset(
         "presence_penalty",
         "timeout",
         "tools",
+        "tool_choice",
         "extra_body",
         "extra_headers",
     }
@@ -146,6 +148,7 @@ class ModelFacade:
         self._client = client
         self._mcp_registry = mcp_registry
         self._usage_stats = ModelUsageStats()
+        self._warned_dropped_kwargs: set[frozenset[str]] = set()
 
     @property
     def model_name(self) -> str:
@@ -187,11 +190,23 @@ class ModelFacade:
         return self._usage_stats
 
     def consolidate_kwargs(self, **kwargs: Any) -> dict[str, Any]:
+        """Combine the model's inference parameters, the call's kwargs, and the provider's extras.
+
+        Call kwargs override inference parameters. ``extra_body`` is merged key by key in the order
+        model config, call, provider, so a later source wins only for the keys it sets; nested values
+        are replaced, not merged. Provider ``extra_headers`` override the call's.
+        """
         # Remove purpose from kwargs to avoid passing it to the model
         kwargs.pop("purpose", None)
-        kwargs = {**self._model_config.inference_parameters.generate_kwargs, **kwargs}
-        if self.model_provider.extra_body:
-            kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
+        inference_kwargs = self._model_config.inference_parameters.generate_kwargs
+        extra_body = {
+            **(inference_kwargs.get("extra_body") or {}),
+            **(kwargs.get("extra_body") or {}),
+            **(self.model_provider.extra_body or {}),
+        }
+        kwargs = {**inference_kwargs, **kwargs}
+        if extra_body:
+            kwargs["extra_body"] = extra_body
         if self.model_provider.extra_headers:
             kwargs["extra_headers"] = {**(kwargs.get("extra_headers") or {}), **self.model_provider.extra_headers}
         # Inject framework-level attribution header when telemetry is enabled.
@@ -337,7 +352,10 @@ class ModelFacade:
             skip_usage_tracking (bool): Whether to skip usage tracking. Default: `False`.
             purpose (str): The purpose of the model usage to show as context in the error message.
                 It is expected to be used by the @catch_llm_exceptions decorator.
-            **kwargs: Additional arguments to pass to the model.
+            **kwargs: Request arguments sent with every completion, such as ``temperature``,
+                ``tool_choice``, or ``extra_body`` (merged as described in ``consolidate_kwargs``).
+                With ``tool_alias``, ``tool_choice`` applies only until the first tool call; a
+                ``tool_choice`` inside ``extra_body`` is sent unchanged on every round.
 
         Returns:
             A tuple containing:
@@ -380,6 +398,10 @@ class ModelFacade:
             completion_kwargs.pop("allow_multiple_choices", None)
             if tool_schemas is not None:
                 completion_kwargs["tools"] = tool_schemas
+            if tool_call_turns > 0:
+                # A forced tool_choice holds only until the first tool call. Repeating it would make every
+                # response a tool call, so the tool loop could never end.
+                completion_kwargs.pop("tool_choice", None)
 
             completion_response = self.completion(
                 messages,
@@ -497,6 +519,10 @@ class ModelFacade:
             completion_kwargs.pop("allow_multiple_choices", None)
             if tool_schemas is not None:
                 completion_kwargs["tools"] = tool_schemas
+            if tool_call_turns > 0:
+                # A forced tool_choice holds only until the first tool call. Repeating it would make every
+                # response a tool call, so the tool loop could never end.
+                completion_kwargs.pop("tool_choice", None)
 
             completion_response = await self.acompletion(
                 messages,
@@ -814,10 +840,15 @@ class ModelFacade:
                 metadata[key] = value
 
         if metadata:
-            logger.debug(
-                "Unknown kwargs %s dropped (not forwarded as model parameters). "
+            # Warn once per set of names; repeats go to DEBUG so a per-record call can't flood the log.
+            dropped = frozenset(metadata)
+            log = logger.debug if dropped in self._warned_dropped_kwargs else logger.warning
+            self._warned_dropped_kwargs.add(dropped)
+            log(
+                "Unknown kwargs %s dropped for model %r (not forwarded as model parameters). "
                 "Use 'extra_body' to pass non-standard parameters to the model.",
-                metadata.keys(),
+                sorted(dropped),
+                self.model_alias,
             )
 
         return ChatCompletionRequest(**request_fields)

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/tests/engine/models/clients/test_anthropic.py

@@ -267,6 +267,57 @@ def test_completion_translates_openai_tool_schema_to_anthropic() -> None:
             "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
         }
     ]
+    assert "tool_choice" not in payload
+
+
+def test_completion_translates_tool_choice_to_anthropic() -> None:
+    sync_mock = make_mock_sync_client(_text_response())
+    client = _make_client(sync_client=sync_mock)
+
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[{"role": "user", "content": "Hi"}],
+        tools=[{"type": "function", "function": {"name": "search"}}],
+        tool_choice={"type": "function", "function": {"name": "search"}},
+    )
+    client.completion(request)
+
+    payload = sync_mock.post.call_args.kwargs["json"]
+    assert payload["tool_choice"] == {"type": "tool", "name": "search"}
+
+
+@pytest.mark.asyncio
+async def test_acompletion_translates_tool_choice_to_anthropic() -> None:
+    async_mock = make_mock_async_client(_text_response())
+    client = _make_client(async_client=async_mock)
+
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[{"role": "user", "content": "Hi"}],
+        tools=[{"type": "function", "function": {"name": "search"}}],
+        tool_choice="required",
+    )
+    await client.acompletion(request)
+
+    payload = async_mock.post.call_args.kwargs["json"]
+    assert payload["tool_choice"] == {"type": "any"}
+
+
+def test_completion_extra_body_tool_choice_overrides_translated_value() -> None:
+    sync_mock = make_mock_sync_client(_text_response())
+    client = _make_client(sync_client=sync_mock)
+
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[{"role": "user", "content": "Hi"}],
+        tools=[{"type": "function", "function": {"name": "search"}}],
+        tool_choice="required",
+        extra_body={"tool_choice": {"type": "auto", "disable_parallel_tool_use": True}},
+    )
+    client.completion(request)
+
+    payload = sync_mock.post.call_args.kwargs["json"]
+    assert payload["tool_choice"] == {"type": "auto", "disable_parallel_tool_use": True}
 
 
 def test_completion_translates_tool_turns_from_chat_messages() -> None:
@@ -666,6 +717,30 @@ def test_http_error_maps_to_provider_error(status_code: int, expected_kind: Prov
     assert exc_info.value.kind == expected_kind
 
 
+def test_http_400_tool_choice_rejection_maps_to_unsupported_params() -> None:
+    error_json = {
+        "error": {
+            "type": "invalid_request_error",
+            "message": 'tool_choice: type "tool" and "any" are not supported for this model.',
+        }
+    }
+    sync_mock = make_mock_sync_client(error_json, status_code=400)
+    client = _make_client(sync_client=sync_mock)
+
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[{"role": "user", "content": "Hi"}],
+        tools=[{"type": "function", "function": {"name": "search"}}],
+        tool_choice="required",
+    )
+    with pytest.raises(ProviderError) as exc_info:
+        client.completion(request)
+
+    assert exc_info.value.kind == ProviderErrorKind.UNSUPPORTED_PARAMS
+    assert "not supported for this model" in exc_info.value.message
+    sync_mock.post.assert_called_once()
+
+
 def test_transport_timeout_raises_provider_error() -> None:
     sync_mock = MagicMock()
     sync_mock.post = MagicMock(side_effect=TimeoutError("timed out"))
@@ -721,6 +796,28 @@ def test_completion_wraps_invalid_tool_schema_as_bad_request() -> None:
     assert "missing a function name" in exc_info.value.message
 
 
+@pytest.mark.parametrize(
+    "tool_choice",
+    [
+        pytest.param("always", id="unknown-string"),
+        pytest.param({"type": []}, id="malformed-type"),
+    ],
+)
+def test_completion_wraps_unsupported_tool_choice_as_bad_request(tool_choice: object) -> None:
+    client = _make_client()
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[{"role": "user", "content": "Hi"}],
+        tool_choice=tool_choice,
+    )
+
+    with pytest.raises(ProviderError) as exc_info:
+        client.completion(request)
+
+    assert exc_info.value.kind == ProviderErrorKind.BAD_REQUEST
+    assert "tool_choice" in exc_info.value.message
+
+
 # --- Unsupported capabilities ---
 
 

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/tests/engine/models/clients/test_anthropic_translation.py

@@ -24,6 +24,7 @@ from data_designer.engine.models.clients.adapters.anthropic_translation import (
     translate_request_messages,
     translate_tool_call,
     translate_tool_calls,
+    translate_tool_choice,
     translate_tool_definition,
     translate_tool_result_content,
     translate_tool_result_message,
@@ -345,6 +346,56 @@ def test_translate_tool_definition_normalizes_supported_shapes(
     assert translate_tool_definition(tool) == expected
 
 
+@pytest.mark.parametrize(
+    ("tool_choice", "expected"),
+    [
+        pytest.param("auto", {"type": "auto"}, id="openai-auto"),
+        pytest.param("required", {"type": "any"}, id="openai-required"),
+        pytest.param("none", {"type": "none"}, id="openai-none"),
+        pytest.param(
+            {"type": "function", "function": {"name": "search"}},
+            {"type": "tool", "name": "search"},
+            id="openai-named-function",
+        ),
+        pytest.param(
+            {"type": "tool", "name": "search"},
+            {"type": "tool", "name": "search"},
+            id="anthropic-named-tool",
+        ),
+        pytest.param(
+            {"type": "any", "disable_parallel_tool_use": True},
+            {"type": "any", "disable_parallel_tool_use": True},
+            id="anthropic-any",
+        ),
+    ],
+)
+def test_translate_tool_choice_normalizes_supported_shapes(
+    tool_choice: object,
+    expected: dict[str, object],
+) -> None:
+    translated = translate_tool_choice(tool_choice)
+
+    assert translated == expected
+    assert translated is not tool_choice
+
+
+@pytest.mark.parametrize(
+    ("tool_choice", "match"),
+    [
+        pytest.param("always", "must be 'auto', 'required', 'none', or a named function", id="unknown-string"),
+        pytest.param({"type": "function", "function": {}}, "missing a function name", id="function-without-name"),
+        pytest.param({"type": "tool"}, "missing a tool name", id="tool-without-name"),
+        pytest.param({"type": "bogus"}, "unsupported type", id="unknown-type"),
+        pytest.param({"type": []}, "unsupported type", id="list-type"),
+        pytest.param({"type": {}}, "unsupported type", id="dict-type"),
+        pytest.param(1, "must be a string or an object", id="wrong-type"),
+    ],
+)
+def test_translate_tool_choice_rejects_unsupported_values(tool_choice: object, match: str) -> None:
+    with pytest.raises(ValueError, match=match):
+        translate_tool_choice(tool_choice)
+
+
 def test_translate_content_blocks_converts_images_and_preserves_other_blocks() -> None:
     blocks = translate_content_blocks(
         [

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/tests/engine/models/clients/test_client_errors.py

@@ -5,16 +5,49 @@ from __future__ import annotations
 
 from typing import Any
 
+import httpx
 import pytest
 
 from data_designer.engine.models.clients.errors import (
     ProviderError,
     ProviderErrorKind,
+    infer_error_kind_from_exception,
     map_http_error_to_provider_error,
     map_http_status_to_provider_error_kind,
 )
 
 
+class PeerDropped(httpx.ReadError):
+    pass
+
+
+@pytest.mark.parametrize(
+    "exception_type,expected_kind",
+    [
+        (httpx.NetworkError, ProviderErrorKind.API_CONNECTION),
+        (httpx.ConnectError, ProviderErrorKind.API_CONNECTION),
+        (httpx.ReadError, ProviderErrorKind.API_CONNECTION),
+        (httpx.WriteError, ProviderErrorKind.API_CONNECTION),
+        (httpx.CloseError, ProviderErrorKind.API_CONNECTION),
+        (httpx.RemoteProtocolError, ProviderErrorKind.API_CONNECTION),
+        (PeerDropped, ProviderErrorKind.API_CONNECTION),
+        (httpx.ConnectTimeout, ProviderErrorKind.TIMEOUT),
+        (httpx.ReadTimeout, ProviderErrorKind.TIMEOUT),
+        (httpx.WriteTimeout, ProviderErrorKind.TIMEOUT),
+        (httpx.PoolTimeout, ProviderErrorKind.TIMEOUT),
+        (httpx.LocalProtocolError, ProviderErrorKind.API_ERROR),
+        (httpx.UnsupportedProtocol, ProviderErrorKind.API_ERROR),
+        (httpx.DecodingError, ProviderErrorKind.API_ERROR),
+        (ConnectionError, ProviderErrorKind.API_CONNECTION),
+        (TimeoutError, ProviderErrorKind.TIMEOUT),
+        (RuntimeError, ProviderErrorKind.API_ERROR),
+    ],
+    ids=lambda value: value.__name__ if isinstance(value, type) else value.value,
+)
+def test_infer_error_kind_from_exception(exception_type: type[Exception], expected_kind: ProviderErrorKind) -> None:
+    assert infer_error_kind_from_exception(exception_type("Server disconnected")) == expected_kind
+
+
 class StubHttpResponse:
     def __init__(
         self,

{data_designer_engine-0.9.4 → data_designer_engine-0.9.5}/tests/engine/models/clients/test_factory.py

@@ -3,9 +3,12 @@
 
 from __future__ import annotations
 
+import asyncio
 from unittest.mock import MagicMock
 
+import httpx
 import pytest
+from pytest_httpx import HTTPXMock
 
 from data_designer.config.models import (
     ChatCompletionInferenceParams,
@@ -17,9 +20,11 @@ from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.clients.adapters.anthropic import AnthropicClient
 from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode
 from data_designer.engine.models.clients.adapters.openai_compatible import OpenAICompatibleClient
+from data_designer.engine.models.clients.errors import ProviderError, ProviderErrorKind
 from data_designer.engine.models.clients.factory import create_model_client
 from data_designer.engine.models.clients.model_request_executor import ModelRequestExecutor
 from data_designer.engine.models.clients.retry import RetryConfig
+from data_designer.engine.models.clients.types import ChatCompletionRequest
 from data_designer.engine.models.request_admission.controller import AdaptiveRequestAdmissionController
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.engine.testing import InMemoryAdmissionEventSink
@@ -247,3 +252,98 @@ def test_no_request_admission_returns_inner_client_directly(
     client = create_model_client(openai_model_config, secret_resolver, openai_registry)
     assert isinstance(client, OpenAICompatibleClient)
     assert not isinstance(client, ModelRequestExecutor)
+
+
+@pytest.mark.asyncio(loop_scope="session")
+@pytest.mark.parametrize("concurrency_mode", [ClientConcurrencyMode.SYNC, ClientConcurrencyMode.ASYNC])
+@pytest.mark.parametrize("exception_type", [httpx.RemoteProtocolError, httpx.ReadError, httpx.WriteError])
+async def test_request_admission_retries_transport_errors(
+    openai_model_config: ModelConfig,
+    secret_resolver: SecretResolver,
+    openai_registry: ModelProviderRegistry,
+    httpx_mock: HTTPXMock,
+    concurrency_mode: ClientConcurrencyMode,
+    exception_type: type[httpx.TransportError],
+) -> None:
+    sink = InMemoryAdmissionEventSink()
+    controller = AdaptiveRequestAdmissionController(event_sink=sink)
+    controller.register(provider_name="openai-prod", model_id="gpt-test", alias="test-model", max_parallel_requests=1)
+    client = create_model_client(
+        openai_model_config,
+        secret_resolver,
+        openai_registry,
+        retry_config=RetryConfig(max_retries=1, backoff_factor=0.0),
+        client_concurrency_mode=concurrency_mode,
+        request_admission=controller,
+        request_event_sink=sink,
+    )
+    httpx_mock.add_exception(exception_type("Server disconnected without sending a response."))
+    httpx_mock.add_response(json={"choices": [{"message": {"role": "assistant", "content": "ok"}}]})
+    request = ChatCompletionRequest(model="gpt-test", messages=[])
+
+    try:
+        response = (
+            await client.acompletion(request)
+            if concurrency_mode == ClientConcurrencyMode.ASYNC
+            else await asyncio.to_thread(client.completion, request)
+        )
+    finally:
+        await client.aclose()
+        client.close()
+
+    assert response.message.content == "ok"
+    assert len(httpx_mock.get_requests()) == 2
+    acquired = [event for event in sink.request_events if event.event_kind == "request_lease_acquired"]
+    released = [event for event in sink.request_events if event.event_kind == "request_lease_released"]
+    assert len(acquired) == len(released) == 2
+    assert len({event.request_lease_id for event in acquired}) == 2
+    assert {event.request_lease_id for event in acquired} == {event.request_lease_id for event in released}
+    completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"]
+    assert [event.diagnostics["outcome"] for event in completed] == ["api_connection", "success"]
+
+
+@pytest.mark.asyncio(loop_scope="session")
+@pytest.mark.parametrize("concurrency_mode", [ClientConcurrencyMode.SYNC, ClientConcurrencyMode.ASYNC])
+@pytest.mark.parametrize("exception_type", [httpx.RemoteProtocolError, httpx.ReadError, httpx.WriteError])
+@pytest.mark.parametrize("max_retries", [0, 2])
+async def test_request_admission_transport_errors_respect_retry_budget(
+    openai_model_config: ModelConfig,
+    secret_resolver: SecretResolver,
+    openai_registry: ModelProviderRegistry,
+    httpx_mock: HTTPXMock,
+    concurrency_mode: ClientConcurrencyMode,
+    exception_type: type[httpx.TransportError],
+    max_retries: int,
+) -> None:
+    controller = AdaptiveRequestAdmissionController()
+    controller.register(provider_name="openai-prod", model_id="gpt-test", alias="test-model", max_parallel_requests=1)
+    client = create_model_client(
+        openai_model_config,
+        secret_resolver,
+        openai_registry,
+        retry_config=RetryConfig(max_retries=max_retries, backoff_factor=0.0),
+        client_concurrency_mode=concurrency_mode,
+        request_admission=controller,
+    )
+    error = exception_type("Server disconnected")
+    httpx_mock.add_exception(error, is_reusable=True)
+    request = ChatCompletionRequest(model="gpt-test", messages=[])
+
+    try:
+        with pytest.raises(ProviderError) as exc_info:
+            if concurrency_mode == ClientConcurrencyMode.ASYNC:
+                await client.acompletion(request)
+            else:
+                await asyncio.to_thread(client.completion, request)
+    finally:
+        await client.aclose()
+        client.close()
+
+    assert exc_info.value.kind == ProviderErrorKind.API_CONNECTION
+    assert exc_info.value.__cause__ is error
+    assert exc_info.value.provider_name == "openai-prod"
+    assert exc_info.value.model_name == "gpt-test"
+    assert len(httpx_mock.get_requests()) == max_retries + 1
+    snapshot = controller.pressure.global_snapshot("openai-prod", "gpt-test")
+    assert snapshot is not None
+    assert snapshot.aggregate_in_flight == 0