ModelRouter/internal/gateway/chat.go

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"`
	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(req *chatRequest) ([]*provider.Provider, string) {
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	cands, effective := g.resolveByModel(model)
	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
}

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
}

// 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()
	}
	cands, effective := g.resolveCands(&req)
	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:           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,
	}
	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
}

// 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 (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 = 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)
}

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{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}
		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.resolveByModel(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,
	})
}