feat: publish sanitized custom.97 source

This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207

No files matched your search

+5 -2
View File
@@ -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)
})
}
}
+14 -48
View File
@@ -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
}
+10 -9
View File
@@ -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 {
+25
View File
@@ -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)
})
}
}