Skip to content

Commit

Permalink
genai: add WithAttribution client option
Browse files Browse the repository at this point in the history
Add an option that lets Google products identify which calls
are made with this client.
  • Loading branch information
jba committed Jul 3, 2024
1 parent f043fb2 commit 4c938bf
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 1 deletion.
10 changes: 9 additions & 1 deletion genai/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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[*attribution](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
}

Expand Down
13 changes: 13 additions & 0 deletions genai/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -864,3 +864,16 @@ func uploadFile(t *testing.T, ctx context.Context, client *Client, filename stri
})
return file
}

func TestAttribution(t *testing.T) {
a := WithAttribution("k", "v")
opts := []option.ClientOption{option.WithAPIKey("x"), a, option.WithEndpoint("e")}
got, ok := optionOfType[*attribution](opts)
if !ok {
t.Fatal("not found")
}
want := a.(*attribution)
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)
}
}
41 changes: 41 additions & 0 deletions genai/option.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// 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"
)

// WithAttribution sets request information identifying the
// product that is calling this client.
func WithAttribution(key, value string) option.ClientOption {
return &attribution{key: key, value: value}
}

type attribution struct {
internaloption.EmbeddableAdapter
key, value string
}

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
}

0 comments on commit 4c938bf

Please sign in to comment.