| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111 |
- package service
- import (
- "context"
- "strings"
- "github.com/Wei-Shaw/sub2api/internal/config"
- )
- // ModelAvailabilityDiagnosis describes whether the requested model can be
- // served by any persistently eligible account in the group (active with its
- // schedulable setting enabled), ignoring transient state such as rate limits,
- // overload, temporary unschedulability, and runtime blocks. Handlers use this
- // on the "no available accounts" error path to distinguish 404
- // model_not_found from 503 service_unavailable.
- type ModelAvailabilityDiagnosis struct {
- // HasAccountsInPool is true if the group has at least one persistently
- // eligible account on the queried platform (or, for Anthropic/Gemini, on
- // the platform plus mixed-scheduled Antigravity accounts).
- HasAccountsInPool bool
- // HasModelSupport is true if at least one account's model mapping admits
- // the requested model.
- HasModelSupport bool
- }
- // ModelAvailabilityDiagnoser is implemented by gateway services that can
- // report whether the requested model is configured to be served by any
- // account. Both *GatewayService and *OpenAIGatewayService implement this so
- // handlers in either package can share a single classifier.
- type ModelAvailabilityDiagnoser interface {
- DiagnoseModelAvailabilityForPlatform(
- ctx context.Context,
- groupID *int64,
- requestedModel string,
- platform string,
- ) ModelAvailabilityDiagnosis
- }
- // DiagnoseModelAvailabilityForPlatform inspects accounts enabled for scheduling
- // by persistent configuration and returns whether the requested model is
- // configured to be served by any of them. The dedicated repository query
- // bypasses scheduler snapshots and deliberately ignores transient rate-limit,
- // overload, temporary-unschedulable, expiry, quota, and runtime-block state.
- //
- // Safe to call on the error path: returns {true,true} on any internal failure
- // or when the inputs preclude meaningful diagnosis (empty model, etc.), so
- // callers stay on the 503 fallback branch.
- func (s *GatewayService) DiagnoseModelAvailabilityForPlatform(
- ctx context.Context,
- groupID *int64,
- requestedModel string,
- platform string,
- ) ModelAvailabilityDiagnosis {
- if s == nil {
- return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}
- }
- requestedModel = strings.TrimSpace(requestedModel)
- if requestedModel == "" {
- // No model specified — cannot decide model_not_found. Caller falls back to 503.
- return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}
- }
- if strings.TrimSpace(platform) == "" {
- // Without a platform we cannot scope the lookup; bail out to the
- // 503 branch rather than make an unscoped scan.
- return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}
- }
- if s.accountRepo == nil {
- return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}
- }
- useMixed := platform == PlatformAnthropic || platform == PlatformGemini
- platforms := []string{platform}
- if useMixed {
- platforms = append(platforms, PlatformAntigravity)
- }
- queryGroupID := groupID
- includeGrouped := false
- if useMixed {
- // Preserve the generic scheduler's scope rules: an explicit group wins
- // for mixed scheduling, while group-less simple mode scans all accounts.
- if groupID == nil && s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
- includeGrouped = true
- }
- } else if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
- queryGroupID = nil
- includeGrouped = true
- }
- accounts, err := s.accountRepo.ListModelAvailabilityCandidates(ctx, queryGroupID, platforms, includeGrouped)
- if err != nil {
- // Conservative fallback: pretend everything is fine so the caller
- // returns 503 (we don't want to flip to 404 just because a lookup
- // hiccup'd).
- return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}
- }
- diag := ModelAvailabilityDiagnosis{}
- for i := range accounts {
- if useMixed && accounts[i].Platform == PlatformAntigravity && !accounts[i].IsMixedSchedulingEnabled() {
- continue
- }
- diag.HasAccountsInPool = true
- if s.isModelSupportedByAccountWithContext(ctx, &accounts[i], requestedModel) {
- diag.HasModelSupport = true
- return diag
- }
- }
- return diag
- }
|