diff --git a/components/model/gemini/README.md b/components/model/gemini/README.md index d90a3cbf0..5d1bd592a 100644 --- a/components/model/gemini/README.md +++ b/components/model/gemini/README.md @@ -298,6 +298,21 @@ if lp := msg.ResponseMeta.LogProbs; lp != nil { > Note: `Logprobs` only takes effect when `ResponseLogprobs` is true. To disable logprobs at call time, use `gemini.WithResponseLogprobs(false)`. +## Labels + +Attach user-defined metadata labels to a request via `gemini.WithLabels(...)`. Labels are key-value pairs used for billing attribution, request tracking, and filtering in Cloud Logging and Monitoring. + +```go +msg, _ := cm.Generate(ctx, input, + gemini.WithLabels(map[string]string{ + "team": "search", + "environment": "production", + }), +) +``` + +> Note: Labels are only supported on the **Vertex AI** backend. The Gemini API backend rejects requests that set labels (`labels parameter is not supported in Gemini API`), so only use this option with a Vertex AI client. + ## Examples See the following examples for more usage: diff --git a/components/model/gemini/README_zh.md b/components/model/gemini/README_zh.md index f8a1b8376..9f1f3a53a 100644 --- a/components/model/gemini/README_zh.md +++ b/components/model/gemini/README_zh.md @@ -293,6 +293,21 @@ if lp := msg.ResponseMeta.LogProbs; lp != nil { > 注意:`Logprobs` 仅在 `ResponseLogprobs=true` 时生效。若想在单次调用中关闭 logprobs,请使用 `gemini.WithResponseLogprobs(false)`。 +## 标签(Labels) + +通过 `gemini.WithLabels(...)` 为请求附加用户自定义的元数据标签。标签是键值对,可用于计费归因、请求追踪以及在 Cloud Logging 和 Monitoring 中进行过滤。 + +```go +msg, _ := cm.Generate(ctx, input, + gemini.WithLabels(map[string]string{ + "team": "search", + "environment": "production", + }), +) +``` + +> 注意:标签仅在 **Vertex AI** 后端受支持。Gemini API 后端会拒绝设置了标签的请求(`labels parameter is not supported in Gemini API`),因此该选项只能与 Vertex AI 客户端一起使用。 + ## 示例 diff --git a/components/model/gemini/gemini.go b/components/model/gemini/gemini.go index 6d9de0b5d..e7b3d2c5a 100644 --- a/components/model/gemini/gemini.go +++ b/components/model/gemini/gemini.go @@ -601,6 +601,13 @@ func (cm *ChatModel) genInputAndConf(input []*schema.Message, opts ...model.Opti m.ImageConfig = geminiOptions.ImageConfig } + if len(geminiOptions.Labels) > 0 { + if cm.cli.ClientConfig().Backend != genai.BackendVertexAI { + return "", nil, nil, nil, fmt.Errorf("labels are only supported on the Vertex AI backend") + } + m.Labels = geminiOptions.Labels + } + if len(geminiOptions.CachedContentName) > 0 { m.CachedContent = geminiOptions.CachedContentName // remove system instruction and tools when using cached content diff --git a/components/model/gemini/gemini_test.go b/components/model/gemini/gemini_test.go index 3010a7f7a..84a9925bb 100644 --- a/components/model/gemini/gemini_test.go +++ b/components/model/gemini/gemini_test.go @@ -1706,6 +1706,39 @@ func TestLogprobsOptionsOverride(t *testing.T) { }) } +func TestWithLabels(t *testing.T) { + ctx := context.Background() + user := []*schema.Message{{Role: schema.User, Content: "hi"}} + labels := map[string]string{"team": "search", "env": "prod"} + + vertexClient, err := genai.NewClient(ctx, &genai.ClientConfig{ + Backend: genai.BackendVertexAI, + APIKey: "fake-key", + }) + assert.NoError(t, err) + + t.Run("labels applied on vertex backend", func(t *testing.T) { + cm := &ChatModel{cli: vertexClient, model: "gemini-x"} + _, _, m, _, err := cm.genInputAndConf(user, WithLabels(labels)) + assert.NoError(t, err) + assert.Equal(t, labels, m.Labels) + }) + + t.Run("labels rejected on gemini api backend", func(t *testing.T) { + cm := &ChatModel{cli: &genai.Client{Models: &genai.Models{}}, model: "gemini-x"} + _, _, _, _, err := cm.genInputAndConf(user, WithLabels(labels)) + assert.Error(t, err) + assert.Contains(t, err.Error(), "Vertex AI") + }) + + t.Run("no labels leaves field unset on gemini api backend", func(t *testing.T) { + cm := &ChatModel{cli: &genai.Client{Models: &genai.Models{}}, model: "gemini-x"} + _, _, m, _, err := cm.genInputAndConf(user) + assert.NoError(t, err) + assert.Nil(t, m.Labels) + }) +} + func TestConvCandidateLogprobs(t *testing.T) { candidate := &genai.Candidate{ Content: &genai.Content{ diff --git a/components/model/gemini/option.go b/components/model/gemini/option.go index e1ac48921..6a0fd37d1 100644 --- a/components/model/gemini/option.go +++ b/components/model/gemini/option.go @@ -30,6 +30,7 @@ type options struct { ResponseModalities []GeminiResponseModality ImageConfig *genai.ImageConfig CachedContentName string + Labels map[string]string ResponseLogprobs *bool Logprobs *int32 } @@ -75,6 +76,20 @@ func WithImageConfig(cfg *genai.ImageConfig) model.Option { }) } +// WithLabels sets user-defined metadata labels on the request. +// Labels are key-value pairs used for billing attribution, request tracking, +// and filtering in Cloud Logging and Monitoring. +// +// Labels are only supported on the Vertex AI backend. The Gemini API backend +// rejects requests that set labels, so this option must only be used with a +// Vertex AI client. +// Optional. +func WithLabels(labels map[string]string) model.Option { + return model.WrapImplSpecificOptFn(func(o *options) { + o.Labels = labels + }) +} + // WithResponseLogprobs sets whether to return the log probabilities of the // tokens chosen by the model at each step. This overrides Config.ResponseLogprobs. func WithResponseLogprobs(enable bool) model.Option {