mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
feeeei
daebb7a80c
Gemini 分组对外提供 AI Studio 兼容协议，分组内账号不一定是 AI Studio。
分组只剩 Vertex 服务账号时，GET /v1beta/models 与 /v1beta/models/{model}
选中服务账号后报 502 unsupported account type: service_account。
- SelectAccountForAIStudioEndpoints 改名 SelectAccountForModelsEndpoints，
Vertex 账号排在所有 AI Studio 兼容账号之后参与选号
- ForwardAIStudioGET 拆为 ForwardListModels / ForwardGetModel，按账号类型分流：
AI Studio 兼容账号原样转发；Vertex 服务账号读取 Google 发布方模型列表并
转换为 AI Studio 格式，模型详情在该模型所用区域的列表中查找；Vertex API Key
账号返回内置列表
- Vertex 发布方模型分页读取抽成 listVertexGeminiPublisherModels，与后台
同步上游模型共用
- 模型接口响应体上限改用 gateway.models_list_read_max_bytes，超出时报错
84 lines
3.2 KiB
Go
84 lines
3.2 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type geminiVertexModelsTokenCache struct{}
|
|
|
|
func (geminiVertexModelsTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
|
return "ya29.vertex-token", nil
|
|
}
|
|
|
|
func (geminiVertexModelsTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
|
return nil
|
|
}
|
|
|
|
func (geminiVertexModelsTokenCache) DeleteAccessToken(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
func (geminiVertexModelsTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
|
return false, nil
|
|
}
|
|
|
|
func (geminiVertexModelsTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
// 分组里只有 Agent Platform (Vertex) 服务账号时，模型列表与模型详情接口返回 Vertex 上的 Gemini 模型。
|
|
func TestGeminiV1BetaModelsEndpoints_VertexServiceAccountOnlyGroup(t *testing.T) {
|
|
groupID := int64(8)
|
|
repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {{
|
|
ID: 18,
|
|
Platform: service.PlatformGemini,
|
|
Type: service.AccountTypeServiceAccount,
|
|
Credentials: map[string]any{
|
|
"location": "global",
|
|
"service_account_json": map[string]any{
|
|
"type": "service_account",
|
|
"project_id": "proj",
|
|
"private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n",
|
|
"client_email": "svc@proj.iam.gserviceaccount.com",
|
|
},
|
|
},
|
|
}}}}}
|
|
upstream := &geminiMixedModelsUpstream{status: http.StatusOK, body: `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}]}`}
|
|
tokenProvider := service.NewGeminiTokenProvider(nil, geminiVertexModelsTokenCache{}, nil)
|
|
h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, tokenProvider, nil, upstream, nil, &config.Config{})}
|
|
apiKey := &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini}}
|
|
|
|
t.Run("list", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
|
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
|
|
|
h.GeminiV1BetaListModels(c)
|
|
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
|
require.JSONEq(t, `{"models":[{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}]}`, rec.Body.String())
|
|
})
|
|
|
|
t.Run("get", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil)
|
|
c.Params = gin.Params{{Key: "model", Value: "gemini-3.8-flash"}}
|
|
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
|
|
|
h.GeminiV1BetaGetModel(c)
|
|
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
|
require.JSONEq(t, `{"name":"models/gemini-3.8-flash","supportedGenerationMethods":["generateContent","streamGenerateContent"]}`, rec.Body.String())
|
|
})
|
|
}
|