feat: add response_format parameter to [params] config#76
Conversation
Add support for the response_format parameter, enabling provider-level
structured output constraints. When set, the LLM is constrained at the
decoding level to produce valid JSON matching the configured format.
Supported types:
- json_object: forces valid JSON output (no schema constraint)
- json_schema: forces JSON conforming to a provided JSON Schema
Provider support:
- OpenAI/OpenAI-compatible: maps to the response_format request field
- Ollama: maps to the format field ("json" or schema object)
- Other providers: field is silently ignored (no breaking change)
TOML usage:
[params.response_format]
type = "json_schema"
[params.response_format.schema]
name = "extraction"
type = "object"
[params.response_format.schema.properties.result]
type = "string"
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a ResponseFormat abstraction and config, validates and parses it from TOML, wires it into CLI requests, and implements provider-specific serialization/validation for OpenAI and Ollama with accompanying unit tests. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~28 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
internal/provider/provider.go (1)
71-74: Clarify thatIsSetmeans “constraint active,” not merely configured.
Type: "text"can be explicitly configured but returns false here. Renaming this to something likeIsConstrained()or adjusting the comment would make the contract less slippery.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/provider/provider.go` around lines 71 - 74, The method ResponseFormat.IsSet() currently returns false for Type == "text" which is confusing because that value can be explicitly configured; update naming/comment to reflect that the method checks whether a constraint is active (not merely configured). Rename the method IsSet to IsConstrained (or change the doc comment) so the contract matches behavior, and update all call sites referencing ResponseFormat.IsSet to use ResponseFormat.IsConstrained (or keep IsSet and change its comment to "IsConstrained / returns true when a constraint is active"); ensure the symbol ResponseFormat and its method rename are updated consistently across the codebase and unit tests.internal/provider/openai_test.go (1)
1520-1647: Add response_format coverage forSendStreamtoo.These tests nicely pin the non-streaming wire format, but
SendStreambuilds its request separately and now has its own response_format branch. A small table over{name, request, want, streaming}would catch regressions in both paths without letting the stream gremlin sneak by wearing sunglasses. As per coding guidelines,**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/provider/openai_test.go` around lines 1520 - 1647, Add table-driven tests mirroring the three existing TestOpenAI_Send_* cases but exercising SendStream instead of Send: create a []struct with name, request (using Request and ResponseFormat), expected presence/type of response_format, and streaming=true; for each case start an httptest.Server that inspects the incoming request body for "response_format" like the existing tests, call NewOpenAI(..., WithOpenAIBaseURL(server.URL)).SendStream(ctx, req), and assert the response_format presence and fields (type, json_schema.strict/name/schema) match expectations; ensure you cover the omit case, json_object case, and json_schema case and follow the table-driven pattern per test file guidelines.internal/provider/ollama_test.go (1)
1198-1301: Exercise the streaming format path as well.These cover
Send, butSendStreamhas a separate request body and also setsformat. A table-driven test that runs the sameResponseFormatcases through both paths would keep the Ollama stream path honest. As per coding guidelines,**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/provider/ollama_test.go` around lines 1198 - 1301, Add a table-driven test that exercises both Send and SendStream for the same ResponseFormat cases (omitted, json_object, json_schema) so the stream path is validated; create a []struct{name string, format ResponseFormat, want...} loop and for each case call o.Send and o.SendStream (or a helper that invokes both) against an httptest.Server handler that inspects the request body for the "format" field (as existing tests do), asserting the expected presence/shape (string "json" for json_object, stripped schema object for json_schema, and omitted when empty); reuse the existing test logic and symbols (Send, SendStream, ResponseFormat) and follow the table-driven pattern used in other *_test.go files.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmd/run.go`:
- Around line 403-408: When populating req.ResponseFormat and when printing the
verbose/dry-run paramsLine, avoid treating the literal "text" as an active
constraint; use the same predicate the provider uses (i.e., the
provider.ResponseFormat.IsSet() semantics) instead of
cfg.Params.ResponseFormat.IsSet() so that "text" is treated as unset/omitted.
Concretely, change the branches that currently check
cfg.Params.ResponseFormat.IsSet() (the block assigning req.ResponseFormat and
the verbose/dry-run paramsLine branches) to check the provider-style predicate
or replicate its logic (ignore Type == "text" as equivalent to unset) before
assigning req.ResponseFormat or appending response_format to the printed
paramsLine.
In `@internal/agent/agent.go`:
- Around line 191-200: The validator currently lets a
params.response_format.schema be provided without a response_format.type and
silently ignores it; update the validation in agent.go to explicitly reject
cases where cfg.Params.ResponseFormat.Schema is non-empty but
cfg.Params.ResponseFormat.Type is empty or missing: add a guard that checks
len(cfg.Params.ResponseFormat.Schema) > 0 && cfg.Params.ResponseFormat.Type ==
"" and return a ValidationError with a clear message (e.g.,
"params.response_format.type is required when params.response_format.schema is
provided"). Keep the existing type whitelist check (switch on
cfg.Params.ResponseFormat.Type) and the json_schema-specific schema-required
check, but ensure the new guard runs early enough to catch schema-only blocks
even when IsSet() is false.
In `@internal/provider/ollama.go`:
- Around line 544-557: The helper buildOllamaFormat currently downgrades
ResponseFormat{Type:"json_schema"} with an empty Schema to "json"; change
buildOllamaFormat to return an error when rf.Type == "json_schema" but rf.Schema
is empty (instead of returning "json") so callers know the input constraint
failed, and keep the existing behavior of returning the schema map when
non-empty; update Send and SendStream to call buildOllamaFormat, detect that
error, and wrap it in a ProviderError with category Input before
marshaling/returning (use the ProviderError type and Input category to follow
the provider error guidelines).
In `@internal/provider/openai.go`:
- Around line 353-374: The buildOpenAIResponseFormat function must validate that
a ResponseFormat with Type == "json_schema" contains a non-empty schema (aside
from an optional "name") and return a categorized ProviderError instead of
producing an invalid payload; change buildOpenAIResponseFormat to return
(*openaiResponseFormat, error), check when rf.Type == "json_schema" that
rf.Schema has at least one key other than "name" (or else return a ProviderError
indicating a malformed response format), and on success construct and return the
openaiResponseFormat (openaiJSONSchemaConfig{Name, Strict:true, Schema:...})
with nil error; update all callers of buildOpenAIResponseFormat to handle the
error and propagate or surface the ProviderError appropriately.
---
Nitpick comments:
In `@internal/provider/ollama_test.go`:
- Around line 1198-1301: Add a table-driven test that exercises both Send and
SendStream for the same ResponseFormat cases (omitted, json_object, json_schema)
so the stream path is validated; create a []struct{name string, format
ResponseFormat, want...} loop and for each case call o.Send and o.SendStream (or
a helper that invokes both) against an httptest.Server handler that inspects the
request body for the "format" field (as existing tests do), asserting the
expected presence/shape (string "json" for json_object, stripped schema object
for json_schema, and omitted when empty); reuse the existing test logic and
symbols (Send, SendStream, ResponseFormat) and follow the table-driven pattern
used in other *_test.go files.
In `@internal/provider/openai_test.go`:
- Around line 1520-1647: Add table-driven tests mirroring the three existing
TestOpenAI_Send_* cases but exercising SendStream instead of Send: create a
[]struct with name, request (using Request and ResponseFormat), expected
presence/type of response_format, and streaming=true; for each case start an
httptest.Server that inspects the incoming request body for "response_format"
like the existing tests, call NewOpenAI(...,
WithOpenAIBaseURL(server.URL)).SendStream(ctx, req), and assert the
response_format presence and fields (type, json_schema.strict/name/schema) match
expectations; ensure you cover the omit case, json_object case, and json_schema
case and follow the table-driven pattern per test file guidelines.
In `@internal/provider/provider.go`:
- Around line 71-74: The method ResponseFormat.IsSet() currently returns false
for Type == "text" which is confusing because that value can be explicitly
configured; update naming/comment to reflect that the method checks whether a
constraint is active (not merely configured). Rename the method IsSet to
IsConstrained (or change the doc comment) so the contract matches behavior, and
update all call sites referencing ResponseFormat.IsSet to use
ResponseFormat.IsConstrained (or keep IsSet and change its comment to
"IsConstrained / returns true when a constraint is active"); ensure the symbol
ResponseFormat and its method rename are updated consistently across the
codebase and unit tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 41dfddec-826e-4cbf-a6e0-91b14a91dc80
📒 Files selected for processing (8)
cmd/run.gointernal/agent/agent.gointernal/agent/agent_test.gointernal/provider/ollama.gointernal/provider/ollama_test.gointernal/provider/openai.gointernal/provider/openai_test.gointernal/provider/provider.go
Address CodeRabbit review feedback: - Align agent-level IsSet() to exclude "text" (matches provider behavior) - Reject schema blocks that omit response_format.type - Return ProviderError from buildOllamaFormat/buildOpenAIResponseFormat instead of silently degrading when schema is missing - Add tests for all new validation paths
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/provider/ollama.go (1)
552-578:buildOllamaFormatlooks good — previous concern resolved.The helper now:
- Maps
json_object→"json".- Requires a non-empty
Schemaforjson_schemaand returns aProviderError{Category: ErrCategoryBadRequest}otherwise (no more sneaky downgrade to"json"🙂).- Strips the
"name"key from the schema map to matchbuildOpenAIResponseFormat's behavior.- Rejects unknown types with an actionable
ProviderError.Error messages are user-actionable and the category routing through
mapProviderErrorincmd/run.gowill land on exit code 1 as intended. As per coding guidelines, "Provider errors must be wrapped in ProviderError with category (Auth, RateLimit, Server, Input, Network, Unknown)" — both error paths here are properly wrapped.One tiny thought for later (not blocking): since
buildOpenAIResponseFormatand this function share the "copy map minusname" dance, a small shared helper (e.g.,stripSchemaName(rf.Schema)) inprovider.gocould keep the two provider mappings from drifting if the stripped-field set ever grows. Totally fine to defer.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/provider/ollama.go` around lines 552 - 578, Extract the repeated "copy schema map minus the 'name' key" logic into a small helper (e.g., stripSchemaName) and use it from buildOllamaFormat and buildOpenAIResponseFormat to avoid duplication and keep behavior in sync; implement a function that accepts map[string]interface{} and returns a new map[string]interface{} without the "name" key, then replace the manual copy loops in buildOllamaFormat and buildOpenAIResponseFormat with calls to that helper.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@internal/provider/ollama.go`:
- Around line 552-578: Extract the repeated "copy schema map minus the 'name'
key" logic into a small helper (e.g., stripSchemaName) and use it from
buildOllamaFormat and buildOpenAIResponseFormat to avoid duplication and keep
behavior in sync; implement a function that accepts map[string]interface{} and
returns a new map[string]interface{} without the "name" key, then replace the
manual copy loops in buildOllamaFormat and buildOpenAIResponseFormat with calls
to that helper.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: cb0eb25c-ee3a-42b2-9955-40643d90d79b
📒 Files selected for processing (6)
internal/agent/agent.gointernal/agent/agent_test.gointernal/provider/ollama.gointernal/provider/ollama_test.gointernal/provider/openai.gointernal/provider/openai_test.go
✅ Files skipped from review due to trivial changes (1)
- internal/agent/agent_test.go
🚧 Files skipped from review as they are similar to previous changes (4)
- internal/provider/ollama_test.go
- internal/provider/openai.go
- internal/agent/agent.go
- internal/provider/openai_test.go
|
@jrswab - for additional context on the motivation here: we are building a pipeline that uses axe to extract structured data from web-harvested content, which then gets ingested into Snowflake for analytics. The response_format parameter lets us enforce valid JSON output at the provider level, which is critical for reliable automated extraction at scale without manual cleanup. |
Adds a custom UnmarshalTOML method so response_format accepts both forms: [params] response_format = "json_object" # string shorthand for the common case [params.response_format] # table form when schema validation is needed type = "json_schema" [params.response_format.schema] name = "my_schema" type = "object" Updates scaffold template to show the string shorthand as the primary example. Adds 5 new test cases for string shorthand parsing and validation. All existing table-format tests continue to pass.
Summary
response_formatsupport to the[params]TOML config, enabling provider-level structured output constraints. When configured, the LLM's token generation is constrained at decoding time to produce valid JSON - a decoding constraint enforced by the provider, not a prompt-level instruction.response_formatrequest field) and Ollama (maps to theformatfield). Providers that do not support this parameter are unaffected - no breaking change.Motivation
Axe agents that need structured data back from the LLM currently rely on prompt instructions ("respond in valid JSON matching this schema"). This is generally reliable, but at scale (thousands of invocations), even a low failure rate produces parse errors that require retry logic and manual intervention.
The OpenAI API and Ollama both support
response_formatnatively - it constrains token sampling so the model cannot produce invalid JSON. This parameter follows the same pattern astemperatureandmax_tokensin[params]: a provider-level API parameter that axe passes through to the request body.Supported types
json_objectjson_schematext(or omitted)TOML usage
Changes
internal/agent/agent.goResponseFormatConfigstruct, validation, scaffold templateinternal/provider/provider.goResponseFormattype on the universalRequeststructcmd/run.gointernal/provider/openai.goresponse_formatfield (Send + SendStream)internal/provider/ollama.goformatfield ("json"or schema object)*_test.go(3 files)Test plan