DeepFabric 4.4.0__py3-none-any.whl

deepfabric/llm/client.py

@@ -0,0 +1,1206 @@
+import asyncio
+import os
+
+from functools import lru_cache
+from typing import Any
+
+import anthropic
+import openai
+import outlines
+
+from google import genai
+from google.genai import types as genai_types
+from pydantic import BaseModel
+
+from ..exceptions import DataSetGeneratorError
+from .errors import handle_provider_error
+from .rate_limit_config import (
+    RateLimitConfig,
+    create_rate_limit_config,
+    get_default_rate_limit_config,
+)
+from .retry_handler import RetryHandler, retry_with_backoff, retry_with_backoff_async
+
+# JSON Schema union type keys that need recursive processing
+_UNION_KEYS = ("anyOf", "oneOf", "allOf")
+
+# Provider-specific parameter mappings
+# Maps generic parameter names to provider-specific equivalents
+# Format: {provider: {generic_name: provider_specific_name}}
+_PROVIDER_PARAM_MAPPINGS: dict[str, dict[str, str]] = {
+    "gemini": {
+        "max_tokens": "max_output_tokens",
+    },
+    # Add other provider mappings as needed
+    # "anthropic": {"some_param": "anthropic_param"},
+}
+
+
+class LLMClient:
+    """Wrapper for Outlines models with retry logic and error handling."""
+
+    def __init__(
+        self,
+        provider: str,
+        model_name: str,
+        rate_limit_config: RateLimitConfig | dict | None = None,
+        **kwargs,
+    ):
+        """Initialize LLM client.
+
+        Args:
+            provider: Provider name
+            model_name: Model identifier
+            rate_limit_config: Rate limiting configuration (None uses provider defaults)
+            **kwargs: Additional client configuration
+        """
+        self.provider = provider
+        self.model_name = model_name
+        self._client_kwargs = kwargs  # Store for lazy async model initialization
+
+        # Initialize rate limiting
+        if isinstance(rate_limit_config, dict):
+            self.rate_limit_config = create_rate_limit_config(provider, rate_limit_config)
+        elif rate_limit_config is None:
+            self.rate_limit_config = get_default_rate_limit_config(provider)
+        else:
+            self.rate_limit_config = rate_limit_config
+
+        self.retry_handler = RetryHandler(self.rate_limit_config, provider)
+
+        self.model: Any = make_outlines_model(provider, model_name, **kwargs)
+        # Lazy-initialize async_model only when needed
+        self._async_model: Any | None = None
+        self._async_model_initialized: bool = False
+        if self.model is None:
+            msg = f"Failed to create model for {provider}/{model_name}"
+            raise DataSetGeneratorError(msg)
+
+    @property
+    def async_model(self) -> Any | None:
+        """Lazily initialize and return the async model.
+
+        The async model is only created when first accessed, reducing memory
+        and connection overhead when only sync operations are used.
+        """
+        if not self._async_model_initialized:
+            self._async_model = make_async_outlines_model(
+                self.provider, self.model_name, **self._client_kwargs
+            )
+            self._async_model_initialized = True
+        return self._async_model
+
+    def generate(self, prompt: str, schema: Any, max_retries: int = 3, **kwargs) -> Any:  # noqa: ARG002
+        """Generate structured output using the provided schema.
+
+        Args:
+            prompt: Input prompt
+            schema: Pydantic model or other schema type
+            max_retries: Maximum number of retry attempts (deprecated, use rate_limit_config)
+            **kwargs: Additional generation parameters
+
+        Returns:
+            Generated output matching the schema
+
+        Raises:
+            DataSetGeneratorError: If generation fails after all retries
+
+        Note:
+            The max_retries parameter is deprecated. Use rate_limit_config in __init__ instead.
+            If provided, it will be ignored in favor of the configured retry handler.
+        """
+        return self._generate_with_retry(prompt, schema, **kwargs)
+
+    # Providers that only support async generation (native API, not Outlines)
+    _ASYNC_ONLY_PROVIDERS = frozenset({"anthropic", "gemini"})
+
+    @retry_with_backoff
+    def _generate_with_retry(self, prompt: str, schema: Any, **kwargs) -> Any:
+        """Internal method that performs actual generation with retry wrapper.
+
+        This method is decorated with retry_with_backoff to handle rate limits
+        and transient errors automatically.
+        """
+        # Check for async-only providers
+        if self.provider in self._ASYNC_ONLY_PROVIDERS:
+            raise DataSetGeneratorError(
+                f"Synchronous generation is not supported for {self.provider}. "
+                f"Use generate_async() instead.",
+                context={"provider": self.provider},
+            )
+
+        # Convert provider-specific parameters
+        kwargs = self._convert_generation_params(**kwargs)
+
+        # For OpenAI and OpenRouter, ensure all properties are in required array
+        generation_schema = schema
+        if (
+            self.provider in ("openai", "openrouter")
+            and isinstance(schema, type)
+            and issubclass(schema, BaseModel)
+        ):
+            generation_schema = _create_openai_compatible_schema(schema)
+
+        # Generate JSON string using the model
+        json_output = self.model(prompt, generation_schema, **kwargs)
+
+        # Parse and validate the JSON response with the ORIGINAL schema
+        # This ensures we still get proper validation
+        try:
+            return schema.model_validate_json(json_output)
+        except Exception as e:
+            raise DataSetGeneratorError(
+                f"Generation validation failed: {e}",
+                context={"raw_content": json_output},
+            ) from e
+
+    async def generate_async(self, prompt: str, schema: Any, max_retries: int = 3, **kwargs) -> Any:  # noqa: ARG002
+        """Asynchronously generate structured output using provider async clients.
+
+        Args:
+            prompt: Input prompt
+            schema: Pydantic model or other schema type
+            max_retries: Maximum number of retry attempts (deprecated, use rate_limit_config)
+            **kwargs: Additional generation parameters
+
+        Returns:
+            Generated output matching the schema
+
+        Raises:
+            DataSetGeneratorError: If generation fails after all retries
+
+        Note:
+            The max_retries parameter is deprecated. Use rate_limit_config in __init__ instead.
+            If provided, it will be ignored in favor of the configured retry handler.
+        """
+        if self.async_model is None:
+            # Fallback to running the synchronous path in a worker thread
+            return await asyncio.to_thread(self.generate, prompt, schema, **kwargs)
+
+        return await self._generate_async_with_retry(prompt, schema, **kwargs)
+
+    @retry_with_backoff_async
+    async def _generate_async_with_retry(self, prompt: str, schema: Any, **kwargs) -> Any:
+        """Internal async method that performs actual generation with retry wrapper.
+
+        This method is decorated with retry_with_backoff_async to handle rate limits
+        and transient errors automatically.
+        """
+        kwargs = self._convert_generation_params(**kwargs)
+
+        # Native providers (Anthropic, Gemini) handle schema transformation internally
+        # For OpenAI and OpenRouter, ensure all properties are in required array
+        generation_schema = schema
+        if self.provider in self._ASYNC_ONLY_PROVIDERS:
+            # Native model classes handle schema transformation internally
+            pass
+        elif (
+            self.provider in ("openai", "openrouter")
+            and isinstance(schema, type)
+            and issubclass(schema, BaseModel)
+        ):
+            generation_schema = _create_openai_compatible_schema(schema)
+
+        # Ensure we have an async model; if not, fall back to running the sync path
+        async_model = self.async_model
+        if async_model is None:
+            # Note: This will raise an error for async-only providers
+            return await asyncio.to_thread(self.generate, prompt, schema, **kwargs)
+
+        # Call the async model (guaranteed non-None by check above)
+        json_output = await async_model(prompt, generation_schema, **kwargs)
+
+        # Validate with original schema to ensure proper validation
+        try:
+            return schema.model_validate_json(json_output)
+        except Exception as e:
+            raise DataSetGeneratorError(
+                f"Async generation validation failed: {e}",
+                context={"raw_content": json_output},
+            ) from e
+
+    async def generate_async_stream(self, prompt: str, schema: Any, max_retries: int = 3, **kwargs):  # noqa: ARG002
+        """Asynchronously generate structured output with streaming text chunks.
+
+        This method streams the LLM's output text as it's generated, then returns
+        the final parsed Pydantic model. It yields tuples of (chunk, result) where:
+        - During streaming: (text_chunk, None)
+        - When complete: (None, final_pydantic_model)
+
+        Args:
+            prompt: Input prompt
+            schema: Pydantic model or other schema type
+            max_retries: Maximum number of retry attempts (deprecated, use rate_limit_config)
+            **kwargs: Additional generation parameters
+
+        Yields:
+            tuple[str | None, Any | None]:
+                - (chunk, None) during streaming
+                - (None, model) when generation is complete
+
+        Raises:
+            DataSetGeneratorError: If generation fails after all retries
+
+        Note:
+            The max_retries parameter is deprecated. Use rate_limit_config in __init__ instead.
+
+        Example:
+            >>> async for chunk, result in client.generate_async_stream(prompt, MyModel):
+            ...     if chunk:
+            ...         print(chunk, end='', flush=True)  # Display streaming text
+            ...     if result:
+            ...         return result  # Final parsed model
+        """
+        # Call streaming generation directly (retry decorator doesn't work with generators)
+        kwargs = self._convert_generation_params(**kwargs)
+
+        # Native providers (Anthropic, Gemini) handle schema transformation internally
+        # For OpenAI and OpenRouter, ensure all properties are in required array
+        generation_schema = schema
+        if self.provider in self._ASYNC_ONLY_PROVIDERS:
+            # Native model classes handle schema transformation internally
+            pass
+        elif (
+            self.provider in ("openai", "openrouter")
+            and isinstance(schema, type)
+            and issubclass(schema, BaseModel)
+        ):
+            generation_schema = _create_openai_compatible_schema(schema)
+
+        # Check if model supports streaming
+        async_model = self.async_model or self.model
+        if not hasattr(async_model, "generate_stream"):
+            # Fallback: no streaming support, yield entire result at once
+            result = await self.generate_async(prompt, schema, **kwargs)
+            yield (None, result)
+            return
+
+        # Stream generation
+        accumulated_text: list[str] = []
+        try:
+            # For sync models used in async context
+            if self.async_model is None:
+                # Use asyncio.Queue for true streaming from sync generator
+                # This yields chunks as they arrive instead of waiting for all
+                queue: asyncio.Queue[str | None] = asyncio.Queue()
+                stream_error: list[Exception] = []
+
+                def _produce_chunks():
+                    """Run sync generator and put chunks in queue."""
+                    try:
+                        for chunk in self.model.generate_stream(
+                            prompt, generation_schema, **kwargs
+                        ):
+                            asyncio.run_coroutine_threadsafe(
+                                queue.put(chunk), asyncio.get_event_loop()
+                            )
+                    except Exception as e:
+                        stream_error.append(e)
+                    finally:
+                        # Signal completion
+                        asyncio.run_coroutine_threadsafe(queue.put(None), asyncio.get_event_loop())
+
+                # Start producer in background thread
+                loop = asyncio.get_event_loop()
+                loop.run_in_executor(None, _produce_chunks)
+
+                # Consume chunks as they arrive
+                while True:
+                    chunk = await queue.get()
+                    if chunk is None:
+                        break
+                    accumulated_text.append(chunk)
+                    yield (chunk, None)
+
+                # Re-raise any error from the producer thread
+                if stream_error:
+                    raise DataSetGeneratorError(
+                        f"Streaming generation failed in producer: {stream_error[0]}",
+                        context={
+                            "raw_content": "".join(accumulated_text) if accumulated_text else None
+                        },
+                    ) from stream_error[0]
+            else:
+                # True async streaming
+                stream = self.async_model.generate_stream(prompt, generation_schema, **kwargs)
+                async for chunk in stream:
+                    accumulated_text.append(chunk)
+                    yield (chunk, None)
+
+            # Parse accumulated JSON with original schema
+            full_text = "".join(accumulated_text)
+            result = schema.model_validate_json(full_text)
+            yield (None, result)
+
+        except Exception as e:
+            # Wrap and raise error with raw content for debugging
+            raw_content = "".join(accumulated_text) if accumulated_text else None
+            raise DataSetGeneratorError(
+                f"Streaming generation failed: {e}",
+                context={"raw_content": raw_content},
+            ) from e
+
+    def _convert_generation_params(self, **kwargs) -> dict:
+        """Convert generic parameters to provider-specific ones.
+
+        Uses the _PROVIDER_PARAM_MAPPINGS static dictionary for extensible
+        parameter conversion across different providers.
+        """
+        mappings = _PROVIDER_PARAM_MAPPINGS.get(self.provider, {})
+        for generic_name, provider_name in mappings.items():
+            if generic_name in kwargs:
+                kwargs[provider_name] = kwargs.pop(generic_name)
+
+        return kwargs
+
+    def __repr__(self) -> str:
+        return f"LLMClient(provider={self.provider}, model={self.model_name})"
+
+
+class GeminiModel:
+    """Gemini API client using native structured outputs.
+
+    Uses Gemini's native response_json_schema parameter which is more reliable
+    than Outlines' wrapper for structured output generation.
+    """
+
+    def __init__(self, client: genai.Client, model_name: str):
+        self.client = client
+        self.model_name = model_name
+
+    def _prepare_json_schema(self, schema: type[BaseModel]) -> dict:
+        """Prepare JSON schema for Gemini by stripping incompatible fields.
+
+        Note: We only strip additionalProperties and incompatible fields here.
+        Ref inlining is intentionally NOT done to avoid excessive nesting depth
+        that can exceed Gemini's schema limits for deeply nested models.
+        """
+        json_schema = schema.model_json_schema()
+        return _strip_additional_properties(json_schema)
+
+    async def __call__(
+        self,
+        prompt: str,
+        schema: type[BaseModel],
+        **kwargs,
+    ) -> str:
+        """Asynchronously generate structured output using Gemini's native JSON schema support.
+
+        Args:
+            prompt: The input prompt
+            schema: Pydantic model class for structured output
+            **kwargs: Additional generation parameters (temperature, max_output_tokens, etc.)
+
+        Returns:
+            JSON string matching the schema
+        """
+        # Use prepared schema with Gemini-compatible transformations
+        prepared_schema = self._prepare_json_schema(schema)
+
+        try:
+            config = genai_types.GenerateContentConfig(
+                response_mime_type="application/json",
+                response_json_schema=prepared_schema,
+                **kwargs,
+            )
+        except Exception as e:
+            raise DataSetGeneratorError(f"Failed to create Gemini config: {e}") from e
+
+        # Call Gemini API directly using async method
+        response = await self.client.aio.models.generate_content(
+            model=self.model_name,
+            contents=prompt,
+            config=config,
+        )
+
+        # Safely check for empty or blocked responses
+        # Some SDK versions may return None or have blocked candidates
+        if not response.candidates:
+            raise DataSetGeneratorError(
+                "Gemini returned empty response",
+                context={"finish_reason": None},
+            )
+
+        first_candidate = response.candidates[0]
+        if first_candidate.content is None or not first_candidate.content.parts:
+            raise DataSetGeneratorError(
+                "Gemini returned empty response",
+                context={"finish_reason": getattr(first_candidate, "finish_reason", None)},
+            )
+
+        if response.text is None:
+            raise DataSetGeneratorError("Gemini returned empty response")
+
+        return response.text
+
+
+class AnthropicModel:
+    """Anthropic API client using native structured outputs.
+
+    Uses Anthropic's beta structured outputs feature (structured-outputs-2025-11-13)
+    which provides guaranteed JSON schema compliance through constrained decoding.
+    """
+
+    # Beta header for structured outputs feature
+    STRUCTURED_OUTPUTS_BETA = "structured-outputs-2025-11-13"
+
+    def __init__(self, client: anthropic.AsyncAnthropic, model_name: str):
+        self.client = client
+        self.model_name = model_name
+
+    def _prepare_json_schema(self, schema: type[BaseModel]) -> dict:
+        """Prepare JSON schema for Anthropic structured outputs.
+
+        Anthropic's structured outputs require:
+        - additionalProperties: false on all objects
+        - All properties in required array
+        """
+        json_schema = schema.model_json_schema()
+        return _ensure_anthropic_strict_mode_compliance(json_schema)
+
+    async def __call__(
+        self,
+        prompt: str,
+        schema: type[BaseModel],
+        **kwargs,
+    ) -> str:
+        """Asynchronously generate structured output using Anthropic's native JSON schema support.
+
+        Args:
+            prompt: The input prompt
+            schema: Pydantic model class for structured output
+            **kwargs: Additional generation parameters (temperature, max_tokens, etc.)
+
+        Returns:
+            JSON string matching the schema
+        """
+        prepared_schema = self._prepare_json_schema(schema)
+
+        # Extract max_tokens with a sensible default (16384 for structured outputs)
+        max_tokens = kwargs.pop("max_tokens", 16384)
+
+        try:
+            response = await self.client.beta.messages.create(
+                model=self.model_name,
+                max_tokens=max_tokens,
+                betas=[self.STRUCTURED_OUTPUTS_BETA],
+                messages=[{"role": "user", "content": prompt}],
+                output_format={
+                    "type": "json_schema",
+                    "schema": prepared_schema,
+                },
+                **kwargs,
+            )
+        except anthropic.BadRequestError as e:
+            raise DataSetGeneratorError(
+                f"Anthropic structured output request failed: {e}",
+                context={"schema": prepared_schema},
+            ) from e
+
+        # Check for refusals
+        if response.stop_reason == "refusal":
+            raise DataSetGeneratorError(
+                "Anthropic refused the request for safety reasons",
+                context={"stop_reason": response.stop_reason},
+            )
+
+        # Check for max_tokens truncation
+        if response.stop_reason == "max_tokens":
+            raise DataSetGeneratorError(
+                "Anthropic response truncated due to max_tokens limit",
+                context={"stop_reason": response.stop_reason},
+            )
+
+        # Extract text from response
+        if not response.content:
+            raise DataSetGeneratorError("Anthropic returned empty response")
+
+        text_block = response.content[0]
+        if not hasattr(text_block, "text") or not text_block.text:
+            raise DataSetGeneratorError("Anthropic returned empty text response")
+
+        return text_block.text
+
+
+def _raise_api_key_error(env_var: str) -> None:
+    """Raise an error for missing API key."""
+    msg = f"{env_var} environment variable not set"
+    raise DataSetGeneratorError(msg)
+
+
+def _get_gemini_api_key() -> str:
+    """Retrieve Gemini API key from environment variables.
+
+    Returns:
+        The API key string
+
+    Raises:
+        DataSetGeneratorError: If no API key is found
+    """
+    for name in ("GOOGLE_API_KEY", "GEMINI_API_KEY"):
+        if api_key := os.getenv(name):
+            return api_key
+    _raise_api_key_error("GOOGLE_API_KEY or GEMINI_API_KEY")
+    # This return is never reached but satisfies type checker
+    raise AssertionError("unreachable")
+
+
+# Provider to environment variable mapping
+PROVIDER_API_KEY_MAP: dict[str, list[str]] = {
+    "openai": ["OPENAI_API_KEY"],
+    "anthropic": ["ANTHROPIC_API_KEY"],
+    "gemini": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
+    "openrouter": ["OPENROUTER_API_KEY"],
+    "ollama": [],  # No API key required
+    # Test providers for unit tests, no API key required
+    "test": [],
+    "override": [],
+}
+
+
+def validate_provider_api_key(provider: str) -> tuple[bool, str | None]:
+    """Validate that the required API key exists for a provider.
+
+    Args:
+        provider: Provider name (openai, anthropic, gemini, ollama, openrouter)
+
+    Returns:
+        Tuple of (is_valid, error_message). If valid, error_message is None.
+    """
+    env_vars = PROVIDER_API_KEY_MAP.get(provider)
+
+    if env_vars is None:
+        return False, f"Unknown provider: {provider}"
+
+    # Ollama doesn't need an API key
+    if not env_vars:
+        return True, None
+
+    # Check if any of the required env vars are set
+    for env_var in env_vars:
+        if os.getenv(env_var):
+            return True, None
+
+    # Build helpful error message
+    if len(env_vars) == 1:
+        return False, f"Missing API key: {env_vars[0]} environment variable is not set"
+    var_list = " or ".join(env_vars)
+    return False, f"Missing API key: Set {var_list} environment variable"
+
+
+def get_required_api_key_env_var(provider: str) -> str | None:
+    """Get the environment variable name(s) required for a provider.
+
+    Args:
+        provider: Provider name
+
+    Returns:
+        Human-readable string of required env var(s), or None if no key required
+    """
+    env_vars = PROVIDER_API_KEY_MAP.get(provider)
+    if not env_vars:
+        return None
+    if len(env_vars) == 1:
+        return env_vars[0]
+    return " or ".join(env_vars)
+
+
+def _raise_unsupported_provider_error(provider: str) -> None:
+    """Raise an error for unsupported provider."""
+    msg = f"Unsupported provider: {provider}"
+    raise DataSetGeneratorError(msg)
+
+
+def _get_openai_client_config(
+    api_key_env_var: str | None,
+    default_base_url: str | None,
+    dummy_key: str | None = None,
+    **kwargs,
+) -> tuple[str | None, dict[str, Any]]:
+    """Extract common configuration for OpenAI-compatible clients.
+
+    Args:
+        api_key_env_var: Environment variable name for API key (None to skip check)
+        default_base_url: Default base URL if not provided in kwargs (None for OpenAI default)
+        dummy_key: Dummy API key to use if api_key_env_var is None (e.g., for Ollama)
+        **kwargs: Additional client configuration (may include base_url override)
+
+    Returns:
+        Tuple of (api_key, client_kwargs) for client initialization
+
+    Raises:
+        DataSetGeneratorError: If required API key is missing
+    """
+    # Get API key from environment or use dummy key
+    if api_key_env_var:
+        api_key = os.getenv(api_key_env_var)
+        if not api_key:
+            _raise_api_key_error(api_key_env_var)
+    else:
+        api_key = dummy_key
+
+    # Set up base_url if provided
+    client_kwargs = {k: v for k, v in kwargs.items() if k != "base_url"}
+    if default_base_url:
+        client_kwargs["base_url"] = kwargs.get("base_url", default_base_url)
+
+    return api_key, client_kwargs
+
+
+def _create_openai_compatible_client(
+    api_key_env_var: str | None,
+    default_base_url: str | None,
+    dummy_key: str | None = None,
+    **kwargs,
+) -> openai.OpenAI:
+    """Create an OpenAI-compatible client for providers that use OpenAI's API format.
+
+    Args:
+        api_key_env_var: Environment variable name for API key (None to skip check)
+        default_base_url: Default base URL if not provided in kwargs (None for OpenAI default)
+        dummy_key: Dummy API key to use if api_key_env_var is None (e.g., for Ollama)
+        **kwargs: Additional client configuration (may include base_url override)
+
+    Returns:
+        Configured OpenAI client instance
+
+    Raises:
+        DataSetGeneratorError: If required API key is missing
+    """
+    api_key, client_kwargs = _get_openai_client_config(
+        api_key_env_var, default_base_url, dummy_key, **kwargs
+    )
+    return openai.OpenAI(api_key=api_key, **client_kwargs)
+
+
+def _create_async_openai_compatible_client(
+    api_key_env_var: str | None,
+    default_base_url: str | None,
+    dummy_key: str | None = None,
+    **kwargs,
+) -> openai.AsyncOpenAI:
+    """Create an async OpenAI-compatible client for providers that use OpenAI's API format.
+
+    Args:
+        api_key_env_var: Environment variable name for API key (None to skip check)
+        default_base_url: Default base URL if not provided in kwargs (None for OpenAI default)
+        dummy_key: Dummy API key to use if api_key_env_var is None (e.g., for Ollama)
+        **kwargs: Additional client configuration (may include base_url override)
+
+    Returns:
+        Configured AsyncOpenAI client instance
+
+    Raises:
+        DataSetGeneratorError: If required API key is missing
+    """
+    api_key, client_kwargs = _get_openai_client_config(
+        api_key_env_var, default_base_url, dummy_key, **kwargs
+    )
+    return openai.AsyncOpenAI(api_key=api_key, **client_kwargs)
+
+
+def _is_incompatible_object(schema: dict) -> bool:
+    """Check if a schema represents an object incompatible with Gemini.
+
+    Gemini rejects objects with no properties defined (like dict[str, Any]).
+
+    Args:
+        schema: JSON schema dictionary
+
+    Returns:
+        True if the schema is an incompatible object type
+    """
+    return schema.get("type") == "object" and "properties" not in schema and "$ref" not in schema
+
+
+def _is_incompatible_array(schema: dict) -> bool:
+    """Check if a schema represents an array with incompatible items.
+
+    Arrays with object items that have no properties (like list[dict[str, Any]])
+    are incompatible with Gemini.
+
+    Args:
+        schema: JSON schema dictionary
+
+    Returns:
+        True if the schema is an incompatible array type
+    """
+    if schema.get("type") != "array" or "items" not in schema:
+        return False
+
+    items = schema["items"]
+    return isinstance(items, dict) and _is_incompatible_object(items)
+
+
+def _inline_refs(schema_dict: dict, defs: dict | None = None) -> dict:
+    """
+    Recursively inline $ref references in a JSON schema.
+
+    Gemini's structured output can be unreliable with $ref. This function
+    resolves all references by inlining the definitions directly.
+
+    Args:
+        schema_dict: JSON schema dictionary
+        defs: The $defs dictionary from the root schema
+
+    Returns:
+        Modified schema dict with all $ref inlined
+    """
+    if not isinstance(schema_dict, dict):
+        return schema_dict
+
+    # Use provided defs or extract from schema (guaranteed to be dict after this)
+    resolved_defs: dict = defs if defs is not None else schema_dict.get("$defs", {})
+
+    # Handle $ref - inline the referenced definition
+    if "$ref" in schema_dict:
+        ref_path = schema_dict["$ref"]
+        # Parse reference like "#/$defs/GraphSubtopic"
+        if ref_path.startswith("#/$defs/"):
+            def_name = ref_path[len("#/$defs/") :]
+            if def_name in resolved_defs:
+                # Return a copy of the definition with refs inlined
+                inlined = _inline_refs(dict(resolved_defs[def_name]), resolved_defs)
+                # Preserve any other properties from the original (like description)
+                for key, value in schema_dict.items():
+                    if key != "$ref":
+                        inlined[key] = value
+                return inlined
+        # If we can't resolve, return as-is
+        return schema_dict
+
+    # Create a copy to modify
+    result: dict[str, Any] = {}
+
+    for key, value in schema_dict.items():
+        if key == "$defs":
+            # Skip $defs since we're inlining them
+            continue
+        if key == "properties" and isinstance(value, dict):
+            result[key] = {
+                prop_name: _inline_refs(prop_schema, resolved_defs)
+                for prop_name, prop_schema in value.items()
+            }
+        elif key == "items" and isinstance(value, dict):
+            result[key] = _inline_refs(value, resolved_defs)
+        elif key in _UNION_KEYS and isinstance(value, list):
+            result[key] = [_inline_refs(variant, resolved_defs) for variant in value]
+        else:
+            result[key] = value
+
+    return result
+
+
+def _strip_additional_properties(schema_dict: dict) -> dict:
+    """
+    Recursively remove additionalProperties from JSON schema and handle dict[str, Any] fields.
+
+    Gemini doesn't support:
+    1. additionalProperties field in JSON schemas
+    2. Objects with no properties defined (e.g., dict[str, Any])
+    3. Arrays whose items are objects with no properties (e.g., list[dict[str, Any]])
+
+    Fields like dict[str, Any] have additionalProperties: true and no properties defined.
+    Gemini requires that object-type fields must have properties, so we exclude these
+    fields from the schema entirely.
+
+    Note: This function preserves $defs and does NOT inline refs. Use _inline_refs
+    separately if ref inlining is needed.
+
+    Args:
+        schema_dict: JSON schema dictionary
+
+    Returns:
+        Modified schema dict without additionalProperties and dict[str, Any] fields
+    """
+    if not isinstance(schema_dict, dict):
+        return schema_dict
+
+    # For Gemini, identify and remove incompatible fields
+    if "properties" in schema_dict:
+        properties_to_remove = []
+        for prop_name, prop_schema in schema_dict["properties"].items():
+            if not isinstance(prop_schema, dict):
+                continue
+
+            # Check for direct incompatibilities
+            if prop_schema.get("additionalProperties") is True:
+                # Remove fields with additionalProperties: true (e.g., dict[str, Any])
+                properties_to_remove.append(prop_name)
+            elif _is_incompatible_object(prop_schema):
+                # Remove objects with no properties
+                properties_to_remove.append(prop_name)
+            elif _is_incompatible_array(prop_schema):
+                # Remove arrays with incompatible items (e.g., list[dict[str, Any]])
+                properties_to_remove.append(prop_name)
+            elif "anyOf" in prop_schema and any(
+                isinstance(variant, dict)
+                and (_is_incompatible_object(variant) or _is_incompatible_array(variant))
+                for variant in prop_schema["anyOf"]
+            ):
+                # Check if anyOf contains incompatible variants
+                properties_to_remove.append(prop_name)
+
+        # Remove incompatible properties from the schema
+        for prop_name in properties_to_remove:
+            del schema_dict["properties"][prop_name]
+
+        # Update required array to exclude removed properties
+        if "required" in schema_dict:
+            schema_dict["required"] = [
+                r for r in schema_dict["required"] if r not in properties_to_remove
+            ]
+
+    # Remove additionalProperties from current level
+    schema_dict.pop("additionalProperties", None)
+
+    # Recursively process nested structures
+    if "$defs" in schema_dict:
+        for def_name, def_schema in schema_dict["$defs"].items():
+            schema_dict["$defs"][def_name] = _strip_additional_properties(def_schema)
+
+    # Process properties recursively
+    if "properties" in schema_dict:
+        for prop_name, prop_schema in schema_dict["properties"].items():
+            schema_dict["properties"][prop_name] = _strip_additional_properties(prop_schema)
+
+    # Process items (for arrays)
+    if "items" in schema_dict:
+        schema_dict["items"] = _strip_additional_properties(schema_dict["items"])
+
+    # Process union types (anyOf, oneOf, allOf)
+    for union_key in _UNION_KEYS:
+        if union_key in schema_dict:
+            schema_dict[union_key] = [
+                _strip_additional_properties(variant) for variant in schema_dict[union_key]
+            ]
+
+    return schema_dict
+
+
+def _ensure_anthropic_strict_mode_compliance(schema_dict: dict) -> dict:
+    """Ensure schema complies with Anthropic's structured outputs requirements.
+
+    Anthropic's structured outputs require:
+    1. For objects, 'additionalProperties' must be explicitly set to false
+    2. ALL properties must be in the 'required' array (no optional fields allowed)
+    3. No fields with additionalProperties: true (incompatible with strict mode)
+
+    This is similar to OpenAI's strict mode requirements.
+
+    Args:
+        schema_dict: JSON schema dictionary
+
+    Returns:
+        Modified schema dict meeting Anthropic structured outputs requirements
+    """
+    # Reuse OpenAI's strict mode compliance function as requirements are similar
+    return _ensure_openai_strict_mode_compliance(schema_dict)
+
+
+def _strip_ref_sibling_keywords(schema_dict: dict) -> dict:
+    """
+    Remove sibling keywords from $ref properties.
+
+    OpenAI's strict mode doesn't allow $ref to have additional keywords like 'description'.
+    When Pydantic generates schemas, it adds 'description' alongside '$ref' for nested models
+    that have Field(description=...). This function strips those extra keywords.
+
+    Args:
+        schema_dict: JSON schema dictionary
+
+    Returns:
+        Modified schema dict with $ref siblings removed
+    """
+    if not isinstance(schema_dict, dict):
+        return schema_dict
+
+    # If this dict has $ref, remove all sibling keywords except $ref itself
+    if "$ref" in schema_dict:
+        return {"$ref": schema_dict["$ref"]}
+
+    # Process properties recursively
+    if "properties" in schema_dict:
+        for prop_name, prop_schema in schema_dict["properties"].items():
+            schema_dict["properties"][prop_name] = _strip_ref_sibling_keywords(prop_schema)
+
+    # Process $defs recursively
+    if "$defs" in schema_dict:
+        for def_name, def_schema in schema_dict["$defs"].items():
+            schema_dict["$defs"][def_name] = _strip_ref_sibling_keywords(def_schema)
+
+    # Process items (for arrays)
+    if "items" in schema_dict:
+        schema_dict["items"] = _strip_ref_sibling_keywords(schema_dict["items"])
+
+    # Process union types (anyOf, oneOf, allOf)
+    for union_key in _UNION_KEYS:
+        if union_key in schema_dict:
+            schema_dict[union_key] = [
+                _strip_ref_sibling_keywords(variant) for variant in schema_dict[union_key]
+            ]
+
+    return schema_dict
+
+
+def _ensure_openai_strict_mode_compliance(schema_dict: dict) -> dict:
+    """
+    Ensure schema complies with OpenAI's strict mode requirements.
+
+    OpenAI's strict mode requires:
+    1. For objects, 'additionalProperties' must be explicitly set to false
+    2. ALL properties must be in the 'required' array (no optional fields allowed)
+    3. No fields with additionalProperties: true (incompatible with strict mode)
+    4. $ref cannot have sibling keywords like 'description'
+
+    Fields like dict[str, Any] have additionalProperties: true and cannot be
+    represented in strict mode, so they are excluded from the schema entirely.
+
+    Args:
+        schema_dict: JSON schema dictionary
+
+    Returns:
+        Modified schema dict meeting OpenAI strict mode requirements
+    """
+    if not isinstance(schema_dict, dict):
+        return schema_dict
+
+    # First, strip sibling keywords from $ref properties
+    schema_dict = _strip_ref_sibling_keywords(schema_dict)
+
+    # For OpenAI strict mode, identify and remove dict[str, Any] fields
+    # These have additionalProperties: true which is incompatible with strict mode
+    if "properties" in schema_dict:
+        properties_to_remove = []
+        for prop_name, prop_schema in schema_dict["properties"].items():
+            # Check for direct additionalProperties: true
+            if isinstance(prop_schema, dict) and prop_schema.get("additionalProperties") is True:
+                # Remove fields with additionalProperties: true (e.g., dict[str, Any])
+                properties_to_remove.append(prop_name)
+            # Check for anyOf containing object variants with additionalProperties: true
+            elif isinstance(prop_schema, dict) and "anyOf" in prop_schema:
+                for variant in prop_schema["anyOf"]:
+                    if isinstance(variant, dict) and variant.get("additionalProperties") is True:
+                        # This anyOf contains an incompatible object variant - remove entire field
+                        properties_to_remove.append(prop_name)
+                        break
+
+        # Remove incompatible properties from the schema
+        for prop_name in properties_to_remove:
+            del schema_dict["properties"][prop_name]
+
+        # Update required array to exclude removed properties
+        if "required" in schema_dict:
+            schema_dict["required"] = [
+                r for r in schema_dict["required"] if r not in properties_to_remove
+            ]
+
+        # After removing incompatible fields, ensure ALL remaining properties are required
+        # OpenAI strict mode doesn't allow optional fields
+        property_keys = list(schema_dict["properties"].keys())
+        schema_dict["required"] = property_keys
+        schema_dict["additionalProperties"] = False
+
+    # For all objects (including those without properties), set additionalProperties to false
+    if schema_dict.get("type") == "object":
+        schema_dict["additionalProperties"] = False
+
+    # Recursively process nested structures
+    if "$defs" in schema_dict:
+        for def_name, def_schema in schema_dict["$defs"].items():
+            schema_dict["$defs"][def_name] = _ensure_openai_strict_mode_compliance(def_schema)
+
+    # Process properties recursively
+    if "properties" in schema_dict:
+        for prop_name, prop_schema in schema_dict["properties"].items():
+            schema_dict["properties"][prop_name] = _ensure_openai_strict_mode_compliance(
+                prop_schema
+            )
+
+    # Process items (for arrays)
+    if "items" in schema_dict:
+        schema_dict["items"] = _ensure_openai_strict_mode_compliance(schema_dict["items"])
+
+    # Process union types (anyOf, oneOf, allOf)
+    # This must be done to handle nested structures like list[dict[str, Any]] | None
+    # where the dict[str, Any] is inside an array variant
+    for union_key in _UNION_KEYS:
+        if union_key in schema_dict:
+            schema_dict[union_key] = [
+                _ensure_openai_strict_mode_compliance(variant) for variant in schema_dict[union_key]
+            ]
+
+    return schema_dict
+
+
+@lru_cache(maxsize=128)
+def _get_cached_openai_schema(schema: type[BaseModel]) -> type[BaseModel]:
+    """
+    Get or create a cached OpenAI-compatible version of a Pydantic schema.
+
+    This function caches transformed schemas to avoid repeated processing
+    of the same Pydantic model during multiple generation calls.
+
+    Args:
+        schema: Original Pydantic model
+
+    Returns:
+        Cached wrapper model that generates OpenAI-compatible schemas
+    """
+
+    # Create a new model class that overrides model_json_schema
+    class OpenAICompatModel(schema):  # type: ignore[misc,valid-type]
+        @classmethod
+        def model_json_schema(cls, **kwargs):
+            # Get the original schema
+            original_schema = super().model_json_schema(**kwargs)
+            # Ensure OpenAI strict mode compliance
+            return _ensure_openai_strict_mode_compliance(original_schema)
+
+    # Set name and docstring
+    OpenAICompatModel.__name__ = f"{schema.__name__}OpenAICompat"
+    OpenAICompatModel.__doc__ = schema.__doc__
+
+    # Rebuild model to resolve forward references (e.g., PendingToolCall in AgentStep)
+    OpenAICompatModel.model_rebuild()
+
+    return OpenAICompatModel
+
+
+def _create_openai_compatible_schema(schema: type[BaseModel]) -> type[BaseModel]:
+    """
+    Create an OpenAI-compatible version of a Pydantic schema.
+
+    OpenAI's strict mode requires that all objects have 'additionalProperties: false'.
+    This function ensures the schema meets those requirements while preserving
+    Pydantic's correct handling of required vs optional fields.
+
+    Uses caching to avoid repeated transformation of the same schema.
+
+    Args:
+        schema: Original Pydantic model
+
+    Returns:
+        Wrapper model that generates OpenAI-compatible schemas
+    """
+    return _get_cached_openai_schema(schema)
+
+
+def make_outlines_model(provider: str, model_name: str, **kwargs) -> Any:
+    """Create an Outlines model for the specified provider and model.
+
+    Args:
+        provider: Provider name (openai, anthropic, gemini, ollama)
+        model_name: Model identifier
+        **kwargs: Additional parameters passed to the client
+
+    Returns:
+        Outlines model instance
+
+    Raises:
+        DataSetGeneratorError: If provider is unsupported or configuration fails
+    """
+    try:
+        if provider == "openai":
+            client = _create_openai_compatible_client(
+                api_key_env_var="OPENAI_API_KEY",
+                default_base_url=None,  # Use OpenAI's default
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "ollama":
+            client = _create_openai_compatible_client(
+                api_key_env_var=None,  # No API key required
+                default_base_url="http://localhost:11434/v1",
+                dummy_key="ollama",
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "openrouter":
+            client = _create_openai_compatible_client(
+                api_key_env_var="OPENROUTER_API_KEY",
+                default_base_url="https://openrouter.ai/api/v1",
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "anthropic":
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key:
+                _raise_api_key_error("ANTHROPIC_API_KEY")
+            # Use native Anthropic structured outputs API
+            client = anthropic.AsyncAnthropic(api_key=api_key, **kwargs)
+            return AnthropicModel(client, model_name)
+
+        if provider == "gemini":
+            api_key = _get_gemini_api_key()
+            # Use direct Gemini API instead of Outlines for better structured output reliability
+            client = genai.Client(api_key=api_key)
+            return GeminiModel(client, model_name)
+
+        _raise_unsupported_provider_error(provider)
+
+    except DataSetGeneratorError:
+        # Re-raise our own errors (like missing API keys)
+        raise
+    except Exception as e:
+        # Use the organized error handler
+        raise handle_provider_error(e, provider, model_name) from e
+
+
+def make_async_outlines_model(provider: str, model_name: str, **kwargs) -> Any | None:
+    """Create an async Outlines model when the provider supports it.
+
+    Returns ``None`` for providers without async-capable clients.
+    """
+
+    try:
+        if provider == "openai":
+            client = _create_async_openai_compatible_client(
+                api_key_env_var="OPENAI_API_KEY",
+                default_base_url=None,  # Use OpenAI's default
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "ollama":
+            client = _create_async_openai_compatible_client(
+                api_key_env_var=None,  # No API key required
+                default_base_url="http://localhost:11434/v1",
+                dummy_key="ollama",
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "openrouter":
+            client = _create_async_openai_compatible_client(
+                api_key_env_var="OPENROUTER_API_KEY",
+                default_base_url="https://openrouter.ai/api/v1",
+                **kwargs,
+            )
+            return outlines.from_openai(client, model_name)
+
+        if provider == "anthropic":
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key:
+                _raise_api_key_error("ANTHROPIC_API_KEY")
+            # Use native Anthropic structured outputs API
+            client = anthropic.AsyncAnthropic(api_key=api_key, **kwargs)
+            return AnthropicModel(client, model_name)
+
+        if provider == "gemini":
+            api_key = _get_gemini_api_key()
+            # Use direct async Gemini API for better structured output reliability
+            client = genai.Client(api_key=api_key)
+            return GeminiModel(client, model_name)
+
+    except DataSetGeneratorError:
+        raise
+    except Exception as e:
+        raise handle_provider_error(e, provider, model_name) from e
+
+    # Outlines does not currently expose async structured generation wrappers
+    # for the remaining providers. Fallback to synchronous execution later.
+    return None