120 lines
3.8 KiB
Go
120 lines
3.8 KiB
Go
package handler
|
||
|
||
import (
|
||
"context"
|
||
"strings"
|
||
"testing"
|
||
|
||
"github.com/qianfree/team-api/relay/common"
|
||
)
|
||
|
||
type modelMappingStub struct {
|
||
models map[string]bool
|
||
calls []string
|
||
}
|
||
|
||
func (s *modelMappingStub) GetModelMapping(_ context.Context, modelName string) (string, string, error) {
|
||
s.calls = append(s.calls, modelName)
|
||
if !s.models[modelName] {
|
||
return "", "", common.ErrModelNotFound
|
||
}
|
||
return modelName, "chat", nil
|
||
}
|
||
|
||
func TestResolveRelayModel(t *testing.T) {
|
||
// 后缀语法已移除：目录名即请求名，不再有「字面模型优先 / 剥后缀回退」两段查找
|
||
t.Run("catalog model resolves literally", func(t *testing.T) {
|
||
provider := &modelMappingStub{models: map[string]bool{"gpt-4o": true}}
|
||
lookup, err := resolveRelayModel(t.Context(), provider, "gpt-4o")
|
||
if err != nil {
|
||
t.Fatalf("resolveRelayModel() error = %v", err)
|
||
}
|
||
if lookup != "gpt-4o" {
|
||
t.Errorf("lookup model = %q, want %q", lookup, "gpt-4o")
|
||
}
|
||
})
|
||
|
||
// 名字带 -max 等字样的真实目录模型按字面命中，不做任何后缀解释
|
||
t.Run("literal model ending in effort-like word", func(t *testing.T) {
|
||
provider := &modelMappingStub{models: map[string]bool{"qwen3.8-max": true}}
|
||
lookup, err := resolveRelayModel(t.Context(), provider, "qwen3.8-max")
|
||
if err != nil {
|
||
t.Fatalf("resolveRelayModel() error = %v", err)
|
||
}
|
||
if lookup != "qwen3.8-max" {
|
||
t.Errorf("lookup model = %q, want %q", lookup, "qwen3.8-max")
|
||
}
|
||
if len(provider.calls) != 1 || provider.calls[0] != "qwen3.8-max" {
|
||
t.Errorf("lookup calls = %v, want exactly one literal lookup", provider.calls)
|
||
}
|
||
})
|
||
|
||
// 目录中不存在即报错，不再尝试剥后缀回退（"o3-max" 不会命中 "o3"）
|
||
t.Run("missing model returns error without suffix fallback", func(t *testing.T) {
|
||
provider := &modelMappingStub{models: map[string]bool{"o3": true}}
|
||
_, err := resolveRelayModel(t.Context(), provider, "o3-max")
|
||
if err == nil {
|
||
t.Fatal("want error for model absent from catalog")
|
||
}
|
||
if len(provider.calls) != 1 {
|
||
t.Errorf("lookup calls = %v, want single literal lookup", provider.calls)
|
||
}
|
||
})
|
||
}
|
||
|
||
// TestEstimateInputTokens 多模态请求体的预扣估算：base64 内联媒体必须按固定成本计入，
|
||
// 不能按字节折算——一张 2MB 的图按 len/4 会被估成约 70 万 token，预扣直接冻穿余额。
|
||
func TestEstimateInputTokens(t *testing.T) {
|
||
// 模拟 2MB 的 base64 图片载荷
|
||
payload := strings.Repeat("A", 2<<20)
|
||
|
||
cases := []struct {
|
||
name string
|
||
body string
|
||
want int
|
||
}{
|
||
{
|
||
"纯文本按字节折算",
|
||
strings.Repeat("x", 400),
|
||
100,
|
||
},
|
||
{
|
||
"空体",
|
||
"",
|
||
0,
|
||
},
|
||
{
|
||
// 文本骨架 40 字节 + 1 个媒体
|
||
"单图不按字节折算",
|
||
`{"url":"data:image/png;base64,` + payload + `"}`,
|
||
(len(`{"url":"data:image/png;base64,`+`"}`))/4 + inlineMediaTokenCost,
|
||
},
|
||
{
|
||
"双图按个数累加",
|
||
`{"a":"data:image/png;base64,` + payload + `","b":"data:image/png;base64,` + payload + `"}`,
|
||
(len(`{"a":"data:image/png;base64,`+`","b":"data:image/png;base64,`+`"}`))/4 + 2*inlineMediaTokenCost,
|
||
},
|
||
{
|
||
// 畸形 data URL（无 base64 标记）退化为纯文本路径
|
||
"畸形 data URL",
|
||
`{"url":"data:image/png,` + strings.Repeat("A", 400) + `"}`,
|
||
len(`{"url":"data:image/png,`+`"}`)/4 + 100,
|
||
},
|
||
}
|
||
|
||
for _, c := range cases {
|
||
t.Run(c.name, func(t *testing.T) {
|
||
if got := estimateInputTokens([]byte(c.body)); got != c.want {
|
||
t.Errorf("estimateInputTokens() = %d, want %d", got, c.want)
|
||
}
|
||
})
|
||
}
|
||
|
||
// 回归保护：2MB 图片的估算必须远低于按字节折算的量级
|
||
naive := len(`{"url":"data:image/png;base64,`+payload+`"}`) / 4
|
||
got := estimateInputTokens([]byte(`{"url":"data:image/png;base64,` + payload + `"}`))
|
||
if got >= naive/10 {
|
||
t.Errorf("2MB 图片估算 %d tokens，与按字节折算的 %d 处于同一量级，预扣仍会冻穿余额", got, naive)
|
||
}
|
||
}
|