Multi-Target Model Aliases and Provider Aggregation #716
Conversation
…appings and update related functionality
Summary of ChangesHello @mcowger, 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 introduces a fundamental upgrade to the model aliasing mechanism, transforming it from a simple one-to-one mapping into a sophisticated multi-target aggregation engine. This change empowers users to define aliases that represent a collection of models, enabling the system to pool providers from all associated targets. The primary impact is increased flexibility and resilience in model routing, as requests can now dynamically resolve to the most suitable available model among a defined set, enhancing the overall reliability and performance of the proxy. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and well-implemented feature for multi-target model aliases, allowing a single alias to map to multiple underlying models and aggregating providers. The core logic changes in model_registry.go are robust and supported by a comprehensive new test suite. The management API has also been updated to support this new alias structure.
My review focuses on improving consistency and robustness in the new management API handlers and simplifying some of the new utility functions. I've identified a couple of potential race conditions/inconsistency issues in the management handlers and opportunities to remove redundant code and clarify comments in provider.go.
| h.cfg.Aliases = *body.Aliases | ||
| h.persist(c) | ||
| // Update global aliases after persisting | ||
| if h.cfg.Aliases != nil { | ||
| util.SetAliases(h.cfg.Aliases) | ||
| } |
There was a problem hiding this comment.
The return value of h.persist(c) is not being checked here. Other handlers in this file, such as PutAlias and DeleteAlias, correctly check this return value and abort if persistence fails. This could lead to this endpoint returning a success status even when the configuration change was not saved.
h.cfg.Aliases = *body.Aliases
if !h.persist(c) {
return
}
// Update global aliases after persisting
if h.cfg.Aliases != nil {
util.SetAliases(h.cfg.Aliases)
}| util.SetAliases(h.cfg.Aliases) | ||
| if !h.persist(c) { | ||
| return | ||
| } |
There was a problem hiding this comment.
The order of operations in DeleteAlias is inconsistent with other handlers like PutAlias. Here, the global in-memory alias state is updated via util.SetAliases before the configuration is persisted to disk. If h.persist(c) fails, the in-memory state will be inconsistent with the persisted configuration. To ensure atomicity, you should persist the changes first and only update the in-memory state upon successful persistence.
if !h.persist(c) {
return
}
util.SetAliases(h.cfg.Aliases)| payload: | ||
| default: # Default rules only set parameters when they are missing in the payload. | ||
| - models: | ||
| - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*") | ||
| protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex | ||
| params: # JSON path (gjson/sjson syntax) -> value | ||
| "generationConfig.thinkingConfig.thinkingBudget": 32768 | ||
| override: # Override rules always set parameters, overwriting any existing values. | ||
| - models: | ||
| - name: "gpt-*" # Supports wildcards (e.g., "gpt-*") | ||
| protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex | ||
| params: # JSON path (gjson/sjson syntax) -> value | ||
| "reasoning.effort": "high" |
There was a problem hiding this comment.
The payload configuration section has been uncommented in this pull request. If this change is unrelated to the multi-target alias feature, it might be better to revert it or move it to a separate PR to keep this one focused on a single concern.
| // GetProviderNameAndModel determines all AI service providers capable of serving a model, | ||
| // and returns the resolved model name if an alias was used. | ||
| // It considers all possible sources: direct model name and global aliases. | ||
| // | ||
| // Parameters: | ||
| // - modelName: The name of the model to identify providers for. | ||
| // | ||
| // Returns: | ||
| // - []string: All provider identifiers capable of serving the model. | ||
| // - string: The resolved model name (original or alias target). |
There was a problem hiding this comment.
The documentation for GetProviderNameAndModel is misleading. It states that the function "returns the resolved model name if an alias was used", but the implementation simply returns the original modelName input. This can cause confusion for future maintainers. Please update the comment to accurately describe the function's behavior.
| // GetProviderNameAndModel determines all AI service providers capable of serving a model, | |
| // and returns the resolved model name if an alias was used. | |
| // It considers all possible sources: direct model name and global aliases. | |
| // | |
| // Parameters: | |
| // - modelName: The name of the model to identify providers for. | |
| // | |
| // Returns: | |
| // - []string: All provider identifiers capable of serving the model. | |
| // - string: The resolved model name (original or alias target). | |
| // GetProviderNameAndModel determines all AI service providers capable of serving a model, | |
| // and returns the original model name that was passed in. | |
| // It considers all possible sources: direct model name and global aliases. | |
| // | |
| // Parameters: | |
| // - modelName: The name of the model to identify providers for. | |
| // | |
| // Returns: | |
| // - []string: All provider identifiers capable of serving the model. | |
| // - string: The original model name. |
| func SetAliases(aliases map[string][]string) { | ||
| if aliases == nil { | ||
| aliases = make(map[string][]string) | ||
| } | ||
| aliasValue.Store(aliases) | ||
| registry.GetGlobalRegistry().SetAliases(aliases) | ||
| } |
There was a problem hiding this comment.
The aliasValue global variable (defined on line 17) and the ResolveAlias function (lines 275-297) appear to be redundant. The alias data is already stored and managed within the ModelRegistry. This atomic.Value creates a duplicate state, and since util.ResolveAlias is unused, it's effectively dead code.
To simplify, you could remove aliasValue and ResolveAlias. SetAliases can also be simplified, as the nil check is already handled inside registry.SetAliases.
func SetAliases(aliases map[string][]string) {
registry.GetGlobalRegistry().SetAliases(aliases)
}1 participant
Overview
This PR introduces a significant enhancement to the model aliasing system, moving from a simple 1:1 mapping on a per provider basis to a robust multi-target aggregation engine. This allows a single alias to point to multiple underlying models, effectively creating a unified "virtual model" that pools providers from all targets.
Key Features
balanced-basicalias can point to bothclaude-3-5-sonnetandgemini-1.5-pro.alias: target) and new sequence mappings (alias: [t1, t2]) using a custom YAML unmarshaler.PUT /aliases/:aliasendpoint.Technical Changes
internal/config/config.go: ChangedAliasestype tomap[string][]stringand implementedUnmarshalYAMLfor backward compatibility.internal/registry/model_registry.go:GetModelProvidersandGetModelCountto aggregate data across all alias targets.ResolveModelForClientto iterate through targets and find the first compatible match.isModelAvailableLockedto check availability across the entire target set.internal/api/handlers/management/config_basic.go: Updated handlers to accepttargets(slice) ortarget(string) in JSON payloads.internal/util/provider.go: Updated global utility functions to handle the new slice-based alias structure.Testing
TestMultiTargetAliasin internal/registry/model_registry_test.go to verify provider aggregation and cross-provider resolution.Example Configuration