mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-19 16:17:56 +00:00
feat(gateway): retrieve models from the visible catalogue
This commit is contained in:
@@ -1116,8 +1116,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// Models handles listing available models
|
||||
// GET /v1/models
|
||||
// Models lists visible models, or retrieves the exact list entry for a model path parameter.
|
||||
// GET /v1/models and /v1/models/:model (also exposed through root aliases)
|
||||
// Returns models based on account configurations (model_mapping whitelist)
|
||||
// Falls back to default models if no whitelist is configured
|
||||
func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
@@ -1173,18 +1173,12 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
|
||||
// Fallback to default models
|
||||
if platform == service.PlatformOpenAI {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": openai.DefaultModels,
|
||||
})
|
||||
writeModelsListResponse(c, openai.DefaultModels)
|
||||
return
|
||||
}
|
||||
|
||||
if platform == service.PlatformGemini {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": geminicli.DefaultModels,
|
||||
})
|
||||
writeModelsListResponse(c, geminicli.DefaultModels)
|
||||
return
|
||||
}
|
||||
if platform == service.PlatformGrok {
|
||||
@@ -1192,10 +1186,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": claude.DefaultModels,
|
||||
})
|
||||
writeModelsListResponse(c, claude.DefaultModels)
|
||||
}
|
||||
|
||||
// CodexModels returns the effective group model list using the manifest shape
|
||||
@@ -1320,10 +1311,7 @@ func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
CreatedAt: "2024-01-01T00:00:00Z",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})
|
||||
writeModelsListResponse(c, models)
|
||||
}
|
||||
|
||||
func writeAllowlistedModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
@@ -1382,10 +1370,7 @@ func writeGrokModelsList(c *gin.Context, modelIDs []string) {
|
||||
models = append(models, item)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})
|
||||
writeModelsListResponse(c, models)
|
||||
}
|
||||
|
||||
func grokModelSupportsConfigurableReasoning(modelID string) bool {
|
||||
@@ -1418,10 +1403,7 @@ func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
|
||||
DisplayName: modelID,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})
|
||||
writeModelsListResponse(c, models)
|
||||
}
|
||||
|
||||
// modelListingSource 汇总模型列表过滤的候选来源:账号映射键(availableModels)
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
middleware2 "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"
|
||||
)
|
||||
|
||||
func requestModelForTest(h *GatewayHandler, group *service.Group, modelID, etag string) *httptest.ResponseRecorder {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Request.Header.Set("If-None-Match", etag)
|
||||
if modelID != "" {
|
||||
c.Params = gin.Params{{Key: "model", Value: modelID}}
|
||||
}
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
|
||||
h.Models(c)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestRetrieveModelMatchesVisibleCatalogue(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformComposite} {
|
||||
for _, mapped := range []bool{false, true} {
|
||||
name := platform + "/fallback"
|
||||
if mapped {
|
||||
name = platform + "/mapped"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
accountPlatform := platform
|
||||
if platform == service.PlatformComposite {
|
||||
accountPlatform = service.PlatformOpenAI
|
||||
}
|
||||
credentials := map[string]any{}
|
||||
if mapped {
|
||||
credentials["model_mapping"] = map[string]any{"custom-model": "upstream-only-model"}
|
||||
}
|
||||
group := &service.Group{ID: 71, Platform: platform}
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
|
||||
group.ID: {{ID: 1, Platform: accountPlatform, Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Credentials: credentials}},
|
||||
}})
|
||||
list := requestModelForTest(h, group, "", "")
|
||||
require.Equal(t, http.StatusOK, list.Code, list.Body.String())
|
||||
var catalog struct {
|
||||
Data []json.RawMessage `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
|
||||
require.NotEmpty(t, catalog.Data)
|
||||
var model struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(catalog.Data[0], &model))
|
||||
retrieved := requestModelForTest(h, group, model.ID, "")
|
||||
require.Equal(t, http.StatusOK, retrieved.Code, retrieved.Body.String())
|
||||
require.JSONEq(t, string(catalog.Data[0]), retrieved.Body.String())
|
||||
require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "unknown-model", "").Code)
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{"absent-from-source"}}
|
||||
hidden := requestModelForTest(h, group, model.ID, "")
|
||||
require.Equal(t, http.StatusNotFound, hidden.Code, hidden.Body.String())
|
||||
require.Contains(t, hidden.Body.String(), `"code":"model_not_found"`)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrievePinnedModelPreservesMetadataFilteringAndErrors(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
status int
|
||||
selected []string
|
||||
wantStatus int
|
||||
}{
|
||||
{name: "metadata", wantStatus: http.StatusOK},
|
||||
{name: "hidden", selected: []string{"other-model"}, wantStatus: http.StatusNotFound},
|
||||
{name: "upstream failure", status: http.StatusServiceUnavailable, wantStatus: http.StatusBadGateway},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
|
||||
2: `{"data":[{"id":"special-model","ID":"not-a-model-id","owned_by":"source-owner","created":123,"extra":{"context":999}}]}`,
|
||||
}, statuses: map[int64]int{}}
|
||||
if tc.status != 0 {
|
||||
upstream.statuses[2] = tc.status
|
||||
}
|
||||
codex := newPinnedCodexTestHandler([]service.Account{newPinnedCodexAccount(2, service.StatusActive, true, false)}, upstream, 3)
|
||||
h := &GatewayHandler{openAIGatewayService: codex.gatewayService, maxAccountSwitches: 3}
|
||||
group := &service.Group{ID: 72, Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{Enabled: len(tc.selected) > 0, Models: tc.selected},
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2}}}
|
||||
got := requestModelForTest(h, group, "special-model", "collection-etag")
|
||||
require.Equal(t, tc.wantStatus, got.Code, got.Body.String())
|
||||
if tc.wantStatus == http.StatusOK {
|
||||
list := requestModelForTest(h, group, "", "")
|
||||
var catalog struct {
|
||||
Data []json.RawMessage `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
|
||||
require.JSONEq(t, string(catalog.Data[0]), got.Body.String())
|
||||
require.Contains(t, got.Body.String(), `"owned_by":"source-owner"`)
|
||||
require.Contains(t, got.Body.String(), `"created":123`)
|
||||
again := requestModelForTest(h, group, "special-model", list.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusOK, again.Code)
|
||||
require.Empty(t, again.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "not-a-model-id", "").Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,9 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
@@ -17,8 +19,12 @@ func (h *GatewayHandler) pinnedOpenAIModels(c *gin.Context, group *service.Group
|
||||
writeOpenAIModelsError(c, http.StatusInternalServerError, "api_error", "OpenAI model discovery is not configured")
|
||||
return
|
||||
}
|
||||
etag := c.GetHeader("If-None-Match")
|
||||
if c.Param("model") != "" {
|
||||
etag = "" // A collection ETag cannot validate a single-model representation.
|
||||
}
|
||||
response, account, err := h.openAIGatewayService.FetchPinnedOpenAIModelsList(
|
||||
c.Request.Context(), group, h.maxAccountSwitches, c.GetHeader("If-None-Match"),
|
||||
c.Request.Context(), group, h.maxAccountSwitches, etag,
|
||||
)
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
@@ -40,6 +46,10 @@ func writeOpenAIModelsError(c *gin.Context, status int, errorType, message strin
|
||||
}
|
||||
|
||||
func writeOpenAIModelsResponse(c *gin.Context, manifest *service.OpenAIModelsResponse) {
|
||||
if c.Param("model") != "" {
|
||||
writeRetrievedModel(c, manifest.Body)
|
||||
return
|
||||
}
|
||||
if manifest.ETag != "" {
|
||||
c.Header("ETag", manifest.ETag)
|
||||
}
|
||||
@@ -50,3 +60,50 @@ func writeOpenAIModelsResponse(c *gin.Context, manifest *service.OpenAIModelsRes
|
||||
}
|
||||
c.Data(http.StatusOK, "application/json", manifest.Body)
|
||||
}
|
||||
|
||||
// Both discovery endpoints consume the same final catalogue, after group/platform
|
||||
// selection and allowlist filtering. Preserve every field on the selected entry.
|
||||
func writeModelsListResponse(c *gin.Context, models any) {
|
||||
response := gin.H{"object": "list", "data": models}
|
||||
if c.Param("model") == "" {
|
||||
c.JSON(http.StatusOK, response)
|
||||
return
|
||||
}
|
||||
body, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
writeOpenAIModelsError(c, http.StatusInternalServerError, "api_error", "Failed to encode model catalogue")
|
||||
return
|
||||
}
|
||||
writeRetrievedModel(c, body)
|
||||
}
|
||||
|
||||
func writeRetrievedModel(c *gin.Context, body []byte) {
|
||||
var catalog struct {
|
||||
Data []json.RawMessage `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &catalog); err != nil {
|
||||
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue")
|
||||
return
|
||||
}
|
||||
modelID := c.Param("model")
|
||||
for _, raw := range catalog.Data {
|
||||
var model map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &model); err != nil {
|
||||
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue entry")
|
||||
return
|
||||
}
|
||||
var id string
|
||||
if err := json.Unmarshal(model["id"], &id); err != nil {
|
||||
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue ID")
|
||||
return
|
||||
}
|
||||
if id == modelID {
|
||||
c.Data(http.StatusOK, "application/json", raw)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{
|
||||
"type": "invalid_request_error", "code": "model_not_found", "param": "model",
|
||||
"message": fmt.Sprintf("Model %q does not exist or is not available for this group", modelID),
|
||||
}})
|
||||
}
|
||||
|
||||
@@ -208,6 +208,8 @@ func RegisterGatewayRoutes(
|
||||
// /models endpoint with a client_version query and expect the ChatGPT
|
||||
// Codex manifest format; other clients keep the OpenAI-style list.
|
||||
gateway.GET("/models", modelsHandler)
|
||||
// Single-model discovery never selects the Codex client_version manifest.
|
||||
gateway.GET("/models/:model", h.Gateway.Models)
|
||||
gateway.GET("/usage", h.Gateway.Usage)
|
||||
gateway.POST("/live", h.OpenAIGateway.Live)
|
||||
gateway.GET("/live/:call_id", h.OpenAIGateway.LiveSideband)
|
||||
@@ -372,6 +374,7 @@ func RegisterGatewayRoutes(
|
||||
h.OpenAIGateway.ResponsesWebSocket(c)
|
||||
})
|
||||
rootRoute(http.MethodGet, "/models", bodyLimit, modelsHandler)
|
||||
rootRoute(http.MethodGet, "/models/:model", bodyLimit, h.Gateway.Models)
|
||||
rootRoute(http.MethodPost, "/messages/count_tokens", bodyLimit, countTokensHandler)
|
||||
codexDirect := r.Group("/backend-api/codex")
|
||||
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)
|
||||
|
||||
@@ -92,3 +92,59 @@ func TestGatewayRoutesPinnedModelsDispatchesOrdinaryAndCodexRequests(t *testing.
|
||||
require.EqualValues(t, 1, upstream.ordinaryCalls.Load())
|
||||
require.EqualValues(t, 1, upstream.codexCalls.Load())
|
||||
}
|
||||
|
||||
func TestGatewayRoutesRetrievePinnedModel(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repo := &pinnedModelsRoutesRepository{account: service.Account{
|
||||
ID: 7, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true,
|
||||
Credentials: map[string]any{"api_key": "test-models-key", "base_url": "https://models.example/v1"},
|
||||
}}
|
||||
upstream := &pinnedModelsRoutesUpstream{}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
s := service.NewOpenAIGatewayService(repo, nil, nil, nil, nil, nil, nil, cfg,
|
||||
nil, nil, nil, nil, nil, upstream, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h := &handler.Handlers{
|
||||
Gateway: handler.NewGatewayHandler(nil, s, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil),
|
||||
OpenAIGateway: handler.NewOpenAIGatewayHandler(s, nil, nil, nil, nil, nil, nil, nil, cfg),
|
||||
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
|
||||
}
|
||||
group := &service.Group{ID: 1, Platform: service.PlatformOpenAI,
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{7}}}
|
||||
router := gin.New()
|
||||
RegisterGatewayRoutes(router, h, servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
|
||||
if c.GetHeader("Authorization") != "Bearer client-key" {
|
||||
c.AbortWithStatus(http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
|
||||
c.Next()
|
||||
}), nil, nil, nil, nil, nil, cfg)
|
||||
request := func(path, key, etag string) *httptest.ResponseRecorder {
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
req.Header.Set("Authorization", key)
|
||||
req.Header.Set("If-None-Match", etag)
|
||||
router.ServeHTTP(w, req)
|
||||
return w
|
||||
}
|
||||
list := request("/v1/models", "Bearer client-key", "")
|
||||
require.Equal(t, http.StatusOK, list.Code)
|
||||
var catalog struct {
|
||||
Data []json.RawMessage `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
|
||||
require.Len(t, catalog.Data, 1)
|
||||
for _, base := range []string{"/v1/models/", "/models/"} {
|
||||
path := base + "ordinary-upstream-model?client_version=ignored"
|
||||
require.Equal(t, http.StatusUnauthorized, request(path, "", "").Code)
|
||||
got := request(path, "Bearer client-key", list.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusOK, got.Code, got.Body.String())
|
||||
require.JSONEq(t, string(catalog.Data[0]), got.Body.String())
|
||||
require.Empty(t, got.Header().Get("ETag"), "a collection validator cannot validate one model")
|
||||
missing := request(base+"unknown-model", "Bearer client-key", "")
|
||||
require.Equal(t, http.StatusNotFound, missing.Code)
|
||||
require.Contains(t, missing.Body.String(), `"error"`)
|
||||
}
|
||||
require.Zero(t, upstream.codexCalls.Load(), "retrieve never dispatches a Codex manifest")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user