303 lines
12 KiB
Go
303 lines
12 KiB
Go
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)
|
|
}
|