-
Notifications
You must be signed in to change notification settings - Fork 47
Implement structured output generation for both LlamaLanguageModel / MLXLanguageModel #75
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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... |
Related to #27