45 lines
919 B
Go
45 lines
919 B
Go
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
|
|
}
|