210 lines
8.0 KiB
Go
210 lines
8.0 KiB
Go
package multimodal
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestSpeechClientSendsMultipartAndReturnsTranscript(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/inference" || r.Method != http.MethodPost {
|
|
t.Fatalf("unexpected speech request: %s %s", r.Method, r.URL.Path)
|
|
}
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
file, _, err := r.FormFile("file")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer file.Close()
|
|
data, _ := io.ReadAll(file)
|
|
if string(data) != "audio" {
|
|
t.Fatalf("audio payload = %q", data)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"text": "распознанная фраза"})
|
|
}))
|
|
defer server.Close()
|
|
text, err := (SpeechClient{BaseURL: server.URL, Client: server.Client()}).Transcribe(context.Background(), "voice.wav", "audio/wav", []byte("audio"))
|
|
if err != nil || text != "распознанная фраза" {
|
|
t.Fatalf("transcript=%q err=%v", text, err)
|
|
}
|
|
}
|
|
|
|
func TestVisionClientSendsImageDataAndInstruction(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var payload struct {
|
|
Messages []struct {
|
|
Content []struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
ImageURL *struct {
|
|
URL string `json:"url"`
|
|
} `json:"image_url"`
|
|
} `json:"content"`
|
|
} `json:"messages"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(payload.Messages) != 1 || len(payload.Messages[0].Content) != 2 || !strings.Contains(payload.Messages[0].Content[0].Text, "Опиши") || payload.Messages[0].Content[1].ImageURL == nil || !strings.HasPrefix(payload.Messages[0].Content[1].ImageURL.URL, "data:image/png;base64,") {
|
|
t.Fatalf("unexpected vision payload: %+v", payload)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]string{"content": "на изображении круг"}}}})
|
|
}))
|
|
defer server.Close()
|
|
text, err := (VisionClient{BaseURL: server.URL + "/v1", Client: server.Client()}).Describe(context.Background(), "vision-model", "image/png", []byte{0, 1, 2}, "Опиши изображение")
|
|
if err != nil || text != "на изображении круг" {
|
|
t.Fatalf("description=%q err=%v", text, err)
|
|
}
|
|
}
|
|
|
|
func TestVisionClientUsesReasoningContentWhenProviderOmitsContent(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{
|
|
"message": map[string]string{"content": "", "reasoning_content": "чёрный квадрат"},
|
|
}}})
|
|
}))
|
|
defer server.Close()
|
|
text, err := (VisionClient{BaseURL: server.URL, Client: server.Client()}).Describe(context.Background(), "vision-model", "image/png", []byte("image"), "Опиши изображение")
|
|
if err != nil || text != "чёрный квадрат" {
|
|
t.Fatalf("description=%q err=%v", text, err)
|
|
}
|
|
}
|
|
|
|
func TestImageGenerationClientSendsPromptAndReturnsBase64(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/images/generations" || r.Method != http.MethodPost {
|
|
t.Fatalf("unexpected image request: %s %s", r.Method, r.URL.Path)
|
|
}
|
|
var payload struct {
|
|
Model string `json:"model"`
|
|
Prompt string `json:"prompt"`
|
|
Size string `json:"size"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if payload.Model != "image-model" || payload.Prompt != "green bird" || payload.Size != "1024x1024" {
|
|
t.Fatalf("unexpected image payload: %+v", payload)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]string{"b64_json": "aW1hZ2U="}}})
|
|
}))
|
|
defer server.Close()
|
|
image, err := (ImageGenerationClient{BaseURL: server.URL, Client: server.Client()}).Generate(context.Background(), "image-model", "green bird", "1024x1024")
|
|
if err != nil || image.DataBase64 != "aW1hZ2U=" || image.ContentType != "image/png" {
|
|
t.Fatalf("image=%+v err=%v", image, err)
|
|
}
|
|
}
|
|
|
|
func TestVideoUnderstandingClientSendsMultipartAndReturnsDescription(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/understand" || r.Method != http.MethodPost {
|
|
t.Fatalf("unexpected video request: %s %s", r.Method, r.URL.Path)
|
|
}
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if r.FormValue("model") != "video-model" || r.FormValue("instruction") != "describe movement" {
|
|
t.Fatalf("unexpected video fields: model=%q instruction=%q", r.FormValue("model"), r.FormValue("instruction"))
|
|
}
|
|
file, _, err := r.FormFile("file")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer file.Close()
|
|
data, _ := io.ReadAll(file)
|
|
if string(data) != "video" {
|
|
t.Fatalf("video payload=%q", data)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"text": "человек идёт"})
|
|
}))
|
|
defer server.Close()
|
|
text, err := (VideoUnderstandingClient{BaseURL: server.URL, Client: server.Client()}).Describe(context.Background(), "video-model", "clip.mp4", "video/mp4", []byte("video"), "describe movement")
|
|
if err != nil || text != "человек идёт" {
|
|
t.Fatalf("description=%q err=%v", text, err)
|
|
}
|
|
}
|
|
|
|
func TestComfyImageGenerationClientRunsWorkflowAndReturnsImage(t *testing.T) {
|
|
polls := 0
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/prompt":
|
|
var payload struct {
|
|
Prompt map[string]any `json:"prompt"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if payload.Prompt["6"].(map[string]any)["inputs"].(map[string]any)["text"] != "red bird" {
|
|
t.Fatalf("prompt was not forwarded: %+v", payload.Prompt["6"])
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "test-prompt"})
|
|
case "/history/test-prompt":
|
|
polls++
|
|
if polls == 1 {
|
|
_, _ = w.Write([]byte(`{}`))
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"test-prompt": map[string]any{"outputs": map[string]any{"9": map[string]any{"images": []any{map[string]string{"filename": "out.png", "subfolder": "", "type": "output"}}}}}})
|
|
case "/view":
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write([]byte("png"))
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
image, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Second}).Generate(context.Background(), "", "red bird", "256x256")
|
|
if err != nil || image.DataBase64 != "cG5n" || image.ContentType != "image/png" {
|
|
t.Fatalf("image=%+v err=%v", image, err)
|
|
}
|
|
}
|
|
|
|
func TestComfyImageGenerationClientDeletesOnlyCancelledPrompt(t *testing.T) {
|
|
var deleted []string
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/prompt":
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "cancel-me"})
|
|
case "/history/cancel-me":
|
|
<-r.Context().Done()
|
|
case "/queue":
|
|
var payload struct {
|
|
Delete []string `json:"delete"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
deleted = append(deleted, payload.Delete...)
|
|
w.WriteHeader(http.StatusNoContent)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Minute}).Generate(ctx, "", "red bird", "256x256")
|
|
done <- err
|
|
}()
|
|
time.Sleep(10 * time.Millisecond)
|
|
cancel()
|
|
if err := <-done; !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("err=%v", err)
|
|
}
|
|
if len(deleted) != 1 || deleted[0] != "cancel-me" {
|
|
t.Fatalf("deleted=%v", deleted)
|
|
}
|
|
}
|