Vectrum
Sign In
120 lines · 3.8 KB · Go
package llm

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

// Model is one entry in the catalogue at GET /v1/models.
type Model struct {
	ID          string `json:"id"`
	DisplayName string `json:"display_name"`
	Deprecated  bool   `json:"deprecated"`
}

// catalogueTimeout is deliberately short. This runs while a person is looking
// at a wizard that cannot show them a list yet, and a gateway that black-holes
// the request must not hold the setup open for the client's ten-minute
// message timeout.
const catalogueTimeout = 10 * time.Second

// ListModels reads the model catalogue.
//
// It is best-effort by design: plenty of endpoints this harness can be pointed
// at speak /v1/messages and nothing else, so every caller is expected to have
// a fallback — KnownModels — rather than to treat an error here as fatal.
func (c *Client) ListModels(ctx context.Context) ([]Model, error) {
	if !c.HasCredentials() {
		return nil, fmt.Errorf("llm: no credentials: set ANTHROPIC_API_KEY, or run `ant auth login`")
	}
	ctx, cancel := context.WithTimeout(ctx, catalogueTimeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.BaseURL+"/v1/models?limit=100", nil)
	if err != nil {
		return nil, err
	}
	c.setHeaders(req)

	resp, err := c.HTTP.Do(req)
	if err != nil {
		return nil, fmt.Errorf("llm: listing models: %w", err)
	}
	defer resp.Body.Close()

	// Capped: an endpoint that answers an unknown route with a whole HTML
	// page should not be read into memory in full to be thrown away.
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, fmt.Errorf("llm: listing models: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, parseAPIError(resp.StatusCode, resp.Header.Get("retry-after"), raw)
	}

	var out struct {
		Data []Model `json:"data"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llm: %s did not return a model list: %w", c.BaseURL, err)
	}

	models := make([]Model, 0, len(out.Data))
	for _, m := range out.Data {
		// A deprecated model still answers requests, which is exactly why it
		// should not be offered as a fresh default.
		if m.ID == "" || m.Deprecated {
			continue
		}
		if m.DisplayName == "" {
			m.DisplayName = m.ID
		}
		models = append(models, m)
	}
	if len(models) == 0 {
		return nil, fmt.Errorf("llm: %s returned no models", c.BaseURL)
	}
	return models, nil
}

// KnownModels is the fallback catalogue: what this build knows how to price.
//
// Pricing is the honest basis for the list. A model the harness cannot price
// is one whose cost would silently read as blank in the status bar, so the set
// it can price is exactly the set it can fully describe.
func KnownModels() []Model {
	out := make([]Model, 0, len(prices))
	for id := range prices {
		out = append(out, Model{ID: id, DisplayName: displayName(id)})
	}
	// Map iteration is random; the wizard needs the same order every time.
	sort.Slice(out, func(i, j int) bool {
		ri, rj := familyRank(out[i].ID), familyRank(out[j].ID)
		if ri != rj {
			return ri < rj
		}
		return out[i].ID > out[j].ID // newer version first within a family
	})
	return out
}

// familyRank orders the families the way someone choosing between them thinks
// about it: the flagship first, the cheap one last.
func familyRank(id string) int {
	for i, fam := range []string{"claude-opus", "claude-sonnet", "claude-haiku", "claude-fable"} {
		if len(id) >= len(fam) && id[:len(fam)] == fam {
			return i
		}
	}
	return 99
}

// displayName stands in for the API's display_name on a model we only know
// statically. It is derived rather than tabulated: a table would be one more
// place to forget to update when a model is added to prices.
func displayName(id string) string { return "Claude " + PrettyModel(id) }