llmgw/internal/sim/server.go

package sim

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

// Server wraps an Engine in an OpenAI-compatible HTTP surface.
type Server struct {
	eng *Engine
	log *slog.Logger
	mux *http.ServeMux
}

// NewServer builds the HTTP surface for an engine.
func NewServer(eng *Engine, log *slog.Logger) *Server {
	if log == nil {
		log = slog.Default()
	}
	s := &Server{eng: eng, log: log, mux: http.NewServeMux()}
	s.mux.HandleFunc("POST /v1/chat/completions", s.handleChat)
	s.mux.HandleFunc("POST /v1/completions", s.handleCompletions)
	s.mux.HandleFunc("GET /v1/models", s.handleModels)
	s.mux.HandleFunc("GET /health", s.handleHealth)
	s.mux.HandleFunc("GET /metrics", s.handleMetrics)
	return s
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.mux.ServeHTTP(w, r) }

// --- request/response types (only the fields we need) ---

type streamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

type chatMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

type chatRequest struct {
	Model               string         `json:"model"`
	Messages            []chatMessage  `json:"messages"`
	MaxTokens           int            `json:"max_tokens"`
	MaxCompletionTokens int            `json:"max_completion_tokens"`
	Stream              bool           `json:"stream"`
	StreamOptions       *streamOptions `json:"stream_options"`
}

type completionRequest struct {
	Model         string          `json:"model"`
	Prompt        json.RawMessage `json:"prompt"`
	MaxTokens     int             `json:"max_tokens"`
	Stream        bool            `json:"stream"`
	StreamOptions *streamOptions  `json:"stream_options"`
}

type usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
	// OpenAI-compatible cache reporting; surfaces the prefix cache to clients.
	PromptTokensDetails struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
}

// text reduces a message's content to plain text.
// Content may be a string or an array of parts; the simulation coerces both
// crudely, since all it needs is a deterministic token count.
func (m chatMessage) text() string {
	var s string
	if err := json.Unmarshal(m.Content, &s); err == nil {
		return s
	}
	return string(m.Content)
}

func promptFromMessages(msgs []chatMessage) string {
	var b strings.Builder
	for _, m := range msgs {
		b.WriteString(m.Role)
		b.WriteString(": ")
		b.WriteString(m.text())
		b.WriteString("\n")
	}
	return b.String()
}

// --- handlers ---

func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
	var req chatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	maxTok := req.MaxTokens
	if maxTok == 0 {
		maxTok = req.MaxCompletionTokens
	}
	includeUsage := req.StreamOptions != nil && req.StreamOptions.IncludeUsage

	s.serve(w, r, servePlan{
		prompt:       Prompt{Text: promptFromMessages(req.Messages), MaxTokens: maxTok},
		model:        s.modelName(req.Model),
		stream:       req.Stream,
		includeUsage: includeUsage,
		object:       "chat.completion",
		chunkObject:  "chat.completion.chunk",
		chat:         true,
	})
}

func (s *Server) handleCompletions(w http.ResponseWriter, r *http.Request) {
	var req completionRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	var prompt string
	if err := json.Unmarshal(req.Prompt, &prompt); err != nil {
		prompt = string(req.Prompt)
	}
	includeUsage := req.StreamOptions != nil && req.StreamOptions.IncludeUsage

	s.serve(w, r, servePlan{
		prompt:       Prompt{Text: prompt, MaxTokens: req.MaxTokens},
		model:        s.modelName(req.Model),
		stream:       req.Stream,
		includeUsage: includeUsage,
		object:       "text_completion",
		chunkObject:  "text_completion",
		chat:         false,
	})
}

func (s *Server) modelName(requested string) string {
	if requested != "" {
		return requested
	}
	return s.eng.Config().Model
}

func (s *Server) handleModels(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{
		"object": "list",
		"data": []map[string]any{{
			"id":       s.eng.Config().Model,
			"object":   "model",
			"created":  time.Now().Unix(),
			"owned_by": "llmgw-sim",
		}},
	})
}

func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
	w.WriteHeader(http.StatusOK)
}

func (s *Server) handleMetrics(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
	s.eng.WriteMetrics(w)
}

// --- shared request flow ---

type servePlan struct {
	prompt       Prompt
	model        string
	stream       bool
	includeUsage bool
	object       string
	chunkObject  string
	chat         bool
}

func (s *Server) serve(w http.ResponseWriter, r *http.Request, p servePlan) {
	ch, err := s.eng.Submit(r.Context(), p.prompt)
	if err != nil {
		code := http.StatusInternalServerError
		kind := "server_error"
		if err == ErrTooLong {
			code, kind = http.StatusBadRequest, "invalid_request_error"
		}
		writeError(w, code, kind, err.Error())
		return
	}

	id := fmt.Sprintf("cmpl-sim-%d", time.Now().UnixNano())
	created := time.Now().Unix()

	if p.stream {
		s.streamResponse(w, r, ch, p, id, created)
		return
	}
	s.bufferedResponse(w, ch, p, id, created)
}

func (s *Server) bufferedResponse(w http.ResponseWriter, ch <-chan Event, p servePlan, id string, created int64) {
	var b strings.Builder
	n := 0
	var res *Result
	for ev := range ch {
		switch ev.Kind {
		case EventToken:
			b.WriteString(ev.Text)
			n++
		case EventError:
			writeError(w, http.StatusInternalServerError, "server_error", ev.Err.Error())
			return
		case EventDone:
			res = ev.Result
		}
	}

	u := s.usage(p.prompt, n, res)
	body := map[string]any{
		"id": id, "object": p.object, "created": created, "model": p.model,
		"usage": u,
	}
	if p.chat {
		body["choices"] = []map[string]any{{
			"index":         0,
			"message":       map[string]string{"role": "assistant", "content": b.String()},
			"finish_reason": "length",
		}}
	} else {
		body["choices"] = []map[string]any{{
			"index": 0, "text": b.String(), "finish_reason": "length",
		}}
	}
	writeJSON(w, http.StatusOK, body)
}

func (s *Server) streamResponse(w http.ResponseWriter, r *http.Request, ch <-chan Event, p servePlan, id string, created int64) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeError(w, http.StatusInternalServerError, "server_error", "streaming not supported")
		return
	}
	h := w.Header()
	h.Set("Content-Type", "text/event-stream")
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)
	flusher.Flush()

	n := 0
	var res *Result
	for {
		select {
		case <-r.Context().Done():
			// Client disconnected; the engine's reap() drops it from the batch.
			return
		case ev, open := <-ch:
			if !open {
				if p.includeUsage {
					s.writeChunk(w, chunk(id, p, created, nil, "", s.usage(p.prompt, n, res)))
					flusher.Flush()
				}
				fmt.Fprint(w, "data: [DONE]\n\n")
				flusher.Flush()
				return
			}
			switch ev.Kind {
			case EventToken:
				n++
				s.writeChunk(w, chunk(id, p, created, &ev.Text, "", nil))
				flusher.Flush()
			case EventDone:
				res = ev.Result
				s.writeChunk(w, chunk(id, p, created, nil, "length", nil))
				flusher.Flush()
			case EventError:
				s.writeChunk(w, map[string]any{"error": map[string]any{
					"message": ev.Err.Error(), "type": "server_error",
				}})
				flusher.Flush()
				return
			}
		}
	}
}

func (s *Server) writeChunk(w http.ResponseWriter, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		s.log.Error("failed to marshal chunk", "err", err)
		return
	}
	fmt.Fprintf(w, "data: %s\n\n", b)
}

func chunk(id string, p servePlan, created int64, text *string, finish string, u any) map[string]any {
	choice := map[string]any{"index": 0}
	if p.chat {
		delta := map[string]any{}
		if text != nil {
			delta["content"] = *text
		}
		choice["delta"] = delta
	} else if text != nil {
		choice["text"] = *text
	} else {
		choice["text"] = ""
	}
	if finish != "" {
		choice["finish_reason"] = finish
	} else {
		choice["finish_reason"] = nil
	}

	out := map[string]any{
		"id": id, "object": p.chunkObject, "created": created, "model": p.model,
	}
	// The usage chunk carries an empty choices array (OpenAI behaviour).
	if u != nil {
		out["choices"] = []map[string]any{}
		out["usage"] = u
	} else {
		out["choices"] = []map[string]any{choice}
	}
	return out
}

// usage prefers the counts the engine reported. res may be nil (cancelled or
// failed request), in which case prompt tokens are derived from the text and
// cached_tokens is unknown.
func (s *Server) usage(p Prompt, completion int, res *Result) usage {
	var u usage
	u.PromptTokens = CountTokens(p.Text)
	u.CompletionTokens = completion
	if res != nil {
		u.PromptTokens = res.PromptTokens
		u.CompletionTokens = res.CompletionTokens
		u.PromptTokensDetails.CachedTokens = res.CachedTokens
	}
	u.TotalTokens = u.PromptTokens + u.CompletionTokens
	return u
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, code int, kind, msg string) {
	writeJSON(w, code, map[string]any{"error": map[string]any{
		"message": msg, "type": kind, "code": nil,
	}})
}