llmgw/internal/router/batches.go

package router

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// The Batch API: a file of requests, served one by one at the lowest
// priority, its answers written to another file. OpenAI's shape, so its
// clients work unchanged.

// batchEndpoints are what a batch may be of.
var batchEndpoints = map[string]bool{
	"/v1/chat/completions": true,
	"/v1/completions":      true,
	"/v1/embeddings":       true,
}

// batchWindow is the only completion window, as OpenAI has it.
const batchWindow = "24h"

func unixOrNil(t time.Time) any {
	if t.IsZero() {
		return nil
	}
	return t.Unix()
}

func strOrNil(s string) any {
	if s == "" {
		return nil
	}
	return s
}

func batchJSON(b store.Batch) map[string]any {
	out := map[string]any{
		"id": b.ID, "object": "batch", "endpoint": b.Endpoint, "input_file_id": b.InputFileID,
		"completion_window": b.CompletionWindow, "status": b.Status,
		"output_file_id": strOrNil(b.OutputFileID), "error_file_id": strOrNil(b.ErrorFileID),
		"created_at": b.CreatedAt.Unix(), "in_progress_at": unixOrNil(b.InProgressAt),
		"expires_at": unixOrNil(b.ExpiresAt), "finalizing_at": unixOrNil(b.FinalizingAt),
		"completed_at": unixOrNil(b.CompletedAt), "failed_at": unixOrNil(b.FailedAt),
		"expired_at": unixOrNil(b.ExpiredAt), "cancelling_at": unixOrNil(b.CancellingAt),
		"cancelled_at":   unixOrNil(b.CancelledAt),
		"request_counts": map[string]int{"total": b.Total, "completed": b.Completed, "failed": b.Failed},
		"errors":         nil, "metadata": nil,
	}
	if b.Errors != "" {
		out["errors"] = json.RawMessage(b.Errors)
	}
	if b.Metadata != "" {
		out["metadata"] = json.RawMessage(b.Metadata)
	}
	return out
}

func (r *Router) handleCreateBatch(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	var in struct {
		InputFileID      string          `json:"input_file_id"`
		Endpoint         string          `json:"endpoint"`
		CompletionWindow string          `json:"completion_window"`
		Metadata         json.RawMessage `json:"metadata"`
	}
	if err := json.NewDecoder(io.LimitReader(req.Body, 1<<16)).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "the request is not valid JSON: "+err.Error())
		return
	}
	if !batchEndpoints[in.Endpoint] {
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			fmt.Sprintf("endpoint %q: a batch is of /v1/chat/completions, /v1/completions or /v1/embeddings", in.Endpoint))
		return
	}
	if in.CompletionWindow != batchWindow {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "completion_window is 24h")
		return
	}
	meta := ""
	if len(in.Metadata) > 0 && string(in.Metadata) != "null" {
		if _, err := requestTags(in.Metadata, ""); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
			return
		}
		meta = string(in.Metadata)
	}
	f, err := r.store.FileByID(req.Context(), in.InputFileID)
	if errors.Is(err, store.ErrNotFound) || err == nil && (f.Owner != ownerOf(p) || f.Purpose != purposeBatch) {
		writeErrorCode(w, http.StatusBadRequest, "invalid_request_error", "file_not_found",
			fmt.Sprintf("no batch file %q", in.InputFileID))
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	now := time.Now()
	b := store.Batch{ID: newID("batch_"), Owner: ownerOf(p), KeyID: p.KeyID, Endpoint: in.Endpoint,
		InputFileID: f.ID, CompletionWindow: batchWindow, Status: store.BatchValidating,
		CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour), Metadata: meta}
	if err := r.store.PutBatch(req.Context(), b); err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	r.wakeBatches()
	writeJSON(w, http.StatusOK, batchJSON(b))
}

// ownedBatch is a batch the caller owns, or a 404 said.
func (r *Router) ownedBatch(w http.ResponseWriter, req *http.Request, owner string) (store.Batch, bool) {
	b, err := r.store.BatchByID(req.Context(), req.PathValue("id"))
	if errors.Is(err, store.ErrNotFound) || err == nil && b.Owner != owner {
		writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "batch_not_found",
			fmt.Sprintf("no batch %q", req.PathValue("id")))
		return store.Batch{}, false
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return store.Batch{}, false
	}
	return b, true
}

func (r *Router) handleGetBatch(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	if b, ok := r.ownedBatch(w, req, ownerOf(p)); ok {
		writeJSON(w, http.StatusOK, batchJSON(b))
	}
}

func (r *Router) handleCancelBatch(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	b, ok := r.ownedBatch(w, req, ownerOf(p))
	if !ok {
		return
	}
	b, err := r.store.CancelBatch(req.Context(), b.ID, time.Now())
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	r.wakeBatches()
	writeJSON(w, http.StatusOK, batchJSON(b))
}

func (r *Router) handleListBatches(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	limit, _ := strconv.Atoi(req.URL.Query().Get("limit"))
	if limit <= 0 || limit > 100 {
		limit = 20
	}
	bs, err := r.store.Batches(req.Context(), ownerOf(p), req.URL.Query().Get("after"), limit+1)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	more := len(bs) > limit
	if more {
		bs = bs[:limit]
	}
	data := make([]map[string]any, 0, len(bs))
	for _, b := range bs {
		data = append(data, batchJSON(b))
	}
	out := map[string]any{"object": "list", "data": data, "has_more": more, "first_id": nil, "last_id": nil}
	if len(bs) > 0 {
		out["first_id"], out["last_id"] = bs[0].ID, bs[len(bs)-1].ID
	}
	writeJSON(w, http.StatusOK, out)
}

// batchLine is one line of a batch's input file.
type batchLine struct {
	CustomID string          `json:"custom_id"`
	Method   string          `json:"method"`
	URL      string          `json:"url"`
	Body     json.RawMessage `json:"body"`
}

// lineError is one reason a batch's input was refused, as OpenAI lists them.
type lineError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
	Line    int    `json:"line,omitempty"`
}

// parseBatchInput reads a batch's input file, or says every line that is
// wrong with it.
func parseBatchInput(content []byte, endpoint string, maxLines int) ([]batchLine, []lineError) {
	var lines []batchLine
	var errs []lineError
	seen := map[string]bool{}
	n := 0
	for raw := range bytes.SplitSeq(content, []byte("\n")) {
		n++
		raw = bytes.TrimSpace(raw)
		if len(raw) == 0 {
			continue
		}
		bad := func(code, msg string) {
			if len(errs) < 100 {
				errs = append(errs, lineError{Code: code, Message: msg, Line: n})
			}
		}
		var l batchLine
		if err := json.Unmarshal(raw, &l); err != nil {
			bad("invalid_json_line", "not a JSON object: "+err.Error())
			continue
		}
		var body map[string]any
		switch {
		case l.CustomID == "":
			bad("missing_custom_id", "custom_id is required")
		case seen[l.CustomID]:
			bad("duplicate_custom_id", fmt.Sprintf("custom_id %q is used twice", l.CustomID))
		case l.Method != http.MethodPost:
			bad("invalid_method", "method is POST")
		case l.URL != endpoint:
			bad("mismatched_endpoint", fmt.Sprintf("url %q is not the batch's endpoint, %s", l.URL, endpoint))
		case json.Unmarshal(l.Body, &body) != nil || body["model"] == nil:
			bad("invalid_body", "body is an object naming a model")
		default:
			seen[l.CustomID] = true
			lines = append(lines, l)
		}
	}
	if len(lines)+len(errs) == 0 {
		errs = append(errs, lineError{Code: "empty_file", Message: "the file has no requests"})
	}
	if len(lines) > maxLines {
		errs = append(errs, lineError{Code: "too_many_requests",
			Message: fmt.Sprintf("%d requests; a batch holds at most %d", len(lines), maxLines)})
	}
	return lines, errs
}