280 lines
8.4 KiB
Go
280 lines
8.4 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"ai-agent-scaffold-go/internal/model"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
// ============================================================
|
|
// 辅助函数
|
|
// ============================================================
|
|
|
|
func newTestClient(url string) *OpenAIClient {
|
|
return NewOpenAIClient(url, "test-key", "gpt-4", 5*time.Second)
|
|
}
|
|
|
|
// ============================================================
|
|
// Generate 测试
|
|
// ============================================================
|
|
|
|
func TestOpenAIClient_Generate_ReturnsContent(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
assert.Equal(t, "POST", r.Method)
|
|
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
|
|
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{
|
|
{"message": map[string]any{"content": "hello world"}},
|
|
},
|
|
})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
reply, err := client.Generate(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "hello world", reply.Content)
|
|
assert.Empty(t, reply.ToolCalls)
|
|
}
|
|
|
|
func TestOpenAIClient_Generate_ToolCalls(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{
|
|
{"message": map[string]any{
|
|
"content": "",
|
|
"tool_calls": []map[string]any{
|
|
{"id": "call_1", "type": "function", "function": map[string]any{
|
|
"name": "search",
|
|
"arguments": `{"query":"test"}`,
|
|
}},
|
|
},
|
|
}},
|
|
},
|
|
})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
reply, err := client.Generate(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "search for test"},
|
|
}, []ToolDef{{Name: "search", Description: "search tool"}})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, reply.ToolCalls, 1)
|
|
assert.Equal(t, "call_1", reply.ToolCalls[0].ID)
|
|
assert.Equal(t, "search", reply.ToolCalls[0].Name)
|
|
assert.Equal(t, `{"query":"test"}`, reply.ToolCalls[0].Arguments)
|
|
}
|
|
|
|
func TestOpenAIClient_Generate_Upstream500_ReturnsError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
fmt.Fprint(w, "internal error")
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
_, err := client.Generate(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "500")
|
|
}
|
|
|
|
func TestOpenAIClient_Generate_EmptyChoices_ReturnsError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{"choices": []any{}})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
_, err := client.Generate(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "no choices")
|
|
}
|
|
|
|
// ============================================================
|
|
// Stream 测试
|
|
// ============================================================
|
|
|
|
func TestOpenAIClient_Stream_ReturnsChunks(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
assert.Equal(t, "POST", r.Method)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
|
w.(http.Flusher).Flush()
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
var texts []string
|
|
for ev := range events {
|
|
if ev.Delta != "" {
|
|
texts = append(texts, ev.Delta)
|
|
}
|
|
if ev.Done {
|
|
break
|
|
}
|
|
}
|
|
// 检查错误通道
|
|
for err := range errs {
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
assert.Equal(t, []string{"hel", "lo"}, texts)
|
|
}
|
|
|
|
func TestOpenAIClient_Stream_ToolCallDeltas_MergesCorrectly(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
// 工具调用分多个 chunk 发送
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"search\"}}]}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"q\"}}]}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"uery\\\":\\\"test\\\"}\"}}]}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
|
w.(http.Flusher).Flush()
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
var toolCalls []model.ChatToolCall
|
|
for ev := range events {
|
|
if len(ev.ToolCalls) > 0 {
|
|
toolCalls = append(toolCalls, ev.ToolCalls...)
|
|
}
|
|
if ev.Done {
|
|
break
|
|
}
|
|
}
|
|
for range errs {
|
|
}
|
|
|
|
assert.Len(t, toolCalls, 1)
|
|
assert.Equal(t, "call_1", toolCalls[0].ID)
|
|
assert.Equal(t, "search", toolCalls[0].Name)
|
|
assert.Equal(t, `{"query":"test"}`, toolCalls[0].Arguments)
|
|
}
|
|
|
|
func TestOpenAIClient_Stream_Upstream500_ReturnsError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
fmt.Fprint(w, "error")
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := newTestClient(srv.URL)
|
|
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil)
|
|
|
|
// 消费通道直到关闭
|
|
var streamErr error
|
|
done := make(chan struct{})
|
|
go func() {
|
|
for range events {
|
|
}
|
|
close(done)
|
|
}()
|
|
for err := range errs {
|
|
streamErr = err
|
|
}
|
|
<-done
|
|
|
|
assert.Error(t, streamErr)
|
|
assert.Contains(t, streamErr.Error(), "500")
|
|
}
|
|
|
|
// ============================================================
|
|
// encodeMessages / encodeTools 测试
|
|
// ============================================================
|
|
|
|
func TestEncodeMessages_WithToolCalls(t *testing.T) {
|
|
msgs := []model.ChatMessage{
|
|
{Role: model.ChatRoleAssistant, ToolCalls: []model.ChatToolCall{
|
|
{ID: "c1", Name: "search", Arguments: `{"q":"x"}`},
|
|
}},
|
|
}
|
|
encoded := encodeMessages(msgs)
|
|
assert.Len(t, encoded, 1)
|
|
assert.Equal(t, "assistant", encoded[0]["role"])
|
|
calls := encoded[0]["tool_calls"].([]map[string]any)
|
|
assert.Len(t, calls, 1)
|
|
assert.Equal(t, "c1", calls[0]["id"])
|
|
}
|
|
|
|
func TestEncodeTools_CorrectFormat(t *testing.T) {
|
|
tools := []ToolDef{{Name: "search", Description: "search tool"}}
|
|
encoded := encodeTools(tools)
|
|
assert.Len(t, encoded, 1)
|
|
assert.Equal(t, "function", encoded[0]["type"])
|
|
fn := encoded[0]["function"].(map[string]any)
|
|
assert.Equal(t, "search", fn["name"])
|
|
assert.Equal(t, "search tool", fn["description"])
|
|
}
|
|
|
|
func TestEncodeTools_EmptyDescription_UsesFallback(t *testing.T) {
|
|
tools := []ToolDef{{Name: "mytool"}}
|
|
encoded := encodeTools(tools)
|
|
fn := encoded[0]["function"].(map[string]any)
|
|
assert.Contains(t, fn["description"], "mytool")
|
|
}
|
|
|
|
// ============================================================
|
|
// truncate 测试
|
|
// ============================================================
|
|
|
|
func TestTruncate_ShortString_Unchanged(t *testing.T) {
|
|
assert.Equal(t, "abc", truncate("abc", 10))
|
|
}
|
|
|
|
func TestTruncate_LongString_Truncated(t *testing.T) {
|
|
result := truncate("abcdefghijk", 5)
|
|
assert.Equal(t, "abcde...", result)
|
|
}
|
|
|
|
// ============================================================
|
|
// buildRequestBody 测试
|
|
// ============================================================
|
|
|
|
func TestBuildRequestBody_Stream(t *testing.T) {
|
|
body, err := buildRequestBody("gpt-4", []model.ChatMessage{
|
|
{Role: model.ChatRoleUser, Content: "hi"},
|
|
}, nil, true)
|
|
assert.NoError(t, err)
|
|
assert.Contains(t, string(body), `"stream":true`)
|
|
assert.Contains(t, string(body), `"model":"gpt-4"`)
|
|
}
|