diff --git a/genai/client.go b/genai/client.go index 7d0304d..52408d0 100644 --- a/genai/client.go +++ b/genai/client.go @@ -87,7 +87,15 @@ Import the option package as "google.golang.org/api/option".`) if err != nil { return nil, fmt.Errorf("creating discovery client: %w", err) } - gc.SetGoogleClientInfo("gccl", "v"+internal.Version, "genai-go", internal.Version) + + kvs := []string{"gccl", "v" + internal.Version, "genai-go", internal.Version} + if a, ok := optionOfType[*clientInfo](opts); ok { + kvs = append(kvs, a.key, a.value) + } + gc.SetGoogleClientInfo(kvs...) + mc.SetGoogleClientInfo(kvs...) + fc.SetGoogleClientInfo(kvs...) + return &Client{gc, mc, fc, cc, ds}, nil } diff --git a/genai/client_test.go b/genai/client_test.go index bec5e35..28f9c32 100644 --- a/genai/client_test.go +++ b/genai/client_test.go @@ -865,3 +865,16 @@ func uploadFile(t *testing.T, ctx context.Context, client *Client, filename stri }) return file } + +func TestAttribution(t *testing.T) { + a := WithClientInfo("k", "v") + opts := []option.ClientOption{option.WithAPIKey("x"), a, option.WithEndpoint("e")} + got, ok := optionOfType[*clientInfo](opts) + if !ok { + t.Fatal("not found") + } + want := a.(*clientInfo) + if got.key != want.key || got.value != want.value { + t.Errorf("got %q, %q, want %q, %q", got.key, got.value, want.key, want.value) + } +} diff --git a/genai/option.go b/genai/option.go new file mode 100644 index 0000000..5527865 --- /dev/null +++ b/genai/option.go @@ -0,0 +1,44 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package genai + +import ( + "google.golang.org/api/option" + "google.golang.org/api/option/internaloption" +) + +// WithClientInfo sets request information identifying the +// product that is calling this client. +func WithClientInfo(key, value string) option.ClientOption { + return &clientInfo{key: key, value: value} +} + +type clientInfo struct { + internaloption.EmbeddableAdapter + key, value string +} + +// optionOfType returns the first value of opts that has type T, +// along with true. If there is no option of that type, it returns +// the zero value for T and false. +func optionOfType[T option.ClientOption](opts []option.ClientOption) (T, bool) { + for _, opt := range opts { + if opt, ok := opt.(T); ok { + return opt, true + } + } + var z T + return z, false +}