ai_old/tests/chat_test.go

package tests

import (
	"apigo.cc/ai/ai"
	_ "apigo.cc/ai/openai"
	_ "apigo.cc/ai/zhipu"
	"fmt"
	"github.com/ssgo/config"
	"github.com/ssgo/u"
	"testing"
)

var agentConf = map[string]ai.AgentConfig{}

func init() {
	_ = config.LoadConfig("agent", &agentConf)
}

func TestAgent(t *testing.T) {
	//testChat(t, "openai", "https://api.keya.pw/v1", keys.Openai)
	// 尚未支持
	//testCode(t, "openai", "https://api.keya.pw/v1", keys.Openai)
	// keya 消耗过大，计费不合理
	//testAskWithImage(t, "openai", "https://api.keya.pw/v1", keys.Openai, "4032.jpg")
	// keya 不支持
	//testMakeImage(t, "openai", "https://api.keya.pw/v1", keys.Openai, "冬天大雪纷飞，一个男人身穿军绿色棉大衣，戴着红色围巾和绿色帽子走在铺面大雪的小镇路上", "")

	testChat(t, "zhipu")
	//testCode(t, "zhipu", "", keys.Zhipu)
	//testSearch(t, "zhipu", "", keys.Zhipu)

	// 测试图片识别
	//testAskWithImage(t, "zhipu", "", keys.Zhipu, "4032.jpg")

	// 视频似乎尚不支持 glm-4v-plus
	//testAskWithVideo(t, "zhipu", "", keys.Zhipu, "glm-4v", "1080.mp4")

	//testMakeImage(t, "zhipu", "", keys.Zhipu, "冬天大雪纷飞，一个男人身穿军绿色棉大衣，戴着红色围巾和绿色帽子走在铺面大雪的小镇路上", "")
	//testMakeVideo(t, "zhipu", "", keys.Zhipu, "大雪纷飞，男人蹦蹦跳跳", "https://aigc-files.bigmodel.cn/api/cogview/20240904133130c4b7121019724aa3_0.png")
}

func testChat(t *testing.T, agent string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := ag.FastAsk(ai.Messages().User().Text("你是什么模型，请给出具体名称、版本号").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testCode(t *testing.T, agent string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := ag.CodeInterpreterAsk(ai.Messages().User().Text("计算[5,10,20,700,99,310,978,100]的平均值和方差。").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testSearch(t *testing.T, agent string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := ag.WebSearchAsk(ai.Messages().User().Text("今天上海的天气怎么样？").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testAskWithImage(t *testing.T, agent, imageFile string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	ask := `请回答：
1、我用的是什么操作系统？
2、现在的日期时间是？请给出年月日。
3、正在用什么软件播放什么歌？谁演唱的？歌曲的大意是？
4、后面的浏览器中正在浏览什么内容？猜测一下我浏览这个网页是想干嘛？
`
	r, usage, err := ag.MultiAsk(ai.Messages().User().Text(ask).Image("data:image/jpeg;base64,"+u.Base64(u.ReadFileBytesN(imageFile))).Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testAskWithVideo(t *testing.T, agent, videoFile string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	ask := `请回答：
1、正在用什么软件播放什么歌？谁演唱的？歌曲的大意是？
4、后面的浏览器中正在浏览什么内容？猜测一下我浏览这个网页是想干嘛？
`

	r, usage, err := ag.MultiAsk(ai.Messages().User().Text(ask).Video("data:video/mp4,"+u.Base64(u.ReadFileBytesN(videoFile))).Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testMakeImage(t *testing.T, agent, prompt, refImage string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	r, err := ag.FastMakeImage(prompt, "1024x1024", "")

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	for i, v := range r {
		fmt.Println("result:", i, v)
	}
}

func testMakeVideo(t *testing.T, agent, prompt, refImage string) {
	ag := ai.GetAgent(agent)

	if ag == nil {
		t.Fatal("agent is nil")
	}

	r, covers, err := ag.FastMakeVideo(prompt, "1280x720", refImage)

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	for i, v := range r {
		fmt.Println("result:", i, v, covers[i])
	}
}