mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-09 08:58:23 +00:00
fix(models): list composite exact route aliases
This commit is contained in:
1 parent
d6adebd22d
commit
1888de135d
4 files changed
+118
-5
No files matched your search
@@ -1162,7 +1162,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if platform == service.PlatformComposite {
|
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() {
|
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||||
source := availableModels
|
source := availableModels
|
||||||
if len(source) == 0 {
|
if len(source) == 0 {
|
||||||
@@ -1257,7 +1257,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
|||||||
platform = group.Platform
|
platform = group.Platform
|
||||||
}
|
}
|
||||||
if platform == service.PlatformComposite {
|
if platform == service.PlatformComposite {
|
||||||
availableModels := h.compositeAvailableModels(ctx, groupID)
|
availableModels := h.compositeAvailableModels(ctx, groupID, service.CompositeRouteEndpointResponses)
|
||||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||||
if group.ModelAllowlistEnabled() {
|
if group.ModelAllowlistEnabled() {
|
||||||
source := availableModels
|
source := availableModels
|
||||||
@@ -1283,7 +1283,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
|||||||
return fallbackModels
|
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 {
|
if h == nil || h.gatewayService == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -1311,6 +1311,16 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
|
|||||||
models = append(models, model)
|
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
|
return models
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package handler
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -22,6 +23,16 @@ type gatewayModelsAccountRepoStub struct {
|
|||||||
byGroup map[int64][]service.Account
|
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 {
|
type gatewayModelsResponseForTest struct {
|
||||||
Object string `json:"object"`
|
Object string `json:"object"`
|
||||||
Data []gatewayModelItemForTest `json:"data"`
|
Data []gatewayModelItemForTest `json:"data"`
|
||||||
@@ -81,16 +92,79 @@ func (s *gatewayModelsAccountRepoStub) ListModelAvailabilityCandidates(ctx conte
|
|||||||
return s.ListSchedulableByGroupID(ctx, *groupID)
|
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{
|
return &GatewayHandler{
|
||||||
gatewayService: service.NewGatewayService(
|
gatewayService: service.NewGatewayService(
|
||||||
repo,
|
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, 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) {
|
func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) {
|
||||||
antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity)
|
antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity)
|
||||||
require.NotEmpty(t, antigravityIDs)
|
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) {
|
func (r *CompositeRouteResolver) Resolve(ctx context.Context, groupID int64, model, endpoint string) (CompositeRouteDecision, error) {
|
||||||
model = strings.TrimSpace(model)
|
model = strings.TrimSpace(model)
|
||||||
endpoint = normalizeCompositeRouteEndpoint(endpoint)
|
endpoint = normalizeCompositeRouteEndpoint(endpoint)
|
||||||
|
|||||||
@@ -1535,6 +1535,14 @@ func explicitModelMappingClaims(account Account, model string) bool {
|
|||||||
return ok && strings.TrimSpace(mapped) != ""
|
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
|
// GetSchedulablePlatforms returns the concrete platforms that currently have
|
||||||
// schedulable accounts in the target group.
|
// schedulable accounts in the target group.
|
||||||
func (s *GatewayService) GetSchedulablePlatforms(ctx context.Context, groupID *int64) map[string]struct{} {
|
func (s *GatewayService) GetSchedulablePlatforms(ctx context.Context, groupID *int64) map[string]struct{} {
|
||||||
|
|||||||
Reference in new issue
Block a user