AutoZero/llm/usage_test.go

package llm

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
)

// 三家 provider 的 token 记账必须归一到同一口径（见 Usage 注释）：
// InputTokens 含缓存段。Anthropic 尤其容易写错——它把命中/写入单列，
// 照抄 input_tokens 就会在缓存命中时把 10 万上下文记成 5 千。

func TestAnthropicUsageIncludesCache(t *testing.T) {
	sse := strings.Join([]string{
		`data: {"type":"message_start","message":{"usage":{"input_tokens":120,"cache_read_input_tokens":90000,"cache_creation_input_tokens":500}}}`,
		"",
		`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		"",
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}`,
		"",
		`data: {"type":"content_block_stop","index":0}`,
		"",
		// message_delta 常常不回缓存明细：必须沿用 message_start 学到的数
		`data: {"type":"message_delta","usage":{"output_tokens":7}}`,
		"",
		`data: {"type":"message_stop"}`,
		"",
	}, "\n")
	var got map[string]any
	srv := captureServer(&got, sse)
	defer srv.Close()

	p := &AnthropicProvider{BaseURL: srv.URL, APIKey: "k", Model: "claude-x"}
	msg, err := p.Stream(context.Background(), "sys",
		[]Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil)
	if err != nil {
		t.Fatal(err)
	}
	u := msg.Usage
	if u == nil {
		t.Fatal("没有 usage")
	}
	if u.InputTokens != 90620 {
		t.Errorf("InputTokens = %d, want 90620（120 + cache_read 90000 + cache_write 500）", u.InputTokens)
	}
	if u.CacheReadTokens != 90000 || u.CacheCreationTokens != 500 {
		t.Errorf("缓存明细 = %d/%d, want 90000/500", u.CacheReadTokens, u.CacheCreationTokens)
	}
	if u.OutputTokens != 7 {
		t.Errorf("OutputTokens = %d, want 7", u.OutputTokens)
	}
	if msg.Content != "ok" {
		t.Errorf("Content = %q", msg.Content)
	}
}

func TestOpenAIUsageDetails(t *testing.T) {
	sse := strings.Join([]string{
		`data: {"choices":[{"delta":{"content":"hi"}}]}`,
		"",
		`data: {"choices":[{"finish_reason":"stop"}],"usage":{"prompt_tokens":1000,"completion_tokens":20,` +
			`"prompt_tokens_details":{"cached_tokens":950},"completion_tokens_details":{"reasoning_tokens":8}}}`,
		"",
		`data: [DONE]`,
		"",
	}, "\n")
	var got map[string]any
	srv := captureServer(&got, sse)
	defer srv.Close()

	p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "gpt-x"}
	msg, err := p.Stream(context.Background(), "sys",
		[]Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil)
	if err != nil {
		t.Fatal(err)
	}
	u := msg.Usage
	if u == nil {
		t.Fatal("没有 usage")
	}
	// prompt_tokens 本身已含缓存：总量不许二次相加。
	if u.InputTokens != 1000 || u.CacheReadTokens != 950 || u.ReasoningTokens != 8 || u.OutputTokens != 20 {
		t.Errorf("usage 归一错误：%+v", u)
	}
}

func TestResponsesUsageDetails(t *testing.T) {
	sse := strings.Join([]string{
		`data: {"type":"response.output_text.delta","delta":"hi"}`,
		"",
		`data: {"type":"response.completed","response":{"status":"completed","usage":` +
			`{"input_tokens":800,"output_tokens":15,"input_tokens_details":{"cached_tokens":700},` +
			`"output_tokens_details":{"reasoning_tokens":5}}}}`,
		"",
	}, "\n")
	var got map[string]any
	srv := captureServer(&got, sse)
	defer srv.Close()

	p := &ResponsesProvider{BaseURL: srv.URL, APIKey: "k", Model: "rsp-x"}
	msg, err := p.Stream(context.Background(), "sys",
		[]Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil)
	if err != nil {
		t.Fatal(err)
	}
	u := msg.Usage
	if u == nil {
		t.Fatal("没有 usage")
	}
	if u.InputTokens != 800 || u.CacheReadTokens != 700 || u.ReasoningTokens != 5 || u.OutputTokens != 15 {
		t.Errorf("usage 归一错误：%+v", u)
	}
}

// TestProviderIdentity：三家都自报协议与当前模型（会话落盘的归因戳）。
// 换了模型（SetModel）之后 Identity 必须跟着变，否则戳是假的。
func TestProviderIdentity(t *testing.T) {
	var cases = []struct {
		p        Provider
		wantFmt  string
		wantAttr string
	}{
		{&OpenAIProvider{Model: "gpt-a"}, FormatChatCompletions, "gpt-a"},
		{&AnthropicProvider{Model: "claude-a"}, FormatAnthropic, "claude-a"},
		{&ResponsesProvider{Model: "rsp-a"}, FormatResponses, "rsp-a"},
	}
	for _, c := range cases {
		id, ok := c.p.(Identifier)
		if !ok {
			t.Fatalf("%T 没实现 Identifier", c.p)
		}
		f, m := id.Identity()
		if f != c.wantFmt || m != c.wantAttr {
			t.Errorf("%T Identity = %q/%q, want %q/%q", c.p, f, m, c.wantFmt, c.wantAttr)
		}
		if sw, ok := c.p.(ModelSwitcher); !ok {
			t.Errorf("%T 应可切模型", c.p)
		} else {
			sw.SetModel("后到")
			if _, m := id.Identity(); m != "后到" {
				t.Errorf("SetModel 后 Identity 仍报 %q", m)
			}
		}
	}
}

// TestUsageSerializesIntoTranscript：新增明细字段要能落进会话文件，
// 同时不给旧消息添噪声（omitempty）。
func TestUsageSerializesIntoTranscript(t *testing.T) {
	old := Message{Role: RoleAssistant, Usage: &Usage{InputTokens: 5, OutputTokens: 2}}
	b, err := json.Marshal(old)
	if err != nil {
		t.Fatal(err)
	}
	for _, noise := range []string{"cache_read_tokens", "reasoning_tokens", "format", "model"} {
		if strings.Contains(string(b), noise) {
			t.Errorf("零值不该序列化出 %q：%s", noise, b)
		}
	}
	withCache := Message{Role: RoleAssistant, Format: FormatAnthropic, Model: "claude-x",
		Usage: &Usage{InputTokens: 90620, OutputTokens: 7, CacheReadTokens: 90000}}
	b2, err := json.Marshal(withCache)
	if err != nil {
		t.Fatal(err)
	}
	for _, want := range []string{"cache_read_tokens", "90000", `"format":"anthropic"`, `"model":"claude-x"`} {
		if !strings.Contains(string(b2), want) {
			t.Errorf("序列化缺 %q：%s", want, b2)
		}
	}
}