mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
807 lines
23 KiB
Go
807 lines
23 KiB
Go
package gateway
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"llmsproxy/internal/config"
|
|
"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"`
|
|
DisableThinking bool `json:"disable_thinking"`
|
|
ExtraBody map[string]interface{} `json:"extra_body,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,omitempty"`
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
|
ToolCalls json.RawMessage `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.
|
|
// toolCalling requests are anchored: they resolve to exactly one provider
|
|
// (highest-priority available) so a tool-call round never switches models.
|
|
func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provider.Provider, string) {
|
|
model := req.Model
|
|
if model == "" {
|
|
model = g.core.DefaultModel()
|
|
}
|
|
cands, effective := g.resolveByModel(model)
|
|
cands = chatOnly(cands)
|
|
allow := g.allowedModels(ctx)
|
|
if allow != nil {
|
|
cands = filterCandsByModels(cands, allow)
|
|
}
|
|
if !toolRequest(req) {
|
|
return cands, effective
|
|
}
|
|
// tool-call request: pin to one provider (no AUTO fallback across models)
|
|
if len(cands) == 0 {
|
|
return nil, effective
|
|
}
|
|
first := cands[0]
|
|
eff := first.ModelFor(model)
|
|
if eff == "" {
|
|
eff = firstModel(first)
|
|
}
|
|
return []*provider.Provider{first}, eff
|
|
}
|
|
|
|
// filterCandsByModels keeps only providers exposing at least one model of the
|
|
// scope (used for user keys with a restricted model scope). An "AUTO" scope
|
|
// entry means the key is allowed to use any model.
|
|
func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
|
|
allowed := make(map[string]bool, len(allow))
|
|
for _, m := range allow {
|
|
if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
|
|
return cands
|
|
}
|
|
allowed[m.Model] = true
|
|
}
|
|
out := make([]*provider.Provider, 0, len(cands))
|
|
for _, p := range cands {
|
|
for _, id := range p.Models() {
|
|
if allowed[id] {
|
|
out = append(out, p)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// intersectModels restricts a model list to the scope (preserving order). An
|
|
// "AUTO" scope entry grants every model.
|
|
func intersectModels(models []string, allow []config.ModelScope) []string {
|
|
allowed := make(map[string]bool, len(allow))
|
|
any := false
|
|
for _, m := range allow {
|
|
if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
|
|
any = true
|
|
break
|
|
}
|
|
allowed[m.Model] = true
|
|
}
|
|
if any {
|
|
return models
|
|
}
|
|
out := make([]string, 0, len(models))
|
|
seen := map[string]bool{}
|
|
for _, m := range models {
|
|
if allowed[m] && !seen[m] {
|
|
seen[m] = true
|
|
out = append(out, m)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// checkModelScope validates the effective model against the key's model scope
|
|
// and token quota. Returns an error message when rejected. A scope entry with
|
|
// model "AUTO" grants all models; its quota caps the key's total tokens.
|
|
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
|
|
allow := g.allowedModels(ctx)
|
|
if allow == nil {
|
|
return ""
|
|
}
|
|
for _, sc := range allow {
|
|
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
|
if sc.TokenQuota > 0 {
|
|
used := g.scopeTokens(ctx, sc)
|
|
if used >= sc.TokenQuota {
|
|
return fmt.Sprintf("token quota exceeded (%d/%d)", used, sc.TokenQuota)
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
}
|
|
for _, sc := range allow {
|
|
if sc.Model != model {
|
|
continue
|
|
}
|
|
if sc.TokenQuota > 0 {
|
|
used := g.scopeTokens(ctx, sc)
|
|
if used >= sc.TokenQuota {
|
|
return fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota)
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
return fmt.Sprintf("model %q is not allowed for this key", model)
|
|
}
|
|
|
|
// scopeTokens returns the tokens a scope entry has consumed within its reset
|
|
// window (total for AUTO / per model otherwise).
|
|
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
|
|
k := keyID(reqKey(ctx))
|
|
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
|
return g.stats.KeyTokens(k)
|
|
}
|
|
win := AutoPeriodSeconds(sc.Period, sc.Hours)
|
|
return g.stats.WindowTokens(sc.Model, win)
|
|
}
|
|
|
|
func hasScopeModel(list []config.ModelScope, s string) bool {
|
|
for _, x := range list {
|
|
if x.Model == s || (x.Model != "" && strings.EqualFold(x.Model, "AUTO")) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
|
|
if isAuto(model) {
|
|
return g.core.Registry().Resolve("AUTO"), ""
|
|
}
|
|
return g.core.Registry().Resolve(model), model
|
|
}
|
|
|
|
// chatOnly keeps providers that expose at least one chat-capable model, so a
|
|
// chat/AUTO request never lands on an image-only source (or borrows its image
|
|
// model id). Explicit image-kind requests stay on the imageOnly path.
|
|
func chatOnly(cands []*provider.Provider) []*provider.Provider {
|
|
out := make([]*provider.Provider, 0, len(cands))
|
|
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
|
|
}
|
|
|
|
// toolRequest reports whether the request participates in a tool-call round.
|
|
func toolRequest(req *chatRequest) bool {
|
|
if len(req.Tools) > 0 || req.ToolChoice != nil {
|
|
return true
|
|
}
|
|
for _, m := range req.Messages {
|
|
if m.Role == "tool" || len(m.ToolCalls) > 0 {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
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()
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(req.Model), "AUTO") {
|
|
plans := g.autoPlans()
|
|
if len(plans) == 0 {
|
|
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot available (quota exhausted or none configured)")
|
|
return
|
|
}
|
|
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
|
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
|
return
|
|
}
|
|
ctx := r.Context()
|
|
done := g.stats.Begin()
|
|
defer done()
|
|
inner := &types.ChatRequest{
|
|
Messages: req.Messages,
|
|
Temperature: req.Temperature,
|
|
MaxTokens: req.MaxTokens,
|
|
Stream: req.Stream,
|
|
Tools: req.Tools,
|
|
ToolChoice: req.ToolChoice,
|
|
DisableThinking: req.DisableThinking,
|
|
ExtraBody: req.ExtraBody,
|
|
}
|
|
rec := &Req{
|
|
Key: keyID(reqKey(ctx)),
|
|
Type: "chat",
|
|
OK: false,
|
|
}
|
|
if req.Stream {
|
|
rec.Type = "stream"
|
|
g.streamChatAuto(w, ctx, plans, inner, rec)
|
|
return
|
|
}
|
|
g.singleChatAuto(w, ctx, plans, inner, rec)
|
|
return
|
|
}
|
|
if !isAuto(model) {
|
|
if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
|
|
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
|
|
return
|
|
}
|
|
}
|
|
cands, effective := g.resolveCands(r.Context(), &req)
|
|
if len(cands) == 0 {
|
|
writeError(w, http.StatusServiceUnavailable, "no_provider", "no LLM source configured")
|
|
return
|
|
}
|
|
if effective == "" {
|
|
effective = firstModel(cands[0])
|
|
}
|
|
if msg := g.checkModelScope(r.Context(), effective); msg != "" {
|
|
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
|
return
|
|
}
|
|
ctx := r.Context()
|
|
done := g.stats.Begin()
|
|
defer done()
|
|
|
|
inner := &types.ChatRequest{
|
|
Model: effective,
|
|
Messages: req.Messages,
|
|
Temperature: req.Temperature,
|
|
MaxTokens: req.MaxTokens,
|
|
Stream: req.Stream,
|
|
Tools: req.Tools,
|
|
ToolChoice: req.ToolChoice,
|
|
DisableThinking: req.DisableThinking,
|
|
ExtraBody: req.ExtraBody,
|
|
}
|
|
rec := &Req{
|
|
Key: keyID(reqKey(ctx)),
|
|
Type: "chat",
|
|
Model: effective,
|
|
Source: firstSource(cands),
|
|
OK: false,
|
|
}
|
|
if req.Stream {
|
|
rec.Type = "stream"
|
|
g.streamChat(w, ctx, cands, inner, effective, rec)
|
|
return
|
|
}
|
|
g.singleChat(w, ctx, cands, inner, effective, rec)
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// toolCallsWire converts unified tool calls to the OpenAI wire format:
|
|
// tool_calls:[{id,type,function:{name,arguments:StringJSON}}]. Clients expect
|
|
// arguments to be a JSON string, not an object.
|
|
func toolCallsWire(tcs []types.ToolCall) json.RawMessage {
|
|
wire := make([]map[string]interface{}, 0, len(tcs))
|
|
for _, tc := range tcs {
|
|
args := "{}"
|
|
if tc.Arguments != nil {
|
|
if b, err := json.Marshal(tc.Arguments); err == nil {
|
|
args = string(b)
|
|
}
|
|
}
|
|
wire = append(wire, map[string]interface{}{
|
|
"id": tc.ID,
|
|
"type": tc.Type,
|
|
"function": map[string]interface{}{
|
|
"name": tc.Name,
|
|
"arguments": args,
|
|
},
|
|
})
|
|
}
|
|
b, _ := json.Marshal(wire)
|
|
return b
|
|
}
|
|
|
|
func firstSource(cands []*provider.Provider) string {
|
|
if len(cands) > 0 {
|
|
return cands[0].Name()
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// effectiveImageModel picks the image model id that will actually be used so
|
|
// quota checks can target it (AUTO resolves to the first image candidate).
|
|
func effectiveImageModel(model string, cands []*provider.Provider) string {
|
|
if !isAuto(model) && model != "" {
|
|
return model
|
|
}
|
|
if len(cands) > 0 {
|
|
for _, id := range cands[0].Models() {
|
|
if m := cands[0].ModelByID(id); m != nil && m.Kind == "image" {
|
|
return id
|
|
}
|
|
}
|
|
}
|
|
return model
|
|
}
|
|
|
|
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
|
rec.LatMs = 0
|
|
t0 := time.Now()
|
|
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry(cands), req)
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
if err != nil {
|
|
rec.OK = false
|
|
rec.Status = http.StatusBadGateway
|
|
rec.Err = err.Error()
|
|
g.writeRec(rec)
|
|
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
|
return
|
|
}
|
|
rec.OK = true
|
|
rec.Status = http.StatusOK
|
|
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
|
rec.Compl = int64(resp.TokenUsage.Completion)
|
|
rec.Source = usedSrc
|
|
rec.Model = usedModel
|
|
g.writeRec(rec)
|
|
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
|
if resp.ReasoningContent != "" {
|
|
msg.ReasoningContent = resp.ReasoningContent
|
|
}
|
|
if len(resp.ToolCalls) > 0 {
|
|
msg.ToolCalls = toolCallsWire(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)
|
|
}
|
|
|
|
// writeRec records a finished request (audit + aggregates).
|
|
func (g *Gateway) writeRec(rec *Req) {
|
|
if rec == nil {
|
|
return
|
|
}
|
|
if rec.Time == 0 {
|
|
rec.Time = time.Now().UnixMilli()
|
|
}
|
|
g.stats.Record(*rec)
|
|
}
|
|
|
|
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
|
rec.LatMs = 0
|
|
t0 := time.Now()
|
|
rec.OK = true
|
|
rec.Status = http.StatusOK
|
|
defer func() {
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
g.writeRec(rec)
|
|
}()
|
|
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry(cands), req)
|
|
if err != nil {
|
|
rec.OK = false
|
|
rec.Status = http.StatusBadGateway
|
|
rec.Err = err.Error()
|
|
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
|
return
|
|
}
|
|
if usedModel != "" {
|
|
rec.Model = usedModel
|
|
}
|
|
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{Role: "assistant", 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}
|
|
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)) / 3
|
|
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()
|
|
}
|
|
}
|
|
|
|
// autoPlan is one schedulable AUTO slot: a model id pinned to its provider
|
|
// with an optional token quota window. Quota-exhausted slots are skipped.
|
|
type autoPlan struct {
|
|
p *provider.Provider
|
|
model string
|
|
quota int64
|
|
win int64
|
|
}
|
|
|
|
// autoPlans builds the schedulable AUTO slots from the persisted rules. A
|
|
// slot is schedulable while its model is available and (when quota > 0) the
|
|
// tokens used within its reset window are below the quota.
|
|
func (g *Gateway) autoPlans() []autoPlan {
|
|
rules := g.core.AutoRules()
|
|
if len(rules) == 0 {
|
|
return nil
|
|
}
|
|
plans := make([]autoPlan, 0, len(rules))
|
|
for _, e := range rules {
|
|
p := g.core.ProviderForModel(e.Model)
|
|
if p == nil {
|
|
continue
|
|
}
|
|
if m := p.ModelByID(e.Model); m != nil && m.Kind == "image" {
|
|
continue
|
|
}
|
|
win := AutoPeriodSeconds(e.Period, e.Hours)
|
|
if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, win) >= e.TokenQuota {
|
|
continue
|
|
}
|
|
plans = append(plans, autoPlan{p: p, model: e.Model, quota: e.TokenQuota, win: win})
|
|
}
|
|
return plans
|
|
}
|
|
|
|
// singleChatAuto runs a non-streaming AUTO request slot by slot: each slot
|
|
// pins its own model; a slot whose provider errors out is skipped. The first
|
|
// slot to answer wins; when every slot fails, the recorded error is from the
|
|
// last one.
|
|
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
|
|
rec.LatMs = 0
|
|
t0 := time.Now()
|
|
var lastErr error
|
|
for _, pl := range plans {
|
|
if !pl.p.Available() {
|
|
continue
|
|
}
|
|
r := *req
|
|
r.Model = pl.model
|
|
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
|
|
if err != nil {
|
|
lastErr = err
|
|
continue
|
|
}
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
rec.OK = true
|
|
rec.Status = http.StatusOK
|
|
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
|
rec.Compl = int64(resp.TokenUsage.Completion)
|
|
rec.Source = usedSrc
|
|
rec.Model = usedModel
|
|
g.writeRec(rec)
|
|
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
|
if resp.ReasoningContent != "" {
|
|
msg.ReasoningContent = resp.ReasoningContent
|
|
}
|
|
if len(resp.ToolCalls) > 0 {
|
|
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
|
|
}
|
|
out := ChatCompletion{
|
|
ID: newID(),
|
|
Object: "chat.completion",
|
|
Created: time.Now().Unix(),
|
|
Model: usedModel,
|
|
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)
|
|
return
|
|
}
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
if lastErr == nil {
|
|
lastErr = fmt.Errorf("no provider available")
|
|
}
|
|
rec.OK = false
|
|
rec.Status = http.StatusBadGateway
|
|
rec.Err = lastErr.Error()
|
|
g.writeRec(rec)
|
|
writeError(w, http.StatusBadGateway, "upstream_error", lastErr.Error())
|
|
}
|
|
|
|
// streamChatAuto streams an AUTO request. It stays pinned to the first slot
|
|
// whose stream begins; a slot that fails to connect is skipped.
|
|
func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
|
|
rec.LatMs = 0
|
|
t0 := time.Now()
|
|
rec.OK = true
|
|
rec.Status = http.StatusOK
|
|
defer func() {
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
g.writeRec(rec)
|
|
}()
|
|
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: "auto",
|
|
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
|
}) {
|
|
return
|
|
}
|
|
var lastErr error
|
|
for _, pl := range plans {
|
|
if !pl.p.Available() {
|
|
continue
|
|
}
|
|
r := *req
|
|
r.Model = pl.model
|
|
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
|
|
if err != nil {
|
|
lastErr = err
|
|
continue
|
|
}
|
|
if usedModel != "" {
|
|
rec.Model = usedModel
|
|
rec.Source = usedSrc
|
|
}
|
|
for ck := range chunks {
|
|
chunk := ChatChunk{
|
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
|
}
|
|
delta := RespMessage{Role: "assistant", 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}
|
|
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)) / 3
|
|
if !send(chunk) {
|
|
return
|
|
}
|
|
}
|
|
return
|
|
}
|
|
if lastErr == nil {
|
|
lastErr = fmt.Errorf("no provider available")
|
|
}
|
|
rec.OK = false
|
|
rec.Status = http.StatusBadGateway
|
|
rec.Err = lastErr.Error()
|
|
stop := "stop"
|
|
send(ChatChunk{
|
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: "auto",
|
|
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()
|
|
}
|
|
if !isAuto(model) {
|
|
if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
|
|
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
|
|
return
|
|
}
|
|
}
|
|
cands, _ := g.resolveByModel(model)
|
|
cands = imageOnly(cands)
|
|
if allow := g.allowedModels(r.Context()); allow != nil {
|
|
cands = filterCandsByModels(cands, allow)
|
|
}
|
|
if len(cands) == 0 {
|
|
writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
|
|
return
|
|
}
|
|
if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
|
|
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
|
return
|
|
}
|
|
done := g.stats.Begin()
|
|
defer done()
|
|
rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, Source: firstSource(cands), OK: false}
|
|
t0 := time.Now()
|
|
resp, usedSrc, err := g.core.Scheduler().Image(r.Context(), scheduler.FromRegistry(cands), &req)
|
|
rec.LatMs = time.Since(t0).Milliseconds()
|
|
if err != nil {
|
|
rec.Status = http.StatusBadGateway
|
|
rec.Err = err.Error()
|
|
g.writeRec(rec)
|
|
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
|
|
return
|
|
}
|
|
if usedSrc != "" {
|
|
rec.Source = usedSrc
|
|
}
|
|
rec.OK = true
|
|
rec.Status = http.StatusOK
|
|
rec.Compl = int64(len(resp.ImageData))
|
|
g.writeRec(rec)
|
|
writeJSON(w, http.StatusOK, types.ImageGenResponse{
|
|
Created: time.Now().Unix(),
|
|
Data: resp.ImageData,
|
|
})
|
|
} |