diff --git a/apps/api/internal/runner/candidate_platform_filter.go b/apps/api/internal/runner/candidate_platform_filter.go new file mode 100644 index 0000000..a67e73a --- /dev/null +++ b/apps/api/internal/runner/candidate_platform_filter.go @@ -0,0 +1,49 @@ +package runner + +import ( + "fmt" + "strings" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +func filterCandidatesByRequestedPlatform( + candidates []store.RuntimeModelCandidate, + body map[string]any, +) ([]store.RuntimeModelCandidate, error) { + platformID := firstNonEmptyString( + stringFromAny(body["platform_id"]), + stringFromAny(body["platformId"]), + ) + platformModelID := firstNonEmptyString( + stringFromAny(body["platform_model_id"]), + stringFromAny(body["platformModelId"]), + ) + if platformID == "" && platformModelID == "" { + return candidates, nil + } + + filtered := make([]store.RuntimeModelCandidate, 0, len(candidates)) + for _, candidate := range candidates { + if platformID != "" && !strings.EqualFold(candidate.PlatformID, platformID) { + continue + } + if platformModelID != "" && !strings.EqualFold(candidate.PlatformModelID, platformModelID) { + continue + } + filtered = append(filtered, candidate) + } + if len(filtered) > 0 { + return filtered, nil + } + + requested := platformID + if platformModelID != "" { + requested = platformModelID + } + return nil, fmt.Errorf( + "%w: requested platform %q is unavailable for this model", + store.ErrNoModelCandidate, + requested, + ) +} diff --git a/apps/api/internal/runner/candidate_platform_filter_test.go b/apps/api/internal/runner/candidate_platform_filter_test.go new file mode 100644 index 0000000..2c2405f --- /dev/null +++ b/apps/api/internal/runner/candidate_platform_filter_test.go @@ -0,0 +1,49 @@ +package runner + +import ( + "errors" + "testing" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +func TestFilterCandidatesByRequestedPlatform(t *testing.T) { + candidates := []store.RuntimeModelCandidate{ + {PlatformID: "platform-a", PlatformModelID: "model-a"}, + {PlatformID: "platform-b", PlatformModelID: "model-b"}, + } + + t.Run("without selection", func(t *testing.T) { + filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{}) + if err != nil || len(filtered) != len(candidates) { + t.Fatalf("unexpected unfiltered candidates: %+v err=%v", filtered, err) + } + }) + + t.Run("snake case platform id", func(t *testing.T) { + filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{ + "platform_id": "platform-b", + }) + if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "platform-b" { + t.Fatalf("unexpected platform filter result: %+v err=%v", filtered, err) + } + }) + + t.Run("camel case platform model id", func(t *testing.T) { + filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{ + "platformModelId": "model-a", + }) + if err != nil || len(filtered) != 1 || filtered[0].PlatformModelID != "model-a" { + t.Fatalf("unexpected platform model filter result: %+v err=%v", filtered, err) + } + }) + + t.Run("selection must be accessible and enabled", func(t *testing.T) { + filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{ + "platformId": "platform-missing", + }) + if len(filtered) != 0 || !errors.Is(err, store.ErrNoModelCandidate) { + t.Fatalf("unexpected unavailable platform result: %+v err=%v", filtered, err) + } + }) +} diff --git a/apps/api/internal/runner/pricing.go b/apps/api/internal/runner/pricing.go index 724acec..6bb6ef3 100644 --- a/apps/api/internal/runner/pricing.go +++ b/apps/api/internal/runner/pricing.go @@ -31,6 +31,10 @@ func (s *Service) Estimate(ctx context.Context, kind string, model string, body if err != nil { return EstimateResult{}, err } + candidates, err = filterCandidatesByRequestedPlatform(candidates, body) + if err != nil { + return EstimateResult{}, err + } candidates, _, err = filterRuntimeCandidatesByRequest(kind, model, modelType, body, candidates) if err != nil { return EstimateResult{}, err diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 0c1e3bb..3eb931d 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -222,6 +222,9 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, CacheAffinityKeys: cacheAffinityKeys.Lookup, CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy, }) + if err == nil { + candidates, err = filterCandidatesByRequestedPlatform(candidates, body) + } if err != nil { if task.Kind == "responses" && responseExecution.PreviousChain != nil { err = responseChainUnavailableError()