Skip to content

feat: add response_format parameter to [params] config#76

Merged
jrswab merged 3 commits into
jrswab:masterfrom
tllongdev:feat/response-format
Apr 21, 2026
Merged

feat: add response_format parameter to [params] config#76
jrswab merged 3 commits into
jrswab:masterfrom
tllongdev:feat/response-format

Conversation

@tllongdev

@tllongdev tllongdev commented Apr 21, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Adds response_format support 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.
  • Implements for OpenAI/OpenAI-compatible APIs (maps to the response_format request field) and Ollama (maps to the format field). 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_format natively - it constrains token sampling so the model cannot produce invalid JSON. This parameter follows the same pattern as temperature and max_tokens in [params]: a provider-level API parameter that axe passes through to the request body.

Supported types

Type Behavior
json_object Constrains output to valid JSON, no schema enforcement
json_schema Constrains output to JSON conforming to a provided JSON Schema
text (or omitted) Default behavior, no constraint

TOML usage

# Constrain output to valid JSON
[params.response_format]
type = "json_object"

# Constrain output to JSON matching a specific schema
[params.response_format]
type = "json_schema"

[params.response_format.schema]
name = "extraction"
type = "object"

[params.response_format.schema.properties.result]
type = "string"

Changes

File What changed
internal/agent/agent.go ResponseFormatConfig struct, validation, scaffold template
internal/provider/provider.go ResponseFormat type on the universal Request struct
cmd/run.go Wires config to request, verbose/dry-run output
internal/provider/openai.go Maps to OpenAI response_format field (Send + SendStream)
internal/provider/ollama.go Maps to Ollama format field ("json" or schema object)
*_test.go (3 files) 10 new tests covering config parsing, validation, and wire format

Test plan

  • All existing tests pass (16 packages, 0 failures)
  • New tests verify: response_format omitted when unset, json_object wire format, json_schema wire format with schema passthrough, TOML parsing for both types, validation rejects invalid types, validation requires schema for json_schema

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"
@coderabbitai

coderabbitai Bot commented Apr 21, 2026

Copy link
Copy Markdown

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 620160f9-1ef5-406e-ac4b-3d727028af9a

📥 Commits

Reviewing files that changed from the base of the PR and between c16cfa3 and 5156f68.

📒 Files selected for processing (2)
  • internal/agent/agent.go
  • internal/agent/agent_test.go
🚧 Files skipped from review as they are similar to previous changes (2)
  • internal/agent/agent_test.go
  • internal/agent/agent.go

📝 Walkthrough

Walkthrough

Adds 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

Cohort / File(s) Summary
Core types & config
internal/provider/provider.go, internal/agent/agent.go, internal/agent/agent_test.go
Introduce ResponseFormat and ResponseFormatConfig types and IsSet(); extend provider.Request and ParamsConfig with response_format; add TOML unmarshalling, validation rules, and unit tests for parsing/validation and scaffold output.
CLI integration
cmd/run.go
Populate req.ResponseFormat from config when set; update verbose and dry-run parameter logging to include response_format=<type> when present.
Ollama provider
internal/provider/ollama.go, internal/provider/ollama_test.go
Add optional format field to Ollama request payload; implement buildOllamaFormat mapping (json_object"json", json_schema→object with name removed); return provider errors for unsupported types or missing schema; add tests covering serialization and error cases.
OpenAI provider
internal/provider/openai.go, internal/provider/openai_test.go
Add OpenAI response_format wire structs and translator: derive/require name (default "response"), remove original name from embedded schema, enforce strict:true; fail early on invalid configs; add tests for serialization and bad-schema errors.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~28 minutes

Poem

A tiny format finds its voice,
TOML whispers what to do,
Providers translate, tests rejoice —
Schemas trimmed and logging true. ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 21.88% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The pull request title 'feat: add response_format parameter to [params] config' clearly and concisely describes the main change: adding a new response_format configuration parameter to the params section.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

🧹 Nitpick comments (3)
internal/provider/provider.go (1)

71-74: Clarify that IsSet means “constraint active,” not merely configured.

Type: "text" can be explicitly configured but returns false here. Renaming this to something like IsConstrained() 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 for SendStream too.

These tests nicely pin the non-streaming wire format, but SendStream builds 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, but SendStream has a separate request body and also sets format. A table-driven test that runs the same ResponseFormat cases 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

📥 Commits

Reviewing files that changed from the base of the PR and between 72df532 and 0dcd919.

📒 Files selected for processing (8)
  • cmd/run.go
  • internal/agent/agent.go
  • internal/agent/agent_test.go
  • internal/provider/ollama.go
  • internal/provider/ollama_test.go
  • internal/provider/openai.go
  • internal/provider/openai_test.go
  • internal/provider/provider.go

Comment thread cmd/run.go
Comment thread internal/agent/agent.go
Comment thread internal/provider/ollama.go
Comment thread internal/provider/openai.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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
internal/provider/ollama.go (1)

552-578: buildOllamaFormat looks good — previous concern resolved.

The helper now:

  • Maps json_object"json".
  • Requires a non-empty Schema for json_schema and returns a ProviderError{Category: ErrCategoryBadRequest} otherwise (no more sneaky downgrade to "json" 🙂).
  • Strips the "name" key from the schema map to match buildOpenAIResponseFormat's behavior.
  • Rejects unknown types with an actionable ProviderError.

Error messages are user-actionable and the category routing through mapProviderError in cmd/run.go will 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 buildOpenAIResponseFormat and this function share the "copy map minus name" dance, a small shared helper (e.g., stripSchemaName(rf.Schema)) in provider.go could 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

📥 Commits

Reviewing files that changed from the base of the PR and between 0dcd919 and c16cfa3.

📒 Files selected for processing (6)
  • internal/agent/agent.go
  • internal/agent/agent_test.go
  • internal/provider/ollama.go
  • internal/provider/ollama_test.go
  • internal/provider/openai.go
  • internal/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

@tllongdev

tllongdev commented Apr 21, 2026

Copy link
Copy Markdown
Contributor Author

@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.
@jrswab jrswab merged commit 289f9ea into jrswab:master Apr 21, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants