lpra/internal/server/ai.go

package server

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

// The assistant: a chat in the panel that sees the workspace and suggests
// commands. The browser sends the talk and what it shows; the server asks
// the model - its address, kind and key are the admins', set in Настройки
// and never sent to the browser - and streams the answer back as text.

const (
	cfgAIType  = "ai_type" // gradio, openai or anthropic
	cfgAIURL   = "ai_url"
	cfgAIKey   = "ai_key"
	cfgAIModel = "ai_model"

	// DefaultAIURL is the proxy to the Qwen space the assistant uses by default.
	DefaultAIURL = "https://agentproxy.minomix.ru"
)

// aiUserKey is the config key of a user's own switch: the assistant is off
// for everyone until they turn it on in their profile.
func aiUserKey(id int64) string { return "ai_user_" + strconv.FormatInt(id, 10) }

type aiConfig struct{ Type, URL, Key, Model string }

type aiMsg struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

func (s *Server) aiConfig() aiConfig {
	get := func(k, def string) string {
		if v, _ := s.db.GetConfig(k); v != "" {
			return v
		}
		return def
	}
	return aiConfig{Type: get(cfgAIType, "gradio"), URL: strings.TrimRight(get(cfgAIURL, DefaultAIURL), "/"), Key: get(cfgAIKey, ""), Model: get(cfgAIModel, "")}
}

var aiClient = &http.Client{Timeout: 10 * time.Minute}

// aiStream asks the model for the answer to msgs (the last one the user's)
// under the system prompt, handing emit the text as it comes.
func aiStream(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	switch cfg.Type {
	case "openai":
		return aiOpenAI(ctx, cfg, system, msgs, emit)
	case "anthropic":
		return aiAnthropic(ctx, cfg, system, msgs, emit)
	default:
		return aiGradio(ctx, cfg, system, msgs, emit)
	}
}

// sseData calls fn with each "data:" payload of an event stream until fn
// says stop or the stream ends.
func sseData(r io.Reader, fn func(data string) (stop bool, err error)) error {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64<<10), 16<<20)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		stop, err := fn(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
		if err != nil || stop {
			return err
		}
	}
	return sc.Err()
}

func aiPost(ctx context.Context, url string, body any, hdr map[string]string) (*http.Response, error) {
	b, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	for k, v := range hdr {
		req.Header.Set(k, v)
	}
	resp, err := aiClient.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 400))
		resp.Body.Close()
		return nil, fmt.Errorf("модель ответила %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
	}
	return resp, nil
}

func aiOpenAI(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	all := append([]aiMsg{{Role: "system", Content: system}}, msgs...)
	hdr := map[string]string{}
	if cfg.Key != "" {
		hdr["Authorization"] = "Bearer " + cfg.Key
	}
	resp, err := aiPost(ctx, cfg.URL+"/chat/completions", map[string]any{"model": cfg.Model, "stream": true, "messages": all}, hdr)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return sseData(resp.Body, func(d string) (bool, error) {
		if d == "[DONE]" {
			return true, nil
		}
		var m struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
			Error *struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		if m.Error != nil {
			return true, errors.New(m.Error.Message)
		}
		for _, c := range m.Choices {
			if c.Delta.Content != "" {
				emit(c.Delta.Content)
			}
		}
		return false, nil
	})
}

func aiAnthropic(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	model := cfg.Model
	if model == "" {
		model = "claude-sonnet-5-5"
	}
	resp, err := aiPost(ctx, cfg.URL+"/v1/messages", map[string]any{"model": model, "max_tokens": 4096, "system": system, "messages": msgs, "stream": true},
		map[string]string{"x-api-key": cfg.Key, "anthropic-version": "2023-06-01"})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return sseData(resp.Body, func(d string) (bool, error) {
		var m struct {
			Type  string `json:"type"`
			Delta struct {
				Text string `json:"text"`
			} `json:"delta"`
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		switch m.Type {
		case "content_block_delta":
			emit(m.Delta.Text)
		case "message_stop":
			return true, nil
		case "error":
			return true, errors.New(m.Error.Message)
		}
		return false, nil
	})
}

// aiGradio talks to a Gradio chat app (chat_predict of the Qwen Omni demo):
// the history and the new text go in through the queue, and the chatbot's
// history comes back - whole, or as diffs - with the answer growing at its
// end.
func aiGradio(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	hdr := map[string]string{}
	if cfg.Key != "" {
		hdr["X-Agent-Key"] = cfg.Key
	}
	get := func(url string) (*http.Response, error) {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
		if err != nil {
			return nil, err
		}
		for k, v := range hdr {
			req.Header.Set(k, v)
		}
		resp, err := aiClient.Do(req)
		if err == nil && resp.StatusCode/100 != 2 {
			resp.Body.Close()
			return nil, fmt.Errorf("модель ответила %d", resp.StatusCode)
		}
		return resp, err
	}
	// Which function is chat_predict, from the app's config.
	resp, err := get(cfg.URL + "/config")
	if err != nil {
		return err
	}
	var conf struct {
		Dependencies []struct {
			ID      int     `json:"id"`
			APIName string  `json:"api_name"`
			Targets [][]any `json:"targets"`
		} `json:"dependencies"`
	}
	err = json.NewDecoder(resp.Body).Decode(&conf)
	resp.Body.Close()
	if err != nil {
		return errors.New("модель: не прочиталась её конфигурация")
	}
	fn, trigger := -1, any(nil)
	for _, d := range conf.Dependencies {
		if d.APIName == "chat_predict" {
			fn = d.ID
			if len(d.Targets) > 0 && len(d.Targets[0]) > 0 {
				trigger = d.Targets[0][0]
			}
		}
	}
	if fn < 0 {
		return errors.New("модель: нет функции chat_predict")
	}
	last := msgs[len(msgs)-1]
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	session := hex.EncodeToString(b)
	jr, err := aiPost(ctx, cfg.URL+"/gradio_api/queue/join", map[string]any{
		"data":     []any{last.Content, nil, nil, nil, msgs[:len(msgs)-1], system, 0.7, 0.8, 20},
		"fn_index": fn, "trigger_id": trigger, "event_data": nil, "session_hash": session,
	}, hdr)
	if err != nil {
		return err
	}
	jr.Body.Close()
	resp, err = get(cfg.URL + "/gradio_api/queue/data?session_hash=" + session)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	sent := "" // what of the answer has gone out
	show := func(text string) {
		if strings.HasPrefix(text, sent) {
			emit(text[len(sent):])
			sent = text
		}
	}
	answer := "" // the answer as the chatbot has it now
	return sseData(resp.Body, func(d string) (bool, error) {
		var m struct {
			Msg     string `json:"msg"`
			Success *bool  `json:"success"`
			Output  struct {
				Data  []json.RawMessage `json:"data"`
				Error string            `json:"error"`
			} `json:"output"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		switch m.Msg {
		case "process_generating", "process_completed":
		case "close_stream":
			return true, nil
		default:
			return false, nil
		}
		if m.Success != nil && !*m.Success {
			msg := m.Output.Error
			if msg == "" {
				msg = "модель не ответила"
			}
			return true, errors.New("модель: " + msg)
		}
		if len(m.Output.Data) > 4 {
			answer = gradioAnswer(m.Output.Data[4], answer)
			show(answer)
		}
		return m.Msg == "process_completed", nil
	})
}

// gradioAnswer is the assistant's last message after one update of the
// chatbot: a whole history, or diffs ("add", "append", "replace") to it.
func gradioAnswer(raw json.RawMessage, cur string) string {
	var whole []struct {
		Role    string `json:"role"`
		Content any    `json:"content"`
	}
	if json.Unmarshal(raw, &whole) == nil {
		if n := len(whole); n > 0 {
			if s, ok := whole[n-1].Content.(string); ok && whole[n-1].Role == "assistant" {
				return s
			}
		}
		return cur
	}
	var diffs [][]json.RawMessage
	if json.Unmarshal(raw, &diffs) != nil {
		return cur
	}
	for _, d := range diffs {
		if len(d) < 3 {
			continue
		}
		var op string
		_ = json.Unmarshal(d[0], &op)
		var path []any
		_ = json.Unmarshal(d[1], &path)
		switch {
		case op == "add" && len(path) == 1:
			var msg struct {
				Role    string `json:"role"`
				Content any    `json:"content"`
			}
			if json.Unmarshal(d[2], &msg) == nil && msg.Role == "assistant" {
				s, _ := msg.Content.(string)
				cur = s
			}
		case (op == "append" || op == "replace") && len(path) == 2 && path[1] == "content":
			var s string
			if json.Unmarshal(d[2], &s) == nil {
				if op == "append" {
					cur += s
				} else {
					cur = s
				}
			}
		}
	}
	return cur
}

const aiSystem = `Ты - ассистент панели LPRA (удалённый доступ к машинам: терминалы, файлы, порты). Отвечай по-русски, коротко и по делу.
Ниже - что пользователь сейчас видит в панели: его машины, вкладки, активная панель, текст терминала или список файлов.
Действия - отдельными блоками, панель выполнит их с правами пользователя:
` + "```run машина\nкоманда\n```" + ` - команда в терминале этой машины (машину можно не указывать - тогда активная);
` + "```read машина:/путь\n```" + ` - прочитать файл;
` + "```ls машина:/путь\n```" + ` - список папки;
` + "```write машина:/путь\nполное новое содержимое файла\n```" + ` - записать файл целиком.
Имя машины - как в списке машин. Одна команда на блок, без пояснений внутри блока. Результат действия придёт следующим сообщением - не выдумывай его.
Опасные команды (удаление, перезапуск, изменение системы) - только с предупреждением. Перед правкой файла сначала прочитай его.`

const aiAgent = `
Режим агента: ты сам ведёшь задачу шаг за шагом. В каждом ответе - коротко, что делаешь, и ровно одно действие. Получив результат, решай следующий шаг.
Когда задача решена (или решить нельзя) - ответь итогом без блоков действий.`

const aiThink = `
Режим размышления: начни ответ строкой ::думаю:: и подумай - разбери задачу, варианты и риски; затем строка ::ответ:: и сам ответ.`

// The talk is kept in the database, the last aiKeep messages of each user.
const aiKeep = 200

func (s *Server) aiAppend(owner int64, role, content string) {
	if content == "" {
		return
	}
	if len(content) > 64<<10 {
		content = utf8Safe(content[:64<<10])
	}
	_, _ = s.db.sql.Exec(`INSERT INTO ai_messages (admin_id, role, content, at) VALUES (?, ?, ?, ?)`, owner, role, content, time.Now().Unix())
	_, _ = s.db.sql.Exec(`DELETE FROM ai_messages WHERE admin_id = ? AND id <= (SELECT id FROM ai_messages WHERE admin_id = ? ORDER BY id DESC LIMIT 1 OFFSET ?)`, owner, owner, aiKeep)
}

// aiHistory is the user's talk, oldest first, the last limit messages.
func (s *Server) aiHistory(owner int64, limit int) []aiMsg {
	rows, err := s.db.sql.Query(`SELECT role, content FROM ai_messages WHERE admin_id = ? ORDER BY id DESC LIMIT ?`, owner, limit)
	out := []aiMsg{}
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var m aiMsg
		if rows.Scan(&m.Role, &m.Content) == nil {
			out = append(out, m)
		}
	}
	for l, r := 0, len(out)-1; l < r; l, r = l+1, r-1 {
		out[l], out[r] = out[r], out[l]
	}
	return out
}

// POST /api/panel/ai/chat - {message | result, context, agent, think}: the
// new words (or an action's result) are kept with the talk, and the answer,
// streamed as plain text, is kept too. The model sees a window of the talk
// (aiWindow) and the panel as it is now. A failure after the start ends the
// text with "[ошибка: …]".
func (s *Server) handleAIChat(w http.ResponseWriter, r *http.Request) {
	a := s.currentAdmin(r)
	if on, _ := s.db.GetConfig(aiUserKey(a.ID)); on != "1" {
		http.Error(w, "ассистент выключен в профиле", http.StatusForbidden)
		return
	}
	var in struct {
		Message string `json:"message"`
		Result  string `json:"result"`
		Context string `json:"context"`
		Agent   bool   `json:"agent"`
		Think   bool   `json:"think"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&in); err != nil || (in.Message == "" && in.Result == "") {
		http.Error(w, "плохой запрос", http.StatusBadRequest)
		return
	}
	s.aiAppend(a.ID, "result", in.Result)
	s.aiAppend(a.ID, "user", in.Message)
	msgs := aiWindow(s.aiHistory(a.ID, 60), 24000)
	if len(in.Context) > 30000 {
		in.Context = utf8Safe(in.Context[len(in.Context)-30000:])
	}
	system := aiSystem
	if in.Agent {
		system += aiAgent
	}
	if in.Think {
		system += aiThink
	}
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("X-Accel-Buffering", "no")
	fl, _ := w.(http.Flusher)
	var answer strings.Builder
	emit := func(t string) {
		if t == "" {
			return
		}
		answer.WriteString(t)
		_, _ = io.WriteString(w, t)
		if fl != nil {
			fl.Flush()
		}
	}
	err := aiStream(r.Context(), s.aiConfig(), system+"\n\n--- Панель сейчас ---\n"+in.Context, msgs, emit)
	s.aiAppend(a.ID, "assistant", stripThink(answer.String()))
	if err != nil && r.Context().Err() == nil {
		if answer.Len() == 0 {
			w.WriteHeader(http.StatusBadGateway)
		}
		_, _ = io.WriteString(w, "\n[ошибка: "+err.Error()+"]")
	}
}

// GET /api/panel/ai/history - the talk to show, oldest first.
func (s *Server) handleAIHistory(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, s.aiHistory(s.currentAdmin(r).ID, 60))
}

// POST /api/panel/ai/clear - a new talk.
func (s *Server) handleAIClear(w http.ResponseWriter, r *http.Request) {
	_, _ = s.db.sql.Exec(`DELETE FROM ai_messages WHERE admin_id = ?`, s.currentAdmin(r).ID)
	writeJSON(w, map[string]any{"ok": true})
}

// POST /api/panel/ai/note - {content}: an action's result the user got by
// hand, kept with the talk for the next question.
func (s *Server) handleAINote(w http.ResponseWriter, r *http.Request) {
	var in struct {
		Content string `json:"content"`
	}
	if json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&in) == nil {
		s.aiAppend(s.currentAdmin(r).ID, "result", in.Content)
	}
	writeJSON(w, map[string]any{"ok": true})
}

// GET /api/panel/ai - whether the assistant is on for this user.
func (s *Server) handleAIState(w http.ResponseWriter, r *http.Request) {
	on, _ := s.db.GetConfig(aiUserKey(s.currentAdmin(r).ID))
	writeJSON(w, map[string]any{"on": on == "1"})
}

// POST /api/panel/ai/toggle - {on}: the user's own switch.
func (s *Server) handleAIToggle(w http.ResponseWriter, r *http.Request) {
	var in struct {
		On bool `json:"on"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<10)).Decode(&in); err != nil {
		writeJSON(w, map[string]any{"ok": false, "error": "плохой запрос"})
		return
	}
	v := ""
	if in.On {
		v = "1"
	}
	_ = s.db.SetConfig(aiUserKey(s.currentAdmin(r).ID), v)
	writeJSON(w, map[string]any{"ok": true})
}

// POST /settings/ai - the model: kind, address, model name and key (an
// empty key keeps the one there).
func (s *Server) handleSetAI(w http.ResponseWriter, r *http.Request) {
	if !checkCSRF(r) {
		http.Error(w, "bad csrf", http.StatusForbidden)
		return
	}
	typ := r.FormValue("ai_type")
	if typ != "openai" && typ != "anthropic" {
		typ = "gradio"
	}
	_ = s.db.SetConfig(cfgAIType, typ)
	_ = s.db.SetConfig(cfgAIURL, strings.TrimSpace(r.FormValue("ai_url")))
	_ = s.db.SetConfig(cfgAIModel, strings.TrimSpace(r.FormValue("ai_model")))
	if k := strings.TrimSpace(r.FormValue("ai_key")); k != "" {
		_ = s.db.SetConfig(cfgAIKey, k)
	}
	if r.FormValue("ai_key_clear") == "1" {
		_ = s.db.SetConfig(cfgAIKey, "")
	}
	http.Redirect(w, r, "/settings#ai", http.StatusSeeOther)
}

// aiWindow is what of the stored talk goes to the model: the newest
// messages within budget characters - the last one always - with older
// action results and answers cut to their head and tail (the newest result
// whole), results given as the user's words, and messages of one role in a
// row merged (some models want them alternating). It starts with the user.
func aiWindow(rows []aiMsg, budget int) []aiMsg {
	cut := func(s string, max int) string {
		if len(s) <= max {
			return s
		}
		head, tail := max/3, max-max/3
		return utf8Safe(s[:head]) + "\n…[сокращено]…\n" + utf8Safe(s[len(s)-tail:])
	}
	lastResult := -1
	for i := len(rows) - 1; i >= 0; i-- {
		if rows[i].Role == "result" {
			lastResult = i
			break
		}
	}
	var picked []aiMsg
	total := 0
	for i := len(rows) - 1; i >= 0; i-- {
		m := rows[i]
		switch {
		case m.Role == "result":
			max := 1500
			if i == lastResult {
				max = 12000
			}
			m = aiMsg{Role: "user", Content: "Результат действия:\n" + cut(m.Content, max)}
		case m.Role == "assistant" && i < len(rows)-2:
			m.Content = cut(m.Content, 3000)
		case m.Role != "assistant":
			m.Role = "user"
		}
		if len(picked) > 0 && total+len(m.Content) > budget {
			break
		}
		total += len(m.Content)
		picked = append(picked, m)
	}
	// Oldest first, from the user's first word.
	for l, r := 0, len(picked)-1; l < r; l, r = l+1, r-1 {
		picked[l], picked[r] = picked[r], picked[l]
	}
	for len(picked) > 0 && picked[0].Role != "user" {
		picked = picked[1:]
	}
	var out []aiMsg
	for _, m := range picked {
		if n := len(out); n > 0 && out[n-1].Role == m.Role {
			out[n-1].Content += "\n\n" + m.Content
			continue
		}
		out = append(out, m)
	}
	return out
}

// utf8Safe trims broken runes off the ends of a cut string.
func utf8Safe(s string) string {
	for len(s) > 0 && !utf8.ValidString(s) {
		if r, _ := utf8.DecodeRuneInString(s); r == utf8.RuneError {
			s = s[1:]
		} else {
			s = s[:len(s)-1]
		}
	}
	return s
}

// stripThink drops the thinking (::думаю:: … ::ответ::, or <think>) before
// an answer is kept: it is not worth its room in the talk.
func stripThink(s string) string {
	t := strings.TrimLeft(s, " \n")
	for _, p := range [][2]string{{"::думаю::", "::ответ::"}, {"<think>", "</think>"}} {
		if strings.HasPrefix(t, p[0]) {
			if i := strings.Index(t, p[1]); i >= 0 {
				return strings.TrimLeft(t[i+len(p[1]):], " \n")
			}
			return ""
		}
	}
	return s
}