package ai

import "context"

type ctxKey int

const (
	ctxMaxTokens ctxKey = iota + 1
	ctxTemperature
)

// WithMaxTokens 限制 completion 輸出上限（靈感短聊用小值可明顯加速）。
func WithMaxTokens(ctx context.Context, n int) context.Context {
	if n <= 0 {
		return ctx
	}
	return context.WithValue(ctx, ctxMaxTokens, n)
}

// MaxTokensFrom returns 0 if unset.
func MaxTokensFrom(ctx context.Context) int {
	if ctx == nil {
		return 0
	}
	v, _ := ctx.Value(ctxMaxTokens).(int)
	return v
}

// WithTemperature sets sampling temperature when > 0.
func WithTemperature(ctx context.Context, t float64) context.Context {
	if t <= 0 {
		return ctx
	}
	return context.WithValue(ctx, ctxTemperature, t)
}

// TemperatureFrom returns 0 if unset (caller uses default).
func TemperatureFrom(ctx context.Context) float64 {
	if ctx == nil {
		return 0
	}
	v, _ := ctx.Value(ctxTemperature).(float64)
	return v
}