wehub-resource-sync
e071084ebe
govulncheck / govulncheck (push) Waiting to run
Harness (E2E) / Harnesses (mock LLM) (push) Waiting to run
Harness (E2E) / Provider harnesses (live LLM conformance) (push) Waiting to run
Lint / golangci-lint (push) Waiting to run
Run Tests / Unit Tests (push) Waiting to run
Run Tests / Etcd Integration Tests (push) Waiting to run
294 lines
9.2 KiB
Go
294 lines
9.2 KiB
Go
package ai_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"go-micro.dev/v6/ai"
|
|
_ "go-micro.dev/v6/ai/anthropic"
|
|
_ "go-micro.dev/v6/ai/atlascloud"
|
|
_ "go-micro.dev/v6/ai/gemini"
|
|
_ "go-micro.dev/v6/ai/groq"
|
|
_ "go-micro.dev/v6/ai/minimax"
|
|
_ "go-micro.dev/v6/ai/mistral"
|
|
_ "go-micro.dev/v6/ai/openai"
|
|
_ "go-micro.dev/v6/ai/together"
|
|
)
|
|
|
|
func TestStreamProvidersConformToOpenAICompatibleSSE(t *testing.T) {
|
|
providers := conformingStreamProviders(t)
|
|
|
|
for _, provider := range providers {
|
|
provider := provider
|
|
t.Run(provider, func(t *testing.T) {
|
|
var sawRequest bool
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sawRequest = true
|
|
if r.URL.Path != "/v1/chat/completions" {
|
|
t.Fatalf("path = %s, want /v1/chat/completions", r.URL.Path)
|
|
}
|
|
if got := r.Header.Get("Accept"); got != "text/event-stream" {
|
|
t.Fatalf("Accept = %q, want text/event-stream", got)
|
|
}
|
|
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
|
t.Fatalf("Authorization = %q, want bearer API key", got)
|
|
}
|
|
|
|
var body map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
if body["model"] == "" {
|
|
t.Fatal("request omitted model")
|
|
}
|
|
if body["stream"] != true {
|
|
t.Fatalf("stream = %#v, want true", body["stream"])
|
|
}
|
|
streamOptions, ok := body["stream_options"].(map[string]any)
|
|
if !ok || streamOptions["include_usage"] != true {
|
|
t.Fatalf("stream_options = %#v, want include_usage=true", body["stream_options"])
|
|
}
|
|
messages, ok := body["messages"].([]any)
|
|
if !ok || len(messages) != 4 {
|
|
t.Fatalf("messages = %#v, want system + history + prompt", body["messages"])
|
|
}
|
|
wantRoles := []string{"system", "user", "assistant", "user"}
|
|
for i, wantRole := range wantRoles {
|
|
message, ok := messages[i].(map[string]any)
|
|
if !ok || message["role"] != wantRole {
|
|
t.Fatalf("message[%d] = %#v, want role %q", i, messages[i], wantRole)
|
|
}
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = w.Write([]byte(": keepalive\n\n"))
|
|
_, _ = w.Write([]byte("event: ignored\n\n"))
|
|
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n"))
|
|
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n"))
|
|
_, _ = w.Write([]byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5}}\n\n"))
|
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
model := ai.New(provider, ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL))
|
|
if model == nil {
|
|
t.Fatalf("ai.New(%q) returned nil", provider)
|
|
}
|
|
stream, err := model.Stream(context.Background(), &ai.Request{
|
|
SystemPrompt: "system",
|
|
Messages: []ai.Message{
|
|
{Role: "user", Content: "previous question"},
|
|
{Role: "assistant", Content: "previous answer"},
|
|
},
|
|
Prompt: "current question",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Stream returned error: %v", err)
|
|
}
|
|
defer stream.Close()
|
|
if !sawRequest {
|
|
t.Fatal("server did not receive stream request")
|
|
}
|
|
|
|
assertStreamReply(t, stream, "hel")
|
|
assertStreamReply(t, stream, "lo")
|
|
usage, err := stream.Recv()
|
|
if err != nil {
|
|
t.Fatalf("usage chunk error: %v", err)
|
|
}
|
|
if usage.Reply != "" || usage.Usage != (ai.Usage{InputTokens: 3, OutputTokens: 2, TotalTokens: 5}) {
|
|
t.Fatalf("usage chunk = %#v", usage)
|
|
}
|
|
if _, err := stream.Recv(); !errors.Is(err, io.EOF) {
|
|
t.Fatalf("final error = %v, want EOF", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamProvidersCloseCancelsInFlightRequest(t *testing.T) {
|
|
for _, provider := range conformingStreamProviders(t) {
|
|
provider := provider
|
|
t.Run(provider, func(t *testing.T) {
|
|
released := make(chan struct{})
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n"))
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
<-r.Context().Done()
|
|
close(released)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
stream, err := ai.New(provider, ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL)).Stream(context.Background(), &ai.Request{Prompt: "Hello"})
|
|
if err != nil {
|
|
t.Fatalf("Stream returned error: %v", err)
|
|
}
|
|
assertStreamReply(t, stream, "hel")
|
|
if err := stream.Close(); err != nil {
|
|
t.Fatalf("Close returned error: %v", err)
|
|
}
|
|
if err := stream.Close(); err != nil {
|
|
t.Fatalf("second Close returned error: %v", err)
|
|
}
|
|
|
|
select {
|
|
case <-released:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("server did not observe canceled stream request")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamProvidersPropagateProviderErrors(t *testing.T) {
|
|
for _, provider := range conformingStreamProviders(t) {
|
|
provider := provider
|
|
t.Run(provider, func(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "upstream quota exhausted", http.StatusTooManyRequests)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
stream, err := ai.New(provider, ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL)).Stream(context.Background(), &ai.Request{Prompt: "Hello"})
|
|
if err == nil {
|
|
_ = stream.Close()
|
|
t.Fatal("Stream returned nil error for provider failure")
|
|
}
|
|
if !strings.Contains(err.Error(), "429") || !strings.Contains(err.Error(), "upstream quota exhausted") {
|
|
t.Fatalf("Stream error = %v, want provider status and body", err)
|
|
}
|
|
if strings.Contains(err.Error(), "test-key") {
|
|
t.Fatal("provider error leaked API key")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamProvidersHonorCanceledContextBeforeRequest(t *testing.T) {
|
|
for _, provider := range conformingStreamProviders(t) {
|
|
provider := provider
|
|
t.Run(provider, func(t *testing.T) {
|
|
var sawRequest bool
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sawRequest = true
|
|
http.Error(w, "unexpected request", http.StatusInternalServerError)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
stream, err := ai.New(provider, ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL)).Stream(ctx, &ai.Request{Prompt: "Hello"})
|
|
if err == nil {
|
|
_ = stream.Close()
|
|
t.Fatal("Stream returned nil error for canceled context")
|
|
}
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Stream error = %v, want context.Canceled", err)
|
|
}
|
|
if sawRequest {
|
|
t.Fatal("provider sent request after context was already canceled")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConfiguredProviderStreamsSkipWithoutCredentials(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
provider string
|
|
keyEnv string
|
|
modelEnv string
|
|
}{
|
|
{provider: "openai", keyEnv: "OPENAI_API_KEY", modelEnv: "OPENAI_MODEL"},
|
|
{provider: "groq", keyEnv: "GROQ_API_KEY", modelEnv: "GROQ_MODEL"},
|
|
{provider: "mistral", keyEnv: "MISTRAL_API_KEY", modelEnv: "MISTRAL_MODEL"},
|
|
{provider: "together", keyEnv: "TOGETHER_API_KEY", modelEnv: "TOGETHER_MODEL"},
|
|
{provider: "atlascloud", keyEnv: "ATLASCLOUD_API_KEY", modelEnv: "ATLASCLOUD_MODEL"},
|
|
{provider: "anthropic", keyEnv: "ANTHROPIC_API_KEY", modelEnv: "ANTHROPIC_MODEL"},
|
|
{provider: "gemini", keyEnv: "GEMINI_API_KEY", modelEnv: "GEMINI_MODEL"},
|
|
} {
|
|
tc := tc
|
|
t.Run(tc.provider, func(t *testing.T) {
|
|
key := os.Getenv(tc.keyEnv)
|
|
if key == "" {
|
|
t.Skipf("%s not set; skipping configured provider stream check", tc.keyEnv)
|
|
}
|
|
|
|
opts := []ai.Option{ai.WithAPIKey(key)}
|
|
if model := os.Getenv(tc.modelEnv); model != "" {
|
|
opts = append(opts, ai.WithModel(model))
|
|
}
|
|
stream, err := ai.New(tc.provider, opts...).Stream(context.Background(), &ai.Request{Prompt: "Reply with exactly: ok"})
|
|
if err != nil {
|
|
t.Fatalf("Stream returned error: %v", err)
|
|
}
|
|
defer stream.Close()
|
|
|
|
deadline := time.After(30 * time.Second)
|
|
for {
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("timed out waiting for provider stream chunk")
|
|
default:
|
|
}
|
|
chunk, err := stream.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
t.Fatal("provider stream ended without content")
|
|
}
|
|
t.Fatalf("Recv returned error: %v", err)
|
|
}
|
|
if chunk.Reply != "" {
|
|
return
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func conformingStreamProviders(t *testing.T) []string {
|
|
t.Helper()
|
|
providers := ai.RegisteredProviders("stream")
|
|
allowed := map[string]struct{}{
|
|
"atlascloud": {},
|
|
"groq": {},
|
|
"minimax": {},
|
|
"mistral": {},
|
|
"openai": {},
|
|
"together": {},
|
|
}
|
|
var out []string
|
|
for _, provider := range providers {
|
|
if _, ok := allowed[provider]; ok {
|
|
out = append(out, provider)
|
|
}
|
|
}
|
|
want := []string{"atlascloud", "groq", "minimax", "mistral", "openai", "together"}
|
|
if !reflect.DeepEqual(out, want) {
|
|
t.Fatalf("conforming stream providers = %#v, want %#v (registered stream providers: %#v)", out, want, providers)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func assertStreamReply(t *testing.T, stream ai.Stream, want string) {
|
|
t.Helper()
|
|
chunk, err := stream.Recv()
|
|
if err != nil {
|
|
t.Fatalf("Recv error = %v, want reply %q", err, want)
|
|
}
|
|
if chunk.Reply != want {
|
|
t.Fatalf("Reply = %q, want %q", chunk.Reply, want)
|
|
}
|
|
}
|