ollie/cmd/olliesrv/internal/fs/cache.go

package fs

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"ollie/cmd/olliesrv/internal/backend"
)

// ModelCache caches the model list from all backends.
// Created by NewRoot and used by the /models handler.
// Also invalidated by root-level /ctl.
type ModelCache struct {
	mu      sync.Mutex
	data    []byte
	fetched bool
}

func NewModelCache() *ModelCache { return &ModelCache{} }

func (c *ModelCache) Get() []byte {
	c.mu.Lock()
	if !c.fetched {
		c.mu.Unlock()
		c.refresh()
		c.mu.Lock()
	}
	defer c.mu.Unlock()
	return c.data
}

func (c *ModelCache) Invalidate() {
	c.mu.Lock()
	c.fetched = false
	c.data = nil
	c.mu.Unlock()
}

func (c *ModelCache) refresh() {
	var sb strings.Builder
	for _, name := range backend.Backends() {
		b, err := backend.NewWithName(name)
		if err != nil {
			continue
		}
		// Use ModelsInfo if the backend implements ModelLister
		if lister, ok := b.(backend.ModelLister); ok {
			for _, m := range lister.ModelsInfo(context.Background()) {
				pricing := m.Pricing
				// Fall back to static pricing if not provided by API
				if pricing == nil {
					pricing = backend.LookupStaticPricing(m.ID)
				}
				writeModelLine(&sb, name, m.ID, pricing)
			}
		} else {
			for _, m := range b.Models(context.Background()) {
				// Try static pricing lookup for backends without ModelsInfo
				pricing := backend.LookupStaticPricing(m)
				writeModelLine(&sb, name, m, pricing)
			}
		}
	}
	c.mu.Lock()
	c.data = []byte(sb.String())
	c.fetched = true
	c.mu.Unlock()
}

// writeModelLine writes a single model line to the builder.
// Format: backend\tmodel[\tin\tout\tcache_read\tcache_write]
// Pricing values are marked with (e) suffix if estimated.
func writeModelLine(sb *strings.Builder, backendName, modelID string, pricing *backend.ModelPricing) {
	sb.WriteString(backendName)
	sb.WriteByte('\t')
	sb.WriteString(modelID)
	if pricing != nil {
		sb.WriteByte('\t')
		sb.WriteString(formatPricingWithEstimate(pricing.Input, pricing.Estimated))
		sb.WriteByte('\t')
		sb.WriteString(formatPricingWithEstimate(pricing.Output, pricing.Estimated))
		sb.WriteByte('\t')
		sb.WriteString(formatPricingWithEstimate(pricing.CacheRead, pricing.Estimated))
		sb.WriteByte('\t')
		sb.WriteString(formatPricingWithEstimate(pricing.CacheWrite, pricing.Estimated))
	}
	sb.WriteByte('\n')
}

// modelPricingBlock renders a model name followed by its per-million-token
// USD pricing, in the format returned by the agent ctl "model" command.
// Resolution mirrors ModelCache.refresh: API-provided pricing (ModelLister)
// first, then the static estimate table. Returns name-only when no pricing
// is known.
func modelPricingBlock(be backend.Backend, modelID string) string {
	var sb strings.Builder
	sb.WriteString(modelID)
	sb.WriteByte('\n')

	var pricing *backend.ModelPricing
	if lister, ok := be.(backend.ModelLister); ok {
		for _, m := range lister.ModelsInfo(context.Background()) {
			if m.ID == modelID {
				pricing = m.Pricing
				break
			}
		}
	}
	if pricing == nil {
		pricing = backend.LookupStaticPricing(modelID)
	}
	if pricing == nil {
		return sb.String()
	}

	fmt.Fprintf(&sb, "In/Mtok: %s\n", formatPricingWithEstimate(pricing.Input, pricing.Estimated))
	fmt.Fprintf(&sb, "Out/Mtok: %s\n", formatPricingWithEstimate(pricing.Output, pricing.Estimated))
	fmt.Fprintf(&sb, "CacheRead/Mtok: %s\n", formatPricingWithEstimate(pricing.CacheRead, pricing.Estimated))
	fmt.Fprintf(&sb, "CacheWrite/Mtok: %s\n", formatPricingWithEstimate(pricing.CacheWrite, pricing.Estimated))
	return sb.String()
}

// formatPricingWithEstimate formats per-token cost as per-million-tokens USD string.
// Returns "-" if the value is zero. Appends "(e)" if estimated is true and value is non-zero.
func formatPricingWithEstimate(perToken float64, estimated bool) string {
	if perToken == 0 {
		return "-"
	}
	perMillion := perToken * 1_000_000
	var s string
	if perMillion >= 1 {
		s = fmt.Sprintf("$%.2f", perMillion)
	} else {
		s = fmt.Sprintf("$%.4f", perMillion)
	}
	if estimated {
		s += "(e)"
	}
	return s
}