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:
root
283
internal/gateway/chat.go
Normal file
283
internal/gateway/chat.go
Normal file
@ -0,0 +1,283 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/provider"
|
||||
"llmsproxy/internal/scheduler"
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// chatRequest mirrors the OpenAI chat completions request the gateway accepts.
|
||||
type chatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []types.ChatMessage `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Tools []interface{} `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
}
|
||||
|
||||
// ChatCompletion is the non-streaming OpenAI response object.
|
||||
type ChatCompletion struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChatChoice `json:"choices"`
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type ChatChoice struct {
|
||||
Index int `json:"index"`
|
||||
Message RespMessage `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type RespMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
ToolCalls []types.ToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
type ChatChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
}
|
||||
|
||||
type ChunkChoice struct {
|
||||
Index int `json:"index"`
|
||||
Delta RespMessage `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
var seq int64
|
||||
|
||||
func newID() string {
|
||||
n := atomic.AddInt64(&seq, 1)
|
||||
return fmt.Sprintf("chatcmpl-%d", n)
|
||||
}
|
||||
|
||||
func isAuto(m string) bool {
|
||||
m = strings.TrimSpace(m)
|
||||
return m == "" || strings.EqualFold(m, "AUTO")
|
||||
}
|
||||
|
||||
// resolveCands picks the ordered candidate providers for a requested model.
|
||||
func (g *Gateway) resolveCands(model string) ([]*provider.Provider, string) {
|
||||
if model == "" || isAuto(model) {
|
||||
return g.core.Registry().Resolve("AUTO"), ""
|
||||
}
|
||||
return g.core.Registry().Resolve(model), model
|
||||
}
|
||||
|
||||
func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
|
||||
return
|
||||
}
|
||||
var req chatRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if len(req.Messages) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "messages is required")
|
||||
return
|
||||
}
|
||||
model := req.Model
|
||||
if model == "" {
|
||||
model = g.core.DefaultModel()
|
||||
}
|
||||
cands, effective := g.resolveCands(model)
|
||||
if len(cands) == 0 {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no LLM source configured")
|
||||
return
|
||||
}
|
||||
if effective == "" {
|
||||
effective = firstModel(cands[0])
|
||||
}
|
||||
ctx := r.Context()
|
||||
|
||||
inner := &types.ChatRequest{
|
||||
Model: normalizeModel(model),
|
||||
Messages: req.Messages,
|
||||
Temperature: req.Temperature,
|
||||
MaxTokens: req.MaxTokens,
|
||||
Stream: req.Stream,
|
||||
Tools: req.Tools,
|
||||
ToolChoice: req.ToolChoice,
|
||||
}
|
||||
if req.Stream {
|
||||
g.streamChat(w, ctx, cands, inner, effective)
|
||||
return
|
||||
}
|
||||
g.singleChat(w, ctx, cands, inner, effective)
|
||||
}
|
||||
|
||||
func normalizeModel(m string) string {
|
||||
if isAuto(m) {
|
||||
return ""
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func firstModel(p *provider.Provider) string {
|
||||
ms := p.Models()
|
||||
if len(ms) > 0 {
|
||||
return ms[0]
|
||||
}
|
||||
return "auto"
|
||||
}
|
||||
|
||||
// imageOnly keeps providers exposing at least one image-kind model.
|
||||
func imageOnly(cands []*provider.Provider) []*provider.Provider {
|
||||
var out []*provider.Provider
|
||||
for _, p := range cands {
|
||||
for _, id := range p.Models() {
|
||||
if m := p.ModelByID(id); m != nil && m.Kind == "image" {
|
||||
out = append(out, p)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string) {
|
||||
resp, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry(cands), req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
}
|
||||
if len(resp.ToolCalls) > 0 {
|
||||
msg.ToolCalls = resp.ToolCalls
|
||||
}
|
||||
out := ChatCompletion{
|
||||
ID: newID(),
|
||||
Object: "chat.completion",
|
||||
Created: time.Now().Unix(),
|
||||
Model: effective,
|
||||
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
out.Usage = &resp.TokenUsage
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string) {
|
||||
chunks, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry(cands), req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
flusher, _ := w.(http.Flusher)
|
||||
id := newID()
|
||||
created := time.Now().Unix()
|
||||
send := func(obj interface{}) bool {
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
|
||||
return false
|
||||
}
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
if !send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
for ck := range chunks {
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
}
|
||||
delta := RespMessage{Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
delta.ReasoningContent = ck.ReasoningContent
|
||||
}
|
||||
if len(ck.ToolCalls) > 0 {
|
||||
delta.ToolCalls = ck.ToolCalls
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
if !send(chunk) {
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
})
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
|
||||
return
|
||||
}
|
||||
var req types.ImageGenRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.Prompt == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "prompt is required")
|
||||
return
|
||||
}
|
||||
model := req.Model
|
||||
if model == "" {
|
||||
model = g.core.DefaultModel()
|
||||
}
|
||||
cands, _ := g.resolveCands(model)
|
||||
cands = imageOnly(cands)
|
||||
if len(cands) == 0 {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
|
||||
return
|
||||
}
|
||||
resp, err := g.core.Scheduler().Image(r.Context(), scheduler.FromRegistry(cands), &req)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, types.ImageGenResponse{
|
||||
Created: time.Now().Unix(),
|
||||
Data: resp.ImageData,
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user