Cyrene_v2/internal/runtime/model/fallback.go

package model

import (
	"context"
	"fmt"

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

// Step binds a model name to an Adapter for one entry in a fallback chain. An
// empty Model keeps the incoming Request.Model.
type Step struct {
	Model   string
	Adapter Adapter
}

// Chain is an Adapter that tries its steps in order and returns the first
// successful response. It implements Adapter, so it can be used anywhere a
// single provider is expected.
type Chain struct {
	steps []Step
	// groupMaxTokens is the active model group's own budget (models.json
	// maxTokens). When > 0 it applies to every step of the group's chain and
	// wins over fallbackMaxTokens.
	groupMaxTokens int
	// fallbackMaxTokens is the last-resort budget (LLM_FALLBACK_MAX_TOKENS). It
	// only applies to the final step of a chain that actually has a fallback.
	fallbackMaxTokens int
}

// NewChain builds a fallback chain from the given steps.
func NewChain(steps ...Step) *Chain {
	return &Chain{steps: steps}
}

// WithTokenBudget sets the chain's token budgets: the active group's own budget
// and the last-resort fallback budget. It must be called before the chain is
// published to concurrent callers (the runtime builds a fresh chain per
// snapshot and swaps it in atomically), so the fields are never mutated in
// flight.
func (c *Chain) WithTokenBudget(groupMaxTokens, fallbackMaxTokens int) *Chain {
	c.groupMaxTokens = groupMaxTokens
	c.fallbackMaxTokens = fallbackMaxTokens
	return c
}

// applyBudget resolves the effective MaxTokens for step idx, honoring the
// precedence: an explicit Request.MaxTokens > the group's MaxTokens > the
// last-resort LLM_FALLBACK_MAX_TOKENS > the adapter default. The last-resort
// budget only applies when the chain really has a fallback step (len > 1) and
// only to that final step; a single-step chain is a primary model, not a
// fallback, so its adapter default is left untouched.
func (c *Chain) applyBudget(req Request, idx int) Request {
	if req.MaxTokens != 0 {
		return req
	}
	if c.groupMaxTokens > 0 {
		req.MaxTokens = c.groupMaxTokens
		return req
	}
	if c.fallbackMaxTokens > 0 && len(c.steps) > 1 && idx == len(c.steps)-1 {
		req.MaxTokens = c.fallbackMaxTokens
	}
	return req
}

// NewOpenAIChatChain builds a fallback chain from a single supplier config: the
// primary model first, then cfg.Fallback (when set and distinct), both served
// by an OpenAIChatAdapter against the same OpenAI-compatible endpoint. It
// delegates to the provider-aware builder so the OpenAI path and the factory stay
// in sync. It returns ErrNoModel when no base URL or primary model is configured.
func NewOpenAIChatChain(cfg config.LLMConfig, opts ...OpenAIOption) (*Chain, error) {
	return newOpenAIChain(cfg, opts...)
}

// Name returns the chain identifier.
func (c *Chain) Name() string { return "fallback" }

// Chat tries each step in order and returns the first success. If every step
// fails, the last error is returned wrapped as ErrUpstream.
func (c *Chain) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
	if len(c.steps) == 0 {
		return Response{}, ErrNoSteps
	}
	var lastErr error
	for i, s := range c.steps {
		stepReq := c.applyBudget(req, i)
		if s.Model != "" {
			stepReq.Model = s.Model
		}
		resp, err := s.Adapter.Chat(ctx, stepReq, opts)
		if err == nil {
			return resp, nil
		}
		lastErr = err
	}
	return Response{}, fmt.Errorf("model: %w: %v", ErrUpstream, lastErr)
}

// ChatStream tries each step in order. If a step fails before delivering any
// event, the chain moves to the next step; once a step has delivered output it
// is committed, so a mid-stream failure surfaces as an error rather than mixing
// content across providers.
func (c *Chain) ChatStream(ctx context.Context, req Request, onEvent func(StreamEvent)) (Response, error) {
	if len(c.steps) == 0 {
		return Response{}, ErrNoSteps
	}
	var lastErr error
	for i, s := range c.steps {
		stepReq := c.applyBudget(req, i)
		if s.Model != "" {
			stepReq.Model = s.Model
		}
		delivered := false
		wrapped := func(ev StreamEvent) {
			delivered = true
			if onEvent != nil {
				onEvent(ev)
			}
		}
		resp, err := s.Adapter.ChatStream(ctx, stepReq, wrapped)
		if err == nil {
			return resp, nil
		}
		if delivered {
			return Response{}, fmt.Errorf("model: %w: stream step %q failed after partial output: %v", ErrUpstream, s.Model, err)
		}
		lastErr = err
	}
	return Response{}, fmt.Errorf("model: %w: %v", ErrUpstream, lastErr)
}