llmgw/internal/router/files.go

package router

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

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

// The Files API, as far as batches need it: a client uploads a file of
// requests, and reads back the files a batch writes.

const (
	purposeBatch       = "batch"
	purposeBatchOutput = "batch_output"
)

func newID(prefix string) string {
	var b [12]byte
	_, _ = rand.Read(b[:])
	return prefix + hex.EncodeToString(b[:])
}

// ownerOf is who a file or a batch belongs to: the key's project, so every
// key of it sees them, or the person for a personal key.
func ownerOf(p auth.Principal) string {
	switch {
	case p.ProjectID != "":
		return "project:" + p.ProjectID
	case p.UserID != "":
		return "user:" + p.UserID
	}
	return "key:" + p.KeyID
}

// batchPrincipal is the caller, or a refusal when there is no store to keep
// files and batches in.
func (r *Router) batchPrincipal(w http.ResponseWriter, req *http.Request) (auth.Principal, bool) {
	p, ok := auth.FromContext(req.Context())
	if !ok || r.store == nil {
		writeError(w, http.StatusNotImplemented, "invalid_request_error",
			"files and batches need authentication and the store, which this gateway runs without")
		return auth.Principal{}, false
	}
	// A batch's files are prompts and answers kept: not under a policy
	// that keeps nothing.
	if r.noStore(guard.Subject{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID,
		TeamID: p.TeamID, OrgID: p.OrgID, Path: req.URL.Path}) {
		writeError(w, http.StatusForbidden, "invalid_request_error",
			"the gateway's policy for this caller keeps nothing, and a batch is prompts and answers kept")
		return auth.Principal{}, false
	}
	return p, true
}

func (r *Router) batchCfg() (maxBytes int64, maxLines, concurrency, retainDays int) {
	b := r.cfg.Batch
	maxBytes, maxLines, concurrency, retainDays = b.MaxFileBytes, b.MaxLines, b.Concurrency, b.RetainDays
	if maxBytes <= 0 {
		maxBytes = 100 << 20
	}
	if maxLines <= 0 {
		maxLines = 50000
	}
	if concurrency <= 0 {
		concurrency = 4
	}
	if retainDays <= 0 {
		retainDays = 30
	}
	return
}

func fileJSON(f store.File) map[string]any {
	out := map[string]any{"id": f.ID, "object": "file", "bytes": f.Bytes, "created_at": f.CreatedAt.Unix(),
		"filename": f.Filename, "purpose": f.Purpose, "status": "processed"}
	if !f.ExpiresAt.IsZero() {
		out["expires_at"] = f.ExpiresAt.Unix()
	}
	return out
}

func (r *Router) handleUploadFile(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	maxBytes, _, _, _ := r.batchCfg()
	req.Body = http.MaxBytesReader(w, req.Body, maxBytes+1<<20)
	if err := req.ParseMultipartForm(32 << 20); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			fmt.Sprintf("a multipart form with purpose and file, at most %d MB: %v", maxBytes>>20, err))
		return
	}
	if purpose := req.FormValue("purpose"); purpose != purposeBatch {
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			fmt.Sprintf("purpose %q: this gateway keeps files for batches only (purpose batch)", purpose))
		return
	}
	file, header, err := req.FormFile("file")
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "no file in the form: "+err.Error())
		return
	}
	defer file.Close()
	content, err := io.ReadAll(io.LimitReader(file, maxBytes+1))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	if int64(len(content)) > maxBytes {
		writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error",
			fmt.Sprintf("the file is larger than %d MB", maxBytes>>20))
		return
	}
	f := store.File{ID: newID("file-"), Owner: ownerOf(p), Purpose: purposeBatch,
		Filename: header.Filename, Bytes: int64(len(content)), CreatedAt: time.Now()}
	if err := r.store.PutFile(req.Context(), f, content); err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	writeJSON(w, http.StatusOK, fileJSON(f))
}

// ownedFile is a file the caller owns, or a 404 said.
func (r *Router) ownedFile(w http.ResponseWriter, req *http.Request, p auth.Principal) (store.File, bool) {
	f, err := r.store.FileByID(req.Context(), req.PathValue("id"))
	if errors.Is(err, store.ErrNotFound) || err == nil && f.Owner != ownerOf(p) {
		writeErrorCode(w, http.StatusNotFound, "invalid_request_error", "file_not_found",
			fmt.Sprintf("no file %q", req.PathValue("id")))
		return store.File{}, false
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return store.File{}, false
	}
	return f, true
}

func (r *Router) handleListFiles(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	fs, err := r.store.Files(req.Context(), ownerOf(p), req.URL.Query().Get("purpose"), 0)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	data := make([]map[string]any, 0, len(fs))
	for _, f := range fs {
		data = append(data, fileJSON(f))
	}
	writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data})
}

func (r *Router) handleGetFile(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	if f, ok := r.ownedFile(w, req, p); ok {
		writeJSON(w, http.StatusOK, fileJSON(f))
	}
}

func (r *Router) handleFileContent(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	f, ok := r.ownedFile(w, req, p)
	if !ok {
		return
	}
	b, err := r.store.FileContent(req.Context(), f.ID)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", f.Filename))
	_, _ = w.Write(b)
}

func (r *Router) handleDeleteFile(w http.ResponseWriter, req *http.Request) {
	p, ok := r.batchPrincipal(w, req)
	if !ok {
		return
	}
	f, ok := r.ownedFile(w, req, p)
	if !ok {
		return
	}
	if _, err := r.store.DeleteFile(req.Context(), f.ID); err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"id": f.ID, "object": "file", "deleted": true})
}