Cyrene_v2/internal/runtime/model/runtime.go

package model

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"sync/atomic"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config/runtimeconfig"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// ModelRouter is the routing contract the agent (and loop) depend on. *Router
// implements it; RuntimeLLM implements it too, so a hot-swappable provider can
// serve as both the adapter and the router without a separate router reference
// to keep in sync.
type ModelRouter interface {
	RouteRequest(req Request) (model string, needsVision bool)
}

// runtimeChainSnapshot is an immutable, self-contained LLM stack snapshot. It is
// stored and swapped atomically: a request that already captured a snapshot
// keeps using that exact chain+router for its whole call, so a reload never
// interrupts an in-flight request.
type runtimeChainSnapshot struct {
	chain  *Chain
	router *Router
	group  runtimeconfig.Group
	groups []runtimeconfig.Group
}

// RuntimeLLM is the hot-swappable LLM provider stack for the runtime. It
// implements Adapter (delegating to the current chain) and ModelRouter
// (delegating to the current router). A panel save reloads models.json, rebuilds
// a fresh chain+router and atomically swaps it in: new calls use the new stack
// while calls already in flight finish on the old one.
type RuntimeLLM struct {
	path    string
	env     config.LLMConfig
	log     *slog.Logger
	current atomic.Value // *runtimeChainSnapshot
}

// NewRuntimeLLM builds the provider stack from models.json at path, falling back
// to env. It loads the initial snapshot so the runtime is usable immediately.
// log may be nil (a discard logger is used).
func NewRuntimeLLM(path string, env config.LLMConfig, log *slog.Logger) *RuntimeLLM {
	if log == nil {
		log = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	m := &RuntimeLLM{path: path, env: env, log: logger.WithSource(log, "model")}
	m.loadInitial()
	return m
}

// Path returns the models.json path this manager reads and reloads from.
func (m *RuntimeLLM) Path() string { return m.path }

// Groups returns a copy of the current loaded groups; the active group (the one
// a new request uses by default) is index 0.
func (m *RuntimeLLM) Groups() []runtimeconfig.Group {
	src := m.load().groups
	out := make([]runtimeconfig.Group, len(src))
	copy(out, src)
	return out
}

// Group returns the active group (index 0).
func (m *RuntimeLLM) Group() runtimeconfig.Group { return m.load().group }

// Current returns the current chain and router. The returned chain stays valid
// for calls already in flight; new calls pick up the latest swap.
func (m *RuntimeLLM) Current() (*Chain, *Router) {
	snap := m.load()
	return snap.chain, snap.router
}

// Reload re-reads models.json (falling back to env) and atomically rebuilds the
// in-use chain+router. In-flight calls keep the previous snapshot.
func (m *RuntimeLLM) Reload() error {
	groups, warn := runtimeconfig.LoadModels(m.path, m.env)
	if warn != nil {
		m.log.Warn("models.json reload fell back to env", "error", warn)
	}
	return m.swap(groups)
}

// Provider returns the active group's provider protocol.
func (m *RuntimeLLM) Provider() string { return m.load().group.Provider }

// Model returns the active group's primary model name.
func (m *RuntimeLLM) Model() string { return m.load().group.Model }

// Name implements Adapter, returning the active group id.
func (m *RuntimeLLM) Name() string { return m.load().group.ID }

// Chat implements Adapter, delegating to the current chain. A request that began
// before a reload completes on the snapshot it captured. When no model is
// configured (nil chain) it returns ErrNoModel rather than panicking.
func (m *RuntimeLLM) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
	chain := m.load().chain
	if chain == nil {
		return Response{}, ErrNoModel
	}
	return chain.Chat(ctx, req, opts)
}

// ChatStream implements Adapter.
func (m *RuntimeLLM) ChatStream(ctx context.Context, req Request, onEvent func(StreamEvent)) (Response, error) {
	chain := m.load().chain
	if chain == nil {
		return Response{}, ErrNoModel
	}
	return chain.ChatStream(ctx, req, onEvent)
}

// RouteRequest implements ModelRouter, delegating to the current router.
func (m *RuntimeLLM) RouteRequest(req Request) (string, bool) {
	r := m.load().router
	if r == nil {
		return "", false
	}
	return r.RouteRequest(req)
}

// loadInitial builds the first snapshot; a build failure leaves a nil chain so
// calls report a clear no-model error instead of panicking.
func (m *RuntimeLLM) loadInitial() {
	groups, warn := runtimeconfig.LoadModels(m.path, m.env)
	if warn != nil {
		m.log.Warn("models.json load fell back to env", "error", warn)
	}
	_ = m.swap(groups)
}

// swap rebuilds and atomically stores a snapshot for groups. A blank slice falls
// back to the env-synthesized group so there is always at least one.
func (m *RuntimeLLM) swap(groups []runtimeconfig.Group) error {
	if len(groups) == 0 {
		groups = runtimeconfig.GroupFromEnv(m.env)
	}
	snap, err := buildSnapshot(groups, m.env)
	if err != nil {
		return err
	}
	m.current.Store(snap)
	return nil
}

// load returns the current snapshot, or an empty (nil chain) one when the
// manager was never initialized.
func (m *RuntimeLLM) load() *runtimeChainSnapshot {
	if v := m.current.Load(); v != nil {
		return v.(*runtimeChainSnapshot)
	}
	return &runtimeChainSnapshot{}
}

// buildSnapshot turns the active group (index 0) into a chain+router. A group
// without a base URL/model yields a nil chain (no model configured), surfaced as
// ErrNoModel by the adapter methods. The Router keeps the env vision model so a
// dedicated vision model survives a reload that only changes the vendor group.
func buildSnapshot(groups []runtimeconfig.Group, env config.LLMConfig) (*runtimeChainSnapshot, error) {
	active := groups[0]
	cfg := active.ToLLMConfig()
	chain, err := NewChatChain(cfg)
	if err != nil && !errors.Is(err, ErrNoModel) {
		return nil, err
	}
	if chain != nil {
		// Budgets are resolved inside the chain: the active group's own
		// maxTokens (models.json) wins, and env.FallbackMaxTokens only caps the
		// last step of a chain that really has a fallback (design/04 §3.1).
		chain = chain.WithTokenBudget(active.MaxTokens, env.FallbackMaxTokens)
	}
	router := NewRouter(DefaultModelForProvider(cfg),
		WithVisionModel(env.VisionModel),
		WithDefaultProvider(cfg.Provider),
	)
	return &runtimeChainSnapshot{chain: chain, router: router, group: active, groups: groups}, nil
}