feat(generation): add support for custom structured output schema#634
feat(generation): add support for custom structured output schema#634lukehinds wants to merge 1 commit into
Conversation
- Introduce `generation.output_schema` and `output.format` parameters to configure direct structured output based on a user-defined JSON schema. - When `output_schema` is provided, the generator now bypasses the standard conversation format and uses constrained decoding to generate output that adheres directly to the specified schema. - Implement a dynamic model creation utility (`make_dynamic_model`) to convert a raw JSON schema dictionary into a Pydantic-compatible class, enabling its use with constrained decoding. - Update the `outlines` and `transformers` dependencies to their latest versions to support these new capabilities. Signed-off-by: Luke Hinds <lukehinds@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for custom structured output schemas, enabling constrained decoding directly into user-defined JSON formats. Key updates include new configuration parameters, a dynamic model factory to bridge raw schemas with the generation pipeline, and a dedicated generation path that bypasses the standard conversation builder. Review feedback highlights that the dynamic model implementation lacks actual schema validation and robust serialization for complex structures. Additionally, the reviewer recommends unifying the retry logic in the generator to eliminate duplication and ensure consistent error handling across all generation modes.
| def model_validate_json(cls, json_data: str) -> "DynamicModel": | ||
| try: | ||
| data = json.loads(json_data) | ||
| except json.JSONDecodeError as exc: | ||
| raise ValueError(f"Invalid JSON from model: {exc}") from exc | ||
| return cls(data) |
There was a problem hiding this comment.
The model_validate_json method only performs a basic json.loads() check and does not verify that the resulting data conforms to the provided schema_dict. Since this class is intended to replace a Pydantic model in the structured output pipeline, it should ideally perform schema validation to ensure downstream components receive data in the expected format. If a full validation library like jsonschema is not available, consider at least documenting this limitation.
| def model_dump(self, exclude_none: bool = False, **_kwargs: Any) -> dict: | ||
| if exclude_none: | ||
| return {k: v for k, v in self._data.items() if v is not None} | ||
| return dict(self._data) |
There was a problem hiding this comment.
The model_dump implementation is shallow and its exclude_none logic is not recursive. Additionally, it assumes self._data is a dictionary (by calling .items()), which will fail if the JSON schema defines a top-level array. Pydantic's model_dump typically handles nested structures and various JSON types. Consider making this implementation more robust to match the expected Pydantic interface.
| if config.output_schema: | ||
| schema_model = self._get_custom_schema_model() | ||
| for attempt in range(max_attempts): | ||
| try: | ||
| result = await self.llm_client.generate_async( | ||
| prompt, | ||
| schema_model, | ||
| max_tokens=config.max_tokens, | ||
| ) | ||
| return True, result | ||
| except Exception as e: # noqa: BLE001 | ||
| last_error = e | ||
| if is_validation_error(e) and attempt < self.config.sample_retries: | ||
| self._emit_retry(sample_idx, attempt, max_attempts, e) | ||
| continue | ||
| return False, last_error | ||
| return False, last_error or Exception("Custom schema generation failed") |
There was a problem hiding this comment.
This block introduces a separate retry loop for custom schema generation that duplicates logic from the standard path. Furthermore, the except block at line 1066 returns immediately for any error that is not a validation error (e.g., transient network issues or rate limits), whereas the standard path likely retries such cases. It is recommended to unify the retry logic to ensure consistent error handling and reduce maintenance overhead.
1 participant
generation.output_schemaandoutput.formatparameters to configure direct structured output based on a user-defined JSON schema.output_schemais provided, the generator now bypasses the standard conversation format and uses constrained decoding to generate output that adheres directly to the specified schema.make_dynamic_model) to convert a raw JSON schema dictionary into a Pydantic-compatible class, enabling its use with constrained decoding.outlinesandtransformersdependencies to their latest versions to support these new capabilities.