Skip to content

Instantly share code, notes, and snippets.

@jba

jba/addheader.go

Last active February 9, 2024 17:23
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save jba/3023830445d77d44708c1c4eae77a3f1 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save jba/3023830445d77d44708c1c4eae77a3f1 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. jba revised this gist Feb 9, 2024. 1 changed file with 2 additions and 1 deletion.
    3 changes: 2 additions & 1 deletion addheader.go
    View file
    Original file line number Diff line number Diff line change
    @@ -15,7 +15,8 @@ import (
    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)
    transport, err := htransport.NewTransport(ctx, http.DefaultTransport,
    option.WithScopes("https://www.googleapis.com/auth/cloud-platform"))
    if err != nil {
    log.Fatal(err)
    }
  2. jba created this gist Feb 9, 2024.
    64 changes: 64 additions & 0 deletions addheader.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,64 @@
    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)
    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)
    }