Last active
February 9, 2024 17:23
-
-
Save jba/3023830445d77d44708c1c4eae77a3f1 to your computer and use it in GitHub Desktop.
add custom headers to a GCP client
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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