feat: vertex veo (#1450)

relay/channel/vertex/adaptor.go

@@ -174,6 +174,7 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel
 		return err
 	}
 	req.Set("Authorization", "Bearer "+accessToken)
+	req.Set("x-goog-user-project", a.AccountCredentials.ProjectID)
 	return nil
 }
 

relay/channel/vertex/relay-vertex.go

@@ -12,7 +12,10 @@ func GetModelRegion(other string, localModelName string) string {
 		if m[localModelName] != nil {
 			return m[localModelName].(string)
 		} else {
-			return m["default"].(string)
+			if v, ok := m["default"]; ok {
+				return v.(string)
+			}
+			return "global"
 		}
 	}
 	return other

relay/channel/vertex/service_account.go

@@ -6,14 +6,15 @@ import (
 	"encoding/json"
 	"encoding/pem"
 	"errors"
-	"github.com/bytedance/gopkg/cache/asynccache"
-	"github.com/golang-jwt/jwt"
 	"net/http"
 	"net/url"
 	relaycommon "one-api/relay/common"
 	"one-api/service"
 	"strings"
 
+	"github.com/bytedance/gopkg/cache/asynccache"
+	"github.com/golang-jwt/jwt"
+
 	"fmt"
 	"time"
 )
@@ -137,3 +138,45 @@ func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (s
 
 	return "", fmt.Errorf("failed to get access token: %v", result)
 }
+
+func AcquireAccessToken(creds Credentials, proxy string) (string, error) {
+	signedJWT, err := createSignedJWT(creds.ClientEmail, creds.PrivateKey)
+	if err != nil {
+		return "", fmt.Errorf("failed to create signed JWT: %w", err)
+	}
+	return exchangeJwtForAccessTokenWithProxy(signedJWT, proxy)
+}
+
+func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, error) {
+	authURL := "https://www.googleapis.com/oauth2/v4/token"
+	data := url.Values{}
+	data.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
+	data.Set("assertion", signedJWT)
+
+	var client *http.Client
+	var err error
+	if proxy != "" {
+		client, err = service.NewProxyHttpClient(proxy)
+		if err != nil {
+			return "", fmt.Errorf("new proxy http client failed: %w", err)
+		}
+	} else {
+		client = service.GetHttpClient()
+	}
+
+	resp, err := client.PostForm(authURL, data)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	var result map[string]interface{}
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return "", err
+	}
+
+	if accessToken, ok := result["access_token"].(string); ok {
+		return accessToken, nil
+	}
+	return "", fmt.Errorf("failed to get access token: %v", result)
+}