87 lines
2.6 KiB
Go
87 lines
2.6 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/mmmy/snapgo/internal/domain"
|
|
)
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
return f(req)
|
|
}
|
|
|
|
func TestVisionClientSummarizeImageUsesChatCompletionImageURL(t *testing.T) {
|
|
var requestBody map[string]any
|
|
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
if r.URL.Path != "/v1/chat/completions" {
|
|
t.Fatalf("expected /v1/chat/completions, got %s", r.URL.Path)
|
|
}
|
|
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
|
t.Fatalf("unexpected auth header %q", got)
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
return &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
|
Body: io.NopCloser(bytes.NewBufferString(`{"choices":[{"message":{"content":"summary ok"}}]}`)),
|
|
}, nil
|
|
})
|
|
|
|
client, err := NewVisionClient(domain.LLMProviderQwen, domain.LLMProviderConfig{
|
|
BaseURL: "https://example.test/v1",
|
|
APIKey: "test-key",
|
|
Model: "qwen-vl-plus",
|
|
MaxTokens: 321,
|
|
Temperature: 0.2,
|
|
TimeoutSecs: 5,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("new client: %v", err)
|
|
}
|
|
client.httpClient.Transport = transport
|
|
|
|
text, err := client.SummarizeImage(context.Background(), "describe", "https://cdn.example.com/a.png")
|
|
if err != nil {
|
|
t.Fatalf("summarize image: %v", err)
|
|
}
|
|
if text != "summary ok" {
|
|
t.Fatalf("expected response text, got %q", text)
|
|
}
|
|
if requestBody["model"] != "qwen-vl-plus" {
|
|
t.Fatalf("expected model forwarded, got %#v", requestBody["model"])
|
|
}
|
|
if requestBody["max_tokens"] != float64(321) {
|
|
t.Fatalf("expected max_tokens forwarded, got %#v", requestBody["max_tokens"])
|
|
}
|
|
|
|
messages := requestBody["messages"].([]any)
|
|
content := messages[0].(map[string]any)["content"].([]any)
|
|
if content[0].(map[string]any)["text"] != "describe" {
|
|
t.Fatalf("expected prompt content, got %#v", content[0])
|
|
}
|
|
image := content[1].(map[string]any)
|
|
if image["type"] != "image_url" {
|
|
t.Fatalf("expected image_url item, got %#v", image)
|
|
}
|
|
url := image["image_url"].(map[string]any)["url"]
|
|
if url != "https://cdn.example.com/a.png" {
|
|
t.Fatalf("expected image url forwarded, got %#v", url)
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsURLAcceptsFullEndpoint(t *testing.T) {
|
|
full := "https://example.com/v1/chat/completions"
|
|
if got := chatCompletionsURL(full); got != full {
|
|
t.Fatalf("expected full endpoint unchanged, got %q", got)
|
|
}
|
|
}
|