feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -152,17 +152,20 @@ func (a *Adaptor) getRequestUrl(info *relaycommon.RelayInfo, modelName, suffix s
|
||||
} else {
|
||||
keyPrefix = "?"
|
||||
}
|
||||
// An API key carries no project. Without a configured project ID the URL
|
||||
// keeps the express-mode form, which has no region in its path.
|
||||
projectID := strings.TrimSpace(info.ChannelOtherSettings.VertexProjectID)
|
||||
if a.RequestMode == RequestModeGemini {
|
||||
return fmt.Sprintf(
|
||||
"%s%skey=%s",
|
||||
BuildGoogleModelURL(info.ChannelBaseUrl, DefaultAPIVersion, "", region, modelName, suffix),
|
||||
BuildGoogleModelURL(info.ChannelBaseUrl, DefaultAPIVersion, projectID, region, modelName, suffix),
|
||||
keyPrefix,
|
||||
info.ApiKey,
|
||||
), nil
|
||||
} else if a.RequestMode == RequestModeClaude {
|
||||
return fmt.Sprintf(
|
||||
"%s%skey=%s",
|
||||
BuildAnthropicModelURL(info.ChannelBaseUrl, DefaultAPIVersion, "", region, modelName, suffix),
|
||||
BuildAnthropicModelURL(info.ChannelBaseUrl, DefaultAPIVersion, projectID, region, modelName, suffix),
|
||||
keyPrefix,
|
||||
info.ApiKey,
|
||||
), nil
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package vertex
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetRequestURLWithAPIKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
projectID string
|
||||
region string
|
||||
stream bool
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "project ID keeps the global region",
|
||||
projectID: "my-project",
|
||||
region: `{"default":"global"}`,
|
||||
want: "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-4-argon:generateContent?key=test-key",
|
||||
},
|
||||
{
|
||||
name: "project ID keeps a regional endpoint when streaming",
|
||||
projectID: "my-project",
|
||||
region: `{"default":"us-central1"}`,
|
||||
stream: true,
|
||||
want: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-4-argon:streamGenerateContent?alt=sse&key=test-key",
|
||||
},
|
||||
{
|
||||
name: "no project ID keeps the express-mode URL",
|
||||
region: `{"default":"global"}`,
|
||||
want: "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-4-argon:generateContent?key=test-key",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{
|
||||
IsStream: tt.stream,
|
||||
OriginModelName: "gemini-4-argon",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ApiVersion: tt.region,
|
||||
ApiKey: "test-key",
|
||||
UpstreamModelName: "gemini-4-argon",
|
||||
ChannelOtherSettings: dto.ChannelOtherSettings{
|
||||
VertexKeyType: dto.VertexKeyTypeAPIKey,
|
||||
VertexProjectID: tt.projectID,
|
||||
},
|
||||
},
|
||||
}
|
||||
adaptor := &Adaptor{}
|
||||
adaptor.Init(info)
|
||||
|
||||
got, err := adaptor.GetRequestURL(info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
@@ -53,7 +52,7 @@ func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) {
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create signed JWT: %w", err)
|
||||
}
|
||||
newToken, err := exchangeJwtForAccessToken(signedJWT, info)
|
||||
newToken, err := exchangeJwtForAccessToken(signedJWT, info.ChannelSetting.Proxy)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to exchange JWT for access token: %w", err)
|
||||
}
|
||||
@@ -104,14 +103,24 @@ func createSignedJWT(email, privateKeyPEM string) (string, error) {
|
||||
return signedToken, nil
|
||||
}
|
||||
|
||||
func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (string, error) {
|
||||
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 exchangeJwtForAccessToken(signedJWT, proxy)
|
||||
}
|
||||
|
||||
// exchangeJwtForAccessToken trades a signed service-account JWT for an access
|
||||
// token at Google's OAuth endpoint. It goes through the channel proxy only and
|
||||
// always verifies the certificate; the channel's TLS setting does not apply.
|
||||
func exchangeJwtForAccessToken(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)
|
||||
|
||||
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
|
||||
client, err := service.GetHttpClientWithProxy(proxy)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("new proxy http client failed: %w", err)
|
||||
}
|
||||
@@ -122,50 +131,7 @@ func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (s
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
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.GetHttpClientWithProxy(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{}
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
@@ -47,17 +47,18 @@ func BuildAPIBaseURL(baseURL, version, projectID, region string) string {
|
||||
}
|
||||
|
||||
region = normalizeVertexRegion(region)
|
||||
host := region + "-aiplatform.googleapis.com"
|
||||
switch region {
|
||||
case "global":
|
||||
host = "aiplatform.googleapis.com"
|
||||
case "us", "eu":
|
||||
// The US and EU multi-regions are served from their own hostnames.
|
||||
host = "aiplatform." + region + ".rep.googleapis.com"
|
||||
}
|
||||
if strings.TrimSpace(projectID) == "" {
|
||||
if region == "global" {
|
||||
return fmt.Sprintf("https://aiplatform.googleapis.com/%s", version)
|
||||
}
|
||||
return fmt.Sprintf("https://%s-aiplatform.googleapis.com/%s", region, version)
|
||||
return fmt.Sprintf("https://%s/%s", host, version)
|
||||
}
|
||||
|
||||
if region == "global" {
|
||||
return fmt.Sprintf("https://aiplatform.googleapis.com/%s/projects/%s/locations/global", version, projectID)
|
||||
}
|
||||
return fmt.Sprintf("https://%s-aiplatform.googleapis.com/%s/projects/%s/locations/%s", region, version, projectID, region)
|
||||
return fmt.Sprintf("https://%s/%s/projects/%s/locations/%s", host, version, projectID, region)
|
||||
}
|
||||
|
||||
func BuildPublisherModelURL(baseURL, version, projectID, region, publisher, modelName, action string) string {
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
package vertex
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestBuildGoogleModelURLUsesTheLocationHost(t *testing.T) {
|
||||
tests := []struct {
|
||||
region string
|
||||
want string
|
||||
}{
|
||||
{"global", "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/publishers/google/models/gemini-4-argon:generateContent"},
|
||||
{"us-central1", "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-4-argon:generateContent"},
|
||||
{"us", "https://aiplatform.us.rep.googleapis.com/v1/projects/my-project/locations/us/publishers/google/models/gemini-4-argon:generateContent"},
|
||||
{"eu", "https://aiplatform.eu.rep.googleapis.com/v1/projects/my-project/locations/eu/publishers/google/models/gemini-4-argon:generateContent"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.region, func(t *testing.T) {
|
||||
got := BuildGoogleModelURL("", DefaultAPIVersion, "my-project", tt.region, "gemini-4-argon", "generateContent")
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user