Skip to content

Commit

Permalink
refactor session data
Browse files Browse the repository at this point in the history
  • Loading branch information
@fegloff
fegloff committed Apr 17, 2024
1 parent 147b336 commit ce88b3a
Show file tree
Hide file tree
Showing 11 changed files with 86 additions and 123 deletions.
4 changes: 2 additions & 2 deletions src/bot.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,7 +333,7 @@ const PayableBots: Record<string, PayableBotConfig> = {
claudeBot: { bot: claudeBot },
vertexBot: { bot: vertexBot },
openAiBot: {
enabled: (ctx: OnMessageContext) => ctx.session.openAi.imageGen.isEnabled,
enabled: (ctx: OnMessageContext) => ctx.session.dalle.isEnabled,
bot: openAiBot
},
oneCountryBot: { bot: oneCountryBot }
Expand Down Expand Up @@ -394,7 +394,7 @@ const onMessage = async (ctx: OnMessageContext): Promise<void> => {
return
}
// Any message interacts with ChatGPT (only for private chats or /ask on enabled on group chats)
if (ctx.update.message.chat && (ctx.chat.type === 'private' || ctx.session.openAi.chatGpt.isFreePromptChatGroups)) {
if (ctx.update.message.chat && (ctx.chat.type === 'private' || ctx.session.chatGpt.isFreePromptChatGroups)) {
await openAiBot.onEvent(ctx, (e) => {
logger.error(e)
})
Expand Down
55 changes: 33 additions & 22 deletions src/helpers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,28 +3,28 @@ import { type BotSessionData } from './modules/types'

export function createInitialSessionData (): BotSessionData {
return {
openAi: {
imageGen: {
numImages: config.openAi.dalle.sessionDefault.numImages,
imgSize: config.openAi.dalle.sessionDefault.imgSize,
isEnabled: config.openAi.dalle.isEnabled,
imgRequestQueue: [],
isProcessingQueue: false,
imageGenerated: [],
isInscriptionLotteryEnabled: config.openAi.dalle.isInscriptionLotteryEnabled,
imgInquiried: []
},
chatGpt: {
model: config.openAi.chatGpt.model,
isEnabled: config.openAi.chatGpt.isEnabled,
isFreePromptChatGroups: config.openAi.chatGpt.isFreePromptChatGroups,
chatConversation: [],
price: 0,
usage: 0,
isProcessingQueue: false,
requestQueue: []
}
},
// openAi: {
// imageGen: {
// numImages: config.openAi.dalle.sessionDefault.numImages,
// imgSize: config.openAi.dalle.sessionDefault.imgSize,
// isEnabled: config.openAi.dalle.isEnabled,
// imgRequestQueue: [],
// isProcessingQueue: false,
// imageGenerated: [],
// isInscriptionLotteryEnabled: config.openAi.dalle.isInscriptionLotteryEnabled,
// imgInquiried: []
// },
// chatGpt: {
// model: config.openAi.chatGpt.model,
// isEnabled: config.openAi.chatGpt.isEnabled,
// isFreePromptChatGroups: config.openAi.chatGpt.isFreePromptChatGroups,
// chatConversation: [],
// price: 0,
// usage: 0,
// isProcessingQueue: false,
// requestQueue: []
// }
// },
oneCountry: { lastDomain: '' },
translate: {
languages: [],
Expand All @@ -50,11 +50,22 @@ export function createInitialSessionData (): BotSessionData {
chatGpt: {
model: config.llms.model,
isEnabled: config.llms.isEnabled,
isFreePromptChatGroups: config.openAi.chatGpt.isFreePromptChatGroups,
chatConversation: [],
price: 0,
usage: 0,
isProcessingQueue: false,
requestQueue: []
},
dalle: {
numImages: config.openAi.dalle.sessionDefault.numImages,
imgSize: config.openAi.dalle.sessionDefault.imgSize,
isEnabled: config.openAi.dalle.isEnabled,
imgRequestQueue: [],
isProcessingQueue: false,
imageGenerated: [],
isInscriptionLotteryEnabled: config.openAi.dalle.isInscriptionLotteryEnabled,
imgInquiried: []
}
}
}
2 changes: 1 addition & 1 deletion src/modules/1country/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -554,7 +554,7 @@ export class OneCountryBot implements PayableBot {
).catch(async (e) => { await this.onError(ctx, e, retryCount - 1) })
ctx.transient.analytics.actualResponseTime = now()
if (method === 'editMessageText') {
ctx.session.openAi.chatGpt.chatConversation.pop() // deletes last prompt
ctx.session.chatGpt.chatConversation.pop() // deletes last prompt
}
await sleep(retryAfter * 1000) // wait retryAfter seconds to enable bot
this.botSuspended = false
Expand Down
2 changes: 1 addition & 1 deletion src/modules/document-handler/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ export class DocumentHandler implements PayableBot {
).catch(async (e) => { await this.onError(ctx, e, retryCount - 1) })
ctx.transient.analytics.actualResponseTime = now()
if (method === 'editMessageText') {
ctx.session.openAi.chatGpt.chatConversation.pop() // deletes last prompt
ctx.session.chatGpt.chatConversation.pop() // deletes last prompt
}
await sleep(retryAfter * 1000) // wait retryAfter seconds to enable bot
} else {
Expand Down
4 changes: 2 additions & 2 deletions src/modules/llms/api/openai.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,7 @@ export const streamChatCompletion = async (
}
}
completion = completion.replaceAll('...', '')
const inputTokens = getTokenNumber(conversation[conversation.length - 1].content as string) + ctx.session.openAi.chatGpt.usage
const inputTokens = getTokenNumber(conversation[conversation.length - 1].content as string) + ctx.session.chatGpt.usage
const outputTokens = getTokenNumber(completion)
await ctx.api
.editMessageText(ctx.chat?.id, msgId, completion)
Expand Down Expand Up @@ -257,7 +257,7 @@ export const streamChatVisionCompletion = async (
}
}
completion = completion.replaceAll('...', '')
const inputTokens = getTokenNumber(prompt) + ctx.session.openAi.chatGpt.usage
const inputTokens = getTokenNumber(prompt) + ctx.session.chatGpt.usage
const outputTokens = getTokenNumber(completion)
await ctx.api
.editMessageText(ctx.chat?.id, msgId, completion)
Expand Down
73 changes: 36 additions & 37 deletions src/modules/llms/dalleBot.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,21 +37,15 @@ import { InlineKeyboard } from 'grammy'

export class DalleBot extends LlmsBase {
constructor (payments: BotPayments) {
super(payments, 'DalleBot', 'chatGpt')
super(payments, 'DalleBot', 'dalle')
if (!config.openAi.dalle.isEnabled) {
this.logger.warn('DALL·E 2 Image Bot is disabled in config')
}
}

public getEstimatedPrice (ctx: any): number {
try {
// if (this.isSupportedImageReply(ctx) && !isNaN(+prompts)) {
// const imageNumber = ctx.message?.caption || ctx.message?.text
// const imageSize = ctx.session.openAi.imageGen.imgSize
// const model = getDalleModel(imageSize)
// const price = getDalleModelPrice(model, true, imageNumber) // cents
// return price * priceAdjustment
// }
const session = this.getSession(ctx)
if (
ctx.hasCommand([
SupportedCommands.dalle,
Expand All @@ -60,8 +54,8 @@ export class DalleBot extends LlmsBase {
SupportedCommands.dalleShorter
])
) {
const imageNumber = ctx.session.openAi.imageGen.numImages
const imageSize = ctx.session.openAi.imageGen.imgSize
const imageNumber = session.numImages
const imageSize = session.imgSize
const model = getDalleModel(imageSize)
const price = getDalleModelPrice(model, true, imageNumber) // cents
return price * PRICE_ADJUSTMENT
Expand All @@ -83,11 +77,11 @@ export class DalleBot extends LlmsBase {
SupportedCommands.dalleShort,
SupportedCommands.dalleShorter
])

const session = this.getSession(ctx)
const photo = ctx.message?.reply_to_message?.photo
if (photo) {
const imgId = photo?.[0].file_unique_id ?? ''
if (ctx.session.openAi.imageGen.imgInquiried.find((i) => i === imgId)) {
if (session.imgInquiried.find((i) => i === imgId)) {
return false
}
}
Expand All @@ -112,8 +106,9 @@ export class DalleBot extends LlmsBase {
}

isSupportedImageReply (ctx: OnMessageContext | OnCallBackQueryData): boolean {
const session = this.getSession(ctx)
const photo = ctx.message?.photo ?? ctx.message?.reply_to_message?.photo
if (photo && ctx.session.openAi.imageGen.isEnabled) {
if (photo && session.isEnabled) {
const prompt = ctx.message?.caption ?? ctx.message?.text
if (
prompt &&
Expand Down Expand Up @@ -165,6 +160,7 @@ export class DalleBot extends LlmsBase {
refundCallback: (reason?: string) => void
): Promise<void> {
ctx.transient.analytics.module = this.module
const session = this.getSession(ctx)
const isSupportedEvent = this.isSupportedEvent(ctx)
if (!isSupportedEvent && ctx.chat?.type !== 'private') {
this.logger.warn(`### unsupported command ${ctx.message?.text}`)
Expand All @@ -175,16 +171,16 @@ export class DalleBot extends LlmsBase {
const photo = ctx.message?.photo ?? ctx.message?.reply_to_message?.photo
const prompt = ctx.message?.caption ?? ctx.message?.text ?? ''
const imgId = photo?.[0].file_unique_id ?? ''
if (!ctx.session.openAi.imageGen.imgInquiried.find((i) => i === imgId)) {
ctx.session.openAi.imageGen.imgRequestQueue.push({
if (!session.imgInquiried.find((i) => i === imgId)) {
session.imgRequestQueue.push({
prompt,
photo,
command: 'vision' // !isNaN(+prompt) ? 'alter' : 'vision'
})
if (!ctx.session.openAi.imageGen.isProcessingQueue) {
ctx.session.openAi.imageGen.isProcessingQueue = true
if (!session.isProcessingQueue) {
session.isProcessingQueue = true
await this.onImgRequestHandler(ctx).then(() => {
ctx.session.openAi.imageGen.isProcessingQueue = false
session.isProcessingQueue = false
})
}
} else {
Expand All @@ -197,15 +193,15 @@ export class DalleBot extends LlmsBase {
const photoUrl = getUrlFromText(ctx)
if (photoUrl) {
const prompt = ctx.match
ctx.session.openAi.imageGen.imgRequestQueue.push({
session.imgRequestQueue.push({
prompt,
photoUrl,
command: 'vision' // !isNaN(+prompt) ? 'alter' : 'vision'
})
if (!ctx.session.openAi.imageGen.isProcessingQueue) {
ctx.session.openAi.imageGen.isProcessingQueue = true
if (!session.isProcessingQueue) {
session.isProcessingQueue = true
await this.onImgRequestHandler(ctx).then(() => {
ctx.session.openAi.imageGen.isProcessingQueue = false
session.isProcessingQueue = false
})
}
}
Expand Down Expand Up @@ -240,14 +236,14 @@ export class DalleBot extends LlmsBase {
refundCallback('Prompt has bad words')
return
}
ctx.session.openAi.imageGen.imgRequestQueue.push({
session.imgRequestQueue.push({
command: 'dalle',
prompt
})
if (!ctx.session.openAi.imageGen.isProcessingQueue) {
ctx.session.openAi.imageGen.isProcessingQueue = true
if (!session.isProcessingQueue) {
session.isProcessingQueue = true
await this.onImgRequestHandler(ctx).then(() => {
ctx.session.openAi.imageGen.isProcessingQueue = false
session.isProcessingQueue = false
})
}
return
Expand All @@ -261,17 +257,18 @@ export class DalleBot extends LlmsBase {
}

async onImgRequestHandler (ctx: OnMessageContext | OnCallBackQueryData): Promise<void> {
while (ctx.session.openAi.imageGen.imgRequestQueue.length > 0) {
const session = this.getSession(ctx)
while (session.imgRequestQueue.length > 0) {
try {
const img = ctx.session.openAi.imageGen.imgRequestQueue.shift()
const minBalance = await getMinBalance(ctx, ctx.session.openAi.chatGpt.model)
const img = session.imgRequestQueue.shift()
const minBalance = await getMinBalance(ctx, ctx.session.chatGpt.model)
if (await this.hasBalance(ctx, minBalance)) {
if (img?.command === 'dalle') {
await this.onGenImgCmd(img?.prompt, ctx)
} else {
await this.onInquiryImage(img?.photo, img?.photoUrl, img?.prompt, ctx)
if (img?.photo?.[0].file_unique_id) {
ctx.session.openAi.imageGen.imgInquiried.push(img?.photo?.[0].file_unique_id)
session.imgInquiried.push(img?.photo?.[0].file_unique_id)
}
}
ctx.chatAction = null
Expand All @@ -286,19 +283,20 @@ export class DalleBot extends LlmsBase {

async onGenImgCmd (prompt: string | undefined, ctx: OnMessageContext | OnCallBackQueryData): Promise<void> {
try {
if (ctx.session.openAi.imageGen.isEnabled && ctx.chat?.id) {
const session = this.getSession(ctx)
if (session.isEnabled && ctx.chat?.id) {
ctx.chatAction = 'upload_photo'
// eslint-disable-next-line @typescript-eslint/naming-convention
const { message_id } = await ctx.reply(
'Generating image via OpenAI\'s DALL·E 3...', { message_thread_id: ctx.message?.message_thread_id }
)
const numImages = ctx.session.openAi.imageGen.numImages
const imgSize = ctx.session.openAi.imageGen.imgSize
const numImages = session.numImages
const imgSize = session.imgSize
const imgs = await postGenerateImg(prompt ?? '', numImages, imgSize)
if (imgs.length > 0) {
await Promise.all(imgs.map(async (img: any) => {
if (ctx.session.openAi.imageGen.isInscriptionLotteryEnabled) {
const inlineKeyboard = new InlineKeyboard().text('Share to enter lottery', `share-payload|${ctx.session.openAi.imageGen.imageGenerated.length}`) // ${imgs[0].url}
if (session.isInscriptionLotteryEnabled) {
const inlineKeyboard = new InlineKeyboard().text('Share to enter lottery', `share-payload|${session.imageGenerated.length}`) // ${imgs[0].url}
const msgExtras = getMessageExtras({
caption: `/dalle ${prompt}\n\n Check [q.country](https://q.country) for general lottery information`,
reply_markup: inlineKeyboard,
Expand All @@ -308,7 +306,7 @@ export class DalleBot extends LlmsBase {
const msg = await ctx.replyWithPhoto(img.url, msgExtras)
const genImg = msg.photo
const fileId = genImg?.pop()?.file_id
ctx.session.openAi.imageGen.imageGenerated.push({ prompt, photoUrl: img.url, photoId: fileId })
session.imageGenerated.push({ prompt, photoUrl: img.url, photoId: fileId })
} else {
const msgExtras = getMessageExtras({ caption: `/dalle ${prompt}` })
await ctx.replyWithPhoto(img.url, msgExtras)
Expand Down Expand Up @@ -340,7 +338,8 @@ export class DalleBot extends LlmsBase {
prompt: string | undefined,
ctx: OnMessageContext | OnCallBackQueryData): Promise<void> => {
try {
if (ctx.session.openAi.imageGen.isEnabled) {
const session = this.getSession(ctx)
if (session.isEnabled) {
// let filePath = ''
let imgList = []
if (photo) {
Expand Down
7 changes: 4 additions & 3 deletions src/modules/llms/llmsBase.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ import {
RequestState,
type BotSessionData,
type LlmsSessionData,
type SubagentResult
type SubagentResult,
type ImageGenSessionData
} from '../types'
import { appText } from '../../utils/text'
import { chatService } from '../../database/services'
Expand Down Expand Up @@ -84,8 +85,8 @@ export abstract class LlmsBase implements PayableBot {
this.subagents = subagents
}

protected getSession (ctx: OnMessageContext | OnCallBackQueryData): LlmsSessionData {
return (ctx.session[this.sessionDataKey as keyof BotSessionData] as LlmsSessionData)
protected getSession (ctx: OnMessageContext | OnCallBackQueryData): LlmsSessionData & ImageGenSessionData {
return (ctx.session[this.sessionDataKey as keyof BotSessionData] as LlmsSessionData & ImageGenSessionData)
}

protected async runSubagents (ctx: OnMessageContext | OnCallBackQueryData, msg: ChatConversation): Promise<void> {
Expand Down
8 changes: 4 additions & 4 deletions src/modules/llms/menu/openaiMenu.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,18 +51,18 @@ const chatGPTimageDefaultOptions = new Menu<BotContext>(MenuIds.CHAT_GPT_MODEL)
function getLabel (m: string, ctx: any): string {
let label = m
console.log(
ctx.session.openAi.chatGpt.model,
ctx.session.chatGpt.model,
m,
ctx.session.openAi.chatGpt.model === m
ctx.session.chatGpt.model === m
)
if (ctx.session.openAi.chatGpt.model === m) {
if (ctx.session.chatGpt.model === m) {
label += ' ✅'
}
return label
}

function setModel (m: string, ctx: any): void {
ctx.session.openAi.chatGpt.model = m
ctx.session.chatGpt.model = m
ctx.menu.back()
}

Expand Down
2 changes: 1 addition & 1 deletion src/modules/llms/openaiBot.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ export class OpenAIBot extends LlmsBase {
return
}

if (ctx.chat?.type === 'private' || ctx.session.openAi.chatGpt.isFreePromptChatGroups) {
if (ctx.chat?.type === 'private' || session.isFreePromptChatGroups) {
await this.onChat(ctx, LlmsModelsEnum.GPT_4, true)
return
}
Expand Down
Loading

0 comments on commit ce88b3a

Please sign in to comment.