mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 08:57:57 +00:00
feat: ModelRouter — unified OpenAI-compatible multi-source LLM gateway
- Lua adapters per upstream (transform_request/response/stream_chunk, build_headers signing hooks) - AUTO priority routing with per-model kind (chat/image), explicit source/model routing - Per-source concurrency caps with queueing, exponential backoff, AUTO failover - OpenAI-compatible API: chat completions, SSE streaming, image generations, models - Gateway key auth, web UI for adapter/source management, runtime persistence - e2e test running the real binary against mocked upstreams
This commit is contained in:
107
internal/scheduler/scheduler.go
Normal file
107
internal/scheduler/scheduler.go
Normal file
@ -0,0 +1,107 @@
|
||||
// Package scheduler implements request scheduling across providers: per-source
|
||||
// concurrency caps (acquire with wait = queuing), AUTO model fallback chains,
|
||||
// and exponential backoff via provider health.
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"llmsproxy/internal/provider"
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// Scheduler drives one chat tool call across the candidate provider chain.
|
||||
type Scheduler struct {
|
||||
// MaxRetries how many fallback providers to try before failing.
|
||||
MaxRetries int
|
||||
}
|
||||
|
||||
func New(maxRetries int) *Scheduler {
|
||||
if maxRetries < 0 {
|
||||
maxRetries = 0
|
||||
}
|
||||
return &Scheduler{MaxRetries: maxRetries}
|
||||
}
|
||||
|
||||
// Provider is the minimal interface the scheduler needs to schedule over.
|
||||
type Provider interface {
|
||||
Name() string
|
||||
Available() bool
|
||||
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
|
||||
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
|
||||
Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error)
|
||||
}
|
||||
|
||||
// FromRegistry converts *provider.Provider slices to the scheduler interface.
|
||||
func FromRegistry(ps []*provider.Provider) []Provider {
|
||||
out := make([]Provider, len(ps))
|
||||
for i, p := range ps {
|
||||
out[i] = p
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Chat runs a chat request across cands, falling back on failure.
|
||||
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, error) {
|
||||
attempts := s.MaxRetries + 1
|
||||
var lastErr error
|
||||
for i := 0; i < attempts && i < len(cands); i++ {
|
||||
p := cands[i]
|
||||
resp, err := p.Chat(ctx, req)
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
if err == nil {
|
||||
return resp, 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")
|
||||
}
|
||||
// should not happen
|
||||
return nil, lastErr
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// ChatStream runs a streaming chat across cands, falling back early on connect errors.
|
||||
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
|
||||
attempts := s.MaxRetries + 1
|
||||
var lastErr error
|
||||
for i := 0; i < attempts && i < len(cands); i++ {
|
||||
p := cands[i]
|
||||
resp, err := p.ChatStream(ctx, req)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
|
||||
}
|
||||
if lastErr == nil {
|
||||
if len(cands) == 0 {
|
||||
return nil, fmt.Errorf("no provider available")
|
||||
}
|
||||
}
|
||||
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) {
|
||||
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
|
||||
}
|
||||
lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
|
||||
}
|
||||
if lastErr == nil && len(cands) == 0 {
|
||||
return nil, fmt.Errorf("no provider available")
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
Reference in New Issue
Block a user