101 lines
5.2 KiB
Go
101 lines
5.2 KiB
Go
package dto
|
|
|
|
import (
|
|
"testing"
|
|
|
|
kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// Agent loops replay tool definitions, tool calls, and tool output on every
|
|
// turn; the prompt estimate must see all of them as the model reads them.
|
|
func TestTokenCountMetaCountsAgentLoopContent(t *testing.T) {
|
|
const toolOutput = "package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"
|
|
const imageData = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
raw string
|
|
request interface{ GetTokenCountMeta() *types.TokenCountMeta }
|
|
contains []string
|
|
files int
|
|
detail string
|
|
toolCount int
|
|
}{
|
|
{
|
|
name: "chat",
|
|
raw: `{"model":"m","messages":[
|
|
{"role":"user","content":[{"type":"text","text":"show main.go"},{"type":"image_url","image_url":{"url":"data:image/png;base64,` + imageData + `"}}]},
|
|
{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"main.go\"}"}}]},
|
|
{"role":"tool","tool_call_id":"c1","content":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"}],
|
|
"tools":[{"type":"function","function":{"name":"read_file","description":"Read a file","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}]}`,
|
|
request: &GeneralOpenAIRequest{},
|
|
contains: []string{"read_file", `{"path":"main.go"}`, toolOutput, "Read a file", `"properties":{"path":{"type":"string"}}`},
|
|
// An omitted detail stays empty so the estimate uses the upstream default (auto).
|
|
files: 1,
|
|
toolCount: 1,
|
|
},
|
|
{
|
|
name: "responses",
|
|
raw: `{"model":"m","instructions":"Be brief.\nUse tools.","input":[
|
|
{"type":"message","role":"user","content":[{"type":"input_text","text":"show main.go"}]},
|
|
{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Reading it now."}]},
|
|
{"type":"function_call","call_id":"c1","name":"read_file","arguments":"{\"path\":\"main.go\"}"},
|
|
{"type":"function_call_output","call_id":"c1","output":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"},
|
|
{"type":"custom_tool_call","call_id":"c2","name":"apply_patch","input":"*** Begin Patch"},
|
|
{"type":"custom_tool_call_output","call_id":"c2","output":[{"type":"input_text","text":"Done!"}]}]}`,
|
|
request: &OpenAIResponsesRequest{},
|
|
contains: []string{"Be brief.\nUse tools.", "Reading it now.", "read_file", `{"path":"main.go"}`, toolOutput, "apply_patch", "*** Begin Patch", "Done!"},
|
|
},
|
|
{
|
|
name: "claude",
|
|
raw: `{"model":"m","max_tokens":1,"messages":[
|
|
{"role":"user","content":"show main.go"},
|
|
{"role":"assistant","content":[{"type":"tool_use","id":"c1","name":"read_file","input":{"path":"main.go"}}]},
|
|
{"role":"user","content":[
|
|
{"type":"tool_result","tool_use_id":"c1","content":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"},
|
|
{"type":"tool_result","tool_use_id":"c2","content":[{"type":"text","text":"screenshot taken"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"` + imageData + `"}}]}]}],
|
|
"tools":[{"name":"read_file","description":"Read a file","input_schema":{"type":"object","properties":{"path":{"type":"string"}}}}]}`,
|
|
request: &ClaudeRequest{},
|
|
contains: []string{"read_file", `{"path":"main.go"}`, toolOutput, "screenshot taken", "Read a file", `"properties":{"path":{"type":"string"}}`},
|
|
files: 1,
|
|
toolCount: 1,
|
|
},
|
|
{
|
|
name: "gemini",
|
|
raw: `{"systemInstruction":{"parts":[{"text":"Be brief."}]},"contents":[
|
|
{"role":"user","parts":[{"text":"show main.go"}]},
|
|
{"role":"model","parts":[{"functionCall":{"name":"read_file","args":{"path":"main.go"}}}]},
|
|
{"role":"user","parts":[{"functionResponse":{"name":"read_file","response":{"output":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}","exit_code":0}}}]},
|
|
{"role":"user","parts":[{"functionResponse":{"name":"screenshot","response":{},"parts":[{"inlineData":{"mimeType":"image/png","data":"` + imageData + `"}}]}}]}],
|
|
"generationConfig":{"mediaResolution":"MEDIA_RESOLUTION_LOW"},
|
|
"tools":[{"functionDeclarations":[{"name":"read_file","description":"Read a file"}]}]}`,
|
|
request: &GeminiChatRequest{},
|
|
contains: []string{"Be brief.", "read_file", `{"path":"main.go"}`, toolOutput, "Read a file"},
|
|
files: 1,
|
|
detail: "MEDIA_RESOLUTION_LOW",
|
|
},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
require.NoError(t, kitutil.Unmarshal([]byte(tc.raw), tc.request))
|
|
meta := tc.request.GetTokenCountMeta()
|
|
for _, want := range tc.contains {
|
|
assert.Contains(t, meta.CombineText, want)
|
|
}
|
|
// Tool output and instructions are counted as raw text, not JSON-escaped,
|
|
// and images stay out of the text.
|
|
assert.NotContains(t, meta.CombineText, `\n`)
|
|
assert.NotContains(t, meta.CombineText, "map[")
|
|
assert.NotContains(t, meta.CombineText, imageData)
|
|
require.Len(t, meta.Files, tc.files)
|
|
for _, file := range meta.Files {
|
|
assert.Equal(t, types.FileTypeImage, file.FileType)
|
|
assert.Equal(t, tc.detail, file.Detail)
|
|
}
|
|
assert.Equal(t, tc.toolCount, meta.ToolsCount)
|
|
})
|
|
}
|
|
}
|