shudao-main/shudao-chat-go/tests/token_calculation_test.go

package tests

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"testing"
)

// TestQwenTokenCalculation 测试千问模型的token计算规则
func TestQwenTokenCalculation(t *testing.T) {
	// 加载测试配置
	config := LoadConfig()
	baseURL := config.Qwen3BaseURL
	model := config.Qwen3Model

	// 构建询问token计算规则的请求
	request := Qwen3ChatRequest{
		Model:  model,
		Stream: false,
		Messages: []ChatMessage{
			{
				Role:    "user",
				Content: "请详细解释一下你的token计算规则。具体来说：1. 中文字符如何计算token？2. 英文字符如何计算token？3. 标点符号如何计算token？4. JSON结构字符（如{、}、[、]、\"、:、,）如何计算token？5. 换行符和空格如何计算token？请给出具体的比例或规则。",
			},
		},
	}

	jsonData, err := json.Marshal(request)
	if err != nil {
		t.Fatalf("序列化请求失败: %v", err)
	}

	t.Logf("发送的请求: %s", string(jsonData))

	// 发送POST请求
	resp, err := http.Post(
		fmt.Sprintf("%s/v1/chat/completions", baseURL),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		t.Fatalf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	t.Logf("响应状态码: %d", resp.StatusCode)

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("请求失败，状态码: %d, 响应: %s", resp.StatusCode, string(body))
	}

	// 读取完整的响应内容
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("读取响应失败: %v", err)
	}

	// 解析JSON响应
	var response Qwen3ChatResponse
	if err := json.Unmarshal(body, &response); err != nil {
		t.Fatalf("解析JSON响应失败: %v", err)
	}

	t.Logf("✅ Token计算规则查询成功！")
	t.Logf("千问模型的回复:")
	t.Logf("==========================================")
	t.Logf("%s", response.Choices[0].Message.Content)
	t.Logf("==========================================")
}

// TestQwenTokenLimit 测试千问模型的具体token限制
func TestQwenTokenLimit(t *testing.T) {
	// 加载测试配置
	config := LoadConfig()
	baseURL := config.Qwen3BaseURL
	model := config.Qwen3Model

	// 构建询问token限制的请求
	request := Qwen3ChatRequest{
		Model:  model,
		Stream: false,
		Messages: []ChatMessage{
			{
				Role:    "user",
				Content: "请告诉我你的最大输入token限制是多少？包括输入和输出的总限制。",
			},
		},
	}

	jsonData, err := json.Marshal(request)
	if err != nil {
		t.Fatalf("序列化请求失败: %v", err)
	}

	t.Logf("发送的请求: %s", string(jsonData))

	// 发送POST请求
	resp, err := http.Post(
		fmt.Sprintf("%s/v1/chat/completions", baseURL),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		t.Fatalf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	t.Logf("响应状态码: %d", resp.StatusCode)

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("请求失败，状态码: %d, 响应: %s", resp.StatusCode, string(body))
	}

	// 读取完整的响应内容
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("读取响应失败: %v", err)
	}

	// 解析JSON响应
	var response Qwen3ChatResponse
	if err := json.Unmarshal(body, &response); err != nil {
		t.Fatalf("解析JSON响应失败: %v", err)
	}

	t.Logf("✅ Token限制查询成功！")
	t.Logf("千问模型的回复:")
	t.Logf("==========================================")
	t.Logf("%s", response.Choices[0].Message.Content)
	t.Logf("==========================================")
}

// TestQwenTokenEstimation 测试具体的token估算
func TestQwenTokenEstimation(t *testing.T) {
	// 加载测试配置
	config := LoadConfig()
	baseURL := config.Qwen3BaseURL
	model := config.Qwen3Model

	// 测试文本
	testTexts := []string{
		"你好世界",                      // 4个中文字符
		"Hello World",               // 11个英文字符
		"你好，Hello World！",           // 混合文本
		`{"name": "张三", "age": 25}`, // JSON格式
		"这是一个测试\n包含换行符\t和制表符", // 包含特殊字符
	}

	for i, testText := range testTexts {
		t.Run(fmt.Sprintf("测试文本%d", i+1), func(t *testing.T) {
			request := Qwen3ChatRequest{
				Model:  model,
				Stream: false,
				Messages: []ChatMessage{
					{
						Role:    "user",
						Content: fmt.Sprintf("请告诉我这段文本有多少个token：\"%s\"。请只回答数字，不要其他解释。", testText),
					},
				},
			}

			jsonData, err := json.Marshal(request)
			if err != nil {
				t.Fatalf("序列化请求失败: %v", err)
			}

			// 发送POST请求
			resp, err := http.Post(
				fmt.Sprintf("%s/v1/chat/completions", baseURL),
				"application/json",
				bytes.NewBuffer(jsonData),
			)
			if err != nil {
				t.Fatalf("发送请求失败: %v", err)
			}
			defer resp.Body.Close()

			if resp.StatusCode != http.StatusOK {
				body, _ := io.ReadAll(resp.Body)
				t.Fatalf("请求失败，状态码: %d, 响应: %s", resp.StatusCode, string(body))
			}

			// 读取完整的响应内容
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				t.Fatalf("读取响应失败: %v", err)
			}

			// 解析JSON响应
			var response Qwen3ChatResponse
			if err := json.Unmarshal(body, &response); err != nil {
				t.Fatalf("解析JSON响应失败: %v", err)
			}

			t.Logf("测试文本: \"%s\"", testText)
			t.Logf("字符长度: %d", len(testText))
			t.Logf("千问估算token数: %s", response.Choices[0].Message.Content)
			t.Logf("---")
		})
	}
}