MiFaZhan
93 lines
2.9 KiB
Go
93 lines
2.9 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
func TestResolveOpenAIBaseURL(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
baseURL string
|
|
endpoint string
|
|
want string
|
|
wantErr bool
|
|
}{
|
|
{name: "base url", baseURL: "https://xueyangai.com/v1/", want: "https://xueyangai.com/v1"},
|
|
{name: "legacy endpoint", endpoint: "https://xueyangai.com/v1/chat/completions", want: "https://xueyangai.com/v1"},
|
|
{name: "base url preferred", baseURL: "https://example.com/v1", endpoint: "https://legacy.example.com/v1/chat/completions", want: "https://example.com/v1"},
|
|
{name: "invalid", baseURL: "xueyangai.com/v1", wantErr: true},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
got, err := ResolveOpenAIBaseURL(test.baseURL, test.endpoint)
|
|
if test.wantErr {
|
|
if err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("ResolveOpenAIBaseURL() error = %v", err)
|
|
}
|
|
if got != test.want {
|
|
t.Fatalf("ResolveOpenAIBaseURL() = %s, want %s", got, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleChatModel(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
|
if request.URL.Path != "/v1/chat/completions" {
|
|
t.Errorf("path = %s, want /v1/chat/completions", request.URL.Path)
|
|
}
|
|
if request.Header.Get("Authorization") != "Bearer test-key" {
|
|
t.Errorf("authorization = %s", request.Header.Get("Authorization"))
|
|
}
|
|
var body map[string]any
|
|
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
|
|
t.Errorf("decode request: %v", err)
|
|
}
|
|
if body["model"] != "test-model" {
|
|
t.Errorf("model = %v, want test-model", body["model"])
|
|
}
|
|
writer.Header().Set("Content-Type", "application/json")
|
|
_, _ = writer.Write([]byte(`{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"test-model","choices":[{"index":0,"message":{"role":"assistant","content":"{\"ok\":true}"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
chatModel, err := NewOpenAICompatibleChatModel(context.Background(), OpenAICompatibleConfig{
|
|
APIKey: "test-key",
|
|
BaseURL: server.URL + "/v1",
|
|
Model: "test-model",
|
|
Timeout: 3 * time.Second,
|
|
MaxTokens: 64,
|
|
JSONOutput: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAICompatibleChatModel() error = %v", err)
|
|
}
|
|
response, err := chatModel.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")})
|
|
if err != nil {
|
|
t.Fatalf("Generate() error = %v", err)
|
|
}
|
|
if response.Content != `{"ok":true}` {
|
|
t.Fatalf("response content = %s", response.Content)
|
|
}
|
|
}
|
|
|
|
func TestNewChatModelRejectsUnsupportedProvider(t *testing.T) {
|
|
_, err := NewChatModel(context.Background(), ChatModelConfig{Provider: "unsupported"})
|
|
if err == nil {
|
|
t.Fatal("expected unsupported provider error")
|
|
}
|
|
}
|