fix: make structured output work with reasoning models and empty responses

altk/core/llm/output_parser.py

@@ -5,6 +5,7 @@
     Any,
     Dict,
     List,
+    Literal,
     Optional,
     Type,
     TypeVar,
@@ -59,25 +60,59 @@ def json_schema_to_pydantic_model(
     }
 
     def _map_object_for_prop(prop_schema: Dict[str, Any]) -> Type:
-        """Return dict/str for a property whose declared type is ``object``.
+        """Return a model/dict/str for a property whose declared type is ``object``.
 
         A property is "free-form" if it has no ``properties`` sub-schema; the
-        OpenAI workaround only applies to those.
+        OpenAI workaround only applies to those. An object *with* properties is
+        recursed into so nested constraints survive the conversion — a bare
+        ``dict`` would erase them and let providers emit output that fails the
+        original JSON Schema.
         """
-        if free_form_object_as_str and "properties" not in prop_schema:
+        if "properties" in prop_schema:
+            return json_schema_to_pydantic_model(
+                prop_schema,
+                model_name=f"{model_name}_{_next_nested_id()}",
+                free_form_object_as_str=free_form_object_as_str,
+            )
+        if free_form_object_as_str:
             return str
         return dict
 
+    _nested_count = [0]
+
+    def _next_nested_id() -> int:
+        _nested_count[0] += 1
+        return _nested_count[0]
+
     def parse_type(
         type_def: Union[str, List[str], None],
         prop_schema: Dict[str, Any],
     ) -> Type[T]:
         def _lookup(t: str) -> Type:
-            return (
-                _map_object_for_prop(prop_schema)
-                if t == "object"
-                else type_mapping.get(t, Any)
-            )
+            if t == "object":
+                return _map_object_for_prop(prop_schema)
+            if t == "array":
+                return _map_array_for_prop(prop_schema)
+            return type_mapping.get(t, Any)
+
+        def _map_array_for_prop(prop_schema: Dict[str, Any]) -> Type:
+            """Preserve ``items`` so array element constraints are not lost."""
+            items = prop_schema.get("items")
+            if not isinstance(items, dict):
+                return list
+            item_type = parse_type(items.get("type"), items)
+            return List[item_type]  # type: ignore[valid-type]
+
+        # ``enum`` becomes a real ``Literal`` type so the choices survive even
+        # inside ``items``, where a Field-level constraint could not reach.
+        enum_values = prop_schema.get("enum")
+        if enum_values and all(
+            isinstance(v, (str, int, bool)) or v is None for v in enum_values
+        ):
+            literal = Literal[tuple(enum_values)]  # type: ignore[valid-type]
+            if isinstance(type_def, list) and "null" in type_def:
+                return Optional[literal]  # type: ignore[return-value]
+            return literal  # type: ignore[return-value]
 
         if isinstance(type_def, list):
             python_types = [_lookup(t) for t in type_def]
@@ -93,14 +128,38 @@ def _lookup(t: str) -> Type:
             return _lookup(type_def)
         return Any  # type: ignore[return-value]
 
+    # JSON Schema keyword -> Pydantic ``Field`` argument. Carrying these over
+    # keeps provider-native structured output faithful to the source schema
+    # (a dropped ``minimum``/``enum`` shows up later as a validation failure).
+    _CONSTRAINT_ARGS = {
+        "minimum": "ge",
+        "maximum": "le",
+        "exclusiveMinimum": "gt",
+        "exclusiveMaximum": "lt",
+        "minLength": "min_length",
+        "maxLength": "max_length",
+        "minItems": "min_length",
+        "maxItems": "max_length",
+        "pattern": "pattern",
+    }
+
     for prop_name, prop_schema in schema.get("properties", {}).items():
         field_type: Any = parse_type(prop_schema.get("type"), prop_schema)
         default = ... if prop_name in required_fields else None
         description = prop_schema.get("description", None)
-        field_args = {"description": description} if description else {}
+        field_args: Dict[str, Any] = {"description": description} if description else {}
+        for json_kw, field_kw in _CONSTRAINT_ARGS.items():
+            if json_kw in prop_schema and field_kw not in field_args:
+                field_args[field_kw] = prop_schema[json_kw]
         fields[prop_name] = (field_type, Field(default, **field_args))
 
-    return create_model(model_name, **fields)  # type: ignore
+    model = create_model(model_name, **fields)  # type: ignore
+    # Mirror ``additionalProperties: false`` — providers with strict structured
+    # output need it, and without it the model may invent extra keys that the
+    # original schema then rejects.
+    if schema.get("additionalProperties") is False:
+        model.model_config["extra"] = "forbid"
+    return model
 
 
 def relax_freeform_object_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
@@ -129,6 +188,26 @@ class OutputValidationError(Exception):
     """Raised when LLM output cannot be validated against the provided schema."""
 
 
+def _is_truncated(raw: Any) -> bool:
+    """Return ``True`` when *raw* was cut off by the token limit.
+
+    A ``finish_reason`` of ``"length"`` means the model never finished writing,
+    so whatever came back cannot be valid JSON. Reasoning ("thinking") models
+    hit this routinely: the reasoning tokens consume the whole budget and the
+    content field arrives empty.
+    """
+    choices = getattr(raw, "choices", None) or (
+        raw.get("choices", []) if isinstance(raw, dict) else []
+    )
+    if not choices:
+        return False
+    first = choices[0]
+    finish = getattr(first, "finish_reason", None) or (
+        first.get("finish_reason") if isinstance(first, dict) else None
+    )
+    return finish == "length"
+
+
 class ValidatingLLMClient(BaseLLMClient, ABC):
     """
     An LLMClient wrapper enforcing output structure via:
@@ -178,6 +257,10 @@ def __init__(
         self.default_generation_kwargs: Dict[str, Any] = dict(
             default_generation_kwargs or {}
         )
+        # Set by the wrapped parser: was the most recent reply cut off by the
+        # token limit? Retries use it to grow ``max_tokens`` instead of
+        # re-asking with a budget already known to be too small.
+        self._last_response_truncated: bool = False
         super().__init__(**base_kwargs)
         # Wrap the subclass's _parse_llm_response so empty / malformed LLM
         # outputs retry gracefully (the retry loop treats "" as invalid)
@@ -204,8 +287,18 @@ def configure_validation(
             self.default_generation_kwargs = dict(default_generation_kwargs)
         return self
 
-    @staticmethod
-    def _build_safe_parse(orig):  # noqa: ANN001, ANN205
+    #: Budget assumed to be in play when the provider applied its own default
+    #: (``max_tokens`` was never passed) and the reply came back truncated.
+    ASSUMED_PROVIDER_MAX_TOKENS: int = 1024
+    #: Ceiling for the retry escalation, so a stuck model cannot grow forever.
+    MAX_TOKENS_ESCALATION_LIMIT: int = 16384
+
+    def _escalate_max_tokens(self, current: Optional[int]) -> int:
+        """Return a larger ``max_tokens`` for the next attempt after truncation."""
+        base = current or self.ASSUMED_PROVIDER_MAX_TOKENS
+        return min(base * 4, self.MAX_TOKENS_ESCALATION_LIMIT)
+
+    def _build_safe_parse(self, orig):  # noqa: ANN001, ANN205
         """Wrap ``_parse_llm_response`` so parse failures become retry-worthy
         empty strings instead of raising. Also surfaces a targeted warning
         when a reasoning-only response exhausted the token budget."""
@@ -214,6 +307,7 @@ def _build_safe_parse(orig):  # noqa: ANN001, ANN205
         _logger = _logging.getLogger("altk.core.llm.output_parser")
 
         def _safe_parse(raw):  # noqa: ANN001, ANN202
+            self._last_response_truncated = _is_truncated(raw)
             try:
                 return orig(raw)
             except (ValueError, KeyError):
@@ -251,6 +345,34 @@ def _safe_parse(raw):  # noqa: ANN001, ANN202
     def provider_class(cls) -> Type[Any]:
         """Return the underlying SDK client class, e.g. openai.OpenAI."""
 
+    def supports_native_structured_output(self) -> bool:
+        """Whether the target model honors a native structured-output kwarg.
+
+        Defaults to ``True`` (previous behavior). Providers that can tell which
+        models support it override this; when it returns ``False`` the schema is
+        injected into the system prompt instead, because a model that ignores
+        ``response_format`` cannot be constrained by it.
+        """
+        return True
+
+    def _render_native_schema(
+        self, schema: Union[Dict[str, Any], Type[BaseModel], Type[Any]]
+    ) -> Any:
+        """Render *schema* into the value this provider expects for its native
+        structured-output kwarg (``schema_field``).
+
+        The default converts a JSON Schema dict into a Pydantic model, which is
+        what litellm accepts. Providers whose SDK rejects a model class
+        override this — see the OpenAI/Azure clients, which need a
+        ``{"type": "json_schema", ...}`` dict for ``chat.completions.create``.
+        """
+        if isinstance(schema, dict):
+            return json_schema_to_pydantic_model(
+                schema,
+                free_form_object_as_str=self.free_form_object_as_str,
+            )
+        return schema
+
     @abstractmethod
     def _register_methods(self) -> None:
         """
@@ -406,19 +528,19 @@ def generate(
         if self.prompt_based_validation:
             include_schema_in_system_prompt = True
             schema_field = None
+        # Models that ignore a native schema kwarg must be steered by the
+        # prompt instead; sending ``response_format`` to them is at best a
+        # no-op and at worst returns empty content.
+        elif schema_field and not self.supports_native_structured_output():
+            include_schema_in_system_prompt = True
+            schema_field = None
         current = prompt
         instr = None
         if include_schema_in_system_prompt:
             instr = self._make_instruction(schema)
             current = self._inject_system(prompt, instr)
         if schema_field:
-            kwargs[schema_field] = schema
-            if isinstance(schema, dict):
-                new_schema = json_schema_to_pydantic_model(
-                    schema,
-                    free_form_object_as_str=self.free_form_object_as_str,
-                )
-                kwargs[schema_field] = new_schema
+            kwargs[schema_field] = self._render_native_schema(schema)
 
         last_error: Optional[str] = None
         for _ in range(1, retries + 1):
@@ -434,13 +556,39 @@ def generate(
                     "include_schema_in_system_prompt",
                 ]
             }
-            raw = super()._generate(**{"prompt": current, **filtered_kwargs})
+            raw = ""
             try:
+                # Inside the try: ``_parse_llm_response`` raises ValueError for a
+                # contentless reply, and that must be retried like any other
+                # invalid output rather than aborting the whole call.
+                raw = super()._generate(**{"prompt": current, **filtered_kwargs})
                 if isinstance(raw, str):
                     return self._validate(raw, schema)
                 return raw
-            except OutputValidationError as e:
+            except (OutputValidationError, ValueError) as e:
+                # ValueError covers providers whose ``_parse_llm_response``
+                # rejects an empty/contentless response ("No content or tool
+                # calls found in response"). Without it, a single blank reply
+                # from the backend aborts the whole call and the configured
+                # ``retries`` are never used.
                 last_error = str(e)
+                # An empty response carries no mistake to correct: it is a
+                # transient backend failure (empty content under load). Retry
+                # the original prompt untouched — appending correction turns
+                # would grow a conversation that several backends answer with
+                # another empty response, burning every remaining attempt.
+                if not (isinstance(raw, str) and raw.strip()):
+                    # Truncated by the token limit? Re-asking with the same
+                    # budget yields the identical truncation, so grow it. This
+                    # is the common failure for reasoning models, whose
+                    # "thinking" tokens can consume a small default budget
+                    # (watsonx defaults to 1024) before any content is emitted.
+                    if self._last_response_truncated:
+                        kwargs["max_tokens"] = self._escalate_max_tokens(
+                            kwargs.get("max_tokens")
+                        )
+                    current = self._inject_system(prompt, instr) if instr else prompt
+                    continue
                 correction = (
                     f"The previous response did not conform: {last_error}\nPlease correct it."
                     " And remember to output ONLY the requested schema, without any additional text."
@@ -484,19 +632,19 @@ async def generate_async(
         if self.prompt_based_validation:
             include_schema_in_system_prompt = True
             schema_field = None
+        # Models that ignore a native schema kwarg must be steered by the
+        # prompt instead; sending ``response_format`` to them is at best a
+        # no-op and at worst returns empty content.
+        elif schema_field and not self.supports_native_structured_output():
+            include_schema_in_system_prompt = True
+            schema_field = None
         current = prompt
         instr = None
         if include_schema_in_system_prompt:
             instr = self._make_instruction(schema)
             current = self._inject_system(prompt, instr)
         if schema_field:
-            kwargs[schema_field] = schema
-            if isinstance(schema, dict):
-                new_schema = json_schema_to_pydantic_model(
-                    schema,
-                    free_form_object_as_str=self.free_form_object_as_str,
-                )
-                kwargs[schema_field] = new_schema
+            kwargs[schema_field] = self._render_native_schema(schema)
 
         last_error: Optional[str] = None
         for _ in range(1, retries + 1):
@@ -512,15 +660,41 @@ async def generate_async(
                     "include_schema_in_system_prompt",
                 ]
             }
-            raw = await super()._generate_async(
-                **{"prompt": current, **filtered_kwargs}
-            )
+            raw = ""
             try:
+                # Inside the try: ``_parse_llm_response`` raises ValueError for a
+                # contentless reply, and that must be retried like any other
+                # invalid output rather than aborting the whole call.
+                raw = await super()._generate_async(
+                    **{"prompt": current, **filtered_kwargs}
+                )
                 if isinstance(raw, str):
                     return self._validate(raw, schema)
                 return raw
-            except OutputValidationError as e:
+            except (OutputValidationError, ValueError) as e:
+                # ValueError covers providers whose ``_parse_llm_response``
+                # rejects an empty/contentless response ("No content or tool
+                # calls found in response"). Without it, a single blank reply
+                # from the backend aborts the whole call and the configured
+                # ``retries`` are never used.
                 last_error = str(e)
+                # An empty response carries no mistake to correct: it is a
+                # transient backend failure (empty content under load). Retry
+                # the original prompt untouched — appending correction turns
+                # would grow a conversation that several backends answer with
+                # another empty response, burning every remaining attempt.
+                if not (isinstance(raw, str) and raw.strip()):
+                    # Truncated by the token limit? Re-asking with the same
+                    # budget yields the identical truncation, so grow it. This
+                    # is the common failure for reasoning models, whose
+                    # "thinking" tokens can consume a small default budget
+                    # (watsonx defaults to 1024) before any content is emitted.
+                    if self._last_response_truncated:
+                        kwargs["max_tokens"] = self._escalate_max_tokens(
+                            kwargs.get("max_tokens")
+                        )
+                    current = self._inject_system(prompt, instr) if instr else prompt
+                    continue
                 correction = (
                     f"The previous response did not conform: {last_error}\nPlease correct it."
                     " And remember to output ONLY the requested schema, without any additional text."