mirror of
https://github.com/songquanpeng/one-api.git
synced 2025-12-27 10:15:58 +08:00
feat: support claude and gemini in vertex ai (#1621)
* feat: support claude and gemini in vertex ai * fix: do not show api key field in channel page when the type is VertexAI * fix: update getToken function to include channelId in cache key
This commit is contained in:
63
relay/adaptor/vertexai/token.go
Normal file
63
relay/adaptor/vertexai/token.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package vertexai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
credentials "cloud.google.com/go/iam/credentials/apiv1"
|
||||
"cloud.google.com/go/iam/credentials/apiv1/credentialspb"
|
||||
"github.com/patrickmn/go-cache"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
|
||||
type ApplicationDefaultCredentials struct {
|
||||
Type string `json:"type"`
|
||||
ProjectID string `json:"project_id"`
|
||||
PrivateKeyID string `json:"private_key_id"`
|
||||
PrivateKey string `json:"private_key"`
|
||||
ClientEmail string `json:"client_email"`
|
||||
ClientID string `json:"client_id"`
|
||||
AuthURI string `json:"auth_uri"`
|
||||
TokenURI string `json:"token_uri"`
|
||||
AuthProviderX509CertURL string `json:"auth_provider_x509_cert_url"`
|
||||
ClientX509CertURL string `json:"client_x509_cert_url"`
|
||||
UniverseDomain string `json:"universe_domain"`
|
||||
}
|
||||
|
||||
|
||||
var Cache = cache.New(50*time.Minute, 55*time.Minute)
|
||||
|
||||
const defaultScope = "https://www.googleapis.com/auth/cloud-platform"
|
||||
|
||||
func getToken(ctx context.Context, channelId int, adcJson string) (string, error) {
|
||||
cacheKey := fmt.Sprintf("vertexai-token-%d", channelId)
|
||||
if token, found := Cache.Get(cacheKey); found {
|
||||
return token.(string), nil
|
||||
}
|
||||
adc := &ApplicationDefaultCredentials{}
|
||||
if err := json.Unmarshal([]byte(adcJson), adc); err != nil {
|
||||
return "", fmt.Errorf("Failed to decode credentials file: %w", err)
|
||||
}
|
||||
|
||||
c, err := credentials.NewIamCredentialsClient(ctx, option.WithCredentialsJSON([]byte(adcJson)))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("Failed to create client: %w", err)
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
req := &credentialspb.GenerateAccessTokenRequest{
|
||||
// See https://pkg.go.dev/cloud.google.com/go/iam/credentials/apiv1/credentialspb#GenerateAccessTokenRequest.
|
||||
Name: fmt.Sprintf("projects/-/serviceAccounts/%s", adc.ClientEmail),
|
||||
Scope: []string{defaultScope},
|
||||
}
|
||||
resp, err := c.GenerateAccessToken(ctx, req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("Failed to generate access token: %w", err)
|
||||
}
|
||||
_ = resp
|
||||
|
||||
Cache.Set(cacheKey, resp.AccessToken, cache.DefaultExpiration)
|
||||
return resp.AccessToken, nil
|
||||
}
|
||||
Reference in New Issue
Block a user