From 4519872d65a4cb1aec4079b863b32932eb5da9ea Mon Sep 17 00:00:00 2001 From: anthonytwh Date: Thu, 2 Jul 2026 23:47:42 -0400 Subject: [PATCH 1/3] feat(models): wire maxOutputTokens for native Gemini and Vertex The native Gemini and GeminiVertexAI providers previously exposed no way to cap a response's output tokens via ModelConfig: GeminiConfig was an empty struct and GeminiVertexAIConfig.MaxOutputTokens was declared but never read by the translator or applied at runtime. Both providers subclass the ADK native Gemini model, which takes generation config from the per-request LlmRequest.config rather than the model definition, so nothing plumbed the ModelConfig-level setting through. Wire it end to end: - add MaxOutputTokens to GeminiConfig (v1alpha2) and regenerate the CRDs - carry it on the adk.Gemini / adk.GeminiVertexAI serialization structs - copy Spec.Gemini/Spec.GeminiVertexAI.MaxOutputTokens in the translator - parse it on the Gemini/GeminiVertexAI pydantic models and build the wrappers with it - apply it in a shared _GeminiGenerationConfigMixin that injects max_output_tokens into GenerateContentConfig without overriding a value the agent already set on the request A per-request value always wins over the model-level default. Adds Go translator tests and Python unit tests for the mixin. Signed-off-by: anthonytwh --- go/api/adk/types.go | 2 + .../crd/bases/kagent.dev_modelconfigs.yaml | 4 + go/api/v1alpha2/modelconfig_types.go | 7 +- .../translator/agent/adk_api_translator.go | 7 ++ .../agent/adk_api_translator_test.go | 75 +++++++++++++++++++ .../templates/kagent.dev_modelconfigs.yaml | 4 + .../src/kagent/adk/models/__init__.py | 3 +- .../src/kagent/adk/models/_gemini.py | 38 +++++++++- .../kagent-adk/src/kagent/adk/types.py | 10 ++- .../tests/unittests/models/test_gemini.py | 60 +++++++++++++++ 10 files changed, 205 insertions(+), 5 deletions(-) create mode 100644 python/packages/kagent-adk/tests/unittests/models/test_gemini.py diff --git a/go/api/adk/types.go b/go/api/adk/types.go index 502a1cde3a..5d4f9587b9 100644 --- a/go/api/adk/types.go +++ b/go/api/adk/types.go @@ -171,6 +171,7 @@ func (a *Anthropic) GetType() string { type GeminiVertexAI struct { BaseModel + MaxOutputTokens *int `json:"max_output_tokens,omitempty"` } func (g *GeminiVertexAI) MarshalJSON() ([]byte, error) { @@ -229,6 +230,7 @@ func (o *Ollama) GetType() string { type Gemini struct { BaseModel + MaxOutputTokens *int `json:"max_output_tokens,omitempty"` } func (g *Gemini) MarshalJSON() ([]byte, error) { diff --git a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml index 4dfdc96bce..94f5297f57 100644 --- a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml +++ b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml @@ -534,6 +534,10 @@ spec: type: object gemini: description: Gemini-specific configuration + properties: + maxOutputTokens: + description: Maximum output tokens to generate for a single response + type: integer type: object geminiVertexAI: description: Gemini Vertex AI-specific configuration diff --git a/go/api/v1alpha2/modelconfig_types.go b/go/api/v1alpha2/modelconfig_types.go index 4e0231fbd1..d67f19ce15 100644 --- a/go/api/v1alpha2/modelconfig_types.go +++ b/go/api/v1alpha2/modelconfig_types.go @@ -240,7 +240,12 @@ type OllamaConfig struct { Options map[string]string `json:"options,omitempty"` } -type GeminiConfig struct{} +// GeminiConfig contains Gemini (AI Studio, API-key) specific configuration options +type GeminiConfig struct { + // Maximum output tokens to generate for a single response + // +optional + MaxOutputTokens int `json:"maxOutputTokens,omitempty"` +} // BedrockConfig contains AWS Bedrock-specific configuration options. type BedrockConfig struct { diff --git a/go/core/internal/controller/translator/agent/adk_api_translator.go b/go/core/internal/controller/translator/agent/adk_api_translator.go index f5b1f18797..7d3c568a3a 100644 --- a/go/core/internal/controller/translator/agent/adk_api_translator.go +++ b/go/core/internal/controller/translator/agent/adk_api_translator.go @@ -641,6 +641,10 @@ func (a *adkApiTranslator) translateModel(ctx context.Context, namespace, modelC populateTLSFields(&gemini.BaseModel, model.Spec.TLS) gemini.APIKeyPassthrough = model.Spec.APIKeyPassthrough + if model.Spec.GeminiVertexAI.MaxOutputTokens > 0 { + gemini.MaxOutputTokens = &model.Spec.GeminiVertexAI.MaxOutputTokens + } + return gemini, modelDeploymentData, secretHashBytes, nil case v1alpha2.ModelProviderAnthropicVertexAI: if model.Spec.AnthropicVertexAI == nil { @@ -727,6 +731,9 @@ func (a *adkApiTranslator) translateModel(ctx context.Context, namespace, modelC } // Populate TLS fields in BaseModel populateTLSFields(&gemini.BaseModel, model.Spec.TLS) + if model.Spec.Gemini != nil && model.Spec.Gemini.MaxOutputTokens > 0 { + gemini.MaxOutputTokens = &model.Spec.Gemini.MaxOutputTokens + } return gemini, modelDeploymentData, secretHashBytes, nil case v1alpha2.ModelProviderBedrock: if model.Spec.Bedrock == nil { diff --git a/go/core/internal/controller/translator/agent/adk_api_translator_test.go b/go/core/internal/controller/translator/agent/adk_api_translator_test.go index 8c4202f663..cb7115ed79 100644 --- a/go/core/internal/controller/translator/agent/adk_api_translator_test.go +++ b/go/core/internal/controller/translator/agent/adk_api_translator_test.go @@ -470,6 +470,81 @@ func Test_AdkApiTranslator_AzureOpenAIParams(t *testing.T) { assert.Equal(t, &maxTokens, m.MaxTokens) } +func Test_AdkApiTranslator_GeminiMaxOutputTokens(t *testing.T) { + scheme := schemev1.Scheme + require.NoError(t, v1alpha2.AddToScheme(scheme)) + + geminiModel := &v1alpha2.ModelConfig{ + ObjectMeta: metav1.ObjectMeta{Name: "gm", Namespace: "ns"}, + Spec: v1alpha2.ModelConfigSpec{ + Model: "gemini-2.5-flash", + Provider: v1alpha2.ModelProviderGemini, + APIKeySecret: "keys", + APIKeySecretKey: "GEMINI_API_KEY", + Gemini: &v1alpha2.GeminiConfig{MaxOutputTokens: 2048}, + }, + } + vertexModel := &v1alpha2.ModelConfig{ + ObjectMeta: metav1.ObjectMeta{Name: "vx", Namespace: "ns"}, + Spec: v1alpha2.ModelConfigSpec{ + Model: "gemini-2.5-pro", + Provider: v1alpha2.ModelProviderGeminiVertexAI, + APIKeySecret: "keys", + APIKeySecretKey: "service-account.json", + GeminiVertexAI: &v1alpha2.GeminiVertexAIConfig{ + BaseVertexAIConfig: v1alpha2.BaseVertexAIConfig{ + ProjectID: "my-project", + Location: "us-central1", + }, + MaxOutputTokens: 1024, + }, + }, + } + + makeAgent := func(model string) *v1alpha2.Agent { + return &v1alpha2.Agent{ + ObjectMeta: metav1.ObjectMeta{Name: "a", Namespace: "ns"}, + Spec: v1alpha2.AgentSpec{ + Type: v1alpha2.AgentType_Declarative, + Declarative: &v1alpha2.DeclarativeAgentSpec{ + SystemMessage: "x", + ModelConfig: model, + }, + }, + } + } + + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "ns"}} + + t.Run("native Gemini", func(t *testing.T) { + agent := makeAgent("gm") + kubeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(ns, geminiModel, agent).Build() + trans := translator.NewAdkApiTranslator(kubeClient, types.NamespacedName{Namespace: "ns", Name: "gm"}, nil, "", nil) + + outputs, err := translator.TranslateAgent(context.Background(), trans, agent) + require.NoError(t, err) + + m, ok := outputs.Config.Model.(*adk.Gemini) + require.True(t, ok, "expected model to be of type Gemini") + require.NotNil(t, m.MaxOutputTokens) + assert.Equal(t, 2048, *m.MaxOutputTokens) + }) + + t.Run("Gemini Vertex AI", func(t *testing.T) { + agent := makeAgent("vx") + kubeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(ns, vertexModel, agent).Build() + trans := translator.NewAdkApiTranslator(kubeClient, types.NamespacedName{Namespace: "ns", Name: "vx"}, nil, "", nil) + + outputs, err := translator.TranslateAgent(context.Background(), trans, agent) + require.NoError(t, err) + + m, ok := outputs.Config.Model.(*adk.GeminiVertexAI) + require.True(t, ok, "expected model to be of type GeminiVertexAI") + require.NotNil(t, m.MaxOutputTokens) + assert.Equal(t, 1024, *m.MaxOutputTokens) + }) +} + func Test_AdkApiTranslator_ServiceAccountNameOverride(t *testing.T) { scheme := schemev1.Scheme require.NoError(t, v1alpha2.AddToScheme(scheme)) diff --git a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml index 4dfdc96bce..94f5297f57 100644 --- a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml +++ b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml @@ -534,6 +534,10 @@ spec: type: object gemini: description: Gemini-specific configuration + properties: + maxOutputTokens: + description: Maximum output tokens to generate for a single response + type: integer type: object geminiVertexAI: description: Gemini Vertex AI-specific configuration diff --git a/python/packages/kagent-adk/src/kagent/adk/models/__init__.py b/python/packages/kagent-adk/src/kagent/adk/models/__init__.py index 3176704cef..bad9f56d26 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/__init__.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/__init__.py @@ -1,7 +1,7 @@ from ._anthropic import KAgentAnthropicLlm from ._bedrock import KAgentBedrockLlm from ._embedding import KAgentEmbedding -from ._gemini import KAgentGeminiLlm +from ._gemini import KAgentGeminiLlm, KAgentGeminiVertexAILlm from ._ollama import KAgentOllamaLlm from ._openai import AzureOpenAI, OpenAI from ._sap_ai_core import KAgentSAPAICoreLlm @@ -12,6 +12,7 @@ "KAgentAnthropicLlm", "KAgentBedrockLlm", "KAgentGeminiLlm", + "KAgentGeminiVertexAILlm", "KAgentOllamaLlm", "KAgentEmbedding", "KAgentSAPAICoreLlm", diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_gemini.py b/python/packages/kagent-adk/src/kagent/adk/models/_gemini.py index 8b5189e7e1..2cd03f70c6 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_gemini.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_gemini.py @@ -20,11 +20,34 @@ def _merge_headers(extra_headers: Optional[dict[str, str]]) -> dict[str, str]: return headers -class KAgentGeminiLlm(KAgentTLSMixin, GeminiLLM): +class _GeminiGenerationConfigMixin: + """Applies model-level generation defaults (e.g. max_output_tokens) to each + request, without overriding any per-request value the agent already set. + + The native Gemini/Vertex ADK model takes generation config from the + per-request LlmRequest.config rather than the model definition, so this + mixin bridges the ModelConfig-level setting onto the request. + """ + + async def generate_content_async(self, llm_request, stream: bool = False): + max_output_tokens = getattr(self, "max_output_tokens", None) + if max_output_tokens is not None: + config = llm_request.config + if config is None: + config = types.GenerateContentConfig() + llm_request.config = config + if config.max_output_tokens is None: + config.max_output_tokens = max_output_tokens + async for response in super().generate_content_async(llm_request, stream=stream): + yield response + + +class KAgentGeminiLlm(KAgentTLSMixin, _GeminiGenerationConfigMixin, GeminiLLM): """Gemini API model that applies kagent TLS and header settings.""" extra_headers: Optional[dict[str, str]] = None api_key_passthrough: Optional[bool] = None + max_output_tokens: Optional[int] = None model_config = {"arbitrary_types_allowed": True} @@ -57,3 +80,16 @@ def _live_api_client(self) -> Client: api_key=os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY"), http_options=self._http_options(api_version=self._live_api_version), ) + + +class KAgentGeminiVertexAILlm(_GeminiGenerationConfigMixin, GeminiLLM): + """Gemini Vertex AI model. + + Auth (project/location/ADC) is handled by the native ADK client via the + environment variables the controller sets, so the client is intentionally + not overridden here. Only the model-level generation config is applied. + """ + + max_output_tokens: Optional[int] = None + + model_config = {"arbitrary_types_allowed": True} diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 766e7136b4..c8d5e67bc9 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -19,7 +19,7 @@ from kagent.adk._remote_a2a_tool import KAgentRemoteA2AToolset from kagent.adk.models._anthropic import KAgentAnthropicLlm from kagent.adk.models._bedrock import KAgentBedrockLlm -from kagent.adk.models._gemini import KAgentGeminiLlm +from kagent.adk.models._gemini import KAgentGeminiLlm, KAgentGeminiVertexAILlm from kagent.adk.models._ollama import create_ollama_llm from kagent.adk.models._openai import AzureOpenAI as OpenAIAzure from kagent.adk.models._openai import OpenAI as OpenAINative @@ -294,6 +294,7 @@ class Anthropic(BaseLLM): class GeminiVertexAI(BaseLLM): + max_output_tokens: int | None = None type: Literal["gemini_vertex_ai"] @@ -307,6 +308,7 @@ class Ollama(BaseLLM): class Gemini(BaseLLM): + max_output_tokens: int | None = None type: Literal["gemini"] @@ -662,7 +664,10 @@ def _create_llm_from_model_config(model_config: ModelUnion): **_transport_kwargs(model_config), ) if model_config.type == "gemini_vertex_ai": - return GeminiLLM(model=model_config.model) + return KAgentGeminiVertexAILlm( + model=model_config.model, + max_output_tokens=model_config.max_output_tokens, + ) if model_config.type == "gemini_anthropic": return ClaudeLLM(model=model_config.model) if model_config.type == "ollama": @@ -685,6 +690,7 @@ def _create_llm_from_model_config(model_config: ModelUnion): return KAgentGeminiLlm( model=model_config.model, extra_headers=extra_headers, + max_output_tokens=model_config.max_output_tokens, **_transport_kwargs(model_config), ) if model_config.type == "bedrock": diff --git a/python/packages/kagent-adk/tests/unittests/models/test_gemini.py b/python/packages/kagent-adk/tests/unittests/models/test_gemini.py new file mode 100644 index 0000000000..db0d6530c5 --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/models/test_gemini.py @@ -0,0 +1,60 @@ +"""Tests for the Gemini model-level generation config wiring.""" + +import pytest +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.genai import types + +from kagent.adk.models._gemini import _GeminiGenerationConfigMixin + + +class _FakeBaseLlm: + """Stand-in for the native ADK GeminiLLM: records the request config it + receives and yields a dummy response, without touching any API.""" + + def __init__(self, max_output_tokens=None): + self.max_output_tokens = max_output_tokens + self.seen_max_output_tokens = "unset" + + async def generate_content_async(self, llm_request, stream: bool = False): + self.seen_max_output_tokens = llm_request.config.max_output_tokens + yield LlmResponse() + + +class _Model(_GeminiGenerationConfigMixin, _FakeBaseLlm): + pass + + +def _request(max_output_tokens=None): + return LlmRequest( + model="gemini-2.5-flash", + config=types.GenerateContentConfig(max_output_tokens=max_output_tokens), + ) + + +@pytest.mark.asyncio +async def test_applies_max_output_tokens_when_unset(): + model = _Model(max_output_tokens=2048) + req = _request() + _ = [r async for r in model.generate_content_async(req, stream=False)] + assert req.config.max_output_tokens == 2048 + assert model.seen_max_output_tokens == 2048 + + +@pytest.mark.asyncio +async def test_does_not_override_per_request_value(): + model = _Model(max_output_tokens=2048) + req = _request(max_output_tokens=512) + _ = [r async for r in model.generate_content_async(req, stream=False)] + # A value the caller/agent already set must win. + assert req.config.max_output_tokens == 512 + assert model.seen_max_output_tokens == 512 + + +@pytest.mark.asyncio +async def test_noop_when_model_has_no_cap(): + model = _Model(max_output_tokens=None) + req = _request() + _ = [r async for r in model.generate_content_async(req, stream=False)] + assert req.config.max_output_tokens is None + assert model.seen_max_output_tokens is None From 75d90bfda77fc3fe7559d7165d052b918562ac35 Mon Sep 17 00:00:00 2001 From: anthonytwh Date: Fri, 3 Jul 2026 20:22:15 -0400 Subject: [PATCH 2/3] validate: reject non-positive Gemini maxOutputTokens Address review feedback: the translator treats maxOutputTokens <= 0 as "unset", so an invalid (0/negative) value was silently ignored rather than rejected. Enforce a positive value at validation time: - add +kubebuilder:validation:Minimum=1 to GeminiConfig and GeminiVertexAIConfig.MaxOutputTokens and regenerate the CRDs (Kubernetes now rejects invalid values on apply) - add pydantic ge=1 to the Gemini / GeminiVertexAI max_output_tokens fields so an invalid ModelConfig fails fast at parse time - add unit tests for the parse-time validation Signed-off-by: anthonytwh --- .../crd/bases/kagent.dev_modelconfigs.yaml | 2 ++ go/api/v1alpha2/modelconfig_types.go | 2 ++ .../templates/kagent.dev_modelconfigs.yaml | 2 ++ .../kagent-adk/src/kagent/adk/types.py | 4 ++-- .../tests/unittests/models/test_gemini.py | 20 +++++++++++++++++++ 5 files changed, 28 insertions(+), 2 deletions(-) diff --git a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml index 94f5297f57..000866ec06 100644 --- a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml +++ b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml @@ -537,6 +537,7 @@ spec: properties: maxOutputTokens: description: Maximum output tokens to generate for a single response + minimum: 1 type: integer type: object geminiVertexAI: @@ -550,6 +551,7 @@ spec: type: string maxOutputTokens: description: Maximum output tokens + minimum: 1 type: integer projectID: description: The project ID diff --git a/go/api/v1alpha2/modelconfig_types.go b/go/api/v1alpha2/modelconfig_types.go index d67f19ce15..6de42a777b 100644 --- a/go/api/v1alpha2/modelconfig_types.go +++ b/go/api/v1alpha2/modelconfig_types.go @@ -74,6 +74,7 @@ type GeminiVertexAIConfig struct { // Maximum output tokens // +optional + // +kubebuilder:validation:Minimum=1 MaxOutputTokens int `json:"maxOutputTokens,omitempty"` // Candidate count @@ -244,6 +245,7 @@ type OllamaConfig struct { type GeminiConfig struct { // Maximum output tokens to generate for a single response // +optional + // +kubebuilder:validation:Minimum=1 MaxOutputTokens int `json:"maxOutputTokens,omitempty"` } diff --git a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml index 94f5297f57..000866ec06 100644 --- a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml +++ b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml @@ -537,6 +537,7 @@ spec: properties: maxOutputTokens: description: Maximum output tokens to generate for a single response + minimum: 1 type: integer type: object geminiVertexAI: @@ -550,6 +551,7 @@ spec: type: string maxOutputTokens: description: Maximum output tokens + minimum: 1 type: integer projectID: description: The project ID diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index c8d5e67bc9..1d7abebcb0 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -294,7 +294,7 @@ class Anthropic(BaseLLM): class GeminiVertexAI(BaseLLM): - max_output_tokens: int | None = None + max_output_tokens: int | None = Field(default=None, ge=1) type: Literal["gemini_vertex_ai"] @@ -308,7 +308,7 @@ class Ollama(BaseLLM): class Gemini(BaseLLM): - max_output_tokens: int | None = None + max_output_tokens: int | None = Field(default=None, ge=1) type: Literal["gemini"] diff --git a/python/packages/kagent-adk/tests/unittests/models/test_gemini.py b/python/packages/kagent-adk/tests/unittests/models/test_gemini.py index db0d6530c5..435bc19dfb 100644 --- a/python/packages/kagent-adk/tests/unittests/models/test_gemini.py +++ b/python/packages/kagent-adk/tests/unittests/models/test_gemini.py @@ -4,8 +4,10 @@ from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.genai import types +from pydantic import ValidationError from kagent.adk.models._gemini import _GeminiGenerationConfigMixin +from kagent.adk.types import Gemini, GeminiVertexAI class _FakeBaseLlm: @@ -58,3 +60,21 @@ async def test_noop_when_model_has_no_cap(): _ = [r async for r in model.generate_content_async(req, stream=False)] assert req.config.max_output_tokens is None assert model.seen_max_output_tokens is None + + +_GEMINI_TYPES = [(Gemini, "gemini"), (GeminiVertexAI, "gemini_vertex_ai")] + + +@pytest.mark.parametrize("model_cls,type_name", _GEMINI_TYPES) +@pytest.mark.parametrize("bad_value", [0, -1]) +def test_rejects_non_positive_max_output_tokens(model_cls, type_name, bad_value): + # The translator treats <= 0 as "unset"; reject it at parse time so an + # invalid config fails fast instead of being silently ignored. + with pytest.raises(ValidationError): + model_cls(type=type_name, model="gemini-2.5-flash", max_output_tokens=bad_value) + + +@pytest.mark.parametrize("model_cls,type_name", _GEMINI_TYPES) +def test_accepts_positive_max_output_tokens(model_cls, type_name): + model = model_cls(type=type_name, model="gemini-2.5-flash", max_output_tokens=1) + assert model.max_output_tokens == 1 From 94a662aeb06533fc98284f7e7b83228da07567a2 Mon Sep 17 00:00:00 2001 From: anthonytwh Date: Wed, 8 Jul 2026 17:02:33 -0400 Subject: [PATCH 3/3] feat(go-adk): apply Gemini maxOutputTokens generation config The ADK Gemini model reads generation config from the per-request LLMRequest.Config rather than the model definition, so the ModelConfig-level maxOutputTokens was never applied by the Go ADK agent (only the Python runtime handled it via _GeminiGenerationConfigMixin). Wrap the native Gemini and Gemini Vertex AI models so max_output_tokens is injected into each request without overriding a per-request value. The wrapper is a no-op when maxOutputTokens is unset or non-positive. Signed-off-by: anthonytwh --- go/adk/pkg/agent/agent.go | 12 +++- go/adk/pkg/models/gemini.go | 53 ++++++++++++++++ go/adk/pkg/models/gemini_test.go | 101 +++++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 go/adk/pkg/models/gemini.go create mode 100644 go/adk/pkg/models/gemini_test.go diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index fa9d633d14..813ebba654 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -228,10 +228,14 @@ func CreateLLM(ctx context.Context, m adk.Model, log logr.Logger) (adkmodel.LLM, if err != nil { return nil, fmt.Errorf("failed to build HTTP client for Gemini: %w", err) } - return adkgemini.NewModel(ctx, modelName, &genai.ClientConfig{ + geminiModel, err := adkgemini.NewModel(ctx, modelName, &genai.ClientConfig{ APIKey: apiKey, HTTPClient: httpClient, }) + if err != nil { + return nil, err + } + return models.WrapGeminiWithGenerationConfig(geminiModel, m.MaxOutputTokens), nil case *adk.GeminiVertexAI: project := os.Getenv("GOOGLE_CLOUD_PROJECT") @@ -246,11 +250,15 @@ func CreateLLM(ctx context.Context, m adk.Model, log logr.Logger) (adkmodel.LLM, if modelName == "" { modelName = DefaultGeminiModel } - return adkgemini.NewModel(ctx, modelName, &genai.ClientConfig{ + geminiModel, err := adkgemini.NewModel(ctx, modelName, &genai.ClientConfig{ Backend: genai.BackendVertexAI, Project: project, Location: location, }) + if err != nil { + return nil, err + } + return models.WrapGeminiWithGenerationConfig(geminiModel, m.MaxOutputTokens), nil case *adk.Anthropic: modelName := m.Model diff --git a/go/adk/pkg/models/gemini.go b/go/adk/pkg/models/gemini.go new file mode 100644 index 0000000000..2eff814db7 --- /dev/null +++ b/go/adk/pkg/models/gemini.go @@ -0,0 +1,53 @@ +// Package models: helpers for the native Gemini / Gemini Vertex AI models. +// +// The ADK Gemini model reads its generation config from the per-request +// LLMRequest.Config rather than from the model definition, so a +// ModelConfig-level setting such as maxOutputTokens is not otherwise applied. +// geminiGenerationConfigModel wraps the ADK model and injects the setting into +// each request, mirroring the Python _GeminiGenerationConfigMixin. +package models + +import ( + "context" + "iter" + + "google.golang.org/adk/model" + "google.golang.org/genai" +) + +// geminiGenerationConfigModel wraps an ADK Gemini model.LLM and applies a +// ModelConfig-level generation config (currently max_output_tokens) to each +// request. A per-request value always wins: the cap is only applied when the +// request does not already set MaxOutputTokens. +type geminiGenerationConfigModel struct { + model.LLM + maxOutputTokens int32 +} + +// WrapGeminiWithGenerationConfig wraps inner so that maxOutputTokens is applied +// to requests that don't already set it. When maxOutputTokens is nil or not +// positive, inner is returned unchanged. +func WrapGeminiWithGenerationConfig(inner model.LLM, maxOutputTokens *int) model.LLM { + if maxOutputTokens == nil || *maxOutputTokens <= 0 { + return inner + } + return &geminiGenerationConfigModel{ + LLM: inner, + maxOutputTokens: int32(*maxOutputTokens), + } +} + +// GenerateContent injects the configured max_output_tokens into the request +// config, without overriding a value the request already set, then delegates to +// the wrapped model. +func (m *geminiGenerationConfigModel) GenerateContent(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + if req != nil && m.maxOutputTokens > 0 { + if req.Config == nil { + req.Config = &genai.GenerateContentConfig{} + } + if req.Config.MaxOutputTokens == 0 { + req.Config.MaxOutputTokens = m.maxOutputTokens + } + } + return m.LLM.GenerateContent(ctx, req, stream) +} diff --git a/go/adk/pkg/models/gemini_test.go b/go/adk/pkg/models/gemini_test.go new file mode 100644 index 0000000000..39b6d17a3c --- /dev/null +++ b/go/adk/pkg/models/gemini_test.go @@ -0,0 +1,101 @@ +package models + +import ( + "context" + "iter" + "testing" + + "google.golang.org/adk/model" + "google.golang.org/genai" +) + +// fakeLLM records the request it was called with and returns an empty stream. +type fakeLLM struct { + name string + gotReq *model.LLMRequest + gotCall bool +} + +func (f *fakeLLM) Name() string { return f.name } + +func (f *fakeLLM) GenerateContent(_ context.Context, req *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + f.gotCall = true + f.gotReq = req + return func(yield func(*model.LLMResponse, error) bool) {} +} + +func drain(seq iter.Seq2[*model.LLMResponse, error]) { + for range seq { //nolint:revive // consume the iterator + } +} + +func intPtr(i int) *int { return &i } + +func TestWrapGeminiWithGenerationConfig_NoOpWhenUnset(t *testing.T) { + inner := &fakeLLM{name: "gemini-2.0-flash"} + for _, tc := range []struct { + name string + max *int + }{ + {"nil", nil}, + {"zero", intPtr(0)}, + {"negative", intPtr(-1)}, + } { + t.Run(tc.name, func(t *testing.T) { + got := WrapGeminiWithGenerationConfig(inner, tc.max) + if got != model.LLM(inner) { + t.Errorf("expected inner model returned unchanged, got wrapper %T", got) + } + }) + } +} + +func TestGeminiGenerationConfig_AppliesCapWhenUnset(t *testing.T) { + inner := &fakeLLM{name: "gemini-2.0-flash"} + wrapped := WrapGeminiWithGenerationConfig(inner, intPtr(1024)) + + req := &model.LLMRequest{Config: &genai.GenerateContentConfig{}} + drain(wrapped.GenerateContent(context.Background(), req, false)) + + if !inner.gotCall { + t.Fatal("wrapped model did not delegate to inner") + } + if got := inner.gotReq.Config.MaxOutputTokens; got != 1024 { + t.Errorf("MaxOutputTokens = %d, want 1024", got) + } +} + +func TestGeminiGenerationConfig_NilConfigInitialized(t *testing.T) { + inner := &fakeLLM{name: "gemini-2.0-flash"} + wrapped := WrapGeminiWithGenerationConfig(inner, intPtr(512)) + + req := &model.LLMRequest{} // Config is nil + drain(wrapped.GenerateContent(context.Background(), req, false)) + + if inner.gotReq.Config == nil { + t.Fatal("Config was not initialized") + } + if got := inner.gotReq.Config.MaxOutputTokens; got != 512 { + t.Errorf("MaxOutputTokens = %d, want 512", got) + } +} + +func TestGeminiGenerationConfig_DoesNotOverridePerRequestValue(t *testing.T) { + inner := &fakeLLM{name: "gemini-2.0-flash"} + wrapped := WrapGeminiWithGenerationConfig(inner, intPtr(1024)) + + req := &model.LLMRequest{Config: &genai.GenerateContentConfig{MaxOutputTokens: 256}} + drain(wrapped.GenerateContent(context.Background(), req, false)) + + if got := inner.gotReq.Config.MaxOutputTokens; got != 256 { + t.Errorf("MaxOutputTokens = %d, want 256 (per-request value must win)", got) + } +} + +func TestGeminiGenerationConfig_NamePassthrough(t *testing.T) { + inner := &fakeLLM{name: "gemini-2.0-flash"} + wrapped := WrapGeminiWithGenerationConfig(inner, intPtr(1024)) + if got := wrapped.Name(); got != "gemini-2.0-flash" { + t.Errorf("Name() = %q, want %q", got, "gemini-2.0-flash") + } +}