Skip to content

Commit

Permalink
feat(vertexai/genai): add WithClientInfo option (googleapis#10535)
Browse files Browse the repository at this point in the history
Add an option that adds a key-value pair to the client info header
of requests from the client.
  • Loading branch information
jba authored Jul 11, 2024
1 parent eb63f0d commit 265963b
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 5 deletions.
6 changes: 5 additions & 1 deletion vertexai/genai/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,11 @@ func setGAPICClient[ClientType sgci](ctx context.Context, pf *ClientType, conf c
if err != nil {
return err
}
c.SetGoogleClientInfo("gccl", internal.Version)
kvs := []string{"gccl", internal.Version}
if conf.ciKey != "" && conf.ciValue != "" {
kvs = append(kvs, conf.ciKey, conf.ciValue)
}
c.SetGoogleClientInfo(kvs...)
*pf = c
return nil
}
Expand Down
27 changes: 23 additions & 4 deletions vertexai/genai/option.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,36 @@ func WithREST() option.ClientOption {
return &withREST{}
}

type withREST struct {
internaloption.EmbeddableAdapter
}

func (w *withREST) applyVertexaiOpt(c *config) {
c.withREST = true
}

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

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

func (w *withClientInfo) applyVertexaiOpt(c *config) {
c.ciKey = w.key
c.ciValue = w.value
}

type config struct {
// withREST tells the client to use REST as the underlying transport.
withREST bool
// key-value pair to add to the Google client info header.
ciKey string
ciValue string
}

// newConfig generates a new config with all the given
Expand All @@ -51,7 +74,3 @@ type vertexaiClientOption interface {
option.ClientOption
applyVertexaiOpt(*config)
}

type withREST struct {
internaloption.EmbeddableAdapter
}

0 comments on commit 265963b

Please sign in to comment.