diff --git a/constant/context_key.go b/constant/context_key.go index ccb8010f9476..82ab7bf7b45c 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -27,6 +27,8 @@ const ( ContextKeyChannelCreateTime ContextKey = "channel_create_time" ContextKeyChannelBaseUrl ContextKey = "base_url" ContextKeyChannelType ContextKey = "channel_type" + ContextKeyRequiredChannelType ContextKey = "required_channel_type" + ContextKeyChannelModels ContextKey = "channel_models" ContextKeyChannelSetting ContextKey = "channel_setting" ContextKeyChannelOtherSetting ContextKey = "channel_other_setting" ContextKeyChannelParamOverride ContextKey = "param_override" diff --git a/controller/channel-test.go b/controller/channel-test.go index fffc59d24d5a..8b25a76c102c 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -70,6 +70,26 @@ func resolveChannelTestUserID(c *gin.Context) (int, error) { } func testChannel(ctx context.Context, channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool) testResult { + return testChannelWithVertexStorageDependencies( + ctx, + channel, + testUserID, + testModel, + endpointType, + isStream, + defaultVertexStorageChannelProbeDependencies(), + ) +} + +func testChannelWithVertexStorageDependencies( + ctx context.Context, + channel *model.Channel, + testUserID int, + testModel string, + endpointType string, + isStream bool, + storageDeps vertexStorageChannelProbeDependencies, +) testResult { if ctx == nil { ctx = context.Background() } @@ -106,6 +126,23 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te } } } + if isVertexStorageChannelTest(channel, testModel) { + c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, relayconstant.VertexStorageRoutePrefix, nil) + c.Set("channel", channel.Type) + c.Set("base_url", channel.GetBaseURL()) + newAPIError := middleware.SetupContextForSelectedChannel(c, channel, testModel) + if newAPIError != nil { + return testResult{ + context: c, + localErr: newAPIError, + newAPIError: newAPIError, + } + } + return testResult{ + context: c, + localErr: testVertexStorageChannel(ctx, c, testModel, storageDeps), + } + } endpointType = normalizeChannelTestEndpoint(channel, endpointType) @@ -660,6 +697,12 @@ func shouldUseStreamForAutomaticChannelTest(channel *model.Channel) bool { return channel != nil && channel.Type == constant.ChannelTypeCodex } +func isVertexStorageChannelTest(channel *model.Channel, testModel string) bool { + return channel != nil && + channel.Type == constant.ChannelTypeVertexAi && + strings.HasPrefix(strings.TrimSpace(testModel), relayconstant.VertexStorageModelPrefix) +} + func detectErrorMessageFromJSONBytes(jsonBytes []byte) string { if len(jsonBytes) == 0 { return "" diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..8924dc8cfc6c 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -181,13 +181,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } }() - retryParam := &service.RetryParam{ - Ctx: c, - TokenGroup: relayInfo.TokenGroup, - ModelName: relayInfo.OriginModelName, - RequestPath: c.Request.URL.Path, - Retry: common.GetPointer(0), - } + retryParam := newRelayRetryParam(c, relayInfo) relayInfo.RetryIndex = 0 relayInfo.LastError = nil @@ -311,6 +305,9 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service AutoBan: &autoBanInt, }, nil } + if retryParam.RequiredChannelType == 0 { + retryParam.RequiredChannelType = requiredChannelTypeForRelay(c) + } channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) if err != nil { return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) @@ -328,6 +325,27 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service return channel, nil } +func newRelayRetryParam(c *gin.Context, relayInfo *relaycommon.RelayInfo) *service.RetryParam { + return &service.RetryParam{ + Ctx: c, + TokenGroup: relayInfo.TokenGroup, + ModelName: relayInfo.OriginModelName, + RequestPath: c.Request.URL.Path, + Retry: common.GetPointer(0), + RequiredChannelType: requiredChannelTypeForRelay(c), + } +} + +func requiredChannelTypeForRelay(c *gin.Context) int { + if requiredChannelType := common.GetContextKeyInt(c, constant.ContextKeyRequiredChannelType); requiredChannelType != 0 { + return requiredChannelType + } + if relayconstant.IsVertexStoragePath(c.Request.URL.Path) { + return constant.ChannelTypeVertexAi + } + return 0 +} + func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { if openaiErr == nil { return false @@ -513,13 +531,7 @@ func RelayTask(c *gin.Context) { } }() - retryParam := &service.RetryParam{ - Ctx: c, - TokenGroup: relayInfo.TokenGroup, - ModelName: relayInfo.OriginModelName, - RequestPath: c.Request.URL.Path, - Retry: common.GetPointer(0), - } + retryParam := newRelayRetryParam(c, relayInfo) for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { var channel *model.Channel diff --git a/controller/vertex_storage_channel_probe.go b/controller/vertex_storage_channel_probe.go new file mode 100644 index 000000000000..aaff9b92d99c --- /dev/null +++ b/controller/vertex_storage_channel_probe.go @@ -0,0 +1,181 @@ +package controller + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/relay/channel/vertex" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +const vertexStorageChannelTestContent = "new-api Vertex AI Storage channel test\n" + +type vertexStorageChannelProbeDependencies struct { + newObjectName func() string + acquireAccessToken func(vertex.CachedAccessTokenRequest) (string, error) + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) +} + +func defaultVertexStorageChannelProbeDependencies() vertexStorageChannelProbeDependencies { + return vertexStorageChannelProbeDependencies{ + newObjectName: func() string { + return ".new-api-channel-test/" + uuid.NewString() + "/test.txt" + }, + acquireAccessToken: vertex.AcquireCachedAccessToken, + doProxy: vertex.DoStorageProxy, + } +} + +func testVertexStorageChannel(ctx context.Context, c *gin.Context, testModel string, deps vertexStorageChannelProbeDependencies) error { + if c == nil { + return errors.New("Vertex storage channel test context is required") + } + if deps.newObjectName == nil || deps.acquireAccessToken == nil || deps.doProxy == nil { + return errors.New("Vertex storage channel test dependencies are incomplete") + } + + modelName := strings.TrimSpace(testModel) + if !strings.HasPrefix(modelName, relayconstant.VertexStorageModelPrefix) { + return errors.New("invalid Vertex storage test model") + } + bucket, err := relayconstant.NormalizeVertexStorageBucket(strings.TrimPrefix(modelName, relayconstant.VertexStorageModelPrefix)) + if err != nil { + return fmt.Errorf("invalid Vertex storage test model: %w", err) + } + if common.GetContextKeyInt(c, constant.ContextKeyChannelType) != constant.ChannelTypeVertexAi { + return errors.New("selected channel is not Vertex AI") + } + if !relayconstant.VertexStorageChannelSupports(common.GetContextKeyStringSlice(c, constant.ContextKeyChannelModels), bucket) { + return fmt.Errorf("selected channel does not allow Cloud Storage bucket %q", bucket) + } + + channelOtherSetting, _ := common.GetContextKeyType[dto.ChannelOtherSettings](c, constant.ContextKeyChannelOtherSetting) + if channelOtherSetting.VertexKeyType == dto.VertexKeyTypeAPIKey { + return errors.New("Vertex storage channel test requires service account JSON") + } + credentials := vertex.Credentials{} + if err := common.Unmarshal([]byte(common.GetContextKeyString(c, constant.ContextKeyChannelKey)), &credentials); err != nil { + return errors.New("selected Vertex AI channel credentials are invalid") + } + + channelSetting, _ := common.GetContextKeyType[dto.ChannelSettings](c, constant.ContextKeyChannelSetting) + accessToken, err := deps.acquireAccessToken(vertex.CachedAccessTokenRequest{ + ChannelID: common.GetContextKeyInt(c, constant.ContextKeyChannelId), + ChannelIsMultiKey: common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey), + ChannelMultiKeyIndex: common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex), + Credentials: credentials, + Proxy: channelSetting.Proxy, + }) + if err != nil { + return errors.New("failed to authorize Vertex storage channel test") + } + + objectName := strings.TrimSpace(deps.newObjectName()) + if objectName == "" { + return errors.New("Vertex storage channel test object name is empty") + } + if ctx == nil { + ctx = context.Background() + } + + var probeErrors []error + uploadQuery := url.Values{} + uploadQuery.Set("uploadType", "media") + uploadQuery.Set("name", objectName) + uploadHeader := make(http.Header) + uploadHeader.Set("Content-Type", "text/plain; charset=utf-8") + _, err = runVertexStorageProbeRequest(ctx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationUpload, + Method: http.MethodPost, + Bucket: bucket, + RawQuery: uploadQuery.Encode(), + Header: uploadHeader, + Body: strings.NewReader(vertexStorageChannelTestContent), + ContentLength: int64(len(vertexStorageChannelTestContent)), + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, 0) + if err != nil { + probeErrors = append(probeErrors, fmt.Errorf("upload temporary object: %w", err)) + } + + downloaded, readErr := runVertexStorageProbeRequest(ctx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationGet, + Method: http.MethodGet, + Bucket: bucket, + Object: objectName, + RawQuery: "alt=media", + Header: make(http.Header), + ContentLength: 0, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, int64(len(vertexStorageChannelTestContent))) + if readErr != nil { + probeErrors = append(probeErrors, fmt.Errorf("read temporary object: %w", readErr)) + } else if !bytes.Equal(downloaded, []byte(vertexStorageChannelTestContent)) { + probeErrors = append(probeErrors, errors.New("read temporary object: content mismatch")) + } + + cleanupCtx, cancelCleanup := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancelCleanup() + _, deleteErr := runVertexStorageProbeRequest(cleanupCtx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationDelete, + Method: http.MethodDelete, + Bucket: bucket, + Object: objectName, + Header: make(http.Header), + ContentLength: 0, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, 0) + if deleteErr != nil { + probeErrors = append(probeErrors, fmt.Errorf("delete temporary object %q manually if necessary: %w", objectName, deleteErr)) + } + + return errors.Join(probeErrors...) +} + +func runVertexStorageProbeRequest( + ctx context.Context, + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error), + input vertex.StorageProxyRequest, + maxResponseBytes int64, +) ([]byte, error) { + response, err := doProxy(ctx, input) + if response != nil && response.Body != nil { + defer service.CloseResponseBodyGracefully(response) + } + if err != nil { + return nil, err + } + if response == nil || response.Body == nil { + return nil, errors.New("Google Cloud Storage returned an empty response") + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return nil, fmt.Errorf("Google Cloud Storage returned status %d", response.StatusCode) + } + if maxResponseBytes <= 0 { + return nil, nil + } + body, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes+1)) + if err != nil { + return nil, err + } + if int64(len(body)) > maxResponseBytes { + return nil, errors.New("Google Cloud Storage returned an oversized test object") + } + return body, nil +} diff --git a/controller/vertex_storage_proxy.go b/controller/vertex_storage_proxy.go new file mode 100644 index 000000000000..d39262ab69e6 --- /dev/null +++ b/controller/vertex_storage_proxy.go @@ -0,0 +1,151 @@ +package controller + +import ( + "context" + "io" + "net/http" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/relay/channel/vertex" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" +) + +type vertexStorageProxyDependencies struct { + acquireAccessToken func(vertex.CachedAccessTokenRequest) (string, error) + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) +} + +func defaultVertexStorageProxyDependencies() vertexStorageProxyDependencies { + return vertexStorageProxyDependencies{ + acquireAccessToken: vertex.AcquireCachedAccessToken, + doProxy: vertex.DoStorageProxy, + } +} + +func RelayVertexStorageUpload(c *gin.Context) { + relayVertexStorageProxy(c, vertex.StorageOperationUpload, defaultVertexStorageProxyDependencies()) +} + +func RelayVertexStorageList(c *gin.Context) { + relayVertexStorageProxy(c, vertex.StorageOperationList, defaultVertexStorageProxyDependencies()) +} + +func RelayVertexStorageObject(c *gin.Context) { + operation := vertex.StorageOperationGet + if c.Request.Method == http.MethodDelete { + operation = vertex.StorageOperationDelete + } + relayVertexStorageProxy(c, operation, defaultVertexStorageProxyDependencies()) +} + +func relayVertexStorageProxy(c *gin.Context, operation vertex.StorageOperation, deps vertexStorageProxyDependencies) { + bucket, err := relayconstant.NormalizeVertexStorageBucket(c.Param("bucket")) + if err != nil { + respondVertexStorageProxyError(c, http.StatusBadRequest, "invalid_bucket", "invalid Cloud Storage bucket") + return + } + if common.GetContextKeyInt(c, constant.ContextKeyChannelType) != constant.ChannelTypeVertexAi { + respondVertexStorageProxyError(c, http.StatusInternalServerError, "channel_type_mismatch", "selected channel is not Vertex AI") + return + } + if !relayconstant.VertexStorageChannelSupports(common.GetContextKeyStringSlice(c, constant.ContextKeyChannelModels), bucket) { + respondVertexStorageProxyError(c, http.StatusForbidden, "bucket_not_allowed", "selected channel does not allow this Cloud Storage bucket") + return + } + + channelOtherSetting, _ := common.GetContextKeyType[dto.ChannelOtherSettings](c, constant.ContextKeyChannelOtherSetting) + if channelOtherSetting.VertexKeyType == dto.VertexKeyTypeAPIKey { + respondVertexStorageProxyError(c, http.StatusBadRequest, "unsupported_key_type", "Cloud Storage access requires Vertex AI service account JSON") + return + } + + credentials := vertex.Credentials{} + if err = common.Unmarshal([]byte(common.GetContextKeyString(c, constant.ContextKeyChannelKey)), &credentials); err != nil { + respondVertexStorageProxyError(c, http.StatusInternalServerError, "invalid_channel_credentials", "selected Vertex AI channel credentials are invalid") + return + } + + object := strings.TrimPrefix(c.Param("object"), "/") + if (operation == vertex.StorageOperationGet || operation == vertex.StorageOperationDelete) && object == "" { + respondVertexStorageProxyError(c, http.StatusBadRequest, "object_required", "Cloud Storage object is required") + return + } + if (operation == vertex.StorageOperationGet || operation == vertex.StorageOperationDelete) && relayconstant.ValidateVertexStorageObjectName(object) != nil { + respondVertexStorageProxyError(c, http.StatusBadRequest, "invalid_object", "Cloud Storage object contains an invalid path segment") + return + } + + channelSetting, _ := common.GetContextKeyType[dto.ChannelSettings](c, constant.ContextKeyChannelSetting) + accessToken, err := deps.acquireAccessToken(vertex.CachedAccessTokenRequest{ + ChannelID: common.GetContextKeyInt(c, constant.ContextKeyChannelId), + ChannelIsMultiKey: common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey), + ChannelMultiKeyIndex: common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex), + Credentials: credentials, + Proxy: channelSetting.Proxy, + }) + if err != nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "access_token_failed", "failed to authorize with Vertex AI service account") + return + } + + response, err := deps.doProxy(c.Request.Context(), vertex.StorageProxyRequest{ + Operation: operation, + Method: c.Request.Method, + Bucket: bucket, + Object: object, + RawQuery: c.Request.URL.RawQuery, + Header: c.Request.Header, + Body: c.Request.Body, + ContentLength: c.Request.ContentLength, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }) + if response != nil && response.Body != nil { + defer service.CloseResponseBodyGracefully(response) + } + if err != nil || response == nil || response.Body == nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "upstream_request_failed", "failed to request Google Cloud Storage") + return + } + + responseHeader := vertex.SanitizeStorageResponseHeader(response.Header) + if operation == vertex.StorageOperationUpload && responseHeader.Get("Location") != "" { + rewrittenLocation, rewriteErr := vertex.RewriteStorageResumableLocation(responseHeader.Get("Location"), system_setting.ServerAddress, bucket) + if rewriteErr != nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "invalid_resumable_location", "Google Cloud Storage returned an invalid resumable upload location") + return + } + responseHeader.Set("Location", rewrittenLocation) + } + if response.StatusCode < http.StatusContinue || response.StatusCode > 599 { + respondVertexStorageProxyError(c, http.StatusBadGateway, "invalid_upstream_status", "Google Cloud Storage returned an invalid response status") + return + } + for name, values := range responseHeader { + if !service.ShouldCopyUpstreamHeader(c, name, values) { + continue + } + for _, value := range values { + c.Writer.Header().Add(name, value) + } + } + c.Status(response.StatusCode) + if _, err = io.Copy(c.Writer, response.Body); err != nil { + logger.LogError(c, "failed to stream Google Cloud Storage response") + } +} + +func respondVertexStorageProxyError(c *gin.Context, status int, code, message string) { + c.JSON(status, gin.H{"error": gin.H{ + "message": message, + "type": "invalid_request_error", + "code": code, + }}) +} diff --git a/docs/openapi/relay.json b/docs/openapi/relay.json index 92a6239074e2..2a0a4cd251c0 100644 --- a/docs/openapi/relay.json +++ b/docs/openapi/relay.json @@ -51,6 +51,9 @@ { "name": "OpenAI音频(Audio)" }, + { + "name": "文件/Vertex AI Cloud Storage" + }, { "name": "重排序(Rerank)" }, @@ -71,6 +74,411 @@ } ], "paths": { + "/vertexai/upload/storage/v1/b/{bucket}/o": { + "post": { + "summary": "上传 Cloud Storage 对象", + "description": "通过 Vertex AI 渠道的服务账号转发 GCS media、multipart 或 resumable 初始化请求。Location 会安全改写为当前网关地址；此接口不计费。", + "operationId": "vertexStorageUploadObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "uploadType", + "in": "query", + "required": true, + "schema": { + "type": "string", + "enum": [ + "media", + "multipart", + "resumable" + ] + } + }, + { + "name": "name", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": false, + "content": { + "application/octet-stream": { + "schema": { + "type": "string", + "format": "binary" + } + }, + "application/json": { + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "multipart/related": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "responses": { + "200": { + "description": "上传成功或 resumable 会话已创建。", + "headers": { + "Location": { + "description": "已改写为当前网关的 resumable 上传地址。", + "schema": { + "type": "string", + "format": "uri" + } + } + } + }, + "201": { + "description": "对象创建成功。" + }, + "400": { + "description": "请求、bucket、上传类型或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth、GCS 请求或 resumable Location 校验失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "put": { + "summary": "续传 Cloud Storage 对象分块", + "description": "使用 resumable 初始化返回的网关 Location 上传分块。请求体与响应流式转发；此接口不计费。", + "operationId": "vertexStorageResumeUploadObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "uploadType", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "upload_id", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "name", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "Content-Range", + "in": "header", + "required": false, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": false, + "content": { + "application/octet-stream": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "responses": { + "200": { + "description": "最终分块上传完成。" + }, + "201": { + "description": "对象创建成功。" + }, + "308": { + "description": "分块已接受，resumable 会话尚未完成。" + }, + "400": { + "description": "请求、bucket 或 Content-Range 无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/vertexai/storage/v1/b/{bucket}/o": { + "get": { + "summary": "列举 Cloud Storage 对象", + "description": "列举渠道允许访问的 bucket 中的对象；查询参数按 GCS JSON API 透传。此接口不计费。", + "operationId": "vertexStorageListObjects", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "prefix", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "delimiter", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "pageToken", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Google Cloud Storage 对象列表。" + }, + "400": { + "description": "请求、bucket 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/vertexai/storage/v1/b/{bucket}/o/{object}": { + "get": { + "summary": "读取或下载 Cloud Storage 对象", + "description": "默认返回对象元数据；设置 alt=media 时流式下载对象内容。对象名中的 / 应编码为 %2F；此接口不计费。", + "operationId": "vertexStorageGetObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "object", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "alt", + "in": "query", + "required": false, + "schema": { + "type": "string", + "enum": [ + "json", + "media" + ] + } + }, + { + "name": "generation", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "Range", + "in": "header", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "对象元数据或完整对象内容。" + }, + "206": { + "description": "Range 下载内容。", + "headers": { + "Content-Range": { + "schema": { + "type": "string" + } + } + } + }, + "400": { + "description": "请求、bucket、object 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "delete": { + "summary": "删除 Cloud Storage 对象", + "description": "删除渠道允许访问的 bucket 中的对象。对象名中的 / 应编码为 %2F；此接口不计费。", + "operationId": "vertexStorageDeleteObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "object", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "generation", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "ifGenerationMatch", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "对象已删除。" + }, + "400": { + "description": "请求、bucket、object 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, "/v1/models": { "get": { "summary": "获取模型列表", diff --git a/docs/vertex-ai-storage.md b/docs/vertex-ai-storage.md new file mode 100644 index 000000000000..e027bd4efa78 --- /dev/null +++ b/docs/vertex-ai-storage.md @@ -0,0 +1,181 @@ +# 如何通过 Vertex AI 渠道管理 Cloud Storage 文件 + +## 目标与边界 + +本指南面向需要让 Vertex AI Gemini 通过 `fileData.fileUri` 使用 Google Cloud Storage（GCS）文件的接入工程师。完成配置后，你可以使用 new-api 的固定 `/vertexai` 路由上传、列举、读取、下载和删除指定 bucket 中的对象。 + +这些路由只代理 GCS JSON API 的五组固定操作，不提供任意 Google API、任意主机或任意 URL 的通配代理，也不负责创建 bucket、修改 IAM、管理 ACL、复制或重写对象。授权粒度是整个 bucket，不支持只授权某个对象名前缀。 + +## 前置 IAM 与服务账号 + +1. 在 Google Cloud 项目中创建供 Vertex AI 渠道使用的服务账号。 +2. 为服务账号授予调用 Vertex AI 所需的权限，例如 `roles/aiplatform.user`。 +3. 在目标 bucket 上授予与实际操作匹配的 GCS 权限。需要完整执行本指南的上传、下载、列举和删除操作时，可授予 `roles/storage.objectUser`；生产环境应按最小权限原则拆分只读或只写权限。 +4. 生成服务账号 JSON 密钥，并将完整 JSON 配置到 Vertex AI 渠道的 Key 中。 + +GCS 文件代理不支持 Vertex AI API Key 模式。不要把服务账号 JSON、私钥或网关 Token 写入源码、日志或公开文档。 + +## 配置渠道 bucket + +在管理后台编辑 Vertex AI 渠道，在“存储桶”字段中填写允许访问的纯 bucket 名称，例如： + +```text +example-bucket +archive-bucket-01 +``` + +每个值必须只有 bucket 名称。不要填写 `gs://example-bucket`、`storage:gs:example-bucket`、`example-bucket/path` 或带查询字符串的值。保存后，系统会在渠道 `models` 中以 `storage:gs:` 形式持久化授权；请求只能访问所选渠道精确配置的 bucket。 + +以下示例统一使用脱敏环境变量： + +```bash +export NEW_API_BASE_URL="https://api.example.com" +export NEW_API_TOKEN="" +export GCS_BUCKET="example-bucket" +export GCS_OBJECT="docs%2Freport.pdf" +``` + +## 使用五组固定路由 + +### 1. 使用 POST 上传对象 + +使用 media upload 将本地文件流式上传到指定对象名： + +```bash +curl --fail-with-body \ + -X POST \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/pdf" \ + --data-binary @./report.pdf \ + "${NEW_API_BASE_URL}/vertexai/upload/storage/v1/b/${GCS_BUCKET}/o?uploadType=media&name=${GCS_OBJECT}" +``` + +此路由也支持 GCS 的 `multipart` 上传和 resumable 初始化。对象名中的 `/` 应编码为 `%2F`。 + +### 2. 使用 PUT 续传分块 + +先按“Resumable 上传”一节初始化会话，再对返回的网关 `Location` 发送分块： + +```bash +curl --fail-with-body \ + -X PUT \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/octet-stream" \ + -H "Content-Range: bytes 0-1048575/2097152" \ + --data-binary @./chunk-000.bin \ + "${RESUMABLE_LOCATION}" +``` + +未完成时 GCS 通常返回 `308 Resume Incomplete`；最终分块完成后返回对象元数据。 + +### 3. 列举对象 + +使用 `prefix` 等 GCS 查询参数筛选列表： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o?prefix=docs%2F" +``` + +渠道授权仍以整个 bucket 为单位，`prefix` 只筛选响应，不会缩小授权范围。 + +### 4. 读取元数据或下载对象 + +省略 `alt=media` 时读取对象元数据： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}" +``` + +设置 `alt=media` 时流式下载对象内容： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -o ./downloaded-report.pdf \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}?alt=media" +``` + +需要断点下载时可增加 `Range: bytes=0-1048575`。网关会保留 GCS 的状态码、`Content-Range`、`ETag` 和响应体。 + +### 5. 删除对象 + +删除前确认对象名和 bucket，删除操作不可由网关恢复： + +```bash +curl --fail-with-body \ + -X DELETE \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}" +``` + +## 在 `fileData.fileUri` 中引用对象 + +上传完成后，在 Vertex AI Gemini 请求中使用原始 GCS URI，而不是 `/vertexai` 网关下载地址： + +```json +{ + "contents": [ + { + "role": "user", + "parts": [ + { + "fileData": { + "mimeType": "application/pdf", + "fileUri": "gs://example-bucket/docs/report.pdf" + } + }, + { + "text": "请总结这份报告。" + } + ] + } + ] +} +``` + +执行推理的 Vertex AI 服务账号必须能够读取该对象。网关 Token 对 GCS 代理路由的访问权限不会替代 Google Cloud IAM。 + +## 使用 Resumable 上传 + +初始化 resumable 会话，并只从响应头读取 `Location`： + +```bash +RESUMABLE_LOCATION="$({ + curl --silent --show-error --fail-with-body \ + -X POST \ + -D - \ + -o /dev/null \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/json; charset=UTF-8" \ + -H "X-Upload-Content-Type: application/octet-stream" \ + --data "{\"name\":\"docs/report.pdf\"}" \ + "${NEW_API_BASE_URL}/vertexai/upload/storage/v1/b/${GCS_BUCKET}/o?uploadType=resumable" +} | awk 'tolower($1) == "location:" { sub(/\r$/, "", $2); print $2 }')" +``` + +返回的 `Location` 必须以 `${NEW_API_BASE_URL}/vertexai/` 开头。网关会校验 Google 返回的 session URL，并使用系统配置的服务地址安全改写；改写失败时返回 `502`，不会把 `storage.googleapis.com` session URL 暴露给客户端。后续每个 `PUT` 都会重新执行 Token 鉴权、限流、渠道分发和 bucket 授权。 + +## 常见错误与非计费说明 + +| 现象 | 常见原因 | 排查方式 | +| --- | --- | --- | +| `400 invalid_bucket` | bucket 包含路径、scheme、查询字符串或完整配置前缀 | 只传纯 bucket 名称 | +| `400 unsupported_key_type` | 渠道使用 Vertex AI API Key | 改用服务账号 JSON | +| `400 object_required` | 读取或删除路由缺少对象名 | 对完整对象名做 URL 编码后放入路径 | +| `403 bucket_not_allowed` | 所选渠道未精确配置目标 bucket | 检查渠道“存储桶”配置及 Token/用户组权限 | +| GCS `403` 响应 | 服务账号缺少 bucket IAM 权限 | 检查 bucket IAM 和服务账号身份 | +| `502 access_token_failed` | 服务账号 JSON、私钥、代理或 Google OAuth 异常 | 检查渠道凭证和 Proxy 配置，不要在日志中输出私钥 | +| `502 invalid_resumable_location` | 系统服务地址缺失，或 Google 返回的 session URL 未通过安全校验 | 配置正确的公开服务地址后重新初始化会话 | + +这些 Storage Proxy 请求采用流式传输，不构造模型 RelayInfo，不执行 quota 预扣或结算，也不写模型消费日志。GCS 本身产生的存储、网络和请求费用仍由你的 Google Cloud 项目承担。 + +## 相关链接 + +- [Google Cloud Storage JSON API](https://cloud.google.com/storage/docs/json_api) +- [Cloud Storage IAM roles](https://cloud.google.com/storage/docs/access-control/iam-roles) +- [Resumable uploads](https://cloud.google.com/storage/docs/performing-resumable-uploads) +- [Vertex AI 文件输入](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/document-understanding) diff --git a/middleware/distributor.go b/middleware/distributor.go index 3f53aa350349..83fcf650ad78 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -31,7 +31,16 @@ type ModelRequest struct { } func Distribute() func(c *gin.Context) { + return distribute(0) +} + +func DistributeByChannelType(requiredChannelType int) func(c *gin.Context) { + return distribute(requiredChannelType) +} + +func distribute(requiredChannelType int) func(c *gin.Context) { return func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyRequiredChannelType, requiredChannelType) var channel *model.Channel channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId) modelRequest, shouldSelectChannel, err := getModelRequest(c) @@ -54,6 +63,14 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } + if requiredChannelType != 0 && channel.Type != requiredChannelType { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return + } + if relayconstant.IsVertexStoragePath(c.Request.URL.Path) && !relayconstant.VertexStorageChannelSupports(channel.GetModels(), c.Param("bucket")) { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return + } } else { // Select a channel for the user // check token model mapping @@ -106,6 +123,7 @@ func Distribute() func(c *gin.Context) { affinityUsable := false preferred, err := model.CacheGetChannel(preferredChannelID) if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && + (requiredChannelType == 0 || preferred.Type == requiredChannelType) && channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) @@ -134,11 +152,12 @@ func Distribute() func(c *gin.Context) { if channel == nil { channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ - Ctx: c, - ModelName: modelRequest.Model, - TokenGroup: usingGroup, - RequestPath: c.Request.URL.Path, - Retry: common.GetPointer(0), + Ctx: c, + ModelName: modelRequest.Model, + TokenGroup: usingGroup, + RequestPath: c.Request.URL.Path, + Retry: common.GetPointer(0), + RequiredChannelType: requiredChannelType, }) if err != nil { showGroup := usingGroup @@ -254,7 +273,14 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { var modelRequest ModelRequest shouldSelectChannel := true var err error - if strings.Contains(c.Request.URL.Path, "/mj/") { + if relayconstant.IsVertexStoragePath(c.Request.URL.Path) { + modelName, err := relayconstant.VertexStorageModelName(c.Param("bucket")) + if err != nil { + return nil, false, err + } + modelRequest.Model = modelName + c.Set("relay_mode", relayconstant.RelayModeVertexStorage) + } else if strings.Contains(c.Request.URL.Path, "/mj/") { relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path) if relayMode == relayconstant.RelayModeMidjourneyTaskFetch || relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition || @@ -446,6 +472,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id) common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name) common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type) + common.SetContextKey(c, constant.ContextKeyChannelModels, channel.GetModels()) common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime) common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting()) common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings()) diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..4822acbd0a72 100644 --- a/model/ability.go +++ b/model/ability.go @@ -60,12 +60,18 @@ func GetAllEnableAbilities() []Ability { return abilities } -func getPriority(group string, model string, retry int) (int, error) { - +func getPriority(group string, model string, retry int, requiredChannelType int) (int, error) { var priorities []int - err := DB.Model(&Ability{}). + priorityQuery := DB.Model(&Ability{}). Select("DISTINCT(priority)"). - Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true). + Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) + if requiredChannelType != 0 { + enabledChannelIDs := DB.Model(&Channel{}). + Select("id"). + Where("status = ? and type = ?", common.ChannelStatusEnabled, requiredChannelType) + priorityQuery = priorityQuery.Where("channel_id IN (?)", enabledChannelIDs) + } + err := priorityQuery. Order("priority DESC"). // 按优先级降序排序 Pluck("priority", &priorities).Error // Pluck用于将查询的结果直接扫描到一个切片中 @@ -90,15 +96,29 @@ func getPriority(group string, model string, retry int) (int, error) { return priorityToUse, nil } -func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { - maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) +func getChannelQuery(group string, model string, retry int, requiredChannelType int) (*gorm.DB, error) { + maxPrioritySubQuery := DB.Model(&Ability{}). + Select("MAX(priority)"). + Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) + var enabledChannelIDs *gorm.DB + if requiredChannelType != 0 { + enabledChannelIDs = DB.Model(&Channel{}). + Select("id"). + Where("status = ? and type = ?", common.ChannelStatusEnabled, requiredChannelType) + maxPrioritySubQuery = maxPrioritySubQuery.Where("channel_id IN (?)", enabledChannelIDs) + } channelQuery := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = (?)", group, model, true, maxPrioritySubQuery) + if enabledChannelIDs != nil { + channelQuery = channelQuery.Where("channel_id IN (?)", enabledChannelIDs) + } if retry != 0 { - priority, err := getPriority(group, model, retry) + priority, err := getPriority(group, model, retry, requiredChannelType) if err != nil { return nil, err - } else { - channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority) + } + channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority) + if enabledChannelIDs != nil { + channelQuery = channelQuery.Where("channel_id IN (?)", enabledChannelIDs) } } @@ -106,10 +126,14 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { } func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetChannelByType(group, model, retry, requestPath, 0) +} + +func GetChannelByType(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { var abilities []Ability var err error = nil - channelQuery, err := getChannelQuery(group, model, retry) + channelQuery, err := getChannelQuery(group, model, retry, requiredChannelType) if err != nil { return nil, err } diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..8e8fd68728b4 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -112,21 +112,30 @@ func SyncChannelCache(frequency int) { } func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetRandomSatisfiedChannelByType(group, model, retry, requestPath, 0) +} + +func GetRandomSatisfiedChannelByType(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry, requestPath) + return GetChannelByType(group, model, retry, requestPath, requiredChannelType) } + return getRandomSatisfiedChannelFromCache(group, model, retry, requestPath, requiredChannelType) +} +func getRandomSatisfiedChannelFromCache(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { channelSyncLock.RLock() defer channelSyncLock.RUnlock() // First, try to find channels with the exact model name. channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model) + channels = filterChannelsByType(channels, requiredChannelType) // If no channels found, try to find channels with the normalized model name. if len(channels) == 0 { normalizedModel := ratio_setting.FormatMatchingModelName(model) channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model) + channels = filterChannelsByType(channels, requiredChannelType) } if len(channels) == 0 { @@ -208,6 +217,20 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat return nil, errors.New("channel not found") } +func filterChannelsByType(channelIDs []int, requiredChannelType int) []int { + if requiredChannelType == 0 { + return channelIDs + } + filtered := make([]int, 0, len(channelIDs)) + for _, channelID := range channelIDs { + channel, ok := channelsIDM[channelID] + if ok && channel.Type == requiredChannelType { + filtered = append(filtered, channelID) + } + } + return filtered +} + // filterChannelsByRequestPathAndModel restricts candidates by request path and // model. Only Advanced Custom (type 58) channels are path-checked: they are kept // only when one of their configured routes matches requestPath and model. All diff --git a/relay/channel/vertex/service_account.go b/relay/channel/vertex/service_account.go index 6d5fee0d94fc..cd9f16ccbf2a 100644 --- a/relay/channel/vertex/service_account.go +++ b/relay/channel/vertex/service_account.go @@ -3,21 +3,20 @@ package vertex import ( "crypto/rsa" "crypto/x509" - "encoding/json" "encoding/pem" "errors" + "fmt" "net/http" "net/url" "strings" + "time" + "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/service" "github.com/bytedance/gopkg/cache/asynccache" "github.com/golang-jwt/jwt/v5" - - "fmt" - "time" ) type Credentials struct { @@ -37,32 +36,59 @@ var Cache = asynccache.NewAsyncCache(asynccache.Options{ }, }) -func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) { - var cacheKey string - if info.ChannelIsMultiKey { - cacheKey = fmt.Sprintf("access-token-%d-%d", info.ChannelId, info.ChannelMultiKeyIndex) - } else { - cacheKey = fmt.Sprintf("access-token-%d", info.ChannelId) +type CachedAccessTokenRequest struct { + ChannelID int + ChannelIsMultiKey bool + ChannelMultiKeyIndex int + Credentials Credentials + Proxy string +} + +func AcquireCachedAccessToken(input CachedAccessTokenRequest) (string, error) { + return acquireCachedAccessToken(input, func(signedJWT string) (string, error) { + return exchangeJwtForAccessTokenWithProxy(signedJWT, input.Proxy) + }) +} + +func acquireCachedAccessToken(input CachedAccessTokenRequest, exchange func(string) (string, error)) (string, error) { + cacheKey := fmt.Sprintf("access-token-%d", input.ChannelID) + if input.ChannelIsMultiKey { + cacheKey = fmt.Sprintf("access-token-%d-%d", input.ChannelID, input.ChannelMultiKeyIndex) } - val, err := Cache.Get(cacheKey) - if err == nil { - return val.(string), nil + if value, err := Cache.Get(cacheKey); err == nil { + if token, ok := value.(string); ok && token != "" { + return token, nil + } } - signedJWT, err := createSignedJWT(a.AccountCredentials.ClientEmail, a.AccountCredentials.PrivateKey) + signedJWT, err := createSignedJWT(input.Credentials.ClientEmail, input.Credentials.PrivateKey) if err != nil { return "", fmt.Errorf("failed to create signed JWT: %w", err) } - newToken, err := exchangeJwtForAccessToken(signedJWT, info) + newToken, err := exchange(signedJWT) if err != nil { return "", fmt.Errorf("failed to exchange JWT for access token: %w", err) } - if err := Cache.SetDefault(cacheKey, newToken); err { - return newToken, nil + value := Cache.GetOrSet(cacheKey, newToken) + if token, ok := value.(string); ok && token != "" { + return token, nil } return newToken, nil } +func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) { + input := CachedAccessTokenRequest{ + ChannelID: info.ChannelId, + ChannelIsMultiKey: info.ChannelIsMultiKey, + ChannelMultiKeyIndex: info.ChannelMultiKeyIndex, + Credentials: a.AccountCredentials, + Proxy: info.ChannelSetting.Proxy, + } + return acquireCachedAccessToken(input, func(signedJWT string) (string, error) { + return exchangeJwtForAccessToken(signedJWT, info) + }) +} + func createSignedJWT(email, privateKeyPEM string) (string, error) { privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "-----BEGIN PRIVATE KEY-----", "") @@ -123,7 +149,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 { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } @@ -166,7 +192,7 @@ func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, defer resp.Body.Close() var result map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } diff --git a/relay/channel/vertex/storage_proxy.go b/relay/channel/vertex/storage_proxy.go new file mode 100644 index 000000000000..605cb0240b26 --- /dev/null +++ b/relay/channel/vertex/storage_proxy.go @@ -0,0 +1,207 @@ +package vertex + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/service" +) + +const vertexStorageHost = "storage.googleapis.com" + +type StorageOperation uint8 + +const ( + StorageOperationUpload StorageOperation = iota + 1 + StorageOperationList + StorageOperationGet + StorageOperationDelete +) + +type StorageProxyRequest struct { + Operation StorageOperation + Method string + Bucket string + Object string + RawQuery string + Header http.Header + Body io.Reader + ContentLength int64 + AccessToken string + Proxy string +} + +func buildStorageRequest(ctx context.Context, input StorageProxyRequest) (*http.Request, error) { + if strings.TrimSpace(input.AccessToken) == "" { + return nil, errors.New("Vertex storage access token is required") + } + bucket, err := relayconstant.NormalizeVertexStorageBucket(input.Bucket) + if err != nil { + return nil, err + } + + escapedBucket := url.PathEscape(bucket) + var escapedPath string + switch input.Operation { + case StorageOperationUpload: + if input.Method != http.MethodPost && input.Method != http.MethodPut { + return nil, fmt.Errorf("unsupported Vertex storage upload method %q", input.Method) + } + escapedPath = "/upload/storage/v1/b/" + escapedBucket + "/o" + case StorageOperationList: + if input.Method != http.MethodGet { + return nil, fmt.Errorf("unsupported Vertex storage list method %q", input.Method) + } + escapedPath = "/storage/v1/b/" + escapedBucket + "/o" + case StorageOperationGet, StorageOperationDelete: + expectedMethod := http.MethodGet + if input.Operation == StorageOperationDelete { + expectedMethod = http.MethodDelete + } + if input.Method != expectedMethod { + return nil, fmt.Errorf("unsupported Vertex storage object method %q", input.Method) + } + if err := relayconstant.ValidateVertexStorageObjectName(input.Object); err != nil { + return nil, err + } + escapedPath = "/storage/v1/b/" + escapedBucket + "/o/" + url.PathEscape(input.Object) + default: + return nil, errors.New("unsupported Vertex storage operation") + } + + path, err := url.PathUnescape(escapedPath) + if err != nil { + return nil, fmt.Errorf("build Vertex storage path: %w", err) + } + target := &url.URL{ + Scheme: "https", + Host: vertexStorageHost, + Path: path, + RawPath: escapedPath, + RawQuery: input.RawQuery, + } + request, err := http.NewRequestWithContext(ctx, input.Method, target.String(), input.Body) + if err != nil { + return nil, err + } + request.Header = sanitizeStorageRequestHeader(input.Header) + request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(input.AccessToken)) + if input.ContentLength >= 0 { + request.ContentLength = input.ContentLength + } + return request, nil +} + +func sanitizeStorageRequestHeader(source http.Header) http.Header { + return sanitizeStorageHeader(source, true) +} + +func SanitizeStorageResponseHeader(source http.Header) http.Header { + result := sanitizeStorageHeader(source, false) + result.Set("Cache-Control", "private, no-store") + result.Del("Expires") + result.Del("Age") + return result +} + +func sanitizeStorageHeader(source http.Header, request bool) http.Header { + result := source.Clone() + if result == nil { + result = make(http.Header) + } + remove := map[string]struct{}{ + "connection": {}, + "keep-alive": {}, + "proxy-authenticate": {}, + "proxy-authorization": {}, + "proxy-connection": {}, + "te": {}, + "trailer": {}, + "transfer-encoding": {}, + "upgrade": {}, + } + if request { + remove["authorization"] = struct{}{} + remove["host"] = struct{}{} + remove["content-length"] = struct{}{} + remove["cookie"] = struct{}{} + remove["x-api-key"] = struct{}{} + remove["x-goog-api-key"] = struct{}{} + } + for name, values := range source { + if !strings.EqualFold(name, "Connection") { + continue + } + for _, value := range values { + for _, nominated := range strings.Split(value, ",") { + if nominated = strings.ToLower(strings.TrimSpace(nominated)); nominated != "" { + remove[nominated] = struct{}{} + } + } + } + } + for name := range result { + if _, ok := remove[strings.ToLower(name)]; ok { + delete(result, name) + } + } + return result +} + +func RewriteStorageResumableLocation(location, gatewayBaseURL, bucket string) (string, error) { + normalizedBucket, err := relayconstant.NormalizeVertexStorageBucket(bucket) + if err != nil { + return "", err + } + upstream, err := url.Parse(location) + if err != nil || upstream.Scheme != "https" || upstream.Host != vertexStorageHost || upstream.User != nil || upstream.Fragment != "" { + return "", errors.New("invalid Vertex storage resumable location") + } + expectedPath := "/upload/storage/v1/b/" + url.PathEscape(normalizedBucket) + "/o" + if upstream.EscapedPath() != expectedPath { + return "", errors.New("Vertex storage resumable location bucket does not match") + } + + gateway, err := url.Parse(strings.TrimSpace(gatewayBaseURL)) + if err != nil || (gateway.Scheme != "http" && gateway.Scheme != "https") || gateway.Host == "" || gateway.User != nil { + return "", errors.New("invalid gateway base URL") + } + escapedGatewayPath := relayconstant.VertexStorageRoutePrefix + expectedPath + gatewayPath, err := url.PathUnescape(escapedGatewayPath) + if err != nil { + return "", fmt.Errorf("build Vertex storage resumable gateway path: %w", err) + } + gateway.Path = gatewayPath + gateway.RawPath = escapedGatewayPath + gateway.RawQuery = upstream.RawQuery + gateway.ForceQuery = upstream.ForceQuery + gateway.Fragment = "" + return gateway.String(), nil +} + +func DoStorageProxy(ctx context.Context, input StorageProxyRequest) (*http.Response, error) { + request, err := buildStorageRequest(ctx, input) + if err != nil { + return nil, err + } + client, err := service.GetHttpClientWithProxy(input.Proxy) + if err != nil { + return nil, fmt.Errorf("new proxy http client failed: %w", err) + } + storageClient := *client + storageClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + } + response, err := storageClient.Do(request) + if err != nil { + return nil, err + } + response.Header = SanitizeStorageResponseHeader(response.Header) + return response, nil +} diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go index 5f4b3be5b96f..6c46c8ec0626 100644 --- a/relay/constant/relay_mode.go +++ b/relay/constant/relay_mode.go @@ -54,6 +54,7 @@ const ( RelayModeResponsesCompact RelayModeAlphaSearch + RelayModeVertexStorage ) func Path2RelayMode(path string) int { diff --git a/relay/constant/vertex_storage.go b/relay/constant/vertex_storage.go new file mode 100644 index 000000000000..7193f3222b5d --- /dev/null +++ b/relay/constant/vertex_storage.go @@ -0,0 +1,63 @@ +package constant + +import ( + "errors" + "strings" +) + +const ( + VertexStorageModelPrefix = "storage:gs:" + VertexStorageRoutePrefix = "/vertexai" + VertexStorageUploadRoute = VertexStorageRoutePrefix + "/upload/storage/v1/b/:bucket/o" + VertexStorageListRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o" + VertexStorageObjectRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o/*object" +) + +func NormalizeVertexStorageBucket(raw string) (string, error) { + bucket := strings.TrimSpace(raw) + if bucket == "" || bucket == "." || bucket == ".." { + return "", errors.New("invalid Vertex storage bucket") + } + if strings.HasPrefix(bucket, VertexStorageModelPrefix) || + strings.Contains(bucket, "://") || strings.ContainsAny(bucket, `/\?#`) { + return "", errors.New("invalid Vertex storage bucket") + } + return bucket, nil +} + +func VertexStorageModelName(raw string) (string, error) { + bucket, err := NormalizeVertexStorageBucket(raw) + if err != nil { + return "", err + } + return VertexStorageModelPrefix + bucket, nil +} + +func VertexStorageChannelSupports(models []string, rawBucket string) bool { + modelName, err := VertexStorageModelName(rawBucket) + if err != nil { + return false + } + for _, offered := range models { + if strings.TrimSpace(offered) == modelName { + return true + } + } + return false +} + +func ValidateVertexStorageObjectName(object string) error { + if object == "" { + return errors.New("Vertex storage object is required") + } + for _, segment := range strings.Split(object, "/") { + if segment == "." || segment == ".." { + return errors.New("invalid Vertex storage object") + } + } + return nil +} + +func IsVertexStoragePath(path string) bool { + return strings.HasPrefix(path, VertexStorageRoutePrefix+"/") +} diff --git a/router/relay-router.go b/router/relay-router.go index b230a5a8084c..8194a19678f3 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -1,10 +1,13 @@ package router import ( + "strings" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/controller" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/relay" + relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relaykit/types" "github.com/gin-gonic/gin" @@ -79,6 +82,23 @@ func SetRelayRouter(router *gin.Engine) { controller.Relay(c, types.RelayFormatOpenAIRealtime) }) } + + vertexStorageRouter := router.Group(relayconstant.VertexStorageRoutePrefix) + vertexStorageRouter.Use(middleware.RouteTag("relay")) + vertexStorageRouter.Use(middleware.SystemPerformanceCheck()) + vertexStorageRouter.Use(middleware.TokenAuth()) + vertexStorageRouter.Use(middleware.ModelRequestRateLimit()) + vertexStorageRouter.Use(middleware.DistributeByChannelType(constant.ChannelTypeVertexAi)) + { + uploadPath := strings.TrimPrefix(relayconstant.VertexStorageUploadRoute, relayconstant.VertexStorageRoutePrefix) + listPath := strings.TrimPrefix(relayconstant.VertexStorageListRoute, relayconstant.VertexStorageRoutePrefix) + objectPath := strings.TrimPrefix(relayconstant.VertexStorageObjectRoute, relayconstant.VertexStorageRoutePrefix) + vertexStorageRouter.POST(uploadPath, controller.RelayVertexStorageUpload) + vertexStorageRouter.PUT(uploadPath, controller.RelayVertexStorageUpload) + vertexStorageRouter.GET(listPath, controller.RelayVertexStorageList) + vertexStorageRouter.GET(objectPath, controller.RelayVertexStorageObject) + vertexStorageRouter.DELETE(objectPath, controller.RelayVertexStorageObject) + } { //http router httpRouter := relayV1Router.Group("") diff --git a/router/relay_router_test.go b/router/relay_router_test.go index 579bd7bffa25..f81168a4de6c 100644 --- a/router/relay_router_test.go +++ b/router/relay_router_test.go @@ -124,3 +124,126 @@ func setupRelayRouterTestDB(t *testing.T) { } }) } + +func TestVertexStorageRoutesAreExact(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + SetRelayRouter(engine) + + want := map[string]bool{ + "POST /vertexai/upload/storage/v1/b/:bucket/o": true, + "PUT /vertexai/upload/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o/*object": true, + "DELETE /vertexai/storage/v1/b/:bucket/o/*object": true, + } + got := make(map[string]bool, len(want)) + for _, route := range engine.Routes() { + if strings.HasPrefix(route.Path, "/vertexai/") { + got[route.Method+" "+route.Path] = true + } + } + + assert.Equal(t, want, got) +} + +func TestVertexStorageOpenAPIContract(t *testing.T) { + openAPIBytes, err := os.ReadFile("../docs/openapi/relay.json") + require.NoError(t, err) + var document struct { + Tags []map[string]any `json:"tags"` + Paths map[string]map[string]any `json:"paths"` + } + require.NoError(t, common.Unmarshal(openAPIBytes, &document)) + + const tag = "文件/Vertex AI Cloud Storage" + tagFound := false + for _, item := range document.Tags { + if item["name"] == tag { + tagFound = true + break + } + } + assert.True(t, tagFound) + + expected := map[string][]string{ + "/vertexai/upload/storage/v1/b/{bucket}/o": {"post", "put"}, + "/vertexai/storage/v1/b/{bucket}/o": {"get"}, + "/vertexai/storage/v1/b/{bucket}/o/{object}": {"get", "delete"}, + } + vertexPathCount := 0 + for path, pathItem := range document.Paths { + if !strings.HasPrefix(path, "/vertexai/") { + continue + } + vertexPathCount++ + methods, ok := expected[path] + require.True(t, ok, "unexpected Vertex Storage path %s", path) + for _, method := range methods { + operation, ok := pathItem[method].(map[string]any) + require.True(t, ok, "%s %s", method, path) + assert.Equal(t, []any{tag}, operation["tags"]) + security, ok := operation["security"].([]any) + require.True(t, ok) + require.NotEmpty(t, security) + securityItem, ok := security[0].(map[string]any) + require.True(t, ok) + assert.Contains(t, securityItem, "BearerAuth") + + parameters, ok := operation["parameters"].([]any) + require.True(t, ok) + assertOpenAPIRequiredPathParameter(t, parameters, "bucket") + if strings.Contains(path, "{object}") { + assertOpenAPIRequiredPathParameter(t, parameters, "object") + } + if strings.Contains(path, "/upload/") { + assertOpenAPIParameter(t, parameters, "uploadType", "query") + assertOpenAPIParameter(t, parameters, "name", "query") + requestBody, ok := operation["requestBody"].(map[string]any) + require.True(t, ok) + content, ok := requestBody["content"].(map[string]any) + require.True(t, ok) + assert.Contains(t, content, "application/octet-stream") + if method == "put" { + assertOpenAPIParameter(t, parameters, "Content-Range", "header") + } + } + if strings.Contains(path, "{object}") && method == "get" { + assertOpenAPIParameter(t, parameters, "alt", "query") + assertOpenAPIParameter(t, parameters, "generation", "query") + assertOpenAPIParameter(t, parameters, "Range", "header") + } + responses, ok := operation["responses"].(map[string]any) + require.True(t, ok) + for _, status := range []string{"400", "403", "502", "default"} { + assert.Contains(t, responses, status) + } + } + } + assert.Equal(t, len(expected), vertexPathCount) +} + +func assertOpenAPIRequiredPathParameter(t *testing.T, parameters []any, name string) { + t.Helper() + for _, raw := range parameters { + parameter, ok := raw.(map[string]any) + require.True(t, ok) + if parameter["name"] == name && parameter["in"] == "path" { + assert.Equal(t, true, parameter["required"]) + return + } + } + assert.Fail(t, "missing required path parameter", name) +} + +func assertOpenAPIParameter(t *testing.T, parameters []any, name, location string) { + t.Helper() + for _, raw := range parameters { + parameter, ok := raw.(map[string]any) + require.True(t, ok) + if parameter["name"] == name && parameter["in"] == location { + return + } + } + assert.Fail(t, "missing OpenAPI parameter", name+" in "+location) +} diff --git a/service/channel_select.go b/service/channel_select.go index 0ab88dc84ff2..f6ed75c3b95c 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -11,12 +11,13 @@ import ( ) type RetryParam struct { - Ctx *gin.Context - TokenGroup string - ModelName string - RequestPath string - Retry *int - resetNextTry bool + Ctx *gin.Context + TokenGroup string + ModelName string + RequestPath string + Retry *int + RequiredChannelType int + resetNextTry bool } func (p *RetryParam) GetRetry() int { @@ -115,7 +116,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) + channel, _ = model.GetRandomSatisfiedChannelByType(autoGroup, param.ModelName, priorityRetry, param.RequestPath, param.RequiredChannelType) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -153,7 +154,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) + channel, err = model.GetRandomSatisfiedChannelByType(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, param.RequiredChannelType) if err != nil { return nil, param.TokenGroup, err } diff --git a/web/src/features/channels/components/data-table-row-actions.tsx b/web/src/features/channels/components/data-table-row-actions.tsx index 7e15463cc97d..6a744de68e36 100644 --- a/web/src/features/channels/components/data-table-row-actions.tsx +++ b/web/src/features/channels/components/data-table-row-actions.tsx @@ -61,9 +61,7 @@ import { useAuthStore } from '@/stores/auth-store' import { MODEL_FETCHABLE_TYPES } from '../constants' import { - channelsQueryKeys, handleDeleteChannel, - handleTestChannel, handleToggleChannelStatus, isChannelEnabled, isMultiKeyChannel, @@ -85,7 +83,6 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { const queryClient = useQueryClient() const currentUser = useAuthStore((s) => s.auth.user) const [deleteConfirmOpen, setDeleteConfirmOpen] = useState(false) - const [isTesting, setIsTesting] = useState(false) const [isTogglingStatus, setIsTogglingStatus] = useState(false) const isEnabled = isChannelEnabled(channel) @@ -106,18 +103,6 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { setOpen('test-channel') } - const handleDirectTest = async (e: React.MouseEvent) => { - e.stopPropagation() - setIsTesting(true) - try { - await handleTestChannel(channel.id, { channelName: channel.name }, () => { - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) - }) - } finally { - setIsTesting(false) - } - } - const handleQueryBalance = () => { setCurrentRow(channel) setOpen('balance-query') @@ -191,42 +176,19 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) {