lpra/internal/server/ai_test.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"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(`{"messages":[{"role":"user","content":"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)
	}
}

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)
		}
	}
}