Files

436 lines
17 KiB
Go

package server
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
)
func collect(t *testing.T, cfg aiConfig) (string, error) {
t.Helper()
var b strings.Builder
err := aiStream(context.Background(), cfg, "sys", []aiMsg{{Role: "user", Content: "раз"}, {Role: "assistant", Content: "два"}, {Role: "user", Content: "три"}}, func(s string) { b.WriteString(s) })
return b.String(), err
}
// Gradio (the Qwen space behind the proxy): history in, the answer streamed
// as diffs of the chatbot, the key in X-Agent-Key.
func TestAIGradio(t *testing.T) {
var joined map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("X-Agent-Key") != "k" {
http.Error(w, "no key", 403)
return
}
switch r.URL.Path {
case "/config":
fmt.Fprint(w, `{"dependencies":[{"id":3,"api_name":"other"},{"id":0,"api_name":"chat_predict","targets":[[21,"click"]]}]}`)
case "/gradio_api/queue/join":
_ = json.NewDecoder(r.Body).Decode(&joined)
fmt.Fprint(w, `{"event_id":"e"}`)
case "/gradio_api/queue/data":
w.Header().Set("Content-Type", "text/event-stream")
for _, m := range []string{
`{"msg":"process_starts"}`,
`{"msg":"process_generating","output":{"data":[null,null,null,null,[{"role":"user","content":"три"}],{},{}]}}`,
`{"msg":"process_generating","output":{"data":[[],[],[],[],[["add",[3],{"role":"assistant","content":"При"}]],[],[]]}}`,
`{"msg":"process_generating","output":{"data":[[],[],[],[],[["append",[3,"content"],"вет"]],[],[]]}}`,
`{"msg":"process_completed","success":true,"output":{"data":[{},{},{},{},[{"role":"user","content":"три"},{"role":"assistant","content":"Привет!"}],{},{}]}}`,
} {
fmt.Fprintf(w, "data: %s\n\n", m)
}
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
got, err := collect(t, aiConfig{Type: "gradio", URL: srv.URL, Key: "k"})
if err != nil || got != "Привет!" {
t.Fatalf("got %q, %v", got, err)
}
d := joined["data"].([]any)
if joined["fn_index"].(float64) != 0 || d[0] != "три" || d[5] != "sys" || len(d[4].([]any)) != 2 {
t.Fatalf("join: %v", joined)
}
}
func TestAIGradioError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/config":
fmt.Fprint(w, `{"dependencies":[{"id":0,"api_name":"chat_predict"}]}`)
case "/gradio_api/queue/join":
fmt.Fprint(w, `{"event_id":"e"}`)
default:
fmt.Fprint(w, "data: {\"msg\":\"process_completed\",\"success\":false,\"output\":{\"error\":\"GPU занят\"}}\n\n")
}
}))
defer srv.Close()
if _, err := collect(t, aiConfig{Type: "gradio", URL: srv.URL}); err == nil || !strings.Contains(err.Error(), "GPU занят") {
t.Fatalf("err: %v", err)
}
}
func TestAIOpenAI(t *testing.T) {
var body map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer k" {
http.Error(w, "bad", 400)
return
}
_ = json.NewDecoder(r.Body).Decode(&body)
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"При\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"вет\"}}]}\n\ndata: [DONE]\n\n")
}))
defer srv.Close()
got, err := collect(t, aiConfig{Type: "openai", URL: srv.URL + "/v1", Key: "k", Model: "m"})
if err != nil || got != "Привет" {
t.Fatalf("got %q, %v", got, err)
}
msgs := body["messages"].([]any)
if body["model"] != "m" || body["stream"] != true || len(msgs) != 4 || msgs[0].(map[string]any)["role"] != "system" {
t.Fatalf("body: %v", body)
}
}
func TestAIOpenAIAttachments(t *testing.T) {
var body map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_ = json.NewDecoder(r.Body).Decode(&body)
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer srv.Close()
msg := aiMsg{Role: "user", Content: "посмотри", Attachments: []aiAttachment{
{Name: "image.png", MIME: "image/png", Data: "AQID"},
{Name: "voice.wav", MIME: "audio/wav", Data: "BAUG"},
{Name: "clip.mp4", MIME: "video/mp4", Data: "BwgJ"},
}}
if err := aiOpenAI(context.Background(), aiConfig{Type: "openai", URL: srv.URL, Model: "m"}, "sys", []aiMsg{msg}, func(string) {}); err != nil {
t.Fatal(err)
}
messages := body["messages"].([]any)
content := messages[1].(map[string]any)["content"].([]any)
if len(content) != 4 || content[3].(map[string]any)["text"] != "посмотри" {
t.Fatalf("content: %#v", content)
}
if got := content[0].(map[string]any)["image_url"].(map[string]any)["url"].(string); !strings.HasPrefix(got, "data:image/png;base64,") {
t.Fatalf("image: %q", got)
}
if got := content[1].(map[string]any)["input_audio"].(map[string]any)["format"]; got != "wav" {
t.Fatalf("audio: %#v", content[1])
}
video := content[2].(map[string]any)
if video["type"] != "video_url" || video["fps"] != float64(1) || !strings.HasPrefix(video["video_url"].(map[string]any)["url"].(string), "data:video/mp4;base64,") {
t.Fatalf("video: %#v", video)
}
}
func TestAIAnthropic(t *testing.T) {
var body map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/messages" || r.Header.Get("x-api-key") != "k" || r.Header.Get("anthropic-version") == "" {
http.Error(w, "bad", 400)
return
}
_ = json.NewDecoder(r.Body).Decode(&body)
fmt.Fprint(w, "event: message_start\ndata: {\"type\":\"message_start\"}\n\n"+
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"При\"}}\n\n"+
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"вет\"}}\n\n"+
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
}))
defer srv.Close()
got, err := collect(t, aiConfig{Type: "anthropic", URL: srv.URL, Key: "k", Model: "m"})
if err != nil || got != "Привет" {
t.Fatalf("got %q, %v", got, err)
}
if body["system"] != "sys" || len(body["messages"].([]any)) != 3 {
t.Fatalf("body: %v", body)
}
}
// The chat is on by default, and an explicit user preference can still turn it off.
func TestAIChatDefaultOnAndCanBeDisabled(t *testing.T) {
s, c := newTestServer(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n")
}))
defer srv.Close()
_ = s.db.SetConfig(cfgAIType, "openai")
_ = s.db.SetConfig(cfgAIURL, srv.URL)
post := func() *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodPost, "/api/panel/ai/chat", strings.NewReader(`{"message":"hi","context":"x"}`))
req.AddCookie(c)
w := httptest.NewRecorder()
s.Handler().ServeHTTP(w, req)
return w
}
a, _ := s.db.AdminByUsername("adm")
state := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/panel/ai", nil)
req.AddCookie(c)
s.handleAIState(state, req)
var got struct {
On bool `json:"on"`
}
if err := json.Unmarshal(state.Body.Bytes(), &got); err != nil || !got.On {
t.Fatalf("default state: %s (%v)", state.Body.String(), err)
}
w := post()
b, _ := io.ReadAll(w.Body)
if w.Code != 200 || string(b) != "ok" {
t.Fatalf("on: %d %q", w.Code, b)
}
// The talk is kept: the question and the answer.
cs := s.aiChats(a.ID)
// Named by the question at once (the model may rename it a moment later).
if len(cs) != 1 || cs[0].Title == "" || w.Header().Get("X-Chat-Id") != strconv.FormatInt(cs[0].ID, 10) {
t.Fatalf("chats: %+v %q", cs, w.Header().Get("X-Chat-Id"))
}
if h := s.aiHistoryOf(a.ID, cs[0].ID, 10); len(h) != 2 || h[0].Content != "hi" || h[1].Content != "ok" {
t.Fatalf("history: %+v", h)
}
off := httptest.NewRecorder()
req = httptest.NewRequest(http.MethodPost, "/api/panel/ai/toggle", strings.NewReader(`{"on":false}`))
req.AddCookie(c)
s.handleAIToggle(off, req)
state = httptest.NewRecorder()
req = httptest.NewRequest(http.MethodGet, "/api/panel/ai", nil)
req.AddCookie(c)
s.handleAIState(state, req)
var after struct {
On bool `json:"on"`
}
if err := json.Unmarshal(state.Body.Bytes(), &after); err != nil || after.On {
t.Fatalf("disabled state: %s (%v)", state.Body.String(), err)
}
if w := post(); w.Code != http.StatusForbidden {
t.Fatalf("explicit off: %d %s", w.Code, w.Body.String())
}
}
func TestAIChatKeepsAttachmentTextOnly(t *testing.T) {
s, c := newTestServer(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var in struct {
Messages []map[string]any `json:"messages"`
}
_ = json.NewDecoder(r.Body).Decode(&in)
content := in.Messages[len(in.Messages)-1]["content"].([]any)
if content[0].(map[string]any)["type"] != "image_url" {
http.Error(w, "image missing", http.StatusBadRequest)
return
}
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer srv.Close()
_ = s.db.SetConfig(cfgAIType, "openai")
_ = s.db.SetConfig(cfgAIURL, srv.URL)
a, _ := s.db.AdminByUsername("adm")
_ = s.db.SetConfig(aiUserKey(a.ID), "1")
chatID := s.aiNewChat(a.ID)
s.aiAppendTo(a.ID, chatID, "user", "предыдущий вопрос")
s.aiAppendTo(a.ID, chatID, "assistant", "ответ")
body := fmt.Sprintf(`{"message":"посмотри","chat":%d,"attachments":[{"name":"x.png","mime":"image/png","data":"AQID"}]}`, chatID)
req := httptest.NewRequest(http.MethodPost, "/api/panel/ai/chat", strings.NewReader(body))
req.AddCookie(c)
w := httptest.NewRecorder()
s.Handler().ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("response: %d %s", w.Code, w.Body.String())
}
chat, _ := strconv.ParseInt(w.Header().Get("X-Chat-Id"), 10, 64)
history := s.aiHistoryOf(a.ID, chat, 10)
if chat != chatID || len(history) != 3 || !strings.Contains(history[2].Content, "x.png") || strings.Contains(history[2].Content, "AQID") {
t.Fatalf("history contains wrong attachment data: %+v", history)
}
}
func TestVersionAtLeast(t *testing.T) {
for _, c := range []struct {
v string
want bool
}{{"0.2.0", true}, {"0.10.1", true}, {"0.1.9", false}, {"", false}, {"dev", false}, {"1.0", true}} {
if got := versionAtLeast(c.v, "0.2.0"); got != c.want {
t.Errorf("%q: %v", c.v, got)
}
}
}
// The window the model sees: newest first within the budget, old results
// and answers cut down, results as the user's words, same roles merged.
func TestAIWindow(t *testing.T) {
big := strings.Repeat("x", 5000)
rows := []aiMsg{
{Role: "user", Content: "старое"},
{Role: "assistant", Content: big},
{Role: "result", Content: big},
{Role: "user", Content: "дальше"},
{Role: "assistant", Content: "ок"},
{Role: "result", Content: "вывод"},
}
w := aiWindow(rows, 100000)
if len(w) != 5 || w[0].Role != "user" || w[len(w)-1].Role != "user" {
t.Fatalf("roles: %+v", w)
}
if len(w[1].Content) > 3200 || len(w[2].Content) > 1700+len("дальше")+40 {
t.Fatalf("old not cut: %d %d", len(w[1].Content), len(w[2].Content))
}
if !strings.Contains(w[4].Content, "вывод") || !strings.HasPrefix(w[4].Content, "Результат") {
t.Fatalf("last result: %q", w[4].Content)
}
// A small budget keeps the newest, starting with the user.
w = aiWindow(rows, 60)
if len(w) == 0 || w[0].Role != "user" || !strings.Contains(w[len(w)-1].Content, "вывод") {
t.Fatalf("small: %+v", w)
}
}
func TestStripThink(t *testing.T) {
if got := stripThink("::думаю::\nхм\n::ответ::\n391"); got != "391" {
t.Fatalf("%q", got)
}
if got := stripThink("просто"); got != "просто" {
t.Fatalf("%q", got)
}
}
// Several talks: each its own messages; a new one is named by its first
// question; deleting one leaves the others.
func TestAIChats(t *testing.T) {
s, _ := newTestServer(t)
a, _ := s.db.AdminByUsername("adm")
one := s.aiNewChat(a.ID)
two := s.aiNewChat(a.ID)
s.aiAppendTo(a.ID, one, "user", "первый вопрос про nginx и его настройки, длинный-длинный")
s.aiAppendTo(a.ID, two, "user", "второй")
if h := s.aiHistoryOf(a.ID, one, 10); len(h) != 1 || h[0].Content[:5] != "перв"[:5] {
t.Fatalf("one: %+v", h)
}
list := s.aiChats(a.ID)
if len(list) != 2 || list[0].ID != two || !strings.HasPrefix(list[1].Title, "первый вопрос") || len([]rune(list[1].Title)) > 41 {
t.Fatalf("list: %+v", list)
}
s.aiDeleteChat(a.ID, one)
if list = s.aiChats(a.ID); len(list) != 1 || len(s.aiHistoryOf(a.ID, one, 10)) != 0 {
t.Fatalf("after delete: %+v", list)
}
// Another user sees none of them and cannot write into them.
if len(s.aiChats(a.ID+1)) != 0 {
t.Fatal("another user sees the talks")
}
s.aiAppendTo(a.ID+1, two, "user", "чужое")
if len(s.aiHistoryOf(a.ID, two, 10)) != 1 {
t.Fatal("another user wrote into the talk")
}
}
// The thinking is kept with the answer (and how long it took), so it shows
// again later; the model never gets it back.
func TestAIThinkingKept(t *testing.T) {
s, c := newTestServer(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fl := w.(http.Flusher)
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"::думаю::\\nхм\\n\"}}]}\n\n")
fl.Flush()
time.Sleep(1100 * time.Millisecond)
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"::ответ::\\n391\"}}]}\n\ndata: [DONE]\n\n")
}))
defer srv.Close()
_ = s.db.SetConfig(cfgAIType, "openai")
_ = s.db.SetConfig(cfgAIURL, srv.URL)
a, _ := s.db.AdminByUsername("adm")
_ = s.db.SetConfig(aiUserKey(a.ID), "1")
req := httptest.NewRequest(http.MethodPost, "/api/panel/ai/chat", strings.NewReader(`{"message":"17*23","think":true}`))
req.AddCookie(c)
w := httptest.NewRecorder()
s.Handler().ServeHTTP(w, req)
chat, _ := strconv.ParseInt(w.Header().Get("X-Chat-Id"), 10, 64)
h := s.aiHistoryOf(a.ID, chat, 10)
if len(h) != 2 || !strings.Contains(h[1].Content, "::думаю::") || h[1].ThinkMS < 1000 {
t.Fatalf("kept: %+v", h)
}
if win := aiWindow(h, 10000); strings.Contains(win[len(win)-1].Content, "хм") {
t.Fatalf("the model gets the thinking back: %+v", win)
}
}
// A new chat is named by the model: its first question in two or three
// words, asked aside so the answer does not wait.
func TestAIChatTitle(t *testing.T) {
s, c := newTestServer(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var in struct {
Messages []aiMsg `json:"messages"`
}
_ = json.NewDecoder(r.Body).Decode(&in)
reply := "ok"
if strings.Contains(in.Messages[0].Content, "название") {
reply = "«Выбор поисковика.»"
}
b, _ := json.Marshal(reply)
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":%s}}]}\n\ndata: [DONE]\n\n", b)
}))
defer srv.Close()
_ = s.db.SetConfig(cfgAIType, "openai")
_ = s.db.SetConfig(cfgAIURL, srv.URL)
a, _ := s.db.AdminByUsername("adm")
_ = s.db.SetConfig(aiUserKey(a.ID), "1")
req := httptest.NewRequest(http.MethodPost, "/api/panel/ai/chat", strings.NewReader(`{"message":"какой поисковик ты используешь сейчас"}`))
req.AddCookie(c)
w := httptest.NewRecorder()
s.Handler().ServeHTTP(w, req)
for i := 0; i < 50; i++ {
if cs := s.aiChats(a.ID); len(cs) == 1 && cs[0].Title == "Выбор поисковика" {
return
}
time.Sleep(50 * time.Millisecond)
}
t.Fatalf("title: %+v", s.aiChats(a.ID))
}
// Qwen Omni's files go where its inputs are (audio before image there), and
// a video as Gradio 5 takes one: {video, subtitles}.
func TestAIGradioMediaSlots(t *testing.T) {
var joined map[string]any
n := 0
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/config":
fmt.Fprint(w, `{"components":[{"id":20,"type":"textbox"},{"id":16,"type":"audio"},{"id":17,"type":"image"},{"id":18,"type":"video"},`+
`{"id":13,"type":"chatbot"},{"id":7,"type":"textbox"},{"id":8,"type":"slider"},{"id":9,"type":"slider"},{"id":10,"type":"slider"}],`+
`"dependencies":[{"id":0,"api_name":"chat_predict","targets":[[21,"click"]],"inputs":[20,16,17,18,13,7,8,9,10]}]}`)
case "/gradio_api/upload":
n++
fmt.Fprintf(w, `["/tmp/f%d"]`, n)
case "/gradio_api/queue/join":
_ = json.NewDecoder(r.Body).Decode(&joined)
fmt.Fprint(w, `{"event_id":"e"}`)
case "/gradio_api/queue/data":
fmt.Fprint(w, "data: {\"msg\":\"process_completed\",\"success\":true,\"output\":{\"data\":[{},{},{},{},[{\"role\":\"assistant\",\"content\":\"ok\"}],{},{}]}}\n\n")
}
}))
defer srv.Close()
enc := base64.StdEncoding.EncodeToString([]byte("x"))
msgs := []aiMsg{{Role: "user", Content: "что тут", Attachments: []aiAttachment{
{Name: "a.png", MIME: "image/png", Data: enc},
{Name: "b.mp3", MIME: "audio/mpeg", Data: enc},
{Name: "c.mp4", MIME: "video/mp4", Data: enc},
}}}
if err := aiStream(context.Background(), aiConfig{Type: "gradio", URL: srv.URL}, "sys", msgs, func(string) {}); err != nil {
t.Fatal(err)
}
d := joined["data"].([]any)
path := func(v any) string { m, _ := v.(map[string]any); p, _ := m["path"].(string); return p }
video, _ := d[3].(map[string]any)
if path(d[1]) != "/tmp/f2" || path(d[2]) != "/tmp/f1" || path(video["video"]) != "/tmp/f3" || video["subtitles"] != nil || d[5] != "sys" {
t.Fatalf("data: %v", d)
}
}