middleware

Documentation

Overview

Package middleware provides composable policies for model calls.

Index

• Variables
• func ApproximateTokenCount(text string) int64
• func ContextWithTokenBudget(ctx context.Context, budget *TokenBudget) context.Context
• func Wrap(base models.Agent, policies ...Middleware) (models.Agent, error)
• type Middleware
• type MiddlewareFunc
  • func (f MiddlewareFunc) Wrap(next models.Agent) (models.Agent, error)
• type RateLimitMode
• type RateLimitPolicy
  • func (p RateLimitPolicy) Wrap(next models.Agent) (models.Agent, error)
• type RetryPolicy
  • func (p RetryPolicy) Wrap(next models.Agent) (models.Agent, error)
• type TimeoutPolicy
  • func (p TimeoutPolicy) Wrap(next models.Agent) (models.Agent, error)
• type TokenBudget
  • func NewTokenBudget(max int64, estimator TokenEstimator) (*TokenBudget, error)
  • func TokenBudgetFromContext(ctx context.Context) (*TokenBudget, bool)
  • func (b *TokenBudget) Max() int64
  • func (b *TokenBudget) Remaining() int64
  • func (b *TokenBudget) Reset()
  • func (b *TokenBudget) Used() int64
• type TokenBudgetError
  • func (e *TokenBudgetError) Error() string
  • func (e *TokenBudgetError) Unwrap() error
• type TokenBudgetPolicy
  • func (p TokenBudgetPolicy) Wrap(next models.Agent) (models.Agent, error)
• type TokenEstimator

Variables

var ErrRateLimitExceeded = errors.New("model rate limit exceeded")

ErrRateLimitExceeded indicates that a reject-mode limiter had no permit.

var ErrTokenBudgetExceeded = errors.New("model token budget exceeded")

ErrTokenBudgetExceeded indicates that a model call would exceed, or has already generated output beyond, its configured estimated token budget.

Functions

func ApproximateTokenCount

func ApproximateTokenCount(text string) int64

ApproximateTokenCount estimates one token per four UTF-8 bytes. It is a portable fallback, not provider billing data.

func ContextWithTokenBudget

func ContextWithTokenBudget(ctx context.Context, budget *TokenBudget) context.Context

ContextWithTokenBudget associates a per-run budget with ctx. It overrides the fallback budget configured on TokenBudgetPolicy.

func Wrap

func Wrap(base models.Agent, policies ...Middleware) (models.Agent, error)

Wrap composes policies around base. The first policy is the outermost: it sees the request first and the response last.

Types

type Middleware

type Middleware interface {
	Wrap(models.Agent) (models.Agent, error)
}

Middleware wraps a model with one policy. Implementations should preserve models.ToolCallingAgent when possible and return models.ErrToolCallingUnsupported when the wrapped model does not support native tool calls.

type MiddlewareFunc

type MiddlewareFunc func(models.Agent) (models.Agent, error)

MiddlewareFunc adapts a function into Middleware.

func (MiddlewareFunc) Wrap

func (f MiddlewareFunc) Wrap(next models.Agent) (models.Agent, error)

Wrap applies f to next.

type RateLimitMode

type RateLimitMode uint8

RateLimitMode controls what happens when no request permit is immediately available.

const (
	// RateLimitWait waits for the next permit and respects context cancellation.
	RateLimitWait RateLimitMode = iota
	// RateLimitReject fails immediately with ErrRateLimitExceeded.
	RateLimitReject
)

type RateLimitPolicy

type RateLimitPolicy struct {
	Requests int
	Per      time.Duration
	Burst    int
	Mode     RateLimitMode
}

RateLimitPolicy limits model request starts. One permit is consumed for each retry attempt when this policy is placed inside RetryPolicy.

func (RateLimitPolicy) Wrap

func (p RateLimitPolicy) Wrap(next models.Agent) (models.Agent, error)

Wrap applies the rate-limit policy.

type RetryPolicy

type RetryPolicy struct {
	MaxAttempts    int
	InitialBackoff time.Duration
	MaxBackoff     time.Duration
	Multiplier     float64
	Jitter         float64
	DisableJitter  bool

	// ShouldRetry overrides the default decision. The default retries all
	// non-cancellation errors except policy rejections and unsupported native
	// tool calling. A deadline created by an inner timeout policy is retryable
	// while expiry of the caller's context is not.
	ShouldRetry func(context.Context, error) bool
}

RetryPolicy retries model calls that fail before producing a result. Stream creation can be retried, but errors emitted after a stream starts cannot be retried safely because doing so could duplicate chunks.

func (RetryPolicy) Wrap

func (p RetryPolicy) Wrap(next models.Agent) (models.Agent, error)

Wrap applies the retry policy.

type TimeoutPolicy

type TimeoutPolicy struct {
	Duration time.Duration
}

TimeoutPolicy bounds a complete model operation. For streams, the deadline remains active until the stream finishes rather than only covering setup. The wrapper can return on time even if a custom provider ignores context, though such a provider may continue its work in the background.

func (TimeoutPolicy) Wrap

func (p TimeoutPolicy) Wrap(next models.Agent) (models.Agent, error)

Wrap applies the timeout policy.

type TokenBudget

type TokenBudget struct {
	// contains filtered or unexported fields
}

TokenBudget is a concurrency-safe, reusable estimated-token allowance. Input charges are rejected before a model call. Output charges are recorded after generation, so Used may exceed Max when a non-streaming provider returns more output than remained.

func NewTokenBudget

func NewTokenBudget(max int64, estimator TokenEstimator) (*TokenBudget, error)

NewTokenBudget creates a token budget. A nil estimator uses ApproximateTokenCount.

func TokenBudgetFromContext

func TokenBudgetFromContext(ctx context.Context) (*TokenBudget, bool)

TokenBudgetFromContext returns the budget associated with ctx.

func (*TokenBudget) Max

func (b *TokenBudget) Max() int64

Max returns the configured allowance.

func (*TokenBudget) Remaining

func (b *TokenBudget) Remaining() int64

Remaining returns the estimated tokens still available, clamped to zero.

func (*TokenBudget) Reset

func (b *TokenBudget) Reset()

Reset clears all charges. It is safe to call concurrently, though callers should normally reset only when no requests using the budget are active.

func (*TokenBudget) Used

func (b *TokenBudget) Used() int64

Used returns the estimated tokens charged so far.

type TokenBudgetError

type TokenBudgetError struct {
	Phase     string
	Max       int64
	Used      int64
	Requested int64
}

TokenBudgetError describes a rejected budget charge.

func (*TokenBudgetError) Error

func (e *TokenBudgetError) Error() string

func (*TokenBudgetError) Unwrap

func (e *TokenBudgetError) Unwrap() error

Unwrap supports errors.Is(err, ErrTokenBudgetExceeded).

type TokenBudgetPolicy

type TokenBudgetPolicy struct {
	Budget *TokenBudget
}

TokenBudgetPolicy enforces the context budget when present, otherwise Budget. A nil fallback is valid and makes the policy context-only.

func (TokenBudgetPolicy) Wrap

func (p TokenBudgetPolicy) Wrap(next models.Agent) (models.Agent, error)

Wrap applies the token-budget policy.

type TokenEstimator

type TokenEstimator func(string) int64

TokenEstimator estimates the number of tokens in text. Provider tokenizers can be supplied when exact accounting is required.