xinmi/internal/models/chat/request_budget.go

package chat

import (
	"context"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
)

// Keep the supported set explicit: adding a factory adapter cannot silently
// bypass a controlled task's actual-send budget.
type requestBudgetChat struct {
	inner     Chat
	supported bool
}

func wrapChatRequestBudget(inner Chat, err error) (Chat, error) {
	if err != nil || inner == nil {
		return inner, err
	}
	supported := false
	switch inner.(type) {
	case *RemoteAPIChat, *AnthropicChat:
		supported = true
	}
	return &requestBudgetChat{inner: inner, supported: supported}, nil
}

func (w *requestBudgetChat) GetModelName() string { return w.inner.GetModelName() }
func (w *requestBudgetChat) GetModelID() string   { return w.inner.GetModelID() }

func (w *requestBudgetChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
	if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
		return nil, err
	}
	return w.inner.Chat(ctx, messages, opts)
}

func (w *requestBudgetChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
	if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
		return nil, err
	}
	return w.inner.ChatStream(ctx, messages, opts)
}