e926e4e6df
checkPrivate validates the configured endpoint literal and validated nothing after it. The client followed redirects, so a 302 from the local llama-server would have sent the photo, as a data URI in the POST body, to whatever the redirect named. "No provider in this repo may upload a blob" was true only of the first hop. Redirects are refused now, and the reply is read through a cap rather than however much the endpoint feels like sending. ValidateEndpoint exports the same check so config can fail at startup on a typo instead of logging once and leaving vision quietly off. Found in review of #72.
269 lines
8.8 KiB
Go
269 lines
8.8 KiB
Go
package vision
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"image"
|
|
"image/png"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/kami/maven/internal/media"
|
|
)
|
|
|
|
func testImage(t *testing.T) media.Image {
|
|
t.Helper()
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 32, 32))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
im, err := media.PrepareImage(buf.Bytes(), "test", 32)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return im
|
|
}
|
|
|
|
func TestDisabledAlwaysRefuses(t *testing.T) {
|
|
_, err := Disabled{}.Describe(context.Background(), testImage(t), "что тут?")
|
|
if !errors.Is(err, ErrDisabled) {
|
|
t.Fatalf("got %v, want ErrDisabled", err)
|
|
}
|
|
}
|
|
|
|
// The whole reason this package has its own HTTP client instead of reusing
|
|
// internal/llm.Client: a vision endpoint that is not on this box is refused.
|
|
func TestNewLocalRefusesNonPrivateEndpoints(t *testing.T) {
|
|
bad := []string{
|
|
"https://api.openai.com",
|
|
"http://8.8.8.8:8080",
|
|
"https://vision.example.com", // a name could resolve anywhere
|
|
"ftp://127.0.0.1:8080", // wrong scheme
|
|
"", // nothing to talk to
|
|
}
|
|
for _, ep := range bad {
|
|
if _, err := NewLocal(Config{Endpoint: ep}); err == nil {
|
|
t.Errorf("NewLocal(%q) was accepted", ep)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNewLocalAcceptsLocalEndpoints(t *testing.T) {
|
|
for _, ep := range []string{"http://127.0.0.1:8081", "http://localhost:8081/", "http://192.168.1.104:8081", "http://[::1]:8081"} {
|
|
p, err := NewLocal(Config{Endpoint: ep})
|
|
if err != nil {
|
|
t.Errorf("NewLocal(%q): %v", ep, err)
|
|
continue
|
|
}
|
|
if strings.HasSuffix(p.Endpoint(), "/") {
|
|
t.Errorf("trailing slash kept: %q", p.Endpoint())
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestDescribeSendsImageAsDataURIAndReturnsText(t *testing.T) {
|
|
var gotBody map[string]any
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/chat/completions" {
|
|
t.Errorf("path = %s", r.URL.Path)
|
|
}
|
|
raw, _ := io.ReadAll(r.Body)
|
|
if err := json.Unmarshal(raw, &gotBody); err != nil {
|
|
t.Errorf("unmarshal request: %v", err)
|
|
}
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":" На картинке кот "}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p, err := NewLocal(Config{Endpoint: srv.URL, Model: "qwen-vl"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
text, err := p.Describe(context.Background(), testImage(t), "кто на фото?")
|
|
if err != nil {
|
|
t.Fatalf("describe: %v", err)
|
|
}
|
|
if text != "На картинке кот" {
|
|
t.Errorf("text = %q (should be trimmed)", text)
|
|
}
|
|
|
|
msgs, ok := gotBody["messages"].([]any)
|
|
if !ok || len(msgs) != 1 {
|
|
t.Fatalf("messages = %#v", gotBody["messages"])
|
|
}
|
|
parts, ok := msgs[0].(map[string]any)["content"].([]any)
|
|
if !ok || len(parts) != 2 {
|
|
t.Fatalf("content parts = %#v", msgs[0])
|
|
}
|
|
if got := parts[0].(map[string]any)["text"]; got != "кто на фото?" {
|
|
t.Errorf("prompt = %v", got)
|
|
}
|
|
url := parts[1].(map[string]any)["image_url"].(map[string]any)["url"].(string)
|
|
if !strings.HasPrefix(url, "data:image/jpeg;base64,") {
|
|
t.Errorf("image not sent as a jpeg data uri: %.40s", url)
|
|
}
|
|
}
|
|
|
|
func TestDescribeUsesDefaultPromptWhenNoQuestion(t *testing.T) {
|
|
var sentPrompt string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var body struct {
|
|
Messages []struct {
|
|
Content []struct {
|
|
Text string `json:"text"`
|
|
} `json:"content"`
|
|
} `json:"messages"`
|
|
}
|
|
_ = json.NewDecoder(r.Body).Decode(&body)
|
|
sentPrompt = body.Messages[0].Content[0].Text
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"ок"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p, err := NewLocal(Config{Endpoint: srv.URL, Prompt: "Опиши по-русски."})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := p.Describe(context.Background(), testImage(t), " "); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if sentPrompt != "Опиши по-русски." {
|
|
t.Errorf("prompt = %q", sentPrompt)
|
|
}
|
|
}
|
|
|
|
// A Thinking model sometimes leaves content empty and puts the answer in
|
|
// reasoning_content; internal/llm has the same fallback and vision needs it too.
|
|
func TestDescribeFallsBackToReasoningContent(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"","reasoning_content":"схема платы"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
p, _ := NewLocal(Config{Endpoint: srv.URL})
|
|
text, err := p.Describe(context.Background(), testImage(t), "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if text != "схема платы" {
|
|
t.Errorf("text = %q", text)
|
|
}
|
|
}
|
|
|
|
func TestDescribeErrors(t *testing.T) {
|
|
t.Run("no choices", func(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
_, _ = w.Write([]byte(`{"choices":[]}`))
|
|
}))
|
|
defer srv.Close()
|
|
p, _ := NewLocal(Config{Endpoint: srv.URL})
|
|
if _, err := p.Describe(context.Background(), testImage(t), ""); !errors.Is(err, ErrEmptyReply) {
|
|
t.Errorf("got %v, want ErrEmptyReply", err)
|
|
}
|
|
})
|
|
t.Run("server error", func(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
}))
|
|
defer srv.Close()
|
|
p, _ := NewLocal(Config{Endpoint: srv.URL})
|
|
if _, err := p.Describe(context.Background(), testImage(t), ""); err == nil {
|
|
t.Error("500 was not an error")
|
|
}
|
|
})
|
|
t.Run("empty image", func(t *testing.T) {
|
|
p, _ := NewLocal(Config{Endpoint: "http://127.0.0.1:1"})
|
|
if _, err := p.Describe(context.Background(), media.Image{}, ""); !errors.Is(err, media.ErrEmpty) {
|
|
t.Errorf("got %v, want media.ErrEmpty", err)
|
|
}
|
|
})
|
|
t.Run("context cancelled", func(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
time.Sleep(200 * time.Millisecond)
|
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"поздно"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
p, _ := NewLocal(Config{Endpoint: srv.URL})
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
|
defer cancel()
|
|
if _, err := p.Describe(ctx, testImage(t), ""); err == nil {
|
|
t.Error("cancelled context returned no error")
|
|
}
|
|
})
|
|
}
|
|
|
|
// checkPrivate validates the configured literal and used to validate nothing
|
|
// else. A 302 from the local llama-server would have sent the photo, as a data
|
|
// URI in the POST body, wherever the redirect named.
|
|
func TestLocalProviderDoesNotFollowARedirect(t *testing.T) {
|
|
var elsewhere int32
|
|
away := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
atomic.AddInt32(&elsewhere, 1)
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"leaked"}}]}`)
|
|
}))
|
|
defer away.Close()
|
|
local := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, away.URL+"/v1/chat/completions", http.StatusFound)
|
|
}))
|
|
defer local.Close()
|
|
|
|
p, err := NewLocal(Config{Endpoint: local.URL})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
im := media.Image{JPEG: []byte{0xFF, 0xD8, 0xFF}}
|
|
if _, err := p.Describe(context.Background(), im, "что это"); err == nil {
|
|
t.Fatal("a redirected describe must fail, not follow")
|
|
}
|
|
if n := atomic.LoadInt32(&elsewhere); n != 0 {
|
|
t.Fatalf("the image was sent to the redirect target %d time(s)", n)
|
|
}
|
|
}
|
|
|
|
// The reply is read through a cap. A stuck endpoint should not cost the daemon
|
|
// its memory.
|
|
func TestLocalProviderCapsTheReply(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"`)
|
|
for written := 0; written < MaxReplyBytes+(1<<20); written += 1 << 16 {
|
|
if _, err := io.WriteString(w, strings.Repeat("a", 1<<16)); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
p, err := NewLocal(Config{Endpoint: srv.URL})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
im := media.Image{JPEG: []byte{0xFF, 0xD8, 0xFF}}
|
|
if _, err := p.Describe(context.Background(), im, ""); err == nil {
|
|
t.Fatal("an unbounded reply must fail rather than being read whole")
|
|
}
|
|
}
|
|
|
|
// ValidateEndpoint is what config calls at startup, and it must agree with the
|
|
// constructor.
|
|
func TestValidateEndpointMatchesTheConstructor(t *testing.T) {
|
|
for _, raw := range []string{"http://127.0.0.1:8081", "http://localhost:8081/"} {
|
|
if err := ValidateEndpoint(raw); err != nil {
|
|
t.Errorf("ValidateEndpoint(%q) = %v", raw, err)
|
|
}
|
|
}
|
|
for _, raw := range []string{"http://8.8.8.8:8081", "http://vision.example.com", "ftp://127.0.0.1"} {
|
|
if err := ValidateEndpoint(raw); err == nil {
|
|
t.Errorf("ValidateEndpoint(%q) accepted a non-private endpoint", raw)
|
|
}
|
|
if _, err := NewLocal(Config{Endpoint: raw}); err == nil {
|
|
t.Errorf("NewLocal(%q) accepted what ValidateEndpoint should refuse", raw)
|
|
}
|
|
}
|
|
}
|