mirror of
https://github.com/coaidev/coai.git
synced 2025-05-19 21:10:18 +09:00
77 lines
2.3 KiB
Go
77 lines
2.3 KiB
Go
package api
|
|
|
|
import (
|
|
"chat/auth"
|
|
"chat/types"
|
|
"chat/utils"
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/spf13/viper"
|
|
"time"
|
|
)
|
|
|
|
func GetImage(prompt string) (string, error) {
|
|
res, err := utils.Post(viper.GetString("openai.image_endpoint")+"/images/generations", map[string]string{
|
|
"Content-Type": "application/json",
|
|
"Authorization": "Bearer " + GetRandomKey(viper.GetString("openai.image")),
|
|
}, types.ChatGPTImageRequest{
|
|
Prompt: prompt,
|
|
Size: "512x512",
|
|
N: 1,
|
|
})
|
|
if err != nil || res == nil {
|
|
return "", err
|
|
}
|
|
|
|
if err, ok := res.(map[string]interface{})["error"]; ok {
|
|
return "", fmt.Errorf(err.(map[string]interface{})["message"].(string))
|
|
}
|
|
data := res.(map[string]interface{})["data"].([]interface{})[0].(map[string]interface{})["url"]
|
|
return data.(string), nil
|
|
}
|
|
|
|
func GetImageWithCache(ctx context.Context, prompt string, cache *redis.Client) (string, error) {
|
|
res, err := cache.Get(ctx, fmt.Sprintf(":image:%s", prompt)).Result()
|
|
if err != nil || len(res) == 0 || res == "" {
|
|
res, err := GetImage(prompt)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
cache.Set(ctx, fmt.Sprintf(":image:%s", prompt), res, time.Hour*6)
|
|
return res, nil
|
|
}
|
|
|
|
return res, nil
|
|
}
|
|
|
|
func GetLimitFormat(id int64) string {
|
|
t := time.Now().Format("2006-01-02")
|
|
return fmt.Sprintf(":imagelimit:%s:%d", t, id)
|
|
}
|
|
|
|
func GetImageWithUserLimit(user *auth.User, prompt string, db *sql.DB, cache *redis.Client) (string, error) {
|
|
// 5 images one day per user (count by cache)
|
|
res, err := cache.Get(context.Background(), GetLimitFormat(user.GetID(db))).Result()
|
|
if err != nil || len(res) == 0 || res == "" {
|
|
cache.Set(context.Background(), GetLimitFormat(user.GetID(db)), "1", time.Hour*24)
|
|
return GetImageWithCache(context.Background(), prompt, cache)
|
|
}
|
|
|
|
if res == "5" {
|
|
if auth.ReduceDalle(db, user) {
|
|
return GetImageWithCache(context.Background(), prompt, cache)
|
|
}
|
|
return "", fmt.Errorf("you have reached your limit of 5 free images per day, please buy more dalle usage or wait until tomorrow")
|
|
} else {
|
|
cache.Set(context.Background(), GetLimitFormat(user.GetID(db)), fmt.Sprintf("%d", utils.ToInt(res)+1), time.Hour*24)
|
|
return GetImageWithCache(context.Background(), prompt, cache)
|
|
}
|
|
}
|
|
|
|
func GetImageMarkdown(url string) string {
|
|
return fmt.Sprintln("![image](", url, ")")
|
|
}
|