ModelRouter/internal/gateway/discover.go

package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
	"time"

	"llmsproxy/internal/config"
)

// Model discovery: turn "here is a URL and a key" into a configured source.
//
// The hard part is not the GET /models call — it is guessing which URL the
// upstream actually serves. This gateway builds request paths as
// base_url + adapter endpoint (openai contributes "/chat/completions"), so a
// base_url that already ends in "/v1" yields "/v1/chat/completions" while a
// base_url of "/v1" plus our default image path yields the broken
// "/v1/v1/images/generations". Vendors disagree on whether the version segment
// belongs to the base or to the path, and there is no header that says which.
//
// So discovery does not trust the URL it is given. It tries a small, ordered
// set of candidates, records which one answered, and reports that back as the
// normalized base_url + adapter endpoint pair. A caller who typed the wrong
// form gets a working source anyway, plus the correction needed to make it
// permanent.
//
// What is deliberately NOT here:
//
//   - No guessing of model capabilities from the name. "vision" in an id does
//     not make a model accept images, and marking an image-only endpoint as
//     kind:"chat" makes every AUTO request fail on it. Kind stays "chat" unless
//     the upstream's own metadata says otherwise (see discoveredModelKind).
//   - No live probe per model. A single-token completion per model would cost
//     real money across a 34-model catalog and takes minutes. /models is the
//     cheap, authoritative-enough signal for "this id is served here"; whether
//     it actually answers stays for the normal request path to discover.
//
// The one thing that IS verified live is reachability, because a source that
// 404s is not worth writing to disk.

const (
	// discoveryTimeout bounds the whole probe. Vendors are routinely slow on
	// the first request to an uncached endpoint; 20s is enough for a healthy
	// one and short enough that the UI is not left hanging.
	discoveryTimeout = 20 * time.Second

	// discoveryModelsPath is the OpenAI-compatible listing endpoint. Nearly
	// every vendor proxies it verbatim, including the ones whose chat path
	// differs.
	discoveryModelsPath = "/v1/models"

	// discoveryMaxBody caps the listing response. A catalog with thousands of
	// models is still far under this; anything larger is a captive portal or
	// an error page, not a model list.
	discoveryMaxBody = 8 << 20
)

// discoveredModel is one entry of a parsed /models response, kept in the
// upstream's own terms. The upstream is the only authority on what it serves.
type discoveredModel struct {
	ID   string
	Kind string // "chat" | "image"
	Meta map[string]interface{}
}

// discoveryResult reports what a probe found, including the case where nothing
// worked — the failure text is what the operator needs, so it must name every
// URL that was tried rather than just the last one.
type discoveryResult struct {
	BaseURL     string   `json:"base_url"`
	Endpoint    string   `json:"endpoint"`
	Adapter     string   `json:"adapter"`
	ModelsPath  string   `json:"models_path"`
	Models      []string `json:"models"`
	ImageModels []string `json:"image_models"`
	Tried       []string `json:"tried"`
	LatencyMS   int64    `json:"latency_ms"`
	Note        string   `json:"note,omitempty"`
}

// MarshalJSON guarantees the two list fields are always JSON arrays, never
// null. The WebUI concatenates them directly, and "null" renders as the string
// "undefined" in every row — a failure that only shows up in a browser and not
// in a JSON schema check.
//
// Relying on the constructor to initialize the slices is not enough: the zero
// value of this struct is a legitimate thing for a caller to build, and a
// future code path that forgets the initializer would ship a UI bug silently.
func (r discoveryResult) MarshalJSON() ([]byte, error) {
	type alias discoveryResult // avoid recursion
	out := alias(r)
	if out.Models == nil {
		out.Models = []string{}
	}
	if out.ImageModels == nil {
		out.ImageModels = []string{}
	}
	if out.Tried == nil {
		out.Tried = []string{}
	}
	return json.Marshal(out)
}

// discoveryCandidate is one (base_url, models-path) guess. Order matters: the
// first candidate is the one reported as normalized when several work, because
// it is the one that keeps base_url free of a version segment and therefore
// composes correctly with the adapter's chat path.
type discoveryCandidate struct {
	baseURL    string
	modelsPath string
	// prefer ranks candidates that are more likely to be the vendor's real
	// path. A 200 from a lower-preference candidate is still accepted, but
	// only after every higher-preference one failed — see DiscoverModels.
	prefer int
}

// discoveryCandidates enumerates the URL forms worth trying for a raw input.
//
// The ranking encodes a real observation rather than a guess: the amr-link
// endpoint serves its catalog at "/v1/models" (the form an SDK example shows)
// from a base the operator pasted WITH the "/v1", while several other vendors
// mount it at the host root. Both are probed; the one that answers is reported
// back as a normalized base/endpoint pair.
//
// What we are really deciding is the base_url to persist. The safest one keeps
// the version segment on the ENDPOINT, because openai's adapter endpoint is
// "/chat/completions" and the two are concatenated: base "https://h" +
// endpoint "/v1/chat/completions" is correct, while leaving "/v1" on the base
// doubles the default image path into "/v1/v1/images/generations" — which is
// exactly the 404 observed on the first vela configuration attempt.
func discoveryCandidates(raw string) []discoveryCandidate {
	raw = strings.TrimRight(strings.TrimSpace(raw), "/")
	if raw == "" {
		return nil
	}
	seen := map[string]bool{}
	var out []discoveryCandidate
	add := func(base, path string, prefer int) {
		if seen[base+"|"+path] {
			return
		}
		seen[base+"|"+path] = true
		out = append(out, discoveryCandidate{baseURL: base, modelsPath: path, prefer: prefer})
	}

	stripped := trimVersionSuffix(raw)
	hasVersion := stripped != raw

	// Preference 0: the URL exactly as typed, catalog at the canonical path.
	// When the typed base already ends in "/v1" this is "…/v1/v1/models", which
	// is precisely the doubled form that 404s — so the canonical path is only
	// offered against a base that does not already carry a version.
	if !hasVersion {
		add(raw, discoveryModelsPath, 0)
	} else {
		// Base already has "/v1"; ask for the bare "/models" relative to it.
		add(raw, "/models", 0)
	}
	// Preference 1: host root + canonical versioned path.
	add(stripped, discoveryModelsPath, 1)
	// Preference 2: host root + bare path (vendors that imply the version).
	if hasVersion {
		add(stripped, "/models", 2)
	}

	// Stable insertion sort by preference; add() above is not ordered.
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].prefer < out[j-1].prefer; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}

// trimVersionSuffix removes a trailing "/v1" (or /v2, /v1beta, ...) so the
// host root can be probed on its own.
func trimVersionSuffix(base string) string {
	i := strings.LastIndex(base, "/")
	if i < 0 {
		return base
	}
	if !isVersionSegment(base[i+1:]) {
		return base
	}
	return base[:i]
}

// versionSegment returns the leading path segment of a models path when that
// segment is a version marker ("/v1/models" → "/v1", "/models" → ""). It is the
// piece that must move between base_url and endpoint so the two compose.
func versionSegment(modelsPath string) string {
	parts := strings.Split(strings.Trim(modelsPath, "/"), "/")
	if len(parts) != 2 || !isVersionSegment(parts[0]) {
		return ""
	}
	return "/" + parts[0]
}

// isVersionSegment reports whether a path segment looks like "v1", "v2",
// "v1beta" — and not like a real path name such as "vision" or a wildcard
// path such as "v1x".
//
// The rule is: a leading "v", then digits, then an OPTIONAL suffix drawn from
// the prerelease markers vendors actually use. A free-form alphabetic tail was
// tried and rejected — it classifies "v1x" as a version, and that is how a
// templated path segment would get stripped off a real base URL.
func isVersionSegment(seg string) bool {
	if len(seg) < 2 || seg[0] != 'v' {
		return false
	}
	i := 1
	digits := 0
	for ; i < len(seg); i++ {
		if seg[i] < '0' || seg[i] > '9' {
			break
		}
		digits++
	}
	if digits == 0 {
		return false
	}
	suffix := strings.ToLower(seg[i:])
	if suffix == "" {
		return true
	}
	for _, marker := range []string{"alpha", "beta", "preview", "dev", "rc"} {
		if suffix == marker || strings.HasPrefix(suffix, "-"+marker) {
			return true
		}
	}
	return false
}

// composeEndpoint derives the chat endpoint to persist, given the models path
// that actually answered, the base_url being persisted, and the base the
// operator typed.
//
// The rule is a single invariant: exactly one version segment must end up in
// base_url + endpoint. The adapter contributes "/chat/completions" and the
// default image path contributes "/v1/images/generations", so a version left on
// BOTH breaks one of them — and one left on NEITHER sends chat to the wrong
// path on vendors that require the segment.
func composeEndpoint(modelsPath, persistedBase, typedBase string) string {
	// The version segment currently present on the persisted base.
	onBase := ""
	if b := trimVersionSuffix(persistedBase); b != persistedBase {
		onBase = persistedBase[len(b):]
	}
	// The version segment present in the models path that answered.
	inPath := versionSegment(modelsPath)

	// Prefer the segment that actually WORKED. When the models path carried
	// one and the base does not, move it onto the endpoint.
	if inPath != "" && onBase == "" {
		return inPath + "/chat/completions"
	}
	// Both carry one (or the path had none because its version lives on the
	// base): keep the base as-is and let the adapter default compose.
	if onBase != "" {
		return ""
	}
	// Neither carries one, but the operator typed a version segment that is
	// being stripped off the base. Restore it on the endpoint so the chat path
	// keeps working.
	if b := trimVersionSuffix(typedBase); b != typedBase && typedBase != "" {
		return typedBase[len(b):] + "/chat/completions"
	}
	return ""
}

// DiscoverModels probes rawURL for a model catalog. apiKey is sent as a bearer
// token; some gateways accept an unauthenticated listing, so an empty key is
// allowed and simply omits the header.
//
// It returns on the first candidate that yields a parseable catalog, so the
// reported latency is that of the working path, not the sum of the failures.
func DiscoverModels(ctx context.Context, rawURL, apiKey string) (*discoveryResult, error) {
	cands := discoveryCandidates(rawURL)
	if len(cands) == 0 {
		return nil, fmt.Errorf("base url required")
	}

	client := &http.Client{Timeout: discoveryTimeout}
	var tried []string
	var firstErr string
	typedBase := strings.TrimRight(strings.TrimSpace(rawURL), "/")

	for _, c := range cands {
		full := c.baseURL + c.modelsPath
		tried = append(tried, full)

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, full, nil)
		if err != nil {
			firstErr = err.Error()
			continue
		}
		req.Header.Set("Accept", "application/json")
		if apiKey != "" {
			req.Header.Set("Authorization", "Bearer "+apiKey)
		}

		start := time.Now()
		resp, err := client.Do(req)
		if err != nil {
			firstErr = err.Error()
			continue
		}
		body, readErr := io.ReadAll(io.LimitReader(resp.Body, discoveryMaxBody))
		resp.Body.Close()
		took := time.Since(start)
		if readErr != nil {
			firstErr = readErr.Error()
			continue
		}
		if resp.StatusCode != http.StatusOK {
			firstErr = fmt.Sprintf("%s returned %d", full, resp.StatusCode)
			continue
		}

		models, err := parseModelCatalog(body)
		if err != nil {
			firstErr = fmt.Sprintf("%s: %v", full, err)
			continue
		}

		res := &discoveryResult{
			BaseURL:     c.baseURL,
			ModelsPath:  c.modelsPath,
			Adapter:     "openai",
			Models:      []string{},
			ImageModels: []string{},
			Tried:       append([]string{}, tried...),
			LatencyMS:   took.Milliseconds(),
		}
		// Persist a base/endpoint pair that composes correctly. The adapter's
		// own endpoint is "/chat/completions", so any version segment has to be
		// carried on the endpoint rather than left on the base (where it doubles
		// the default image path into "/v1/v1/images/generations" — the exact
		// 404 seen on the first vela configuration attempt).
		res.Endpoint = composeEndpoint(c.modelsPath, c.baseURL, typedBase)

		// If the operator typed a version segment and we are persisting a base
		// that no longer has it, the endpoint above is what restores it. If the
		// base kept the segment, the adapter default is already correct.

		seen := map[string]bool{}
		for _, m := range models {
			if m.ID == "" || seen[m.ID] {
				continue
			}
			seen[m.ID] = true
			if m.Kind == "image" {
				res.ImageModels = append(res.ImageModels, m.ID)
			} else {
				res.Models = append(res.Models, m.ID)
			}
		}
		sort.Strings(res.Models)
		sort.Strings(res.ImageModels)
		if len(res.Models) == 0 && len(res.ImageModels) == 0 {
			firstErr = fmt.Sprintf("%s: catalog contained no usable model ids", full)
			continue
		}
		if c.baseURL != typedBase {
			res.Note = fmt.Sprintf("base url normalized from %s to %s; endpoint set to %q so it composes with the adapter",
				typedBase, c.baseURL, res.Endpoint)
		}
		return res, nil
	}
	return nil, fmt.Errorf("no model catalog found; tried %s; last error: %s", strings.Join(tried, ", "), firstErr)
}

// upstreamModel is the subset of a /models entry this code reads. Vendors add
// fields freely and rename the nesting; only id is universal, and the modality
// hints below are read opportunistically.
type modelCaps struct {
	ImageGeneration bool `json:"image_generation"`
	Vision          bool `json:"vision"`
}

type upstreamModel struct {
	ID       string `json:"id"`
	Object   string `json:"object"`
	Kind     string `json:"kind"`
	Type     string `json:"type"`
	Modality string `json:"modality"`
	Architecture struct {
		InputModalities  []string `json:"input_modalities"`
		OutputModalities []string `json:"output_modalities"`
	} `json:"architecture"`
	Capabilities modelCaps `json:"capabilities"`

	// Metadata is kept as raw JSON on purpose. The real amr-link catalog nests
	// cost/capability/context_limit here, and a typed struct would silently
	// discard every key this build does not know about — the operator would
	// then lose the context limit their vendor published. Raw also lets
	// input_modalities/output_modalities be read without a second type.
	Metadata json.RawMessage `json:"metadata"`
}

// metadataFields lists the keys worth surfacing in the UI. They are copied
// through verbatim and nothing is inferred from them; an upstream that sends
// no context limit must leave that absent rather than have this code write 0.
var metadataFields = []string{
	"context_length", "context_window", "max_context_length",
	"context_limit", "context_budget", "cost", "capability",
}

// parseModelCatalog reads the OpenAI list shape {"data":[...]} and the bare
// array shape some gateways return. A response that is valid JSON but has
// neither shape is an error, not an empty catalog: reporting "0 models" for a
// login page or a rate-limit JSON would silently create an empty source.
func parseModelCatalog(body []byte) ([]discoveredModel, error) {
	// "null" unmarshals into a nil slice without error, and an HTML login page
	// fails to unmarshal at all. Both must be errors.
	trimmed := strings.TrimSpace(string(body))
	if trimmed == "" || trimmed == "null" {
		return nil, fmt.Errorf("empty response")
	}

	var wrapper struct {
		Data   []upstreamModel `json:"data"`
		Models []upstreamModel `json:"models"`
	}
	if err := json.Unmarshal(body, &wrapper); err == nil {
		if len(wrapper.Data) > 0 {
			return toDiscovered(wrapper.Data), nil
		}
		if len(wrapper.Models) > 0 {
			return toDiscovered(wrapper.Models), nil
		}
		// A well-formed but empty catalog is a legitimate answer from a key
		// that has no models provisioned. Distinguish it from "not a catalog"
		// by checking the shape actually parsed.
		var probe map[string]json.RawMessage
		if json.Unmarshal(body, &probe) == nil {
			if _, ok := probe["data"]; ok {
				return nil, nil
			}
			if _, ok := probe["models"]; ok {
				return nil, nil
			}
		}
	}
	var bare []upstreamModel
	if err := json.Unmarshal(body, &bare); err == nil {
		return toDiscovered(bare), nil
	}
	return nil, fmt.Errorf("response is not a model catalog")
}

func toDiscovered(in []upstreamModel) []discoveredModel {
	out := make([]discoveredModel, 0, len(in))
	for _, m := range in {
		out = append(out, discoveredModel{
			ID:   m.ID,
			Kind: discoveredModelKind(m),
			Meta: modelMeta(m),
		})
	}
	return out
}

// discoveredModelKind trusts the upstream's own modality fields and nothing
// else. It never infers from the model NAME: ids like "gpt-image-1" or
// "ideogram-v2" appear on catalogs that also serve chat, and a wrong kind
// poisons the AUTO chain (an image-only id selected for a chat request fails
// every tier behind it).
func discoveredModelKind(m upstreamModel) string {
	modIn := append(append([]string{}, m.Architecture.InputModalities...),
		modalityList(m.Metadata, "input_modalities")...)
	modOut := append(append([]string{}, m.Architecture.OutputModalities...),
		modalityList(m.Metadata, "output_modalities")...)

	has := func(list []string, want string) bool {
		for _, v := range list {
			if strings.EqualFold(strings.TrimSpace(v), want) {
				return true
			}
		}
		return false
	}
	if m.Capabilities.ImageGeneration || has(modOut, "image") {
		return "image"
	}
	if has(modIn, "image") || has(modIn, "file") || has(modIn, "audio") || m.Capabilities.Vision {
		return "chat" // multimodal input still means a chat model
	}
	// Some catalogs encode modality as a top-level string instead of a list.
	// Only the image-ish spellings count: "chat" here means the upstream said
	// chat, and treating that as "unknown" would ignore its statement.
	switch strings.ToLower(strings.TrimSpace(m.Kind + " " + m.Modality + " " + m.Type)) {
	case "image", "text-to-image", "image_generation", "text_to_image":
		return "image"
	}
	return "chat"
}

// modalityList reads one array-of-strings field out of a raw JSON object,
// tolerating both absence and a non-array value.
func modalityList(raw json.RawMessage, field string) []string {
	if len(raw) == 0 {
		return nil
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil
	}
	v, ok := obj[field]
	if !ok {
		return nil
	}
	var out []string
	if err := json.Unmarshal(v, &out); err != nil {
		return nil
	}
	return out
}

// modelMeta carries through the fields an operator wants to see in the UI
// without this code inventing them. The upstream nests most of them under
// "metadata" (the real amr-link catalog does) while some put the same keys at
// the top level, so both are read and top level wins on a tie. Only keys the
// upstream actually sent are copied, so an absent context limit stays absent
// instead of becoming 0.
func modelMeta(m upstreamModel) map[string]interface{} {
	meta := map[string]interface{}{}

	copyFields := func(raw json.RawMessage) {
		if len(raw) == 0 {
			return
		}
		var obj map[string]json.RawMessage
		if err := json.Unmarshal(raw, &obj); err != nil {
			return
		}
		for _, k := range metadataFields {
			if v, ok := obj[k]; ok {
				if dec := rawOrString(v); dec != nil {
					meta[k] = dec
				}
			}
		}
	}

	copyFields(m.Metadata)      // nested first
	copyFields(mustJSON(m))     // top level wins
	if len(meta) == 0 {
		return nil
	}
	return meta
}

// rawOrString keeps numeric metadata as a number and everything else as its
// decoded string, so the UI can sort on a cost or context length instead of
// receiving "1048000" and having to coerce it.
func rawOrString(raw json.RawMessage) interface{} {
	var num json.Number
	if err := json.Unmarshal(raw, &num); err == nil {
		if f, err := num.Float64(); err == nil {
			return f
		}
	}
	var s string
	if err := json.Unmarshal(raw, &s); err == nil {
		return s
	}
	var any interface{}
	if err := json.Unmarshal(raw, &any); err == nil {
		return any
	}
	return nil
}

func mustJSON(v interface{}) []byte {
	b, _ := json.Marshal(v)
	return b
}

// modelsToConfig converts discovery output into config.Model entries, applying
// the same priority convention the UI uses: chat models the operator is
// likely to want first get a lower number, and image models sort last because
// they are unreachable from /v1/chat/completions.
func modelsToConfig(chat, image []string) []config.Model {
	out := make([]config.Model, 0, len(chat)+len(image))
	for _, id := range chat {
		out = append(out, config.Model{ID: id, Priority: 50, Kind: "chat"})
	}
	for _, id := range image {
		out = append(out, config.Model{ID: id, Priority: 70, Kind: "image"})
	}
	return out
}