feature(ampcode): Improves AMP model mapping with alias support#1314
feature(ampcode): Improves AMP model mapping with alias support#1314
Conversation
Summary of ChangesHello @soilSpoon, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the AMP model mapping capabilities by introducing a sophisticated fallback mechanism. It allows the system to intelligently switch to alternative models, defined as aliases, when the initially selected model encounters issues like quota limits. This improvement ensures greater resilience and reliability in model selection. Additionally, it includes a critical fix to prevent Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Pull request overview
This PR enhances the AMP model mapping functionality to support fallback mechanisms when primary models fail due to issues like quota exhaustion. It enables the system to attempt alternative models (aliases) by integrating OAuth model alias configuration into the model mapper.
Changes:
- Implements fallback logic in authentication manager to retry requests with alternative models when the primary fails
- Extends model mapper to support multiple fallback models via OAuth model aliases
- Adds context keys to pass fallback model information between handlers
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/cliproxy/auth/conductor.go | Refactors execution logic to support fallback model attempts through new executeWithFallback helper |
| internal/api/server.go | Updates config change detection to include OAuth model aliases |
| internal/api/modules/amp/model_mapping.go | Adds OAuth alias lookup and multi-model fallback support to model mapper |
| internal/api/modules/amp/fallback_handlers.go | Updates handler to store and use fallback models from context |
| internal/api/modules/amp/amp.go | Initializes OAuth model alias configuration in AMP module |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| tried := make(map[string]struct{}) | ||
| var lastErr error | ||
|
|
||
| // Track fallback models from context (provided by Amp module fallback_models key) |
There was a problem hiding this comment.
Corrected spelling of 'recieve' to 'receive' in variable name context.
|
|
||
| // Reset tried set for the new model and find its providers | ||
| tried = make(map[string]struct{}) | ||
| providers = util.GetProviderName(thinking.ParseSuffix(routeModel).ModelName) |
There was a problem hiding this comment.
Consider extracting the base model name into a variable before calling util.GetProviderName to improve readability and avoid repeating the thinking.ParseSuffix(routeModel).ModelName pattern.
| providers = util.GetProviderName(thinking.ParseSuffix(routeModel).ModelName) | |
| baseModel := thinking.ParseSuffix(routeModel).ModelName | |
| providers = util.GetProviderName(baseModel) |
There was a problem hiding this comment.
Code Review
This pull request introduces a significant enhancement to the AMP model mapping by adding support for fallback models via aliases. This allows the system to gracefully handle failures like quota exhaustion by trying alternative models. The changes span across model mapping logic, fallback handlers, and the authentication execution flow. I appreciate the refactoring in conductor.go, which greatly improves code clarity and reduces duplication. My review includes suggestions to further improve maintainability by refactoring some duplicated code, correcting a minor logging inaccuracy, and enhancing the way context keys are shared between packages to align with best practices.
| if mappedModels, mappedProviders := resolveMappedModels(); len(mappedModels) > 0 { | ||
| // Mapping found and provider available - rewrite the model in request body | ||
| bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel) | ||
| bodyBytes = rewriteModelInRequest(bodyBytes, mappedModels[0]) | ||
| c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes)) | ||
| // Store mapped model in context for handlers that check it (like gemini bridge) | ||
| c.Set(MappedModelContextKey, mappedModel) | ||
| resolvedModel = mappedModel | ||
| // Store mapped model and fallbacks in context for handlers | ||
| c.Set(MappedModelContextKey, mappedModels[0]) | ||
| if len(mappedModels) > 1 { | ||
| c.Set(FallbackModelsContextKey, mappedModels[1:]) | ||
| } | ||
| resolvedModel = mappedModels[0] | ||
| usedMapping = true | ||
| providers = mappedProviders | ||
| } |
There was a problem hiding this comment.
This block of code is identical to the one at lines 249-261. To improve maintainability and avoid code duplication, consider extracting this logic into a local helper function.
For example:
applyMapping := func(mappedModels []string, mappedProviders []string) {
bodyBytes = rewriteModelInRequest(bodyBytes, mappedModels[0])
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
// Store mapped model and fallbacks in context for handlers
c.Set(MappedModelContextKey, mappedModels[0])
if len(mappedModels) > 1 {
c.Set(FallbackModelsContextKey, mappedModels[1:])
}
resolvedModel = mappedModels[0]
usedMapping = true
providers = mappedProviders
}
if forceMappings {
if mappedModels, mappedProviders := resolveMappedModels(); len(mappedModels) > 0 {
applyMapping(mappedModels, mappedProviders)
}
// ...
} else {
// ...
if len(providers) == 0 {
if mappedModels, mappedProviders := resolveMappedModels(); len(mappedModels) > 0 {
applyMapping(mappedModels, mappedProviders)
}
}
}Note that bodyBytes, resolvedModel, usedMapping, and providers would need to be handled as they are modified by this helper. Using a closure that captures these variables would be a clean way to implement this.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 9 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Track fallback models from context (provided by Amp module fallback_models key) | ||
| var fallbacks []string | ||
| if v := ctx.Value("fallback_models"); v != nil { | ||
| if fs, ok := v.([]string); ok { | ||
| fallbacks = fs | ||
| } | ||
| } |
There was a problem hiding this comment.
Using string literals as context keys is error-prone and can lead to conflicts. Consider using a typed context key instead. The codebase already has internal/routing/ctxkeys/keys.go defining typed context keys like FallbackModels. Use that constant here instead of the string literal "fallback_models".
| } | ||
| }() | ||
|
|
||
| // Use requestCtx as base if available to preserve amp context values (fallback_models, etc.) |
There was a problem hiding this comment.
Corrected capitalization of 'Amp' to 'AMP' for consistency with product naming.
| // Use requestCtx as base if available to preserve amp context values (fallback_models, etc.) | |
| // Use requestCtx as base if available to preserve AMP context values (fallback_models, etc.) |
| func TestRouter_Resolve_ModelMappings(t *testing.T) { | ||
| registry := NewRegistry() | ||
|
|
||
| // Add a provider |
There was a problem hiding this comment.
Missing test coverage for error cases in router resolution, such as when mappings reference non-existent models or when all providers are unavailable. Consider adding negative test cases.
| // TODO: implement based on error type | ||
| // For now, all errors are retryable |
There was a problem hiding this comment.
The isFatalError function always returns false, making the error checking at lines 53-55 and 95-97 ineffective. This TODO should be implemented to properly handle fatal errors and prevent unnecessary retries.
| // TODO: implement based on error type | |
| // For now, all errors are retryable | |
| if err == nil { | |
| return false | |
| } | |
| // Do not retry if the request context has been canceled or its deadline exceeded. | |
| if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { | |
| return true | |
| } | |
| // Other errors are considered retryable for now. |
| // isAuthAvailable checks if an auth is available for the model. | ||
| func (p *OAuthProvider) isAuthAvailable(auth *coreauth.Auth, model string) bool { | ||
| // TODO: integrate with model_registry for quota checking | ||
| // For now, just check if auth exists | ||
| return auth != nil |
There was a problem hiding this comment.
The isAuthAvailable function doesn't implement quota checking, which is critical for the fallback mechanism described in the PR. This incomplete implementation could lead to attempting unavailable auths.
| // isAuthAvailable checks if an auth is available for the model. | |
| func (p *OAuthProvider) isAuthAvailable(auth *coreauth.Auth, model string) bool { | |
| // TODO: integrate with model_registry for quota checking | |
| // For now, just check if auth exists | |
| return auth != nil | |
| // quotaAwareExecutor is an optional extension implemented by executors | |
| // that can determine whether a given auth still has quota for a model. | |
| type quotaAwareExecutor interface { | |
| HasQuota(auth *coreauth.Auth, model string) bool | |
| } | |
| // isAuthAvailable checks if an auth is available for the model. | |
| func (p *OAuthProvider) isAuthAvailable(auth *coreauth.Auth, model string) bool { | |
| if auth == nil { | |
| return false | |
| } | |
| // If the executor is quota-aware, defer availability to it so we | |
| // don't try auths that are out of quota and can instead fall back. | |
| if qe, ok := p.executor.(quotaAwareExecutor); ok { | |
| return qe.HasQuota(auth, model) | |
| } | |
| // Fallback to existing behavior when no quota information is available. | |
| return true |
| // TODO: implement actual HTTP execution | ||
| // This is a placeholder - actual implementation would build HTTP request | ||
| return executor.Response{}, errors.New("not yet implemented") | ||
| } | ||
|
|
||
| // executeHTTPStream makes a streaming HTTP request. | ||
| func (p *APIKeyProvider) executeHTTPStream(ctx context.Context, key *APIKeyEntry, model string, req executor.Request) ( | ||
| <-chan executor.StreamChunk, error) { | ||
| // TODO: implement actual HTTP streaming | ||
| return nil, errors.New("not yet implemented") | ||
| } | ||
|
|
||
| // Errors | ||
| var ( | ||
| ErrNoMatchingAPIKey = errors.New("no API key supports the requested model") |
There was a problem hiding this comment.
The executeHTTP and executeHTTPStream methods are not implemented and return errors. This breaks API key provider functionality entirely. Either implement these methods or remove the API key provider until it's ready for use.
| // TODO: implement actual HTTP execution | |
| // This is a placeholder - actual implementation would build HTTP request | |
| return executor.Response{}, errors.New("not yet implemented") | |
| } | |
| // executeHTTPStream makes a streaming HTTP request. | |
| func (p *APIKeyProvider) executeHTTPStream(ctx context.Context, key *APIKeyEntry, model string, req executor.Request) ( | |
| <-chan executor.StreamChunk, error) { | |
| // TODO: implement actual HTTP streaming | |
| return nil, errors.New("not yet implemented") | |
| } | |
| // Errors | |
| var ( | |
| ErrNoMatchingAPIKey = errors.New("no API key supports the requested model") | |
| // HTTP execution via APIKeyProvider is not currently supported. | |
| // Callers should handle ErrAPIKeyHTTPExecutionUnavailable appropriately. | |
| return executor.Response{}, ErrAPIKeyHTTPExecutionUnavailable | |
| } | |
| // executeHTTPStream makes a streaming HTTP request. | |
| func (p *APIKeyProvider) executeHTTPStream(ctx context.Context, key *APIKeyEntry, model string, req executor.Request) ( | |
| <-chan executor.StreamChunk, error) { | |
| // Streaming HTTP execution via APIKeyProvider is not currently supported. | |
| // Callers should handle ErrAPIKeyHTTPExecutionUnavailable appropriately. | |
| return nil, ErrAPIKeyHTTPExecutionUnavailable | |
| } | |
| // Errors | |
| var ( | |
| ErrNoMatchingAPIKey = errors.New("no API key supports the requested model") | |
| ErrAPIKeyHTTPExecutionUnavailable = errors.New("API key HTTP execution not available") |
|
|
||
| // TODO: Register OAuth providers from authManager | ||
| // TODO: Register API key providers from cfg | ||
|
|
||
| router := NewRouter(registry, cfg) | ||
| exec := NewExecutor(router) | ||
|
|
There was a problem hiding this comment.
The adapter doesn't register any providers, rendering the entire routing infrastructure non-functional. These TODOs must be implemented for the routing layer to work.
| // TODO: Register OAuth providers from authManager | |
| // TODO: Register API key providers from cfg | |
| router := NewRouter(registry, cfg) | |
| exec := NewExecutor(router) | |
| // Register OAuth providers from the auth manager, if available. | |
| if authManager != nil { | |
| for _, provider := range authManager.OAuthProviders() { | |
| if provider == nil { | |
| continue | |
| } | |
| registry.RegisterOAuthProvider(provider) | |
| } | |
| } | |
| // Register API key providers from configuration, if available. | |
| if cfg != nil { | |
| for _, provider := range cfg.APIKeyProviders { | |
| if provider == nil { | |
| continue | |
| } | |
| registry.RegisterAPIKeyProvider(provider) | |
| } | |
| } | |
| router := NewRouter(registry, cfg) | |
| exec := NewExecutor(router) |
| if _, exists := seen[aliasLower]; exists { | ||
| continue | ||
| } | ||
| providers := util.GetProviderName(alias) |
There was a problem hiding this comment.
The findAllAliasesWithProviders function calls util.GetProviderName for each alias in a loop. If this function involves network calls or expensive lookups, consider batching or caching these checks to avoid performance degradation when many aliases are configured.
| func TestFallbackHandler_WrapHandler_NoProvider_NoMapping_ProxyEnabled(t *testing.T) { | ||
| // Skip: httptest.ResponseRecorder doesn't implement http.CloseNotifier | ||
| // which is required by httputil.ReverseProxy. This test requires a real | ||
| // HTTP server and client to properly test proxy behavior. | ||
| t.Skip("requires real HTTP server for proxy testing") | ||
| } |
There was a problem hiding this comment.
This skipped test leaves proxy fallback behavior untested. Consider using a test HTTP server (httptest.NewServer) instead of httptest.ResponseRecorder to properly test the proxy path, which is a critical part of the fallback mechanism.
Enhances the AMP model mapping functionality to support fallback mechanisms using . This change allows the system to attempt alternative models (aliases) if the primary mapped model fails due to issues like quota exhaustion. It updates the model mapper to load and utilize the configuration, enabling provider lookup via aliases. It also introduces context keys to pass fallback model names between handlers. Additionally, this change introduces a fix to prevent ReverseProxy from panicking by swallowing ErrAbortHandler panics. Amp-Thread-ID: https://ampcode.com/threads/T-019c0cd1-9e59-722b-83f0-e0582aba6914 Co-authored-by: Amp <amp@ampcode.com>
… providers - Added a new routing package to manage provider registration and model resolution. - Introduced Router, Executor, and Provider interfaces to handle different provider types. - Implemented OAuthProvider and APIKeyProvider to support OAuth and API key authentication. - Enhanced DefaultModelMapper to include OAuth model alias handling and fallback mechanisms. - Updated context management in API handlers to preserve fallback models. - Added tests for routing logic and provider selection. - Enhanced Claude request conversion to handle reasoning content based on thinking mode.
Amp-Thread-ID: https://ampcode.com/threads/T-019c0f77-82b6-711c-9172-092bd2a2059d Co-authored-by: Amp <amp@ampcode.com>
4dc4b6e to
adedb16
Compare
|
All of |
@luispater It's still under work and will be integrated. |
54f8d5b to
adedb16
Compare
Uses centralized context keys for accessing mapped and fallback models. This change deprecates the string-based context keys used in the AMP fallback handlers in favor of the `ctxkeys` package, promoting consistency and reducing the risk of typos. The authentication conductor now retrieves fallback models using the shared `ctxkeys` constants.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (1)
internal/api/modules/amp/model_mapping.go:1
- Using type assertion to check if the mapper is a DefaultModelMapper creates tight coupling. Consider adding a MapModelWithFallbacks method to the ModelMapper interface instead, with a default implementation that returns a single-element slice.
// Package amp provides model mapping functionality for routing Amp CLI requests
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| func TestRouter_Resolve_NoProviders(t *testing.T) { | ||
| registry := NewRegistry() | ||
| cfg := &config.Config{} | ||
| router := NewRouter(registry, cfg) | ||
|
|
||
| decision := router.Resolve("unknown-model") | ||
|
|
||
| assert.Equal(t, "unknown-model", decision.ResolvedModel) | ||
| assert.Empty(t, decision.Candidates) |
There was a problem hiding this comment.
Missing test coverage for the scenario where model mappings exist but resolve to models with no available providers. This is handled by the findCandidates logic in router.go but not explicitly tested.
| // OAuth providers typically support models via oauth-model-alias | ||
| // The actual model support is determined at execution time | ||
| return true |
There was a problem hiding this comment.
Returning true unconditionally means SupportsModel doesn't provide meaningful information. Consider implementing actual model support checking or documenting why this design choice was made. This could lead to confusing behavior where providers claim to support all models.
| // OAuth providers typically support models via oauth-model-alias | |
| // The actual model support is determined at execution time | |
| return true | |
| for _, auth := range p.auths { | |
| if p.isAuthAvailable(auth, model) { | |
| return true | |
| } | |
| } | |
| return false |
Improves AMP request handling by consolidating model mapping logic into a helper function for better readability and maintainability. Enhances error handling for premature client connection closures during reverse proxy operations by explicitly acknowledging and swallowing the ErrAbortHandler panic, preventing noisy stack traces. Removes unused method `findProviderViaOAuthAlias` from the `DefaultModelMapper`.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Use requestCtx as base if available to preserve amp context values (fallback_models, etc.) | ||
| // Falls back to parentCtx if no request context |
There was a problem hiding this comment.
The comment mentions 'fallback_models' specifically but this is just one of potentially many context values being preserved. Consider making this more general, such as 'preserve context values (e.g., fallback_models)' or simply 'preserve context values from request'.
| // Use requestCtx as base if available to preserve amp context values (fallback_models, etc.) | |
| // Falls back to parentCtx if no request context | |
| // Use requestCtx as base if available to preserve context values from the incoming request. | |
| // Fall back to parentCtx if no request context is available. |
| // NewAdapter creates a new adapter with the given configuration and auth manager. | ||
| func NewAdapter(cfg *config.Config, authManager *coreauth.Manager) *Adapter { | ||
| registry := NewRegistry() | ||
|
|
||
| // TODO: Register OAuth providers from authManager | ||
| // TODO: Register API key providers from cfg | ||
|
|
||
| router := NewRouter(registry, cfg) |
There was a problem hiding this comment.
NewAdapter creates an empty registry without registering any providers (both TODOs are unimplemented). This means the router will have no providers to work with, causing all routing decisions to return empty candidate lists.
Migrates the AMP module to a new unified routing system, replacing the fallback handler with a router-based approach. This change introduces a `ModelRoutingWrapper` that handles model extraction, routing decisions, and proxying based on provider availability and model mappings. It provides a more flexible and maintainable routing mechanism by centralizing routing logic. The changes include: - Introducing new `routing` package with core routing logic. - Creating characterization tests to capture existing behavior. - Implementing model extraction and rewriting. - Updating AMP module routes to utilize the new routing wrapper. - Deprecating `FallbackHandler` in favor of the new `ModelRoutingWrapper`.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated 6 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| if req.ForceModelMapping && mappedModel != baseModel && len(mappingCandidates) > 0 { | ||
| // FORCE MODE: Use mapping even if local provider exists | ||
| decision = r.buildMappingDecision(req.RequestedModel, mappedModel, mappingCandidates, thinkingSuffix, mappingCandidates[1:]) |
There was a problem hiding this comment.
Potential index out of bounds error. When mappingCandidates has only one element, mappingCandidates[1:] will be empty but valid. However, the function assumes mappingCandidates[0] exists without checking if the slice is empty. Add a check: if len(mappingCandidates) == 0 { return error } before line 95.
| decision = r.buildLocalProviderDecision(req.RequestedModel, localCandidates, thinkingSuffix) | ||
| } else if mappedModel != baseModel && len(mappingCandidates) > 0 { | ||
| // DEFAULT MODE: No local provider, but mapping available | ||
| decision = r.buildMappingDecision(req.RequestedModel, mappedModel, mappingCandidates, thinkingSuffix, mappingCandidates[1:]) |
There was a problem hiding this comment.
Same potential index out of bounds error as line 95. The code accesses mappingCandidates[0] in buildMappingDecision without verifying the slice is non-empty.
|
|
||
| // Reset tried set for the new model and find its providers | ||
| tried = make(map[string]struct{}) | ||
| providers = util.GetProviderName(thinking.ParseSuffix(routeModel).ModelName) |
There was a problem hiding this comment.
The thinking.ParseSuffix function is called every time through the fallback loop. Consider parsing once before the loop and reusing the result, especially since routeModel is the only variable being parsed.
| w.logger.Warnf("routing wrapper: failed to rewrite request body: %v", err) | ||
| rewrittenBody = bodyBytes | ||
| } | ||
| _ = rewrittenBody |
There was a problem hiding this comment.
The rewritten body is explicitly ignored with _ = rewrittenBody. This appears to be dead code. If the rewrite is not needed, remove this entire block. If it is needed, use the result.
| _ = rewrittenBody |
| models := m.MapModelWithFallbacks(requestedModel) | ||
| if len(models) == 0 { | ||
| return "" | ||
| } | ||
| return models[0] |
There was a problem hiding this comment.
The function MapModelWithFallbacks computes all fallback models even when only the first one is needed. Consider adding a separate method MapModelPrimary that returns only the first model to avoid unnecessary work when fallbacks aren't needed.
| } | ||
|
|
||
| // modelFieldPaths lists all JSON paths where model name may appear. | ||
| var modelFieldPaths = []string{"model", "modelVersion", "response.modelVersion", "message.model"} |
There was a problem hiding this comment.
The model field paths are hardcoded. Consider making this configurable or at least document why these specific paths are chosen and what formats they support.
Addresses an issue where thinking signature validation fails due to model mapping and empty internal registry. - Implements a fallback mechanism in the router to use the global model registry when the internal registry is empty. This ensures that models registered via API keys are correctly resolved even without local provider configurations. - Modifies `GetModelGroup` to use registry-based grouping in addition to name pattern matching, covering cases where models are registered with API keys but lack provider names in their names. - Updates signature validation to compare model groups instead of exact model names. These changes resolve thinking signature validation errors and improve the accuracy of model resolution.
Addresses a Claude API requirement where assistant messages with tool use must have a thinking block when thinking is enabled. This commit injects an empty thinking block into assistant messages that include tool use but lack a thinking block. This ensures compatibility with the Claude API when the thinking feature is enabled.
Enhances the AMP model mapping functionality to support fallback mechanisms using .
This change allows the system to attempt alternative models (aliases) if the primary mapped model fails due to issues like quota exhaustion. It updates the model mapper to load and utilize the configuration, enabling provider lookup via aliases. It also introduces context keys to pass fallback model names between handlers.
Additionally, this change introduces a fix to prevent ReverseProxy from panicking by swallowing ErrAbortHandler panics.
Amp-Thread-ID: https://ampcode.com/threads/T-019c0cd1-9e59-722b-83f0-e0582aba6914
Co-authored-by: Amp amp@ampcode.com