@@ -2,6 +2,7 @@ package channel
22
33import (
44 "bytes"
5+ "encoding/json"
56 "fmt"
67 "gpt-load/internal/models"
78 "gpt-load/internal/types"
@@ -12,6 +13,7 @@ import (
1213 "strings"
1314 "sync"
1415
16+ "github.com/sirupsen/logrus"
1517 "gorm.io/datatypes"
1618)
1719
@@ -33,9 +35,11 @@ type BaseChannel struct {
3335 upstreamLock sync.Mutex
3436
3537 // Cached fields from the group for stale check
36- channelType string
37- groupUpstreams datatypes.JSON
38- effectiveConfig * types.SystemSettings
38+ channelType string
39+ groupUpstreams datatypes.JSON
40+ effectiveConfig * types.SystemSettings
41+ modelRedirectRules datatypes.JSONMap
42+ modelRedirectStrict bool
3943}
4044
4145// getUpstreamURL selects an upstream URL using a smooth weighted round-robin algorithm.
@@ -107,6 +111,13 @@ func (b *BaseChannel) IsConfigStale(group *models.Group) bool {
107111 if ! reflect .DeepEqual (b .effectiveConfig , & group .EffectiveConfig ) {
108112 return true
109113 }
114+ // Check for model redirect rules changes
115+ if ! reflect .DeepEqual (b .modelRedirectRules , group .ModelRedirectRules ) {
116+ return true
117+ }
118+ if b .modelRedirectStrict != group .ModelRedirectStrict {
119+ return true
120+ }
110121 return false
111122}
112123
@@ -119,3 +130,143 @@ func (b *BaseChannel) GetHTTPClient() *http.Client {
119130func (b * BaseChannel ) GetStreamClient () * http.Client {
120131 return b .StreamClient
121132}
133+
134+ // ApplyModelRedirect applies model redirection based on the group's redirect rules.
135+ func (b * BaseChannel ) ApplyModelRedirect (req * http.Request , bodyBytes []byte , group * models.Group ) ([]byte , error ) {
136+ if len (group .ModelRedirectMap ) == 0 || len (bodyBytes ) == 0 {
137+ return bodyBytes , nil
138+ }
139+
140+ var requestData map [string ]any
141+ if err := json .Unmarshal (bodyBytes , & requestData ); err != nil {
142+ return bodyBytes , nil
143+ }
144+
145+ modelValue , exists := requestData ["model" ]
146+ if ! exists {
147+ return bodyBytes , nil
148+ }
149+
150+ model , ok := modelValue .(string )
151+ if ! ok {
152+ return bodyBytes , nil
153+ }
154+
155+ // Direct match without any prefix processing
156+ if targetModel , found := group .ModelRedirectMap [model ]; found {
157+ requestData ["model" ] = targetModel
158+
159+ // Log the redirection for audit
160+ logrus .WithFields (logrus.Fields {
161+ "group" : group .Name ,
162+ "original_model" : model ,
163+ "target_model" : targetModel ,
164+ "channel" : "json_body" ,
165+ }).Debug ("Model redirected" )
166+
167+ return json .Marshal (requestData )
168+ }
169+
170+ if group .ModelRedirectStrict {
171+ return nil , fmt .Errorf ("model '%s' is not configured in redirect rules" , model )
172+ }
173+
174+ return bodyBytes , nil
175+ }
176+
177+ // TransformModelList transforms the model list response based on redirect rules.
178+ func (b * BaseChannel ) TransformModelList (req * http.Request , bodyBytes []byte , group * models.Group ) (map [string ]any , error ) {
179+ var response map [string ]any
180+ if err := json .Unmarshal (bodyBytes , & response ); err != nil {
181+ logrus .WithError (err ).Debug ("Failed to parse model list response, returning empty" )
182+ return nil , err
183+ }
184+
185+ dataInterface , exists := response ["data" ]
186+ if ! exists {
187+ return response , nil
188+ }
189+
190+ upstreamModels , ok := dataInterface .([]any )
191+ if ! ok {
192+ return response , nil
193+ }
194+
195+ // Build configured source models list (common logic for both modes)
196+ configuredModels := buildConfiguredModels (group .ModelRedirectMap )
197+
198+ // Strict mode: return only configured models (whitelist)
199+ if group .ModelRedirectStrict {
200+ response ["data" ] = configuredModels
201+
202+ logrus .WithFields (logrus.Fields {
203+ "group" : group .Name ,
204+ "model_count" : len (configuredModels ),
205+ "strict_mode" : true ,
206+ }).Debug ("Model list returned (strict mode - configured models only)" )
207+
208+ return response , nil
209+ }
210+
211+ // Non-strict mode: merge upstream + configured models (upstream priority)
212+ merged := mergeModelLists (upstreamModels , configuredModels )
213+ response ["data" ] = merged
214+
215+ logrus .WithFields (logrus.Fields {
216+ "group" : group .Name ,
217+ "upstream_count" : len (upstreamModels ),
218+ "configured_count" : len (configuredModels ),
219+ "merged_count" : len (merged ),
220+ "strict_mode" : false ,
221+ }).Debug ("Model list merged (non-strict mode)" )
222+
223+ return response , nil
224+ }
225+
226+ // buildConfiguredModels builds a list of models from redirect rules
227+ func buildConfiguredModels (redirectMap map [string ]string ) []any {
228+ if len (redirectMap ) == 0 {
229+ return []any {}
230+ }
231+
232+ models := make ([]any , 0 , len (redirectMap ))
233+ for sourceModel := range redirectMap {
234+ models = append (models , map [string ]any {
235+ "id" : sourceModel ,
236+ "object" : "model" ,
237+ "created" : 0 ,
238+ "owned_by" : "system" ,
239+ })
240+ }
241+ return models
242+ }
243+
244+ // mergeModelLists merges upstream and configured model lists
245+ func mergeModelLists (upstream []any , configured []any ) []any {
246+ // Create set of upstream model IDs
247+ upstreamIDs := make (map [string ]bool )
248+ for _ , item := range upstream {
249+ if modelObj , ok := item .(map [string ]any ); ok {
250+ if modelID , ok := modelObj ["id" ].(string ); ok {
251+ upstreamIDs [modelID ] = true
252+ }
253+ }
254+ }
255+
256+ // Start with all upstream models
257+ result := make ([]any , len (upstream ))
258+ copy (result , upstream )
259+
260+ // Add configured models that don't exist in upstream
261+ for _ , item := range configured {
262+ if modelObj , ok := item .(map [string ]any ); ok {
263+ if modelID , ok := modelObj ["id" ].(string ); ok {
264+ if ! upstreamIDs [modelID ] {
265+ result = append (result , item )
266+ }
267+ }
268+ }
269+ }
270+
271+ return result
272+ }
0 commit comments