45 lines
1.3 KiB
Go
45 lines
1.3 KiB
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)
|
||
|
|
}
|