Skip to content

feat: expose max_output_tokens in GenerationConfig across providers - #159

Open
nstuivenberg wants to merge 6 commits into
flock-community:mainfrom
nstuivenberg:feat/max-output-tokens-config
Open

feat: expose max_output_tokens in GenerationConfig across providers#159
nstuivenberg wants to merge 6 commits into
flock-community:mainfrom
nstuivenberg:feat/max-output-tokens-config

Conversation

@nstuivenberg

@nstuivenberg nstuivenberg commented Aug 21, 2026

Copy link
Copy Markdown

Problem

None of the providers (Gemini, VertexAI, OpenAI) let callers configure the LLM's max output token limit — GenerationSettings/GenerationConfig never exposed it. When max_output_tokens is omitted, provider defaults kick in silently; for Gemini this is 8,192 tokens regardless of the model's actual ceiling (confirmed by community reports of the identical pattern, e.g. google-gemini/gemini-cli#23081, and by Google's own docs).

This caused a real production incident: a Kotlin backend using this library to extract structured price-matrix data via gemini-3.1-flash-lite intermittently failed with a JSON parse error ("Expected quotation mark, had EOF") on large documents — the model's structured-output response was silently truncated mid-JSON at the undocumented 8K-token default, and the failure only surfaced three layers downstream as an opaque deserialization error.

Example

A simplified, generic version of the real-world case that motivated this PR: a structured-output agent extracting a price matrix out of a supplier document. Small documents fit inside Gemini's implicit 8K default; a large or combined document's matrix does not, and before this PR there was no way to raise the ceiling.

@AigenticParameter
data class PriceMatrixEntry(
    val kilometersPerYear: Int,
    val durationInMonths: Int,
    val pricePerMonth: Double,
)

@AigenticParameter
data class PriceMatrix(val entries: List<PriceMatrixEntry>)

val agent = agent<String, PriceMatrix> {
    geminiModel {
        apiKey(apiKey)
        modelIdentifier(GeminiModelIdentifier.Gemini3_1FlashLite)
        generationConfig {
            maxOutputTokens(65536)
        }
    }
    task("Extract the price matrix from the document") {
        addInstruction("Return one entry per cell, with kilometers per year as rows and contract duration in months as columns")
    }
}

val run = agent.start(documentAsCsv)

Fix

  • GenerationSettings (core) gains an optional maxOutputTokens: Int? = null.
  • GenerationConfig DSL gains a matching maxOutputTokens(Int) builder function, mirroring the existing thinkingBudget(Int) pattern.
  • Wired through to all three request mappers:
    • Gemini — set on GenerationConfig, serialized as max_output_tokens.
    • VertexAI — applied to GenerateContentConfig.Builder.maxOutputTokens(...), only when configured.
    • OpenAI — see the parameter branching below.
  • OpenAI max_tokens vs max_completion_tokens: which parameter is sent is now decided per model identifier. Reasoning and GPT-5-family identifiers (o1, o1-pro, o3, o3-pro, o3-mini, o4-mini, gpt-5*) get max_completion_tokens, since they reject the older max_tokens with a 400. Classic chat identifiers (gpt-4o, gpt-4o-mini, gpt-4.1*, gpt-4-turbo, gpt-3.5-turbo, the search-preview models) and OpenAIModelIdentifier.Custom get max_tokens. The branch is exhaustive over the sealed identifier hierarchy, so a newly added identifier won't compile until it's classified. Ollama reuses OpenAIModel with a plain ModelIdentifier, which falls into the max_tokens branch.
  • Gemini truncation is no longer silent: when a candidate comes back with finishReason == MAX_TOKENS, GeminiResponseMapper now fails fast with an explicit, actionable error pointing at generationConfig { maxOutputTokens(...) }, instead of letting a truncated payload reach the caller (or the JSON parser). Raising the ceiling reduces how often this happens; this change makes the remaining cases diagnosable.
  • Extracted GeminiClient's Json config to a shared internal geminiJson value, reused by the new serialization tests, so "field omitted when null" is verified against the actual production serializer rather than a reimplementation.
  • Platform run traces: ConfigDto in src/platform/wirespec/gateway.ws gains an optional maxOutputTokens: Integer?, populated by RequestMapper, so the configured limit is visible in published run traces during diagnosis (alongside the existing thinkingBudget).
  • Fixed a pre-existing, unrelated gap found while testing this: gemini, openai, and vertexai all declared kotest-runner-junit5 but never called useJUnitPlatform() on their jvmTest task, so their Kotest specs compiled but silently never executed (0 tests run, green build). Added the missing configuration so all three modules' existing and new specs actually run.

Backward compatibility

maxOutputTokens defaults to null everywhere (matching this PR's requirement of not getting in the way of existing users of the API) — no existing caller needs to change anything, and request payloads are byte-for-byte unchanged unless a caller opts in via the new DSL function. Verified (not assumed) that an unset value is omitted from the serialized Gemini request rather than sent as an explicit null, and that neither OpenAI token parameter is emitted when unset.

Note: GenerationSettings gaining a constructor parameter is source-compatible (default value) but not binary-compatible for precompiled JVM/Android consumers who construct it positionally or via copy() without recompiling against this version. Given the project is pre-1.0 (0.11.0-SNAPSHOT) this seems like an acceptable trade-off, flagging it in case maintainers track binary compatibility more strictly than that.

The ConfigDto change adds an optional field to an outgoing payload sent to the Aigentic platform gateway. It should be compatible with a backend that ignores unknown/absent optional fields, but I can't verify that against the live gateway from here — maintainers should confirm the backend side accepts it before release.

Known limitations / follow-ups (out of scope for this PR)

  • OpenAI reasoning models (o-series, gpt-5.x) can still hit an opaque "unknown type" exception in the OpenAI message mapper if maxCompletionTokens is set too low and the model exhausts the budget on internal reasoning before producing visible content — this PR doesn't address that failure mode.
  • No cross-validation between thinkingBudget and maxOutputTokens, even though on Gemini/VertexAI they can share the same token budget when thinking is enabled.
  • Ollama now receives max_tokens (the parameter its OpenAI-compatible layer documents) rather than max_completion_tokens, and there's a unit test asserting that request shape for Custom identifiers — but this has not been verified end-to-end against a live Ollama instance.
  • :src:platform has the same missing useJUnitPlatform() gap that the three provider modules had: it declares kotest-runner-junit5 for jvmTest but never configures the JUnit platform, so its Kotest specs — including the RequestMapperTest added in this PR — compile but don't execute. I left the fix out of this PR because adding useJUnitPlatform() there surfaces 5 pre-existing, unrelated failures in TestExecutorTest that look like MockK drift rather than anything caused by this change. Flagging it for maintainers as a separate cleanup.

Testing

New specs, all running on the JUnit platform after the build fix:

  • coreAgentConfigTest: the DSL builds maxOutputTokens (alongside thinkingBudget), and defaults to null.
  • geminiGeminiRequestMapperKtTest: the value reaches generationConfig, is omitted when unset, and the serialized body contains "max_output_tokens":65536 when set / no max_output_tokens key when unset (asserted against the production geminiJson serializer).
  • geminiGeminiResponseMapperKtTest: a MAX_TOKENS candidate raises an error mentioning truncation; a STOP candidate still maps to a Message.StructuredOutput as before.
  • openaiChatCompletionsRequestTest: max_completion_tokens for O1 and GPT5, max_tokens for GPT4O and Custom, and neither field set when maxOutputTokens is unset.
  • vertexaiRequestMapperTest: maxOutputTokens() is Optional.of(65536) when configured and Optional.empty() when not.
  • platformRequestMapperTest asserts ConfigDto.maxOutputTokens is populated / left null. Caveat: as noted above, this module doesn't run its Kotest specs yet, so this test compiles but does not currently execute in CI.

./gradlew build passes for core, gemini, openai, and vertexai, with the JUnit platform fix confirmed to make the provider specs actually execute (previously silently skipped). The :src:platform module has a pre-existing, task-ordering-dependent Gradle validation failure on the full build task that reproduces identically on an unmodified main — unrelated to this change.

nstuivenberg added 4 commits August 21, 2026 15:12
…roviders

Adds an optional maxOutputTokens setting to GenerationSettings and a matching maxOutputTokens(Int) DSL function on GenerationConfig, wired through to the Gemini, VertexAI and OpenAI request mappers. Ollama inherits it via OpenAIModel. Without it providers fall back to their own default output limits, which silently truncates large structured responses.
…platform

These modules declare kotest-runner-junit5 and contain specs, but never configure useJUnitPlatform() on their jvmTest task, so the specs compiled and were silently never executed. Only :src:core and :src:tools:openapi configured it.
…the max output token limit

A candidate that hits MAX_TOKENS still carries partial content, which flowed into the mapper and surfaced downstream as an opaque JSON parse error on structured output. Also extracts the client Json configuration into geminiJson so the serialization tests assert against the production configuration instead of a reimplementation.
nstuivenberg added 2 commits August 21, 2026 16:06
…ng OpenAI models

Reasoning models reject max_tokens and require max_completion_tokens, while classic OpenAI models, custom identifiers, Ollama and other OpenAI-compatible servers understand max_tokens and silently ignore max_completion_tokens. The classification is exhaustive over OpenAIModelIdentifier so a new identifier fails to compile until it is classified.
Adds an optional maxOutputTokens field to ConfigDto in gateway.ws and maps it from the agent generation settings when publishing a run, so the configured limit travels with the run trace.
@nstuivenberg
nstuivenberg marked this pull request as ready for review August 21, 2026 14:16
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