mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
fix(models): list composite exact route aliases
This commit is contained in:
@@ -1162,7 +1162,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
}
|
||||
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, "")
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
@@ -1257,7 +1257,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
platform = group.Platform
|
||||
}
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID)
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID, service.CompositeRouteEndpointResponses)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||
if group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
@@ -1283,7 +1283,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
return fallbackModels
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64) []string {
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, endpoint string) []string {
|
||||
if h == nil || h.gatewayService == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -1311,6 +1311,16 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
// A route can expose a public ID that no account model mapping contains.
|
||||
// On lookup failure, retain the existing account-derived catalog only.
|
||||
if routeModels, err := h.gatewayService.GetCompositeRouteModels(ctx, groupID, endpoint); err == nil {
|
||||
for _, model := range routeModels {
|
||||
if _, ok := seen[model]; !ok {
|
||||
seen[model] = struct{}{}
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -22,6 +23,16 @@ type gatewayModelsAccountRepoStub struct {
|
||||
byGroup map[int64][]service.Account
|
||||
}
|
||||
|
||||
type gatewayModelsRouteRepoStub struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *gatewayModelsRouteRepoStub) ListByGroup(_ context.Context, _ int64, _ bool) ([]service.CompositeModelRoute, error) {
|
||||
return s.routes, s.err
|
||||
}
|
||||
|
||||
type gatewayModelsResponseForTest struct {
|
||||
Object string `json:"object"`
|
||||
Data []gatewayModelItemForTest `json:"data"`
|
||||
@@ -81,16 +92,79 @@ func (s *gatewayModelsAccountRepoStub) ListModelAvailabilityCandidates(ctx conte
|
||||
return s.ListSchedulableByGroupID(ctx, *groupID)
|
||||
}
|
||||
|
||||
func newGatewayModelsHandlerForTest(repo service.AccountRepository) *GatewayHandler {
|
||||
func newGatewayModelsHandlerForTest(repo service.AccountRepository, routeRepo ...service.CompositeModelRouteRepository) *GatewayHandler {
|
||||
var resolver *service.CompositeRouteResolver
|
||||
if len(routeRepo) > 0 {
|
||||
resolver = service.NewCompositeRouteResolver(routeRepo[0])
|
||||
}
|
||||
return &GatewayHandler{
|
||||
gatewayService: service.NewGatewayService(
|
||||
repo,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, resolver, nil, nil,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayCompositeRouteModelsAppearInBothCatalogs(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 7784
|
||||
accounts := &gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
|
||||
groupID: {{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{"mapped-model": "upstream"},
|
||||
}}},
|
||||
}}
|
||||
routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{
|
||||
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "mapped-model", MatchType: service.CompositeRouteMatchExact, Enabled: true},
|
||||
{PublicModel: "image-only", MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true},
|
||||
{PublicModel: "prefix-", MatchType: service.CompositeRouteMatchPrefix, Enabled: true},
|
||||
{PublicModel: "disabled", MatchType: service.CompositeRouteMatchExact, Enabled: false},
|
||||
}}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
repo service.CompositeModelRouteRepository
|
||||
allowlist []string
|
||||
want []string
|
||||
}{
|
||||
{"routes", routes, nil, []string{"mapped-model", "route-only", "image-only"}},
|
||||
{"allowlist", routes, []string{"route-only", "disabled", "unknown"}, []string{"route-only"}},
|
||||
{"lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, nil, []string{"mapped-model"}},
|
||||
{"no route repo", nil, nil, []string{"mapped-model"}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
h := newGatewayModelsHandlerForTest(accounts, tc.repo)
|
||||
group := &service.Group{ID: groupID, Platform: service.PlatformComposite}
|
||||
if tc.allowlist != nil {
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: tc.allowlist}
|
||||
}
|
||||
for _, endpoint := range []string{"models", "codex"} {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group})
|
||||
if endpoint == "models" {
|
||||
h.Models(c)
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, tc.want, modelIDsForTest(got.Data))
|
||||
} else {
|
||||
h.CodexModels(c)
|
||||
var got codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
want := tc.want
|
||||
if tc.name == "routes" {
|
||||
want = []string{"mapped-model", "route-only"}
|
||||
}
|
||||
require.ElementsMatch(t, want, codexModelSlugsForTest(got.Models))
|
||||
}
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) {
|
||||
antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity)
|
||||
require.NotEmpty(t, antigravityIDs)
|
||||
|
||||
@@ -22,6 +22,27 @@ func (r *CompositeRouteResolver) SetModelOwnershipResolver(resolver CompositeMod
|
||||
}
|
||||
}
|
||||
|
||||
// ListExactPublicModels returns enabled, concrete route IDs suitable for a model catalog.
|
||||
func (r *CompositeRouteResolver) ListExactPublicModels(ctx context.Context, groupID int64, endpoint string) ([]string, error) {
|
||||
if r == nil || r.repo == nil || groupID <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
routes, err := r.repo.ListByGroup(ctx, groupID, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
models := make([]string, 0, len(routes))
|
||||
for _, route := range routes {
|
||||
if route.Enabled && route.MatchType == CompositeRouteMatchExact &&
|
||||
(endpoint == "" || normalizeCompositeRouteEndpoint(route.Endpoint) == CompositeRouteEndpointAny || normalizeCompositeRouteEndpoint(route.Endpoint) == endpoint) {
|
||||
if model := strings.TrimSpace(route.PublicModel); model != "" {
|
||||
models = append(models, model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (r *CompositeRouteResolver) Resolve(ctx context.Context, groupID int64, model, endpoint string) (CompositeRouteDecision, error) {
|
||||
model = strings.TrimSpace(model)
|
||||
endpoint = normalizeCompositeRouteEndpoint(endpoint)
|
||||
|
||||
@@ -1535,6 +1535,14 @@ func explicitModelMappingClaims(account Account, model string) bool {
|
||||
return ok && strings.TrimSpace(mapped) != ""
|
||||
}
|
||||
|
||||
// GetCompositeRouteModels returns public IDs from enabled exact composite routes.
|
||||
func (s *GatewayService) GetCompositeRouteModels(ctx context.Context, groupID *int64, endpoint string) ([]string, error) {
|
||||
if s == nil || s.compositeResolver == nil || groupID == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return s.compositeResolver.ListExactPublicModels(ctx, *groupID, endpoint)
|
||||
}
|
||||
|
||||
// GetSchedulablePlatforms returns the concrete platforms that currently have
|
||||
// schedulable accounts in the target group.
|
||||
func (s *GatewayService) GetSchedulablePlatforms(ctx context.Context, groupID *int64) map[string]struct{} {
|
||||
|
||||
Reference in New Issue
Block a user