Skip to content

Instantly share code, notes, and snippets.

@jba
Last active February 9, 2024 17:23
Show Gist options
  • Save jba/3023830445d77d44708c1c4eae77a3f1 to your computer and use it in GitHub Desktop.
Save jba/3023830445d77d44708c1c4eae77a3f1 to your computer and use it in GitHub Desktop.
add custom headers to a GCP client
package main
import (
"context"
"fmt"
"log"
"net/http"
"os"
"cloud.google.com/go/vertexai/genai"
"google.golang.org/api/option"
htransport "google.golang.org/api/transport/http"
)
func main() {
ctx := context.Background()
// Make a transport that has all the necessary authentication and other features for Google Cloud.
transport, err := htransport.NewTransport(ctx, http.DefaultTransport,
option.WithScopes("https://www.googleapis.com/auth/cloud-platform"))
if err != nil {
log.Fatal(err)
}
// Wrap it in a headerTransport (see definition below), which can be used to add any headers
// to a request. (It will also overwrite existing headers, so use with caution.)
headers := http.Header{}
headers.Set("My-Header", "17")
transport = &headerTransport{
headers: headers,
base: transport,
}
// Now make an http.Client with our transport and pass it as an option.
hc := &http.Client{Transport: transport}
client, err := genai.NewClient(ctx, os.Getenv("PROJECT"), "us-central1", genai.WithREST(), option.WithHTTPClient(hc))
if err != nil {
log.Fatal(err)
}
defer client.Close()
model := client.GenerativeModel("gemini-pro")
res, err := model.GenerateContent(ctx, genai.Text("My prompt."))
if err != nil {
log.Fatal(err)
}
fmt.Println(res.Candidates[0].Content.Parts)
}
type headerTransport struct {
headers http.Header
base http.RoundTripper
}
func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
rt := t.base
newReq := *req
newReq.Header = make(http.Header)
for k, vv := range req.Header {
newReq.Header[k] = vv
}
for k, v := range t.headers {
newReq.Header[k] = v
}
return rt.RoundTrip(&newReq)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment