Skip to content

Commit

Permalink
imagine
Browse files Browse the repository at this point in the history
  • Loading branch information
radiantspace committed May 4, 2024
1 parent 803af34 commit 95ad786
Show file tree
Hide file tree
Showing 7 changed files with 85 additions and 28 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ in Slack:
- [x] Voice response (OpenAI TTS)
- [x] Threads, i.e. context awareness and memory (OpenAI Assistant Threads)
- [ ] Image recognition
- [ ] Image generation
- [x] Image generation


## Telegram Features
Expand Down
53 changes: 36 additions & 17 deletions backend/app/ai/openai/image.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"io"
"net/http"
"talk2robots/m/v2/app/config"
"talk2robots/m/v2/app/lib"
"talk2robots/m/v2/app/models"
"talk2robots/m/v2/app/payments"
"time"
Expand All @@ -19,7 +20,7 @@ const (
DALLE3_HD = 0.08
)

func CreateImage(ctx context.Context, prompt string) (string, error) {
func CreateImage(ctx context.Context, prompt string) (string, string, error) {
timeNow := time.Now()

// Set the request body
Expand All @@ -28,39 +29,46 @@ func CreateImage(ctx context.Context, prompt string) (string, error) {
Prompt string `json:"prompt"`
N int `json:"n"`
Size string `json:"size"`
Quality string `json:"quality"`
Quality string `json:"quality,omitempty"`
}{
Model: string(models.DallE3),
Prompt: prompt,
N: 1,
Size: "1024x1024",
Quality: "hd",
Model: string(models.DallE3),
Prompt: prompt,
N: 1,
Size: "1024x1024",
}

if lib.IsUserBasic(ctx) {
requestBody.Quality = "hd"
}

// Convert the request body to JSON
requestBodyJSON, err := json.Marshal(requestBody)
if err != nil {
return "", err
return "", "", err
}

// Create the HTTP request
req, err := http.NewRequest("POST", "https://api.openai.com/v1/images/generations", bytes.NewBuffer(requestBodyJSON))
if err != nil {
return "", err
return "", "", err
}

// Set the request headers
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+config.CONFIG.OpenAIAPIKey)

usage := models.CostAndUsage{
Engine: models.DallE3,
PricePerInputUnit: DALLE3_HD,
Cost: 0,
Engine: models.DallE3,
ImagePrice: DALLE3_S,
Usage: models.Usage{
PromptTokens: 1,
ImagesCount: 1,
},
}

if lib.IsUserBasic(ctx) {
usage.ImagePrice = DALLE3_HD
}

status := fmt.Sprintf("status:%d", 0)
defer func() {
config.CONFIG.DataDogClient.Timing("openai.image.latency", time.Since(timeNow), []string{status, "model:" + string(models.DallE3)}, 1)
Expand All @@ -73,21 +81,32 @@ func CreateImage(ctx context.Context, prompt string) (string, error) {
status = fmt.Sprintf("status:%d", resp.StatusCode)
}
if err != nil {
return "", err
return "", "", err
}

// Read the response body
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return "", err
return "", "", err
}

// "{\n \"created\": 1714859364,\n \"data\": [\n {\n \"revised_prompt\": \"Create an image of a fluffy Siberian cat. It has a multi-colored coat in hues of cream, brown, and black. The cat is perched on a vintage bicycle with a shiny silver finish. The bicycle is stationary and placed against a backdrop of a vivid sunset, casting a warm orange glow that reflects off of the spokes of the bicycle wheels.\",\n \"url\": \"https://oaidalleapiprodscus.blob.core.windows.net/private/org-mXLEStBZvlLX16KcOWmBQxRi/user-FolUm62ffE794657USBZWAuZ/img-fu7BaN8wBqO93SpduqbPrjzL.png?st=2024-05-04T20%3A49%3A24Z\u0026se=2024-05-04T22%3A49%3A24Z\u0026sp=r\u0026sv=2021-08-06\u0026sr=b\u0026rscd=inline\u0026rsct=image/png\u0026skoid=6aaadede-4fb3-4698-a8f6-684d7786b067\u0026sktid=a48cca56-e6da-484e-a814-9c849652bcb3\u0026skt=2024-05-04T21%3A45%3A40Z\u0026ske=2024-05-05T21%3A45%3A40Z\u0026sks=b\u0026skv=2021-08-06\u0026sig=XTWxaLMaxX%2Bt8sg8QIbax0jy2cZ0TDcr7W9/i83Yca4%3D\"\n }\n ]\n}\n"
responseJson := struct {
Data []struct {
RevisedPrompt string `json:"revised_prompt"`
URL string `json:"url"`
} `json:"data"`
}{}
if err := json.Unmarshal(responseBody, &responseJson); err != nil {
return "", "", err
}

// Check the response status code
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("unexpected status code: %d, response: %s", resp.StatusCode, responseBody)
return "", "", fmt.Errorf("unexpected status code: %d, response: %s", resp.StatusCode, responseBody)
}

go payments.Bill(ctx, usage)

return string(responseBody), nil
return responseJson.Data[0].URL, responseJson.Data[0].RevisedPrompt, nil
}
4 changes: 4 additions & 0 deletions backend/app/lib/payments.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,10 @@ func UserTotalAudioMinutesKey(user string) string {
return user + ":total_audio_minutes"
}

func UserTotalImagesKey(user string) string {
return user + ":total_images"
}

func UserCurrentThreadPromptKey(user string, topic string) string {
if topic != "" && topic != "0" {
return user + ":" + topic + ":current-thread-prompt-tokens"
Expand Down
2 changes: 1 addition & 1 deletion backend/app/models/openai.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ type CostAndUsage struct {
Engine Engine `json:"engine"`
PricePerInputUnit float64 `json:"price_per_input_unit"`
PricePerOutputUnit float64 `json:"price_per_output_unit"`
ImagePrice float64 `json:"image_price,omitempty"`
Cost float64 `json:"cost"`
Usage Usage `json:"usage"`
User string `json:"user"`
Expand Down Expand Up @@ -173,7 +174,6 @@ type Usage struct {
TotalTokens int `json:"total_tokens"`
AudioDuration float64 `json:"audio_duration"` // only for Whisper API
ImagesCount int `json:"images_count,omitempty"`
ImagePrice float64 `json:"image_price,omitempty"`
}

type ThreadRunRequest struct {
Expand Down
15 changes: 14 additions & 1 deletion backend/app/payments/billing.go
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,8 @@ func Bill(originalContext context.Context, usage models.CostAndUsage) models.Cos
usage.Cost =
float64(usage.Usage.PromptTokens)*usage.PricePerInputUnit +
float64(usage.Usage.CompletionTokens)*usage.PricePerOutputUnit +
usage.Usage.AudioDuration*usage.PricePerInputUnit
usage.Usage.AudioDuration*usage.PricePerInputUnit +
float64(usage.Usage.ImagesCount)*usage.ImagePrice
_, err := redis.RedisClient.IncrByFloat(ctx, "system_totals:cost", usage.Cost).Result()
if err != nil {
log.Errorf("[billing] error incrementing system cost: %v", err)
Expand Down Expand Up @@ -149,6 +150,18 @@ func Bill(originalContext context.Context, usage models.CostAndUsage) models.Cos
config.CONFIG.DataDogClient.Distribution("billing.audio_minutes", usage.Usage.AudioDuration, []string{"engine:" + string(usage.Engine), "user_type:" + userType}, 1)
}

if usage.Usage.ImagesCount > 0 {
_, err = redis.RedisClient.IncrBy(context.Background(), lib.UserTotalImagesKey(usage.User), int64(usage.Usage.ImagesCount)).Result()
if err != nil {
log.Errorf("[billing] error incrementing user total images: %v", err)
}
_, err = redis.RedisClient.IncrBy(context.Background(), "system_totals:images", int64(usage.Usage.ImagesCount)).Result()
if err != nil {
log.Errorf("[billing] error incrementing system total images: %v", err)
}
config.CONFIG.DataDogClient.Distribution("billing.images", float64(usage.Usage.ImagesCount), []string{"engine:" + string(usage.Engine), "user_type:" + userType}, 1)
}

userTotalCost, err := redis.RedisClient.IncrByFloat(context.Background(), lib.UserTotalCostKey(usage.User), usage.Cost).Result()
if err != nil {
log.Errorf("[billing] error getting user total cost: %s", err)
Expand Down
7 changes: 4 additions & 3 deletions backend/app/telegram/commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,12 @@ var VASILISA_BIRTHDAY = time.Date(2007, 12, 13, 23, 45, 0, 0, time.FixedZone("UT
var ONBOARDING_TEXT = `Hi, I'm a bot powered by AI! I can:
- Default 🧠: chat with or answer any questions (/chatgpt)
- ✨ New feature 🎙️: talk to AI using voice messages (/voicegpt)
- ✨ New feature 🖼️: draw, just ask me to draw something ('create/draw/picture an image of a fish riding a bicycle')
- Correct grammar: (/grammar)
- Explain grammar and mistakes: (/teacher)
- ✨ New feature: remember context in /chatgpt and /voicegpt modes (use /clear to clear current thread)
- ✨ New feature: transcribe voice/audio/video messages (/transcribe)
- ✨ New feature: summarize text/voice/audio/video messages (/summarize)
- Remember context in /chatgpt and /voicegpt modes (use /clear to clear current thread)
- Transcribe voice/audio/video messages (/transcribe)
- Summarize text/voice/audio/video messages (/summarize)
Enjoy and let me know if any /support is needed!`

Expand Down
30 changes: 25 additions & 5 deletions backend/app/telegram/telegram.go
Original file line number Diff line number Diff line change
Expand Up @@ -242,17 +242,29 @@ func handleMessage(bot *telego.Bot, message telego.Message) {
}

if IsCreateImageCommand(message.Text) {
url, err := openai.CreateImage(ctx, message.Text)
config.CONFIG.DataDogClient.Incr("telegram.create_image_received", []string{"channel_type:" + message.Chat.Type}, 1)
log.Infof("Generating image in a chat %s..", chatIDString)
sendImageAction(bot, &message)
url, revisedPrompt, err := openai.CreateImage(ctx, message.Text)
if err != nil {
log.Errorf("Error creating image in chat %s: %v", chatIDString, err)
return
}
log.Infof("Sending image to chat %s: %s", chatIDString, url)
log.Infof("Sending image to chat %s", chatIDString)
if len(revisedPrompt) > 1000 {
revisedPrompt = revisedPrompt[:997] + "..."
}
if url != "" {
bot.SendPhoto(&telego.SendPhotoParams{
ChatID: chatID,
Photo: telego.InputFile{URL: url},
_, err := bot.SendPhoto(&telego.SendPhotoParams{
ChatID: chatID,
Photo: telego.InputFile{URL: url},
Caption: revisedPrompt,
MessageThreadID: message.MessageThreadID,
})

if err != nil {
log.Errorf("Error sending image to chat %s: %v", chatIDString, err)
}
}
return
}
Expand Down Expand Up @@ -684,6 +696,14 @@ func sendTypingAction(bot *telego.Bot, message *telego.Message) {
}
}

func sendImageAction(bot *telego.Bot, message *telego.Message) {
chatID := message.Chat.ChatID()
err := bot.SendChatAction(&telego.SendChatActionParams{ChatID: chatID, Action: telego.ChatActionUploadPhoto, MessageThreadID: message.MessageThreadID})
if err != nil {
log.Errorf("Failed to send chat action: %v", err)
}
}

func sendAudioAction(bot *telego.Bot, message *telego.Message) {
chatID := message.Chat.ChatID()
err := bot.SendChatAction(&telego.SendChatActionParams{ChatID: chatID, Action: telego.ChatActionRecordVoice, MessageThreadID: message.MessageThreadID})
Expand Down

0 comments on commit 95ad786

Please sign in to comment.