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"`
}

func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
	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,
		}
		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", "")
	}
}
// 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 := r.URL.Query().Get("key")
	if reqRole(r.Context()) != "admin" {
		// user keys may only see their own usage - records and aggregates are
		// keyed by the masked id (keyID), so filter on that masked form.
		key = keyID(reqKey(r.Context()))
	}
	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()
		}
		w.Header().Set("Content-Type", "text/csv; charset=utf-8")
		w.Header().Set("Content-Disposition", "attachment; filename=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", "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),
				rec.Err,
			})
		}
		cw.Flush()
		return
	}
if r.URL.Query().Get("export") == "keys-csv" {
		w.Header().Set("Content-Type", "text/csv; charset=utf-8")
		w.Header().Set("Content-Disposition", "attachment; filename=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"})
		// Use the same key filtering as the JSON API
		exportKey := r.URL.Query().Get("key")
		if reqRole(r.Context()) != "admin" {
			exportKey = keyID(reqKey(r.Context()))
		}
		// 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, exportKey)
		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, ",")
			}
			reqs := row.Stat.Reqs
			success := row.Stat.OK
			errCount := row.Stat.Err
			prompt := row.Stat.Prompt
			compl := row.Stat.Compl
			tokens := row.Stat.Prompt + row.Stat.Compl
			latSum := row.Stat.LatSum
			latMax := row.Stat.LatMax
			avgLat := int64(0)
			if reqs > 0 {
				avgLat = latSum / 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(reqs, 10),
				strconv.FormatInt(success, 10),
				strconv.FormatInt(errCount, 10),
				strconv.FormatInt(prompt, 10),
				strconv.FormatInt(compl, 10),
				strconv.FormatInt(tokens, 10),
				strconv.FormatInt(avgLat, 10),
				strconv.FormatInt(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)
}