Implement structured output generation for both LlamaLanguageModel / MLXLanguageModel #75
Conversation
There was a problem hiding this comment.
Pull request overview
This PR implements structured output generation for LlamaLanguageModel and MLXLanguageModel by adding constrained token sampling to generate JSON that conforms to a schema. The implementation includes comprehensive tests covering various data types and structures.
Key changes:
- Added
ConstrainedJSONGeneratorthat uses token-level sampling to generate schema-conformant JSON - Implemented
TokenBackendprotocol with adapters for both Llama and MLX models - Enhanced
GenerationGuideto store constraint values for min/max on numbers and arrays - Extended
GenerationSchemawith character validation and schema prompt generation
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| Tests/AnyLanguageModelTests/StructuredGenerationTests.swift | Comprehensive test suite covering simple types, nested structs, enums, arrays, and optionals across all supported model types |
| Tests/AnyLanguageModelTests/GenerableMacroTests.swift | Added round-trip tests for enums, nested structs, and arrays |
| Sources/AnyLanguageModelMacros/GenerableMacro.swift | Refactored guide extraction to use a structured Constraints type and properly parse numeric ranges and array count constraints |
| Sources/AnyLanguageModel/StructuredGeneration.swift | New file implementing token-level constrained JSON generation with TokenBackend protocol and ConstrainedJSONGenerator |
| Sources/AnyLanguageModel/Models/SystemLanguageModel.swift | Updated to use schema-based generation for non-String types and added conversion to FoundationModels.DynamicGenerationSchema |
| Sources/AnyLanguageModel/Models/MLXLanguageModel.swift | Implemented MLXTokenBackend and structured JSON generation with proper token sampling and repetition penalty handling |
| Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift | Implemented LlamaTokenBackend and structured JSON generation with batch-based decoding and sampler integration |
| Sources/AnyLanguageModel/GenerationSchema.swift | Added schemaPrompt() method, character validation for JSON strings, improved node equality checking, and support for constraint propagation |
| Sources/AnyLanguageModel/GenerationGuide.swift | Made GenerationGuide store actual constraint values (min/max, minCount/maxCount) for use during schema generation |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@eastriverlee Thank you for your contribution! And thank you for your patience. I'll have a chance to look a this soon. |
|
@eastriverlee Thanks again for your patience. I just rebased, resolving the conflicts as best I could. I recently merged #59, which takes a slightly different approach for schema conversion. I'm working to harmonize these implementations now... |
|
Nice! We've got CI passing, and everything looking good. Just going to give this one last automated PR review, and we should be good to merge. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 8 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Extract Character extension to separate file
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…e, .count, .minimumCount, .maximumCount instead of GenerationGuide(minimum:maximum:) Updated GenerableMacro to implement Float/Decimal guide methods to set min/max bounds so constraints are preserved
…inators during structured JSON generation, so the sampler can’t select EOS/EOT mid-structure and return a partial/non-object
…lti-string structures
|
Decided to let this one cook a bit more to make sure the implementation was correct. Kicking off another review after making some changes. Let's see what comes back this time... |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 18 changed files in this pull request and generated 21 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@eastriverlee Merged! Thanks for your work on this and your patience with my review. Really excited to support guided generation / structured output for MLX and Llama. |
2 participants
Related to #27