jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-21 01:17:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(keys): role-based gateway keys with admin management UI and per-user model scope

Browse Source
This commit is contained in:
root
2026-08-09 10:01:40 +08:00
parent dec03238dd
commit 3408c9cb1f
10 changed files with 901 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
internal/scheduler/scheduler.go
View File

@ -47,7 +47,10 @@ func FromRegistry(ps []*provider.Provider) []Provider {
// candidate receives a request pinned to its own model (ModelFor), so an AUTO
// chain fallback switches the model id per provider instead of reusing the
// first candidate's model name.
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, error) {
//
// On success it returns the response together with the name of the provider
// and the exact model id that actually served the request (used for stats).
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, string, string, error) {
attempts := s.MaxRetries + 1
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
@ -56,27 +59,25 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
r.Model = p.ModelFor(req.Model)
resp, err := p.Chat(ctx, &r)
if ctx.Err() != nil {
return nil, ctx.Err()
return nil, "", "", ctx.Err()
}
if err == nil {
return resp, nil
return resp, p.Name(), r.Model, nil
}
lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
}
if lastErr == nil {
// if loop couldn't run because cands was short but no error recorded yet
if len(cands) == 0 {
return nil, fmt.Errorf("no provider available")
lastErr = fmt.Errorf("no provider available")
}
// should not happen
return nil, lastErr
}
return nil, lastErr
return nil, "", "", lastErr
}
// ChatStream runs a streaming chat across cands, falling back early on connect
// errors. The request model is pinned per candidate like Chat.
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
// errors. The request model is pinned per candidate like Chat. On success it
// returns the chunk channel plus the serving provider name and model id.
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, string, string, error) {
attempts := s.MaxRetries + 1
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
@ -85,32 +86,31 @@ func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types
r.Model = p.ModelFor(req.Model)
resp, err := p.ChatStream(ctx, &r)
if err == nil {
return resp, nil
return resp, p.Name(), r.Model, nil
}
lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
}
if lastErr == nil {
if len(cands) == 0 {
return nil, fmt.Errorf("no provider available")
}
if lastErr == nil && len(cands) == 0 {
lastErr = fmt.Errorf("no provider available")
}
return nil, lastErr
return nil, "", "", lastErr
}
// Image runs an image-generation request across cands.
func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
// Image runs an image-generation request across cands; returns the used
// provider name on success.
func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.ImageGenRequest) (*types.UnifiedResponse, string, error) {
attempts := s.MaxRetries + 1
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
p := cands[i]
resp, err := p.Image(ctx, req)
if err == nil {
return resp, nil
return resp, p.Name(), nil
}
lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
}
if lastErr == nil && len(cands) == 0 {
return nil, fmt.Errorf("no provider available")
lastErr = fmt.Errorf("no provider available")
}
return nil, lastErr
return nil, "", lastErr
}