Skip to content

Commit

Permalink
Merge pull request #371 from harmony-one/message-sliding-window
Browse files Browse the repository at this point in the history
add sliding window logic + add timestamp field
  • Loading branch information
fegloff authored Nov 4, 2024
2 parents 96eb425 + 36a20e7 commit 6b075d9
Show file tree
Hide file tree
Showing 12 changed files with 224 additions and 28 deletions.
13 changes: 13 additions & 0 deletions src/bot.ts
Original file line number Diff line number Diff line change
Expand Up @@ -509,6 +509,19 @@ const logErrorHandler = (ex: any): void => {
logger.error(ex)
}

// bot.command('testcleanup', async (ctx) => {
// await openAiBot.testCleanup(ctx as OnMessageContext)
// })

bot.command('new', async (ctx) => {
writeCommandLog(ctx as OnMessageContext).catch(logErrorHandler)
await openAiBot.onStop(ctx as OnMessageContext)
return await ctx.reply('Chat history reseted', {
parse_mode: 'Markdown',
message_thread_id: ctx.message?.message_thread_id
})
})

bot.command('more', async (ctx) => {
writeCommandLog(ctx as OnMessageContext).catch(logErrorHandler)
return await ctx.reply(commandsHelpText.more, {
Expand Down
7 changes: 5 additions & 2 deletions src/helpers.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import config from './config'
import { conversationManager } from './modules/llms/utils/conversationManager'
import { LlmModelsEnum } from './modules/llms/utils/llmModelsManager'
import { type DalleImageSize } from './modules/llms/utils/types'
import { type BotSessionData } from './modules/types'
Expand Down Expand Up @@ -26,7 +27,8 @@ export function createInitialSessionData (): BotSessionData {
price: 0,
usage: 0,
isProcessingQueue: false,
requestQueue: []
requestQueue: [],
cleanupState: conversationManager.initializeCleanupTimes()
},
chatGpt: {
model: config.llms.model,
Expand All @@ -36,7 +38,8 @@ export function createInitialSessionData (): BotSessionData {
price: 0,
usage: 0,
isProcessingQueue: false,
requestQueue: []
requestQueue: [],
cleanupState: conversationManager.initializeCleanupTimes()
},
dalle: {
numImages: config.openAi.dalle.sessionDefault.numImages,
Expand Down
15 changes: 10 additions & 5 deletions src/modules/llms/api/athropic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,8 @@ export const anthropicCompletion = async (
completion: {
content: completion[0].text,
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: totalOutputTokens + totalInputTokens,
price: 0,
Expand Down Expand Up @@ -92,7 +93,8 @@ export const xaiCompletion = async (
completion: {
content: completion[0].text,
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: totalOutputTokens + totalInputTokens,
price: 0,
Expand Down Expand Up @@ -202,7 +204,8 @@ export const anthropicStreamCompletion = async (
completion: {
content: completion,
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: parseInt(totalOutputTokens, 10) + parseInt(totalInputTokens, 10),
price: 0,
Expand Down Expand Up @@ -252,7 +255,8 @@ export const toolsChatCompletion = async (
completion: {
content: completion[0].text,
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: totalOutputTokens + totalInputTokens,
price: 0,
Expand All @@ -264,7 +268,8 @@ export const toolsChatCompletion = async (
completion: {
content: 'Timeout error',
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: 0,
price: 0
Expand Down
7 changes: 4 additions & 3 deletions src/modules/llms/api/llmApi.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import axios from 'axios'
import config from '../../../config'
import { type ChatConversation } from '../../types'
import { type ChatConversationWithoutTimestamp, type ChatConversation } from '../../types'
import pino from 'pino'
import { type ChatModel } from '../utils/types'
import { headers } from './helper'
Expand Down Expand Up @@ -36,7 +36,7 @@ interface LlmAddUrlDocument {
interface QueryUrlDocument {
collectioName: string
prompt: string
conversation?: ChatConversation[]
conversation?: ChatConversationWithoutTimestamp[]
}

export const getChatModel = (modelName: string): ChatModel | undefined => {
Expand Down Expand Up @@ -130,7 +130,8 @@ export const llmCompletion = async (
completion: {
content: completion[0].message?.content,
role: 'system',
model
model,
timestamp: Date.now()
},
usage: totalOutputTokens + totalInputTokens,
price: 0
Expand Down
15 changes: 10 additions & 5 deletions src/modules/llms/api/openai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,9 @@ export async function alterGeneratedImg (
}
}

const prepareConversation = (conversation: ChatConversation[], model: string): ChatConversation[] => {
type ConversationOutput = Omit<ChatConversation, 'timestamp' | 'model' | 'id' | 'author' | 'numSubAgents'>

const prepareConversation = (conversation: ChatConversation[], model: string): ConversationOutput[] => {
const messages = conversation.filter(c => c.model === model).map(m => { return { content: m.content, role: m.role } })
if (messages.length !== 1 || model === LlmModelsEnum.O1) {
return messages
Expand Down Expand Up @@ -125,7 +127,8 @@ export async function chatCompletion (
return {
completion: {
content: response.choices[0].message?.content ?? 'Error - no completion available',
role: 'assistant'
role: 'assistant',
timestamp: Date.now()
},
usage: response.usage?.total_tokens, // 2010
price: price * config.openAi.chatGpt.priceAdjustment,
Expand Down Expand Up @@ -215,7 +218,8 @@ export const streamChatCompletion = async (
return {
completion: {
content: completion,
role: 'assistant'
role: 'assistant',
timestamp: Date.now()
},
usage: outputTokens + inputTokens,
price: 0,
Expand Down Expand Up @@ -308,7 +312,8 @@ export const streamChatVisionCompletion = async (
return {
completion: {
content: completion,
role: 'assistant'
role: 'assistant',
timestamp: Date.now()
},
usage: outputTokens + inputTokens,
price: 0,
Expand All @@ -319,7 +324,7 @@ export const streamChatVisionCompletion = async (

export async function improvePrompt (promptText: string, model: string): Promise<string> {
const prompt = `Improve this picture description using max 100 words and don't add additional text to the image: ${promptText} `
const conversation = [{ role: 'user', content: prompt }]
const conversation = [{ role: 'user', content: prompt, timestamp: Date.now() }]
const response = await chatCompletion(conversation, model)
return response.completion?.content as string ?? ''
}
Expand Down
3 changes: 2 additions & 1 deletion src/modules/llms/api/pdfHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ export const handlePdf = async (prompt: string): Promise<PdfCompletion> => {
return {
completion: {
content: response.data.response,
role: 'system'
role: 'system',
timestamp: Date.now()
},
prompt,
price: response.data.cost
Expand Down
10 changes: 6 additions & 4 deletions src/modules/llms/api/vertex.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import axios, { type AxiosResponse } from 'axios'
import config from '../../../config'
import { type OnMessageContext, type ChatConversation, type OnCallBackQueryData } from '../../types'
import { type OnMessageContext, type ChatConversation, type OnCallBackQueryData, type ChatConversationWithoutTimestamp } from '../../types'
import { type LlmCompletion } from './llmApi'
import { type Readable } from 'stream'
import { GrammyError } from 'grammy'
Expand Down Expand Up @@ -29,7 +29,7 @@ export const vertexCompletion = async (
stream: false,
messages: conversation.filter(c => c.model === model)
.map((msg) => {
const msgFiltered: ChatConversation = { content: msg.content, model: msg.model }
const msgFiltered: ChatConversationWithoutTimestamp = { content: msg.content, model: msg.model }
if (model === LlmModelsEnum.CHAT_BISON) {
msgFiltered.author = msg.role
} else {
Expand All @@ -48,7 +48,8 @@ export const vertexCompletion = async (
completion: {
content: response.data._prediction_response[0][0].candidates[0].content,
role: 'bot', // role replace to author attribute will be done later
model
model,
timestamp: Date.now()
},
usage: totalOutputTokens + totalInputTokens,
price: 0
Expand Down Expand Up @@ -145,7 +146,8 @@ export const vertexStreamCompletion = async (
completion: {
content: completion,
role: 'assistant',
model
model,
timestamp: Date.now()
},
usage: parseInt(totalOutputTokens, 10) + parseInt(totalInputTokens, 10),
price: 0,
Expand Down
60 changes: 55 additions & 5 deletions src/modules/llms/llmsBase.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import {
type LLMModelsManager,
type ModelVersion
} from './utils/llmModelsManager'
import { conversationManager } from './utils/conversationManager'

export abstract class LlmsBase implements PayableBot {
public module: string
Expand Down Expand Up @@ -205,7 +206,8 @@ export abstract class LlmsBase implements PayableBot {
id: ctx.message?.message_id,
model,
content: await preparePrompt(ctx, prompt as string),
numSubAgents: 0
numSubAgents: 0,
timestamp: Date.now()
})
if (!session.isProcessingQueue) {
session.isProcessingQueue = true
Expand All @@ -218,7 +220,8 @@ export abstract class LlmsBase implements PayableBot {
id: ctx.message?.message_id ?? ctx.message?.message_thread_id ?? 0,
model,
content: prompt as string ?? '', // await preparePrompt(ctx, prompt as string),
numSubAgents: supportedAgents
numSubAgents: supportedAgents,
timestamp: Date.now()
}
await this.runSubagents(ctx, msg, stream, usesTools) // prompt as string)
}
Expand All @@ -230,6 +233,7 @@ export abstract class LlmsBase implements PayableBot {

async onChatRequestHandler (ctx: OnMessageContext | OnCallBackQueryData, stream: boolean, usesTools: boolean): Promise<void> {
const session = this.getSession(ctx)
session.chatConversation = conversationManager.manageConversationWindow(session.chatConversation, ctx, this.sessionDataKey)
while (session.requestQueue.length > 0) {
try {
const msg = session.requestQueue.shift()
Expand Down Expand Up @@ -272,7 +276,8 @@ export abstract class LlmsBase implements PayableBot {
const chat: ChatConversation = {
content: enhancedPrompt || prompt,
role: 'user',
model: modelVersion
model: modelVersion,
timestamp: Date.now()
}
chatConversation.push(chat)
const payload = {
Expand Down Expand Up @@ -358,7 +363,8 @@ export abstract class LlmsBase implements PayableBot {
conversation.push({
role: 'assistant',
content: completion.completion?.content ?? '',
model
model,
timestamp: Date.now()
})
return {
price: price.price,
Expand All @@ -371,7 +377,8 @@ export abstract class LlmsBase implements PayableBot {
conversation.push({
role: 'assistant',
content: response.completion?.content ?? '',
model
model,
timestamp: Date.now()
})
return {
price: response.price,
Expand Down Expand Up @@ -470,6 +477,49 @@ export abstract class LlmsBase implements PayableBot {
session.price = 0
}

async testCleanup (ctx: OnMessageContext | OnCallBackQueryData): Promise<void> {
const session = this.getSession(ctx)
// Force cleanup times for testing
const now = new Date()
const forcedCleanupTime = new Date(now)
forcedCleanupTime.setHours(2, 59, 0, 0) // Set to 2:59 AM
session.cleanupState = {
nextCleanupTime: forcedCleanupTime.getTime() + (60 * 1000), // 3 AM
lastCleanupTime: forcedCleanupTime.getTime() - (24 * 60 * 60 * 1000) // Yesterday 2:59 AM
}
console.log('Testing cleanup with forced times:', {
nextCleanup: new Date(session.cleanupState.nextCleanupTime).toLocaleString(),
lastCleanup: new Date(session.cleanupState.lastCleanupTime).toLocaleString(),
currentTime: now.toLocaleString()
})
// Add some test messages with various timestamps
if (session.chatConversation.length === 0) {
const yesterday = new Date(now)
yesterday.setDate(yesterday.getDate() - 1)
session.chatConversation = [
{
role: 'user',
content: 'Message from 2 days ago',
model: 'test',
timestamp: yesterday.getTime() - (24 * 60 * 60 * 1000)
},
{
role: 'assistant',
content: 'Message from yesterday',
model: 'test',
timestamp: yesterday.getTime()
},
{
role: 'user',
content: 'Message from today',
model: 'test',
timestamp: now.getTime()
}
]
}
await this.onChatRequestHandler(ctx, false, false)
}

async onError (
ctx: OnMessageContext | OnCallBackQueryData,
e: any,
Expand Down
Loading

0 comments on commit 6b075d9

Please sign in to comment.