feat: pictures for the assistant - the models make and edit them by themselves through the Vertex API (a new picture, or the user's attached photo, or the last picture made), with an API key in Настройки and a switch for the prompt improver; folded descriptions of pictures and searches; the name of an attached file opens or saves it; a deleted chat is wiped from the database files, attachments included
This commit is contained in:
redtea committed
1 parent
c72fcfebf3
commit
dd1244a169
13 files changed
+925
-20
No files matched your search
@@ -0,0 +1,305 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeVertex is the image service: it records what it was asked and answers
|
||||
// like the real one (a job to poll, a picture when done).
|
||||
type fakeVertex struct {
|
||||
srv *httptest.Server
|
||||
mu sync.Mutex
|
||||
auth []string
|
||||
gen map[string]any
|
||||
edit struct {
|
||||
prompt string
|
||||
image []byte
|
||||
name string
|
||||
}
|
||||
polls int
|
||||
failAs string // an error_code the job ends with
|
||||
}
|
||||
|
||||
func newFakeVertex(t *testing.T) *fakeVertex {
|
||||
f := &fakeVertex{}
|
||||
f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.auth = append(f.auth, r.Header.Get("Authorization"))
|
||||
// The pictures are open by their address, as at the real service.
|
||||
if r.URL.Path == "/uploads/userphoto/ai/t.png" {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write([]byte("\x89PNG-original"))
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.Header.Get("Authorization") != "Bearer vx_good" {
|
||||
w.WriteHeader(401)
|
||||
fmt.Fprint(w, `{"ok":false,"error_code":"unauthorized","error":"Неверный ключ"}`)
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case r.URL.Path == "/api/vertex/v1/generate":
|
||||
_ = json.NewDecoder(r.Body).Decode(&f.gen)
|
||||
w.WriteHeader(202)
|
||||
fmt.Fprint(w, `{"ok":true,"job_id":"job_gen1","status":"queued","cost":5}`)
|
||||
case r.URL.Path == "/api/vertex/v1/edit":
|
||||
if err := r.ParseMultipartForm(32 << 20); err != nil {
|
||||
w.WriteHeader(400)
|
||||
fmt.Fprint(w, `{"ok":false,"error_code":"invalid_image","error":"нет фото"}`)
|
||||
return
|
||||
}
|
||||
f.edit.prompt = r.FormValue("prompt")
|
||||
file, hdr, err := r.FormFile("image")
|
||||
if err == nil {
|
||||
f.edit.image, _ = io.ReadAll(file)
|
||||
f.edit.name = hdr.Filename
|
||||
}
|
||||
w.WriteHeader(202)
|
||||
fmt.Fprint(w, `{"ok":true,"job_id":"job_edit1","status":"queued","cost":2}`)
|
||||
case strings.HasPrefix(r.URL.Path, "/api/vertex/v1/jobs/"):
|
||||
f.polls++
|
||||
if f.failAs != "" {
|
||||
fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"failed","error_code":%q,"error":"сбой"}`, f.failAs)
|
||||
return
|
||||
}
|
||||
if f.polls < 2 {
|
||||
fmt.Fprint(w, `{"ok":true,"job_id":"x","status":"running"}`)
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"done","photo_url":"%s/uploads/userphoto/ai/t.png","cost":5,"balance":42}`, f.srv.URL)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
}
|
||||
|
||||
// setImage points the panel at the fake service with a key.
|
||||
func setImage(t *testing.T, s *Server, f *fakeVertex, key string) {
|
||||
t.Helper()
|
||||
_ = s.db.SetConfig(cfgImageURL, f.srv.URL)
|
||||
_ = s.db.SetConfig(cfgImageKey, key)
|
||||
}
|
||||
|
||||
func apiCall(t *testing.T, s *Server, c *http.Cookie, method, path, body string) string {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
req.AddCookie(c)
|
||||
w := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(w, req)
|
||||
return w.Body.String()
|
||||
}
|
||||
|
||||
// A picture from nothing: the prompt, model and ratio go to the service with
|
||||
// the key; the job's progress and its picture come back, the key never.
|
||||
func TestImageGenerate(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"cat in space","model":"artist","ratio":"3:4"}`)
|
||||
if !strings.Contains(got, `"ok":true`) || !strings.Contains(got, `"job_id":"job_gen1"`) || strings.Contains(got, "vx_good") {
|
||||
t.Fatalf("create: %s", got)
|
||||
}
|
||||
if f.gen["prompt"] != "cat in space" || f.gen["model"] != "artist" || f.gen["ratio"] != "3:4" || f.gen["enhance"] != true {
|
||||
t.Fatalf("what the service got: %v", f.gen)
|
||||
}
|
||||
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", ""); !strings.Contains(got, `"status":"running"`) {
|
||||
t.Fatalf("first poll: %s", got)
|
||||
}
|
||||
got = apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", "")
|
||||
if !strings.Contains(got, `"status":"done"`) || !strings.Contains(got, "/uploads/userphoto/ai/t.png") || !strings.Contains(got, `"balance":42`) || strings.Contains(got, "vx_good") {
|
||||
t.Fatalf("done: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Unknown models and ratios are not sent on: the defaults go instead.
|
||||
func TestImageGenerateDefaults(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","model":"gpt","ratio":"7:1"}`)
|
||||
if f.gen["model"] != "fast" || f.gen["ratio"] != "1:1" {
|
||||
t.Fatalf("defaults: %v", f.gen)
|
||||
}
|
||||
}
|
||||
|
||||
// Editing sends the user's photo as a file with the instruction.
|
||||
func TestImageEditFromAttachedPhoto(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
photo := []byte("\xff\xd8\xffJPEGDATA")
|
||||
b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "сделай закат",
|
||||
"image": map[string]string{"mime": "image/jpeg", "data": base64.StdEncoding.EncodeToString(photo)}})
|
||||
got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b))
|
||||
if !strings.Contains(got, `"job_id":"job_edit1"`) {
|
||||
t.Fatalf("edit: %s", got)
|
||||
}
|
||||
if f.edit.prompt != "сделай закат" || !bytes.Equal(f.edit.image, photo) || !strings.HasSuffix(f.edit.name, ".jpg") {
|
||||
t.Fatalf("what the service got: %q %q %q", f.edit.prompt, f.edit.image, f.edit.name)
|
||||
}
|
||||
// A format the service does not take is told so, not sent.
|
||||
b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image": map[string]string{"mime": "image/gif", "data": "R0lG"}})
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "png") {
|
||||
t.Fatalf("gif: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A picture made before can be edited again: fetched from the service's own
|
||||
// address only - never from another.
|
||||
func TestImageEditFromEarlierPicture(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "темнее", "image_url": f.srv.URL + "/uploads/userphoto/ai/t.png"})
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":true`) {
|
||||
t.Fatalf("by url: %s", got)
|
||||
}
|
||||
if !bytes.Equal(f.edit.image, []byte("\x89PNG-original")) || !strings.HasSuffix(f.edit.name, ".png") {
|
||||
t.Fatalf("the picture sent: %q %q", f.edit.image, f.edit.name)
|
||||
}
|
||||
b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image_url": "http://127.0.0.1:9/secret.png"})
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) {
|
||||
t.Fatalf("a foreign address was fetched: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// What goes wrong is told in words: no key, a wrong key, the service's own
|
||||
// codes, a job that failed.
|
||||
func TestImageErrors(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "ключ") {
|
||||
t.Fatalf("no key: %s", got)
|
||||
}
|
||||
setImage(t, s, f, "vx_bad")
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "ключ") {
|
||||
t.Fatalf("wrong key: %s", got)
|
||||
}
|
||||
setImage(t, s, f, "vx_good")
|
||||
f.failAs = "insufficient_credits"
|
||||
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1", ""); !strings.Contains(got, `"status":"failed"`) || !strings.Contains(got, "кредит") {
|
||||
t.Fatalf("failed job: %s", got)
|
||||
}
|
||||
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/..%2Fme", ""); !strings.Contains(got, `"ok":false`) {
|
||||
t.Fatalf("a job id out of shape: %s", got)
|
||||
}
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":""}`); !strings.Contains(got, `"ok":false`) {
|
||||
t.Fatalf("empty prompt: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The assistant off in the profile means no pictures either.
|
||||
func TestImageNeedsAssistantOn(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
a, _ := s.db.AdminByUsername("adm")
|
||||
_ = s.db.SetConfig(aiUserKey(a.ID), "0")
|
||||
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "выключен") {
|
||||
t.Fatalf("off: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The key is saved from Настройки, shown back to the admin, and the chat
|
||||
// learns only that pictures are on.
|
||||
func TestImageKeySetting(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
postAIModels(t, s, c, [][7]string{{"", "M", "gradio", "", "", "", ""}})
|
||||
f := url.Values{"csrf": {"tok"}, "m_type": {"gradio"}, "ai_image_key": {"vx_secret"}}
|
||||
req := httptest.NewRequest(http.MethodPost, "/settings/ai", strings.NewReader(f.Encode()))
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.AddCookie(c)
|
||||
req.AddCookie(&http.Cookie{Name: csrfCookie, Value: "tok"})
|
||||
s.Handler().ServeHTTP(httptest.NewRecorder(), req)
|
||||
if _, k := s.imageConf(); k != "vx_secret" {
|
||||
t.Fatalf("key not saved: %q", k)
|
||||
}
|
||||
if body := get(t, s, c, "/settings?frag=1").Body.String(); !strings.Contains(body, `value="vx_secret"`) {
|
||||
t.Fatal("the settings page does not show the key")
|
||||
}
|
||||
state := get(t, s, c, "/api/panel/ai").Body.String()
|
||||
if strings.Contains(state, "vx_secret") || !strings.Contains(state, `"image":true`) {
|
||||
t.Fatalf("state: %s", state)
|
||||
}
|
||||
}
|
||||
|
||||
// With a key the model is told it can draw; without, it is not.
|
||||
func TestImagePromptOnlyWithKey(t *testing.T) {
|
||||
s, _ := newTestServer(t)
|
||||
if strings.Contains(s.imagePrompt(), "```image") {
|
||||
t.Fatal("told about pictures with no key")
|
||||
}
|
||||
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
||||
p := s.imagePrompt()
|
||||
for _, want := range []string{"```image\n", "```image edit", "fast", "artist", "0.2plus", "16:9"} {
|
||||
if !strings.Contains(p, want) {
|
||||
t.Fatalf("the prompt lacks %q:\n%s", want, p)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The chat gives the model the image instructions when pictures are on.
|
||||
func TestAIChatMentionsImageTool(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
var system string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var b map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
system, _ = b["messages"].([]any)[0].(map[string]any)["content"].(string)
|
||||
fmt.Fprint(w, openAIChunk("content", "ok")+"data: [DONE]\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
postAIModels(t, s, c, [][7]string{{"", "M", "openai", srv.URL, "m", "", ""}})
|
||||
ask := func() {
|
||||
apiCall(t, s, c, "POST", "/api/panel/ai/chat", `{"message":"нарисуй кота","context":"x"}`)
|
||||
}
|
||||
ask()
|
||||
if strings.Contains(system, "```image") {
|
||||
t.Fatal("image tool offered with no key")
|
||||
}
|
||||
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
||||
ask()
|
||||
if !strings.Contains(system, "```image") {
|
||||
t.Fatalf("image tool not offered with a key:\n%s", system)
|
||||
}
|
||||
}
|
||||
|
||||
// The description is improved by the service unless the model asks for it
|
||||
// as it is.
|
||||
func TestImageEnhanceSwitch(t *testing.T) {
|
||||
s, c := newTestServer(t)
|
||||
f := newFakeVertex(t)
|
||||
setImage(t, s, f, "vx_good")
|
||||
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`)
|
||||
if f.gen["enhance"] != true {
|
||||
t.Fatalf("by default: %v", f.gen)
|
||||
}
|
||||
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":false}`)
|
||||
if f.gen["enhance"] != false {
|
||||
t.Fatalf("switched off: %v", f.gen)
|
||||
}
|
||||
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":true}`)
|
||||
if f.gen["enhance"] != true {
|
||||
t.Fatalf("switched on: %v", f.gen)
|
||||
}
|
||||
{
|
||||
p := s.imagePrompt()
|
||||
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
||||
if !strings.Contains(s.imagePrompt(), "noenhance") {
|
||||
t.Fatalf("the model is not told of the switch:\n%s", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user