diff --git a/backend/src/bot/handlers/message/handler.py b/backend/src/bot/handlers/message/handler.py index 902193c..68d7634 100644 --- a/backend/src/bot/handlers/message/handler.py +++ b/backend/src/bot/handlers/message/handler.py @@ -3,10 +3,17 @@ import base64 import contextlib import io import time +from collections.abc import Awaitable, Callable +from typing import Any -from aiogram import Bot, F, Router, html, types +from aiogram import BaseMiddleware, Bot, F, Router, html, types from aiogram.enums import ChatAction -from aiogram.types import KeyboardButton, ReplyKeyboardMarkup, ReplyKeyboardRemove +from aiogram.types import ( + KeyboardButton, + ReplyKeyboardMarkup, + ReplyKeyboardRemove, + TelegramObject, +) from convex import ConvexInt64 from bot.modules.ai import ( @@ -23,6 +30,45 @@ from utils.convex import ConvexClient router = Router() convex = ConvexClient(env.convex_url) +ALBUM_COLLECT_DELAY = 0.5 + + +class AlbumMiddleware(BaseMiddleware): + def __init__(self) -> None: + self.albums: dict[str, list[types.Message]] = {} + self.scheduled: set[str] = set() + + async def __call__( + self, + handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]], + event: TelegramObject, + data: dict[str, Any], + ) -> Any: # noqa: ANN401 + if not isinstance(event, types.Message) or not event.media_group_id: + return await handler(event, data) + + album_id = event.media_group_id + if album_id not in self.albums: + self.albums[album_id] = [] + self.albums[album_id].append(event) + + if album_id in self.scheduled: + return None + + self.scheduled.add(album_id) + await asyncio.sleep(ALBUM_COLLECT_DELAY) + + messages = self.albums.pop(album_id, []) + self.scheduled.discard(album_id) + + if messages: + data["album"] = messages + return await handler(messages[0], data) + return None + + +router.message.middleware(AlbumMiddleware()) + EDIT_THROTTLE_SECONDS = 1.0 TELEGRAM_MAX_LENGTH = 4096 @@ -398,6 +444,74 @@ async def on_text_message(message: types.Message, bot: Bot) -> None: await process_message(message.from_user.id, message.text, bot, message.chat.id) +@router.message(F.media_group_id, F.photo) +async def on_album_message( + message: types.Message, bot: Bot, album: list[types.Message] +) -> None: + if not message.from_user: + return + + await convex.mutation( + "users:getOrCreate", + { + "telegramId": ConvexInt64(message.from_user.id), + "telegramChatId": ConvexInt64(message.chat.id), + }, + ) + + user = await convex.query( + "users:getByTelegramId", {"telegramId": ConvexInt64(message.from_user.id)} + ) + + if not user or not user.get("activeChatId"): + await message.answer("Use /new first to create a chat.") + return + + caption = message.caption or "Process the images according to your task" + + images_base64: list[str] = [] + images_media_types: list[str] = [] + + for msg in album: + if not msg.photo: + continue + + photo = msg.photo[-1] + file = await bot.get_file(photo.file_id) + if not file.file_path: + continue + + buffer = io.BytesIO() + await bot.download_file(file.file_path, buffer) + image_bytes = buffer.getvalue() + images_base64.append(base64.b64encode(image_bytes).decode()) + + ext = file.file_path.rsplit(".", 1)[-1].lower() + media_type = f"image/{ext}" if ext in ("png", "gif", "webp") else "image/jpeg" + images_media_types.append(media_type) + + if not images_base64: + await message.answer("Failed to get photos.") + return + + active_chat_id = user["activeChatId"] + await convex.mutation( + "messages:create", + { + "chatId": active_chat_id, + "role": "user", + "content": caption, + "source": "telegram", + "imagesBase64": images_base64, + "imagesMediaTypes": images_media_types, + }, + ) + + await process_message( + message.from_user.id, caption, bot, message.chat.id, skip_user_message=True + ) + + @router.message(F.photo) async def on_photo_message(message: types.Message, bot: Bot) -> None: if not message.from_user or not message.photo: diff --git a/frontend/src/lib/convex/_generated/api.d.ts b/frontend/src/lib/convex/_generated/api.d.ts index 1f51849..391ba27 100644 --- a/frontend/src/lib/convex/_generated/api.d.ts +++ b/frontend/src/lib/convex/_generated/api.d.ts @@ -8,22 +8,18 @@ * @module */ -import type * as chats from "../chats.js"; -import type * as messages from "../messages.js"; -import type * as pendingGenerations from "../pendingGenerations.js"; -import type * as users from "../users.js"; +import type * as chats from '../chats.js'; +import type * as messages from '../messages.js'; +import type * as pendingGenerations from '../pendingGenerations.js'; +import type * as users from '../users.js'; -import type { - ApiFromModules, - FilterApi, - FunctionReference, -} from "convex/server"; +import type { ApiFromModules, FilterApi, FunctionReference } from 'convex/server'; declare const fullApi: ApiFromModules<{ - chats: typeof chats; - messages: typeof messages; - pendingGenerations: typeof pendingGenerations; - users: typeof users; + chats: typeof chats; + messages: typeof messages; + pendingGenerations: typeof pendingGenerations; + users: typeof users; }>; /** @@ -34,10 +30,7 @@ declare const fullApi: ApiFromModules<{ * const myFunctionReference = api.myModule.myFunction; * ``` */ -export declare const api: FilterApi< - typeof fullApi, - FunctionReference ->; +export declare const api: FilterApi>; /** * A utility for referencing Convex functions in your app's internal API. @@ -47,9 +40,6 @@ export declare const api: FilterApi< * const myFunctionReference = internal.myModule.myFunction; * ``` */ -export declare const internal: FilterApi< - typeof fullApi, - FunctionReference ->; +export declare const internal: FilterApi>; export declare const components: {}; diff --git a/frontend/src/lib/convex/_generated/api.js b/frontend/src/lib/convex/_generated/api.js index 44bf985..24593c7 100644 --- a/frontend/src/lib/convex/_generated/api.js +++ b/frontend/src/lib/convex/_generated/api.js @@ -8,7 +8,7 @@ * @module */ -import { anyApi, componentsGeneric } from "convex/server"; +import { anyApi, componentsGeneric } from 'convex/server'; /** * A utility for referencing Convex functions in your app's API. diff --git a/frontend/src/lib/convex/_generated/dataModel.d.ts b/frontend/src/lib/convex/_generated/dataModel.d.ts index f97fd19..5428df6 100644 --- a/frontend/src/lib/convex/_generated/dataModel.d.ts +++ b/frontend/src/lib/convex/_generated/dataModel.d.ts @@ -9,13 +9,13 @@ */ import type { - DataModelFromSchemaDefinition, - DocumentByName, - TableNamesInDataModel, - SystemTableNames, -} from "convex/server"; -import type { GenericId } from "convex/values"; -import schema from "../schema.js"; + DataModelFromSchemaDefinition, + DocumentByName, + TableNamesInDataModel, + SystemTableNames +} from 'convex/server'; +import type { GenericId } from 'convex/values'; +import schema from '../schema.js'; /** * The names of all of your Convex tables. @@ -27,10 +27,7 @@ export type TableNames = TableNamesInDataModel; * * @typeParam TableName - A string literal type of the table name (like "users"). */ -export type Doc = DocumentByName< - DataModel, - TableName ->; +export type Doc = DocumentByName; /** * An identifier for a document in Convex. @@ -45,8 +42,7 @@ export type Doc = DocumentByName< * * @typeParam TableName - A string literal type of the table name (like "users"). */ -export type Id = - GenericId; +export type Id = GenericId; /** * A type describing your Convex data model. diff --git a/frontend/src/lib/convex/_generated/server.d.ts b/frontend/src/lib/convex/_generated/server.d.ts index bec05e6..1cc047e 100644 --- a/frontend/src/lib/convex/_generated/server.d.ts +++ b/frontend/src/lib/convex/_generated/server.d.ts @@ -9,17 +9,17 @@ */ import { - ActionBuilder, - HttpActionBuilder, - MutationBuilder, - QueryBuilder, - GenericActionCtx, - GenericMutationCtx, - GenericQueryCtx, - GenericDatabaseReader, - GenericDatabaseWriter, -} from "convex/server"; -import type { DataModel } from "./dataModel.js"; + ActionBuilder, + HttpActionBuilder, + MutationBuilder, + QueryBuilder, + GenericActionCtx, + GenericMutationCtx, + GenericQueryCtx, + GenericDatabaseReader, + GenericDatabaseWriter +} from 'convex/server'; +import type { DataModel } from './dataModel.js'; /** * Define a query in this Convex app's public API. @@ -29,7 +29,7 @@ import type { DataModel } from "./dataModel.js"; * @param func - The query function. It receives a {@link QueryCtx} as its first argument. * @returns The wrapped query. Include this as an `export` to name it and make it accessible. */ -export declare const query: QueryBuilder; +export declare const query: QueryBuilder; /** * Define a query that is only accessible from other Convex functions (but not from the client). @@ -39,7 +39,7 @@ export declare const query: QueryBuilder; * @param func - The query function. It receives a {@link QueryCtx} as its first argument. * @returns The wrapped query. Include this as an `export` to name it and make it accessible. */ -export declare const internalQuery: QueryBuilder; +export declare const internalQuery: QueryBuilder; /** * Define a mutation in this Convex app's public API. @@ -49,7 +49,7 @@ export declare const internalQuery: QueryBuilder; * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument. * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible. */ -export declare const mutation: MutationBuilder; +export declare const mutation: MutationBuilder; /** * Define a mutation that is only accessible from other Convex functions (but not from the client). @@ -59,7 +59,7 @@ export declare const mutation: MutationBuilder; * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument. * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible. */ -export declare const internalMutation: MutationBuilder; +export declare const internalMutation: MutationBuilder; /** * Define an action in this Convex app's public API. @@ -72,7 +72,7 @@ export declare const internalMutation: MutationBuilder; * @param func - The action. It receives an {@link ActionCtx} as its first argument. * @returns The wrapped action. Include this as an `export` to name it and make it accessible. */ -export declare const action: ActionBuilder; +export declare const action: ActionBuilder; /** * Define an action that is only accessible from other Convex functions (but not from the client). @@ -80,7 +80,7 @@ export declare const action: ActionBuilder; * @param func - The function. It receives an {@link ActionCtx} as its first argument. * @returns The wrapped function. Include this as an `export` to name it and make it accessible. */ -export declare const internalAction: ActionBuilder; +export declare const internalAction: ActionBuilder; /** * Define an HTTP action. diff --git a/frontend/src/lib/convex/_generated/server.js b/frontend/src/lib/convex/_generated/server.js index bf3d25a..a18aa28 100644 --- a/frontend/src/lib/convex/_generated/server.js +++ b/frontend/src/lib/convex/_generated/server.js @@ -9,14 +9,14 @@ */ import { - actionGeneric, - httpActionGeneric, - queryGeneric, - mutationGeneric, - internalActionGeneric, - internalMutationGeneric, - internalQueryGeneric, -} from "convex/server"; + actionGeneric, + httpActionGeneric, + queryGeneric, + mutationGeneric, + internalActionGeneric, + internalMutationGeneric, + internalQueryGeneric +} from 'convex/server'; /** * Define a query in this Convex app's public API. diff --git a/frontend/src/lib/convex/messages.ts b/frontend/src/lib/convex/messages.ts index 6ab3ac9..6f486cd 100644 --- a/frontend/src/lib/convex/messages.ts +++ b/frontend/src/lib/convex/messages.ts @@ -35,6 +35,8 @@ export const create = mutation({ source: v.union(v.literal('telegram'), v.literal('web')), imageBase64: v.optional(v.string()), imageMediaType: v.optional(v.string()), + imagesBase64: v.optional(v.array(v.string())), + imagesMediaTypes: v.optional(v.array(v.string())), followUpOptions: v.optional(v.array(v.string())), isStreaming: v.optional(v.boolean()) }, @@ -47,6 +49,8 @@ export const create = mutation({ source: args.source, imageBase64: args.imageBase64, imageMediaType: args.imageMediaType, + imagesBase64: args.imagesBase64, + imagesMediaTypes: args.imagesMediaTypes, followUpOptions: args.followUpOptions, createdAt: Date.now(), isStreaming: args.isStreaming @@ -166,11 +170,24 @@ export const getChatImages = query({ .withIndex('by_chat_id', (q) => q.eq('chatId', args.chatId)) .collect(); - return messages - .filter((m) => m.imageBase64 && m.imageMediaType) - .map((m) => ({ - base64: m.imageBase64!, - mediaType: m.imageMediaType! - })); + const images: Array<{ base64: string; mediaType: string }> = []; + + for (const m of messages) { + if (m.imagesBase64 && m.imagesMediaTypes) { + for (let i = 0; i < m.imagesBase64.length; i++) { + images.push({ + base64: m.imagesBase64[i], + mediaType: m.imagesMediaTypes[i] + }); + } + } else if (m.imageBase64 && m.imageMediaType) { + images.push({ + base64: m.imageBase64, + mediaType: m.imageMediaType + }); + } + } + + return images; } }); diff --git a/frontend/src/lib/convex/schema.ts b/frontend/src/lib/convex/schema.ts index 64c9114..a199fc7 100644 --- a/frontend/src/lib/convex/schema.ts +++ b/frontend/src/lib/convex/schema.ts @@ -26,6 +26,8 @@ export default defineSchema({ imageBase64: v.optional(v.string()), imageMediaType: v.optional(v.string()), imageStorageId: v.optional(v.id('_storage')), + imagesBase64: v.optional(v.array(v.string())), + imagesMediaTypes: v.optional(v.array(v.string())), followUpOptions: v.optional(v.array(v.string())), source: v.union(v.literal('telegram'), v.literal('web')), createdAt: v.number(), diff --git a/frontend/src/routes/[mnemonic]/+page.svelte b/frontend/src/routes/[mnemonic]/+page.svelte index 9724624..2cdbd06 100644 --- a/frontend/src/routes/[mnemonic]/+page.svelte +++ b/frontend/src/routes/[mnemonic]/+page.svelte @@ -180,7 +180,7 @@ {#if showScrollButton}