ModelRouter/internal/gateway/api.go

package gateway

import (
	"encoding/csv"
	"encoding/json"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"

	"llmsproxy/internal/config"
)

type adapterPayload struct {
	Name string `json:"name"`
	Code string `json:"code"`
}

func (g *Gateway) handleAdaptersAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/adapters")
	path = strings.Trim(path, "/")

	switch r.Method {
	case http.MethodGet:
		if path == "" {
			writeJSON(w, http.StatusOK, map[string]interface{}{"adapters": g.core.ListAdapters()})
			return
		}
		writeError(w, http.StatusNotFound, "not_found", "adapter code not exposed; edit in UI")
	case http.MethodPost:
		var p adapterPayload
		if err := json.NewDecoder(r.Body).Decode(&p); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := g.core.UploadAdapter(p.Name, p.Code); err != nil {
			writeError(w, http.StatusBadRequest, "adapter_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": p.Name})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "adapter name required")
			return
		}
		if err := g.core.RemoveAdapter(path); err != nil {
			writeError(w, http.StatusBadRequest, "adapter_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// sourcePayload mirrors config.Source for JSON web UI editing.
type sourcePayload struct {
	Name          string                 `json:"name"`
	BaseURL       string                 `json:"base_url"`
	APIKey        string                 `json:"api_key"`
	Adapter       string                 `json:"adapter"`
	Endpoint      string                 `json:"endpoint"`
	ImageEndpoint string                 `json:"image_endpoint"`
	Models        []config.Model         `json:"models"`
	Headers       map[string]string      `json:"headers"`
	Meta          map[string]interface{} `json:"meta"`
	Temperature   float64                `json:"temperature"`
	MaxTokens     int                    `json:"max_tokens"`
	MaxConcurrent int                    `json:"max_concurrent"`
	RPM           int                    `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited
}

func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/sources")
	path = strings.Trim(path, "/")

	switch r.Method {
	case http.MethodGet:
		writeJSON(w, http.StatusOK, map[string]interface{}{"sources": g.core.Sources()})
	case http.MethodPost:
		var p sourcePayload
		body, _ := io.ReadAll(r.Body)
		if err := json.Unmarshal(body, &p); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		src := config.Source{
			Name:          p.Name,
			BaseURL:       p.BaseURL,
			APIKey:        p.APIKey,
			Adapter:       p.Adapter,
			Endpoint:      p.Endpoint,
			ImageEndpoint: p.ImageEndpoint,
			Models:        p.Models,
			Headers:       p.Headers,
			Meta:          p.Meta,
			Temperature:   p.Temperature,
			MaxTokens:     p.MaxTokens,
			MaxConcurrent: p.MaxConcurrent,
			RPM:           p.RPM,
		}
		if err := g.core.AddSource(src); err != nil {
			writeError(w, http.StatusBadRequest, "source_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "source name required")
			return
		}
		if err := g.core.RemoveSource(path); err != nil {
			writeError(w, http.StatusBadRequest, "source_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// csvHeaders stamps the shared download headers for CSV exports.
func csvHeaders(w http.ResponseWriter, filename string) {
	w.Header().Set("Content-Type", "text/csv; charset=utf-8")
	w.Header().Set("Content-Disposition", "attachment; filename="+filename)
}

// handleSourceTemplatesAPI manages reusable source templates (a Source minus
// name and api_key) stored in the runtime file so the WebUI can spin up
// multiple key-bearing sources from one shared template.
func (g *Gateway) handleSourceTemplatesAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/source_templates")
	path = strings.Trim(path, "/")
	switch r.Method {
	case http.MethodGet:
		writeJSON(w, http.StatusOK, map[string]interface{}{"templates": g.core.Templates()})
	case http.MethodPost:
		var t config.SourceTemplate
		if err := json.NewDecoder(r.Body).Decode(&t); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := g.core.SaveTemplate(t); err != nil {
			writeError(w, http.StatusBadRequest, "template_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "templates": g.core.Templates()})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "template name required")
			return
		}
		if err := g.core.RemoveTemplate(path); err != nil {
			writeError(w, http.StatusBadRequest, "template_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// exportKey resolves which masked key id an export covers: admins may pass
// any key filter, user keys are always scoped to themselves (records and
// aggregates are keyed by the masked keyID form).
func exportKey(r *http.Request) string {
	k := r.URL.Query().Get("key")
	if reqRole(r.Context()) != "admin" {
		k = keyID(reqKey(r.Context()))
	}
	return k
}

// handleStatsAPI returns per-key / per-model / per-source usage aggregates and
// the recent request audit trail.
func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	limit := 500
	if v := r.URL.Query().Get("limit"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 20000 {
			limit = n
		}
	}
	key := exportKey(r)
	if r.URL.Query().Get("export") == "csv" {
		from, _ := strconv.ParseInt(r.URL.Query().Get("from"), 10, 64)
		to, _ := strconv.ParseInt(r.URL.Query().Get("to"), 10, 64)
		if to == 0 {
			to = time.Now().UnixMilli()
		}
		csvHeaders(w, "llmsproxy-requests.csv")
		cw := csv.NewWriter(w)
		names := map[string]string{}
		for _, k := range g.core.ListKeys() {
			names[keyID(k.Key)] = k.Name
		}
		_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "first_byte_ms", "cache_hit_tokens", "cache_miss_tokens", "error"})
		for _, rec := range g.stats.AuditRecords(from, to, key) {
			_ = cw.Write([]string{
				time.UnixMilli(rec.Time).Format(time.RFC3339),
				rec.Key,
				names[rec.Key],
				rec.Type,
				rec.Model,
				rec.Source,
				strconv.Itoa(rec.Status),
				strconv.FormatBool(rec.OK),
				strconv.FormatInt(rec.Prompt, 10),
				strconv.FormatInt(rec.Compl, 10),
				strconv.FormatInt(rec.LatMs, 10),
				strconv.FormatInt(rec.FirstByteMs, 10),
				strconv.FormatInt(rec.CacheHit, 10),
				strconv.FormatInt(rec.CacheMiss, 10),
				rec.Err,
			})
		}
		cw.Flush()
		return
	}
	if r.URL.Query().Get("export") == "keys-csv" {
		csvHeaders(w, "llmsproxy-keys.csv")
		cw := csv.NewWriter(w)
		_ = cw.Write([]string{"key", "key_name", "role", "models", "total_requests", "success_requests", "failed_requests", "prompt_tokens", "completion_tokens", "total_tokens", "avg_latency_ms", "max_latency_ms", "created_at"})
		keyFilter := exportKey(r)
		// by_key rows are keyed by the masked id (keyID); build a masked-id ->
		// record map so names/roles/models resolve for the export.
		byMasked := map[string]config.GWKey{}
		for _, k := range g.core.ListKeys() {
			byMasked[keyID(k.Key)] = k
		}
		snap := g.stats.Snapshot(0, keyFilter)
		byKeyRaw, ok := snap["by_key"].([]StatsRow)
		if !ok {
			byKeyRaw = []StatsRow{}
		}
		for _, row := range byKeyRaw {
			keyInfo, found := byMasked[row.Name]
			name, role, models := "", "", ""
			if found {
				name = keyInfo.Name
				role = keyInfo.Role
				modelNames := make([]string, 0, len(keyInfo.Models))
				for _, m := range keyInfo.Models {
					modelNames = append(modelNames, m.Model)
				}
				models = strings.Join(modelNames, ",")
			}
			avgLat := int64(0)
			if row.Stat.Reqs > 0 {
				avgLat = row.Stat.LatSum / row.Stat.Reqs
			}
			createdAt := ""
			if found && keyInfo.CreatedAt > 0 {
				createdAt = time.Unix(keyInfo.CreatedAt, 0).Format(time.RFC3339)
			}
			_ = cw.Write([]string{
				row.Name, name, role, models,
				strconv.FormatInt(row.Stat.Reqs, 10),
				strconv.FormatInt(row.Stat.OK, 10),
				strconv.FormatInt(row.Stat.Err, 10),
				strconv.FormatInt(row.Stat.Prompt, 10),
				strconv.FormatInt(row.Stat.Compl, 10),
				strconv.FormatInt(row.Stat.Prompt+row.Stat.Compl, 10),
				strconv.FormatInt(avgLat, 10),
				strconv.FormatInt(row.Stat.LatMax, 10),
				createdAt,
			})
		}
		cw.Flush()
		return
	}
	snap := g.stats.Snapshot(limit, key)
	keyNames := map[string]string{}
	for _, k := range g.core.ListKeys() {
		keyNames[keyID(k.Key)] = k.Name
	}
	snap["key_names"] = keyNames
	writeJSON(w, http.StatusOK, snap)
}