feat(*): add multiple image support

This commit is contained in:
h
2026-01-21 21:32:46 +01:00
parent 87aa974084
commit cb188f7cc0
9 changed files with 188 additions and 69 deletions

View File

@@ -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:

View File

@@ -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<any, "public">
>;
export declare const api: FilterApi<typeof fullApi, FunctionReference<any, 'public'>>;
/**
* 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<any, "internal">
>;
export declare const internal: FilterApi<typeof fullApi, FunctionReference<any, 'internal'>>;
export declare const components: {};

View File

@@ -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.

View File

@@ -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<DataModel>;
*
* @typeParam TableName - A string literal type of the table name (like "users").
*/
export type Doc<TableName extends TableNames> = DocumentByName<
DataModel,
TableName
>;
export type Doc<TableName extends TableNames> = DocumentByName<DataModel, TableName>;
/**
* An identifier for a document in Convex.
@@ -45,8 +42,7 @@ export type Doc<TableName extends TableNames> = DocumentByName<
*
* @typeParam TableName - A string literal type of the table name (like "users").
*/
export type Id<TableName extends TableNames | SystemTableNames> =
GenericId<TableName>;
export type Id<TableName extends TableNames | SystemTableNames> = GenericId<TableName>;
/**
* A type describing your Convex data model.

View File

@@ -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<DataModel, "public">;
export declare const query: QueryBuilder<DataModel, 'public'>;
/**
* 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<DataModel, "public">;
* @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<DataModel, "internal">;
export declare const internalQuery: QueryBuilder<DataModel, 'internal'>;
/**
* Define a mutation in this Convex app's public API.
@@ -49,7 +49,7 @@ export declare const internalQuery: QueryBuilder<DataModel, "internal">;
* @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<DataModel, "public">;
export declare const mutation: MutationBuilder<DataModel, 'public'>;
/**
* 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<DataModel, "public">;
* @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<DataModel, "internal">;
export declare const internalMutation: MutationBuilder<DataModel, 'internal'>;
/**
* Define an action in this Convex app's public API.
@@ -72,7 +72,7 @@ export declare const internalMutation: MutationBuilder<DataModel, "internal">;
* @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<DataModel, "public">;
export declare const action: ActionBuilder<DataModel, 'public'>;
/**
* 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<DataModel, "public">;
* @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<DataModel, "internal">;
export declare const internalAction: ActionBuilder<DataModel, 'internal'>;
/**
* Define an HTTP action.

View File

@@ -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.

View File

@@ -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;
}
});

View File

@@ -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(),

View File

@@ -180,7 +180,7 @@
{#if showScrollButton}
<button
onclick={scrollToLastMessage}
class="fixed bottom-12 right-3 z-50 flex h-8 w-8 items-center justify-center rounded-full bg-blue-600 text-white shadow-lg animate-pulse"
class="fixed right-3 bottom-12 z-50 flex h-8 w-8 animate-pulse items-center justify-center rounded-full bg-blue-600 text-white shadow-lg"
>
</button>