package server import ( "context" "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "strconv" "strings" "testing" ) 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 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 route: only for a user who turned the assistant on. func TestAIChatNeedsUserOn(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 } if w := post(); w.Code != http.StatusForbidden { t.Fatalf("off: %d %s", w.Code, w.Body.String()) } a, _ := s.db.AdminByUsername("adm") _ = s.db.SetConfig(aiUserKey(a.ID), "1") 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) if len(cs) != 1 || cs[0].Title != "hi" || 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) } } 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") } }