Vectrum
Sign In
265 lines · 8.0 KB · Go
package tui

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/charmbracelet/bubbles/textarea"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"

	"harness/internal/llm"
)

// modelFetchTimeout bounds the wait. Someone is sitting in front of a picker
// that says it is asking; a minute of that is not a picker, it is a hang.
const modelFetchTimeout = 20 * time.Second

// modelsMsg carries the answer to GET /v1/models back into the UI.
type modelsMsg struct {
	models []llm.Model
	err    error
}

// fetchModels asks the endpoint what it serves. It runs off the UI goroutine
// like any other command, and reports failure as a value rather than by
// leaving the picker waiting forever.
func fetchModels(c *llm.Client) tea.Cmd {
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), modelFetchTimeout)
		defer cancel()
		models, err := c.ListModels(ctx)
		return modelsMsg{models: models, err: err}
	}
}

// openModelPicker puts the picker up and starts asking.
func (m *model) openModelPicker() tea.Cmd {
	m.phase = phaseModel
	m.models, m.modelIdx, m.modelsLoading = nil, 0, true
	m.closeMenu()
	return tea.Batch(m.sync(), m.spin.Tick, fetchModels(m.opts.Client))
}

// onModels lands the answer. An answer nobody is waiting for any more is
// dropped: the picker was escaped, and reopening it over whatever the user
// went on to do would be the UI arguing with them.
func (m *model) onModels(msg modelsMsg) tea.Cmd {
	if m.phase != phaseModel {
		return nil
	}
	m.modelsLoading = false
	m.phase = phaseIdle

	switch {
	case msg.err != nil:
		m.tr.addError("could not list models from " + endpointLabel(m.opts.Client) + ": " +
			msg.err.Error() + " — this endpoint may not serve /v1/models; -model or /model <id> still works")
	case len(msg.models) == 0:
		m.tr.addError(endpointLabel(m.opts.Client) + " returned no models — -model or /model <id> still works")
	default:
		m.phase = phaseModel
		m.models = msg.models
		m.modelIdx = modelIndex(msg.models, m.opts.Client.Model())
	}
	return tea.Batch(m.sync(), textarea.Blink)
}

// modelIndex finds the model in force, so the picker opens where the session
// already is rather than at the top of a list.
func modelIndex(models []llm.Model, current string) int {
	for i, mo := range models {
		if mo.ID == current {
			return i
		}
	}
	return 0
}

// onModelKey drives the picker. While the fetch is in flight there is nothing
// to move through, so only the ways out are live.
func (m *model) onModelKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	key := msg.String()
	if key == "esc" || key == "ctrl+c" {
		m.phase = phaseIdle
		m.models, m.modelsLoading = nil, false
		return m, tea.Batch(m.sync(), textarea.Blink)
	}
	if m.modelsLoading || len(m.models) == 0 {
		return m, nil
	}

	last := len(m.models) - 1
	switch key {
	case "up", "ctrl+p":
		m.modelIdx = maxInt(0, m.modelIdx-1)
	case "down", "ctrl+n":
		m.modelIdx = minInt(last, m.modelIdx+1)
	case "home":
		m.modelIdx = 0
	case "end":
		m.modelIdx = last
	case "enter", "tab":
		id := m.models[m.modelIdx].ID
		m.phase = phaseIdle
		m.models = nil
		m.applyModel(id)
		return m, tea.Batch(m.sync(), textarea.Blink)
	default:
		if len(key) == 1 && key[0] >= '1' && key[0] <= '9' {
			if i := int(key[0] - '1'); i <= last {
				m.modelIdx = i
			}
		}
	}
	return m, m.sync()
}

// useModel switches the client. The session keeps its history: the
// conversation so far is replayed to the new model on the next turn, which is
// the point of switching mid-session.
//
// It is separate from applyModel because the setup wizard changes several
// settings at once and says so in one line, where the picker changes this one
// and announces it. (setModel is taken: that is the /model command.)
func (m *model) useModel(id string) {
	if id == "" {
		return
	}
	m.opts.Client.SetModel(id)
	m.refreshSplash()
}

// applyModel switches the client and says so.
func (m *model) applyModel(id string) {
	m.useModel(id)
	m.tr.addNotice("model: " + prettyModel(id) + " (" + id + ") — from the next turn on")
}

// modelView is the picker: a heading, where the list came from, and a row per
// model with what it is called, its id, and what it costs.
func (m *model) modelView() string {
	width := maxInt(1, m.width)
	head := " " + styHeading.Render("Select model")
	from := styDim.Render("from " + endpointLabel(m.opts.Client))

	if m.modelsLoading {
		return head + "\n" + truncateStyled(" "+m.spin.View()+" "+from+styFaint.Render(" · esc to cancel"), width)
	}

	rows := m.modelRows()
	shown := rows
	if len(m.models) > rows {
		shown = maxInt(1, rows-1)
	}
	start := 0
	if m.modelIdx >= shown {
		start = m.modelIdx - shown + 1
	}
	end := minInt(len(m.models), start+shown)

	// One column for the names, so the ids line up under each other.
	names := make([]string, len(m.models))
	col := 0
	for i, mo := range m.models {
		names[i] = mo.DisplayName
		if names[i] == "" {
			names[i] = prettyModel(mo.ID)
		}
		if i >= start && i < end {
			col = maxInt(col, lipgloss.Width(names[i]))
		}
	}
	col = minInt(col, maxInt(10, width/3))

	var b strings.Builder
	b.WriteString(head + "\n")
	b.WriteString(truncateStyled(" "+from+styFaint.Render(" · ↑↓ move · enter apply · esc cancel"), width) + "\n")
	for i := start; i < end; i++ {
		mo := m.models[i]
		marker, name := "  ", styDim.Render(pad(names[i], col))
		if i == m.modelIdx {
			marker, name = styKey.Render("❯ "), styKey.Render(pad(names[i], col))
		}
		// An endpoint that gave no display name leaves the id as the name;
		// repeating it in the detail column says nothing twice.
		detail := ""
		if names[i] != mo.ID {
			detail = mo.ID
		}
		if mo.MaxInputTokens > 0 {
			if detail != "" {
				detail += " · "
			}
			detail += human(mo.MaxInputTokens) + " in"
		}
		if p, ok := llm.PriceOf(mo.ID); ok {
			if detail != "" {
				detail += " · "
			}
			detail += fmt.Sprintf("$%g / $%g per Mtok", p.InputPerMTok, p.OutputPerMTok)
		}
		line := marker + fmt.Sprintf("%d. ", i+1) + name + " " + styFaint.Render(detail)
		if mo.ID == m.opts.Client.Model() {
			line += styOK.Render(" ✓")
		}
		b.WriteString(truncateStyled(" "+line, width) + "\n")
	}
	if n := len(m.models) - (end - start); n > 0 {
		b.WriteString(truncateStyled("   "+styFaint.Render(fmt.Sprintf("… +%d more", n)), width) + "\n")
	}
	return strings.TrimRight(b.String(), "\n")
}

// modelRows is how many model rows fit: the same arithmetic the completion
// menu does, minus the heading and the line saying where the list came from.
func (m *model) modelRows() int {
	room := m.height - 6
	return maxInt(1, minInt(modelMaxRows, room))
}

// modelMaxRows caps the list. Anthropic serves a couple of dozen models and
// nobody scrolls a picker looking for the twenty-fourth.
const modelMaxRows = 10

// ---- describing a model -------------------------------------------------
//
// The picker and the setup wizard show the same models and must say the same
// things about them, so what a model is called and what it costs is decided
// once, here.

// pricedModels is the list to fall back on when the endpoint will not answer
// GET /v1/models. The picker treats that as an error and says so — there is a
// session to go back to. The wizard cannot: it is the first run, and a model
// has to be chosen before there is anything to go back to.
func pricedModels() []llm.Model {
	ids := llm.PricedModels()
	out := make([]llm.Model, 0, len(ids))
	for _, id := range ids {
		out = append(out, llm.Model{ID: id, DisplayName: prettyModel(id)})
	}
	return out
}

// priceLabel is a model's list pricing, or "" when this build cannot price it.
// A blank says "not known"; a zero would say "free", which is worse than
// saying nothing.
func priceLabel(id string) string {
	p, ok := llm.PriceOf(id)
	if !ok {
		return ""
	}
	return fmt.Sprintf("$%g / $%g per Mtok", p.InputPerMTok, p.OutputPerMTok)
}

// modelDetail is the line under a highlighted model: its id, and its context
// window when the endpoint reported one.
func modelDetail(m llm.Model) string {
	detail := m.ID
	if m.MaxInputTokens > 0 {
		detail += " · " + human(m.MaxInputTokens) + " token context"
	}
	return detail
}