refactor: Decouple ModelFacade from LiteLLM via ModelClient adapter (#373)

* plans for model facade overhaul

* update plan

* add review

* address feedback + add more details after several self reviews

* update plan doc

* address nits

* Add cannonical objects

* self-review feedback + address

* add LiteLLMRouter protocol to strongly type bridge router param

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* simplify some things

* add a protol for http response like object

* move HttpResponse

* update PR-1 architecture notes for lifecycle and router protocol

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Address PR #359 feedback: exception wrapping, shared parsing, test improvements

- Wrap all LiteLLM router calls in try/except to normalize raw exceptions
  into canonical ProviderError at the bridge boundary (blocking review item)
- Extract reusable response-parsing helpers into clients/parsing.py for
  shared use across future native adapters
- Add async image parsing path using httpx.AsyncClient to avoid blocking
  the event loop in agenerate_image
- Add retry_after field to ProviderError for future retry engine support
- Fix _to_int_or_none to parse numeric strings from providers
- Create test conftest.py with shared mock_router/bridge_client fixtures
- Parametrize duplicate image generation and error mapping tests
- Add tests for exception wrapping across all bridge methods

* Use contextlib to dry out some code

* Address Greptile feedback: HTTP-date retry-after parsing, docstring clarity

- Parse RFC 7231 HTTP-date strings in Retry-After header (used by
  Azure and Anthropic during rate-limiting) in addition to numeric
  delay-seconds
- Clarify collect_non_none_optional_fields docstring explaining why
  f.default is None is the correct check for optional field forwarding
- Add tests for HTTP-date and garbage Retry-After values

* Address Greptile feedback: FastAPI detail parsing, comment fixes

- Fix misleading comment about prompt field defaults in _IMAGE_EXCLUDE
- Handle list-format detail arrays in _extract_structured_message for
  FastAPI/Pydantic validation errors
- Document scope boundary for vision content in collect_raw_image_candidates

* add PR-2 architecture notes for model facade overhaul

* save progress on pr2

* small refactor

* address feedback

* Address greptile comment in pr1

* refactor ProviderError from dataclass to regular Exception

- Replace @DataClass + __post_init__ with explicit __init__ that calls
  super().__init__ properly, avoiding brittle field-ordering dependency
- Store cause via __cause__ only, removing the redundant .cause attr
- Update match pattern in handle_llm_exceptions for non-dataclass type
- Rename shadowed local `fields` to `optional_fields` in TransportKwargs

* Address greptile feedback

* PR feedback

* track usage tracking in finally block for images

* pr feedback

* wrap facade close in try/catch

* clean up stray params

* fix stray inclusion of metadata

* small regression fix

* address more feedback

---------

Co-authored-by: Johnny Greco <jogreco@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

.gitignore

@@ -107,3 +107,6 @@ packages/data-designer/README.md
 .cursor/rules/cerebro.mdc
 .cursor/mcp.json
 .claude/rules/cerebro.md
+
+# Claude worktrees
+.claude/worktrees/

packages/data-designer-engine/src/data_designer/engine/mcp/facade.py

@@ -10,6 +10,7 @@
     map_http_error_to_provider_error,
     map_http_status_to_provider_error_kind,
 )
+from data_designer.engine.models.clients.factory import create_model_client
 from data_designer.engine.models.clients.types import (
     AssistantMessage,
     ChatCompletionRequest,
@@ -25,12 +26,12 @@
 )
 
 __all__ = [
-    "HttpResponse",
     "AssistantMessage",
     "ChatCompletionRequest",
     "ChatCompletionResponse",
     "EmbeddingRequest",
     "EmbeddingResponse",
+    "HttpResponse",
     "ImageGenerationRequest",
     "ImageGenerationResponse",
     "ImagePayload",
@@ -39,6 +40,7 @@
     "ProviderErrorKind",
     "ToolCall",
     "Usage",
+    "create_model_client",
     "map_http_error_to_provider_error",
     "map_http_status_to_provider_error_kind",
 ]

packages/data-designer-engine/src/data_designer/engine/models/clients/__init__.py

@@ -12,13 +12,13 @@
 from data_designer.engine.models.clients.errors import (
     ProviderError,
     ProviderErrorKind,
+    extract_message_from_exception_string,
     map_http_status_to_provider_error_kind,
 )
 from data_designer.engine.models.clients.parsing import (
     aextract_images_from_chat_response,
     aextract_images_from_image_response,
     aparse_chat_completion_response,
-    collect_non_none_optional_fields,
     extract_embedding_vector,
     extract_images_from_chat_response,
     extract_images_from_image_response,
@@ -32,6 +32,7 @@
     EmbeddingResponse,
     ImageGenerationRequest,
     ImageGenerationResponse,
+    TransportKwargs,
 )
 
 logger = logging.getLogger(__name__)
@@ -75,57 +76,67 @@ def supports_image_generation(self) -> bool:
         return True
 
     def completion(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
+        transport = TransportKwargs.from_request(request)
         with _handle_non_provider_errors(self.provider_name):
             response = self._router.completion(
                 model=request.model,
                 messages=request.messages,
-                **collect_non_none_optional_fields(request),
+                extra_headers=transport.headers or None,
+                **transport.body,
             )
         return parse_chat_completion_response(response)
 
     async def acompletion(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
+        transport = TransportKwargs.from_request(request)
         with _handle_non_provider_errors(self.provider_name):
             response = await self._router.acompletion(
                 model=request.model,
                 messages=request.messages,
-                **collect_non_none_optional_fields(request),
+                extra_headers=transport.headers or None,
+                **transport.body,
             )
         return await aparse_chat_completion_response(response)
 
     def embeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        transport = TransportKwargs.from_request(request)
         with _handle_non_provider_errors(self.provider_name):
             response = self._router.embedding(
                 model=request.model,
                 input=request.inputs,
-                **collect_non_none_optional_fields(request),
+                extra_headers=transport.headers or None,
+                **transport.body,
             )
         vectors = [extract_embedding_vector(item) for item in getattr(response, "data", [])]
         return EmbeddingResponse(vectors=vectors, usage=extract_usage(getattr(response, "usage", None)), raw=response)
 
     async def aembeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        transport = TransportKwargs.from_request(request)
         with _handle_non_provider_errors(self.provider_name):
             response = await self._router.aembedding(
                 model=request.model,
                 input=request.inputs,
-                **collect_non_none_optional_fields(request),
+                extra_headers=transport.headers or None,
+                **transport.body,
             )
         vectors = [extract_embedding_vector(item) for item in getattr(response, "data", [])]
         return EmbeddingResponse(vectors=vectors, usage=extract_usage(getattr(response, "usage", None)), raw=response)
 
     def generate_image(self, request: ImageGenerationRequest) -> ImageGenerationResponse:
-        image_kwargs = collect_non_none_optional_fields(request, exclude=self._IMAGE_EXCLUDE)
+        transport = TransportKwargs.from_request(request, exclude=self._IMAGE_EXCLUDE)
         with _handle_non_provider_errors(self.provider_name):
             if request.messages is not None:
                 response = self._router.completion(
                     model=request.model,
                     messages=request.messages,
-                    **image_kwargs,
+                    extra_headers=transport.headers or None,
+                    **transport.body,
                 )
             else:
                 response = self._router.image_generation(
                     prompt=request.prompt,
                     model=request.model,
-                    **image_kwargs,
+                    extra_headers=transport.headers or None,
+                    **transport.body,
                 )
 
         if request.messages is not None:
@@ -137,19 +148,21 @@ def generate_image(self, request: ImageGenerationRequest) -> ImageGenerationResp
         return ImageGenerationResponse(images=images, usage=usage, raw=response)
 
     async def agenerate_image(self, request: ImageGenerationRequest) -> ImageGenerationResponse:
-        image_kwargs = collect_non_none_optional_fields(request, exclude=self._IMAGE_EXCLUDE)
+        transport = TransportKwargs.from_request(request, exclude=self._IMAGE_EXCLUDE)
         with _handle_non_provider_errors(self.provider_name):
             if request.messages is not None:
                 response = await self._router.acompletion(
                     model=request.model,
                     messages=request.messages,
-                    **image_kwargs,
+                    extra_headers=transport.headers or None,
+                    **transport.body,
                 )
             else:
                 response = await self._router.aimage_generation(
                     prompt=request.prompt,
                     model=request.model,
-                    **image_kwargs,
+                    extra_headers=transport.headers or None,
+                    **transport.body,
                 )
 
         if request.messages is not None:
@@ -183,7 +196,7 @@ def _handle_non_provider_errors(provider_name: str) -> Iterator[None]:
 
         raise ProviderError(
             kind=kind,
-            message=str(exc),
+            message=extract_message_from_exception_string(str(exc)),
             status_code=status_code if isinstance(status_code, int) else None,
             provider_name=provider_name,
             cause=exc,

packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/litellm_bridge.py

@@ -5,8 +5,8 @@
 
 import calendar
 import email.utils
+import json
 import time
-from dataclasses import dataclass
 from enum import Enum
 
 from data_designer.engine.models.clients.types import HttpResponse
@@ -28,20 +28,26 @@ class ProviderErrorKind(str, Enum):
     UNSUPPORTED_CAPABILITY = "unsupported_capability"
 
 
-@dataclass
 class ProviderError(Exception):
-    kind: ProviderErrorKind
-    message: str
-    status_code: int | None = None
-    provider_name: str | None = None
-    model_name: str | None = None
-    retry_after: float | None = None
-    cause: Exception | None = None
-
-    def __post_init__(self) -> None:
-        Exception.__init__(self, self.message)
-        if self.cause is not None:
-            self.__cause__ = self.cause
+    def __init__(
+        self,
+        kind: ProviderErrorKind,
+        message: str,
+        status_code: int | None = None,
+        provider_name: str | None = None,
+        model_name: str | None = None,
+        retry_after: float | None = None,
+        cause: Exception | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.kind = kind
+        self.message = message
+        self.status_code = status_code
+        self.provider_name = provider_name
+        self.model_name = model_name
+        self.retry_after = retry_after
+        if cause is not None:
+            self.__cause__ = cause
 
     def __str__(self) -> str:
         return self.message
@@ -118,6 +124,31 @@ def map_http_error_to_provider_error(
     )
 
 
+def extract_message_from_exception_string(raw: str) -> str:
+    """Extract a human-readable message from a stringified LiteLLM exception.
+
+    LiteLLM often formats errors as ``"Error code: 400 - {json}"``.  This
+    mirrors the structured-key lookup in ``_extract_structured_message`` but
+    operates on a raw string instead of an ``HttpResponse``.
+    """
+    json_start = raw.find("{")
+    if json_start != -1:
+        try:
+            payload = json.loads(raw[json_start:])
+        except (json.JSONDecodeError, ValueError):
+            return raw
+        if isinstance(payload, dict):
+            for key in ("message", "error", "detail"):
+                value = payload.get(key)
+                if isinstance(value, str) and value.strip():
+                    return value.strip()
+                if isinstance(value, dict):
+                    nested = value.get("message")
+                    if isinstance(nested, str) and nested.strip():
+                        return nested.strip()
+    return raw
+
+
 def _extract_response_text(response: HttpResponse) -> str:
     # Try structured JSON extraction first — most providers return structured error
     # bodies and we want the human-readable message, not raw JSON.

packages/data-designer-engine/src/data_designer/engine/models/clients/errors.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import data_designer.lazy_heavy_imports as lazy
+from data_designer.config.models import ModelConfig
+from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.engine.models.clients.adapters.litellm_bridge import LiteLLMBridgeClient
+from data_designer.engine.models.clients.base import ModelClient
+from data_designer.engine.models.litellm_overrides import CustomRouter, LiteLLMRouterDefaultKwargs
+from data_designer.engine.secret_resolver import SecretResolver
+
+
+def create_model_client(
+    model_config: ModelConfig,
+    secret_resolver: SecretResolver,
+    model_provider_registry: ModelProviderRegistry,
+) -> ModelClient:
+    """Create a ModelClient for the given model configuration.
+
+    Resolves the provider, API key, and constructs a LiteLLM router wrapped in
+    a LiteLLMBridgeClient adapter.
+
+    Args:
+        model_config: The model configuration to create a client for.
+        secret_resolver: Resolver for secrets referenced in provider configs.
+        model_provider_registry: Registry of model provider configurations.
+
+    Returns:
+        A ModelClient instance ready for use.
+    """
+    provider = model_provider_registry.get_provider(model_config.provider)
+    api_key = None
+    if provider.api_key:
+        api_key = secret_resolver.resolve(provider.api_key)
+    api_key = api_key or "not-used-but-required"
+
+    litellm_params = lazy.litellm.LiteLLM_Params(
+        model=f"{provider.provider_type}/{model_config.model}",
+        api_base=provider.endpoint,
+        api_key=api_key,
+        max_parallel_requests=model_config.inference_parameters.max_parallel_requests,
+    )
+    deployment = {
+        "model_name": model_config.model,
+        "litellm_params": litellm_params.model_dump(),
+    }
+    router = CustomRouter([deployment], **LiteLLMRouterDefaultKwargs().model_dump())
+    return LiteLLMBridgeClient(provider_name=provider.name, router=router)

packages/data-designer-engine/src/data_designer/engine/models/clients/factory.py

@@ -5,9 +5,9 @@
 
 from __future__ import annotations
 
-import dataclasses
 import json
 import logging
+import uuid
 from typing import Any
 
 from data_designer.config.utils.image_helpers import (
@@ -206,7 +206,7 @@ def extract_tool_calls(raw_tool_calls: Any) -> list[ToolCall]:
 
     normalized_tool_calls: list[ToolCall] = []
     for raw_tool_call in raw_tool_calls:
-        tool_call_id = get_value_from(raw_tool_call, "id") or ""
+        tool_call_id = get_value_from(raw_tool_call, "id") or uuid.uuid4().hex
         function = get_value_from(raw_tool_call, "function")
         name = get_value_from(function, "name") or ""
         arguments_value = get_value_from(function, "arguments")
@@ -333,17 +333,3 @@ def get_first_value_or_none(values: Any) -> Any | None:
     if isinstance(values, list) and values:
         return values[0]
     return None
-
-
-def collect_non_none_optional_fields(request: Any, *, exclude: frozenset[str] = frozenset()) -> dict[str, Any]:
-    """Extract non-None optional fields from a request dataclass, skipping *exclude*.
-
-    The ``f.default is None`` check intentionally targets fields whose default is
-    ``None`` — i.e. truly optional kwargs the caller may or may not set.  Fields with
-    non-``None`` defaults are not "optional" in this forwarding sense and are excluded.
-    """
-    return {
-        f.name: v
-        for f in dataclasses.fields(request)
-        if f.name not in exclude and f.default is None and (v := getattr(request, f.name)) is not None
-    }