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,302 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Pictures for the assistant: the model asks for one in its answer - a
|
||||
// ```image block, a new picture or an edit of a photo - and the panel makes
|
||||
// it through the Vertex API (https://s2.x-grk.com/api/vertex/docs). The
|
||||
// work is a job there: it is put, then asked about until it is done. The
|
||||
// API key is the admins', set in Настройки and never sent to the browser.
|
||||
|
||||
const (
|
||||
cfgImageKey = "ai_image_key"
|
||||
cfgImageURL = "ai_image_url" // not in the settings: the service's address, for tests
|
||||
|
||||
// DefaultImageURL is where the Vertex API is.
|
||||
DefaultImageURL = "https://s2.x-grk.com"
|
||||
|
||||
imageMaxBytes = 20 << 20 // what the service takes of one photo
|
||||
)
|
||||
|
||||
// imageConf is the service's address and the key (none: pictures are off).
|
||||
func (s *Server) imageConf() (base, key string) {
|
||||
base, _ = s.db.GetConfig(cfgImageURL)
|
||||
if base == "" {
|
||||
base = DefaultImageURL
|
||||
}
|
||||
key, _ = s.db.GetConfig(cfgImageKey)
|
||||
return strings.TrimRight(base, "/"), strings.TrimSpace(key)
|
||||
}
|
||||
|
||||
// imagePrompt tells the model it can make pictures, and how - when there
|
||||
// is a key.
|
||||
func (s *Server) imagePrompt() string {
|
||||
if _, key := s.imageConf(); key == "" {
|
||||
return ""
|
||||
}
|
||||
return "\n\nКартинки. Ты умеешь рисовать и править фото: сам, по смыслу разговора, без просьбы «вызови инструмент» и без лишних вопросов, " +
|
||||
"когда пользователь просит нарисовать, сгенерировать, изобразить, сделать картинку, арт, логотип, обои, иллюстрацию " +
|
||||
"или изменить, дорисовать, убрать, заменить что-то на фото.\n" +
|
||||
"```image\nописание картинки\n``` - создать новую. Опиши сцену, стиль, свет, детали; язык описания любой.\n" +
|
||||
"```image edit\nчто изменить на фото\n``` - править фото пользователя: то, что приложено к его сообщению " +
|
||||
"(приложено несколько - после edit назови файл: ```image edit имя.jpg```); фото нет - последнюю созданную картинку.\n" +
|
||||
"В первой строке блока можно задать формат: 1:1 (по умолчанию), 16:9, 9:16, 3:2, 2:3, 4:3, 3:4, auto - и качество: " +
|
||||
"fast (по умолчанию, 1 кредит), artist (5, лучше), 0.2plus (10, лучшее); повышай качество, только если пользователь просит. " +
|
||||
"Например: ```image 16:9 artist```.\n" +
|
||||
"Сервис по умолчанию сам улучшает описание (дописывает детали и стиль) - это хорошо для короткого запроса. " +
|
||||
"Если пользователь дал готовое, точное описание (промпт целиком, логотип «ровно как сказано», текст на картинке, строгая композиция) " +
|
||||
"или просит ничего не добавлять - добавь noenhance: ```image noenhance 3:2```. Для правки фото это не применяется.\n" +
|
||||
"Картинку покажет сама панель: ссылку в ответ не вставляй, скажи коротко, что получилось, и предложи поправить. " +
|
||||
"Делается она до пары минут. Одно действие за ответ."
|
||||
}
|
||||
|
||||
// imageError says in Russian what an error code of the service means.
|
||||
func imageError(code, msg string) string {
|
||||
switch code {
|
||||
case "unauthorized":
|
||||
return "сервис не принял API-ключ - проверьте его в Настройки › ИИ-ассистент › Генерация изображений"
|
||||
case "banned":
|
||||
return "аккаунт сервиса картинок заблокирован"
|
||||
case "insufficient_credits":
|
||||
return "на аккаунте сервиса картинок не хватает кредитов"
|
||||
case "daily_limit":
|
||||
return "исчерпан дневной лимит картинок"
|
||||
case "job_busy":
|
||||
return "уже делается другая картинка - дождитесь её"
|
||||
case "invalid_prompt":
|
||||
return "сервис не принял описание (пустое или длиннее 4000 символов)"
|
||||
case "invalid_image":
|
||||
return "сервис не принял фото (нужен png, jpg или webp)"
|
||||
case "image_too_large":
|
||||
return "фото больше 20 МБ"
|
||||
case "tokens_exhausted":
|
||||
return "сервис картинок перегружен - попробуйте позже"
|
||||
case "generation_failed", "edit_failed":
|
||||
return "сервис не справился с картинкой (кредиты вернутся)"
|
||||
case "source_missing":
|
||||
return "исходное фото потерялось на сервисе - загрузите его заново"
|
||||
}
|
||||
if msg == "" {
|
||||
msg = "сервис картинок ответил ошибкой"
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
var imageClient = &http.Client{Timeout: 90 * time.Second}
|
||||
|
||||
// imageCall sends the request to the service with the key and reads its
|
||||
// JSON; an error is words for the user.
|
||||
func imageCall(ctx context.Context, req *http.Request, key string, v any) error {
|
||||
req = req.WithContext(ctx)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := imageClient.Do(req)
|
||||
if err != nil {
|
||||
return errors.New("сервис картинок не отвечает: " + err.Error())
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
var e struct {
|
||||
OK *bool `json:"ok"`
|
||||
ErrorCode string `json:"error_code"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(b, &e) != nil {
|
||||
return errors.New("сервис картинок ответил " + resp.Status)
|
||||
}
|
||||
if (e.OK != nil && !*e.OK) || resp.StatusCode/100 != 2 {
|
||||
return errors.New(imageError(e.ErrorCode, e.Error))
|
||||
}
|
||||
return json.Unmarshal(b, v)
|
||||
}
|
||||
|
||||
var (
|
||||
imageRatios = map[string]bool{"1:1": true, "16:9": true, "9:16": true, "3:2": true, "2:3": true, "4:3": true, "3:4": true, "auto": true}
|
||||
imageModels = map[string]bool{"fast": true, "artist": true, "0.2plus": true}
|
||||
imageJobRe = regexp.MustCompile(`^[A-Za-z0-9_-]{1,64}$`)
|
||||
imageExt = map[string]string{"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg", "image/webp": "webp"}
|
||||
)
|
||||
|
||||
// POST /api/panel/ai/image - {mode: generate|edit, prompt, model, ratio,
|
||||
// image: {mime, data (base64)} | image_url}: puts the job, answers its id.
|
||||
func (s *Server) handleAIImage(w http.ResponseWriter, r *http.Request) {
|
||||
fail := func(msg string) { writeJSON(w, map[string]any{"ok": false, "error": msg}) }
|
||||
if !s.aiOn(s.currentAdmin(r).ID) {
|
||||
fail("ассистент выключен в профиле")
|
||||
return
|
||||
}
|
||||
base, key := s.imageConf()
|
||||
if key == "" {
|
||||
fail("не задан API-ключ сервиса картинок (Настройки › ИИ-ассистент › Генерация изображений)")
|
||||
return
|
||||
}
|
||||
var in struct {
|
||||
Mode string `json:"mode"`
|
||||
Prompt string `json:"prompt"`
|
||||
Model string `json:"model"`
|
||||
Ratio string `json:"ratio"`
|
||||
Enhance *bool `json:"enhance"` // the service improves the description; on unless false
|
||||
ImageURL string `json:"image_url"`
|
||||
Image struct {
|
||||
MIME string `json:"mime"`
|
||||
Data string `json:"data"`
|
||||
} `json:"image"`
|
||||
}
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 32<<20)).Decode(&in); err != nil {
|
||||
fail("плохой запрос")
|
||||
return
|
||||
}
|
||||
in.Prompt = strings.TrimSpace(in.Prompt)
|
||||
if in.Prompt == "" {
|
||||
fail("пустое описание картинки")
|
||||
return
|
||||
}
|
||||
if r := []rune(in.Prompt); len(r) > 4000 {
|
||||
in.Prompt = string(r[:4000])
|
||||
}
|
||||
var job struct {
|
||||
JobID string `json:"job_id"`
|
||||
Cost int `json:"cost"`
|
||||
}
|
||||
ctx := r.Context()
|
||||
if in.Mode != "edit" {
|
||||
if !imageModels[in.Model] {
|
||||
in.Model = "fast"
|
||||
}
|
||||
if !imageRatios[in.Ratio] {
|
||||
in.Ratio = "1:1"
|
||||
}
|
||||
enhance := in.Enhance == nil || *in.Enhance
|
||||
b, _ := json.Marshal(map[string]any{"prompt": in.Prompt, "model": in.Model, "ratio": in.Ratio, "enhance": enhance})
|
||||
req, _ := http.NewRequest(http.MethodPost, base+"/api/vertex/v1/generate", bytes.NewReader(b))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if err := imageCall(ctx, req, key, &job); err != nil {
|
||||
fail(err.Error())
|
||||
return
|
||||
}
|
||||
} else {
|
||||
mime, data, err := s.imageSource(ctx, base, in.Image.MIME, in.Image.Data, in.ImageURL)
|
||||
if err != nil {
|
||||
fail(err.Error())
|
||||
return
|
||||
}
|
||||
var body bytes.Buffer
|
||||
mw := multipart.NewWriter(&body)
|
||||
h := textproto.MIMEHeader{}
|
||||
h.Set("Content-Disposition", `form-data; name="image"; filename="photo.`+imageExt[mime]+`"`)
|
||||
h.Set("Content-Type", mime)
|
||||
part, _ := mw.CreatePart(h)
|
||||
_, _ = part.Write(data)
|
||||
_ = mw.WriteField("prompt", in.Prompt)
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequest(http.MethodPost, base+"/api/vertex/v1/edit", &body)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
if err := imageCall(ctx, req, key, &job); err != nil {
|
||||
fail(err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
if job.JobID == "" {
|
||||
fail("сервис картинок не дал номер задачи")
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]any{"ok": true, "job_id": job.JobID, "cost": job.Cost})
|
||||
}
|
||||
|
||||
// imageSource is the photo to edit: sent by the browser, or a picture made
|
||||
// before, taken from the service's own address (and from no other).
|
||||
func (s *Server) imageSource(ctx context.Context, base, mime, b64, picURL string) (string, []byte, error) {
|
||||
if picURL != "" {
|
||||
u, err := url.Parse(picURL)
|
||||
bu, _ := url.Parse(base)
|
||||
if err != nil || bu == nil || u.Host != bu.Host || (u.Scheme != "https" && u.Scheme != "http") {
|
||||
return "", nil, errors.New("прежнюю картинку можно править только с сервиса картинок")
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
resp, err := imageClient.Do(req)
|
||||
if err != nil || resp.StatusCode/100 != 2 {
|
||||
if resp != nil {
|
||||
resp.Body.Close()
|
||||
}
|
||||
return "", nil, errors.New("прежняя картинка больше не открывается")
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
data, _ := io.ReadAll(io.LimitReader(resp.Body, imageMaxBytes+1))
|
||||
mime = strings.ToLower(strings.SplitN(resp.Header.Get("Content-Type"), ";", 2)[0])
|
||||
if imageExt[mime] == "" {
|
||||
mime = "image/png"
|
||||
}
|
||||
return mime, data, nil
|
||||
}
|
||||
mime = strings.ToLower(mime)
|
||||
if imageExt[mime] == "" {
|
||||
return "", nil, errors.New("править можно png, jpg и webp")
|
||||
}
|
||||
data, err := base64.StdEncoding.DecodeString(b64)
|
||||
if err != nil || len(data) == 0 {
|
||||
return "", nil, errors.New("нет фото для правки: прикрепите его к сообщению")
|
||||
}
|
||||
if len(data) > imageMaxBytes {
|
||||
return "", nil, errors.New("фото больше 20 МБ")
|
||||
}
|
||||
return mime, data, nil
|
||||
}
|
||||
|
||||
// GET /api/panel/ai/image/{job}?wait=1..45 - the job's state, waiting up to
|
||||
// wait seconds for it to end.
|
||||
func (s *Server) handleAIImageJob(w http.ResponseWriter, r *http.Request) {
|
||||
fail := func(msg string) { writeJSON(w, map[string]any{"ok": false, "error": msg}) }
|
||||
if !s.aiOn(s.currentAdmin(r).ID) {
|
||||
fail("ассистент выключен в профиле")
|
||||
return
|
||||
}
|
||||
base, key := s.imageConf()
|
||||
id := r.PathValue("job")
|
||||
if key == "" || !imageJobRe.MatchString(id) {
|
||||
fail("нет такой задачи")
|
||||
return
|
||||
}
|
||||
wait, _ := strconv.Atoi(r.URL.Query().Get("wait"))
|
||||
if wait < 0 {
|
||||
wait = 0
|
||||
}
|
||||
if wait > 45 {
|
||||
wait = 45
|
||||
}
|
||||
req, _ := http.NewRequest(http.MethodGet, base+"/api/vertex/v1/jobs/"+id+"?wait="+strconv.Itoa(wait), nil)
|
||||
var st struct {
|
||||
Status string `json:"status"`
|
||||
PhotoURL string `json:"photo_url"`
|
||||
Cost int `json:"cost"`
|
||||
Balance int `json:"balance"`
|
||||
ErrorCode string `json:"error_code"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := imageCall(r.Context(), req, key, &st); err != nil {
|
||||
fail(err.Error())
|
||||
return
|
||||
}
|
||||
out := map[string]any{"ok": true, "status": st.Status}
|
||||
switch st.Status {
|
||||
case "done":
|
||||
out["photo_url"], out["cost"], out["balance"] = st.PhotoURL, st.Cost, st.Balance
|
||||
case "failed":
|
||||
out["error"] = imageError(st.ErrorCode, st.Error)
|
||||
}
|
||||
writeJSON(w, out)
|
||||
}
|
||||
Reference in new issue
Block a user