feat: web UI chat render, panels, presence + analytics
This commit is contained in:
@@ -11,12 +11,16 @@ from starlette.applications import Starlette
|
||||
|
||||
from api.auth import BearerAuthMiddleware
|
||||
from api.mcp.server import mcp
|
||||
from api.realtime import hub
|
||||
from api.routers import (
|
||||
accounts,
|
||||
analytics,
|
||||
annotations,
|
||||
avatars,
|
||||
backfill,
|
||||
chats,
|
||||
custom_emoji,
|
||||
events,
|
||||
folders,
|
||||
media,
|
||||
peers,
|
||||
@@ -39,7 +43,10 @@ mcp_app = mcp.http_app(path="/")
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
|
||||
pool = await container.get(asyncpg.Pool)
|
||||
await hub.start(pool)
|
||||
yield
|
||||
await hub.stop()
|
||||
await app_.state.dishka_container.close()
|
||||
|
||||
|
||||
@@ -59,6 +66,7 @@ async def health(pool: FromDishka[asyncpg.Pool]) -> dict[str, bool]:
|
||||
|
||||
|
||||
app.include_router(accounts.router)
|
||||
app.include_router(analytics.router)
|
||||
app.include_router(policy.router)
|
||||
app.include_router(folders.router)
|
||||
app.include_router(backfill.router)
|
||||
@@ -66,8 +74,10 @@ app.include_router(search.router)
|
||||
app.include_router(chats.router)
|
||||
app.include_router(media.router)
|
||||
app.include_router(avatars.router)
|
||||
app.include_router(custom_emoji.router)
|
||||
app.include_router(social.router)
|
||||
app.include_router(presence.router)
|
||||
app.include_router(events.router)
|
||||
app.include_router(peers.router)
|
||||
app.include_router(annotations.router)
|
||||
app.include_router(watches.router)
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.env import env
|
||||
from utils.events import BG_EVENTS_CHANNEL
|
||||
from utils.read import chats as chats_read
|
||||
from utils.read import presence as presence_read
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
QUEUE_MAXSIZE = 256
|
||||
|
||||
|
||||
class Subscriber:
|
||||
def __init__(self, account_id: int, chat_id: int | None) -> None:
|
||||
self.account_id = account_id
|
||||
self.chat_id = chat_id
|
||||
self.queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
|
||||
|
||||
|
||||
class EventHub:
|
||||
def __init__(self) -> None:
|
||||
self._subscribers: set[Subscriber] = set()
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
self._conn: asyncpg.Connection | None = None
|
||||
self._tasks: set[asyncio.Task] = set()
|
||||
|
||||
def subscribe(self, account_id: int, chat_id: int | None) -> Subscriber:
|
||||
sub = Subscriber(account_id, chat_id)
|
||||
self._subscribers.add(sub)
|
||||
return sub
|
||||
|
||||
def unsubscribe(self, sub: Subscriber) -> None:
|
||||
self._subscribers.discard(sub)
|
||||
|
||||
async def start(self, pool: asyncpg.Pool) -> None:
|
||||
self._pool = pool
|
||||
conn = await asyncpg.connect(dsn=env.db.connection_url)
|
||||
await conn.add_listener(BG_EVENTS_CHANNEL, self._on_notify)
|
||||
self._conn = conn
|
||||
logger.info("Realtime hub listening on %s", BG_EVENTS_CHANNEL)
|
||||
|
||||
async def stop(self) -> None:
|
||||
for task in self._tasks:
|
||||
task.cancel()
|
||||
if self._conn is not None:
|
||||
await self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
def _on_notify(
|
||||
self, _conn: asyncpg.Connection, _pid: int, _channel: str, payload: str
|
||||
) -> None:
|
||||
task = asyncio.create_task(self._dispatch(payload))
|
||||
self._tasks.add(task)
|
||||
task.add_done_callback(self._tasks.discard)
|
||||
|
||||
async def _dispatch(self, payload: str) -> None:
|
||||
try:
|
||||
event = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
return
|
||||
account_id = event.get("account_id")
|
||||
chat_id = event.get("chat_id")
|
||||
targets = [
|
||||
sub
|
||||
for sub in self._subscribers
|
||||
if sub.account_id == account_id
|
||||
and (sub.chat_id is None or sub.chat_id == chat_id)
|
||||
]
|
||||
if not targets:
|
||||
return
|
||||
frame = await self._build_frame(event)
|
||||
if frame is None:
|
||||
return
|
||||
for sub in targets:
|
||||
try:
|
||||
sub.queue.put_nowait(frame)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning("Dropping event for slow subscriber")
|
||||
|
||||
async def _build_frame( # noqa: PLR0911
|
||||
self, event: dict[str, Any]
|
||||
) -> dict[str, Any] | None:
|
||||
if self._pool is None:
|
||||
return None
|
||||
kind = event.get("kind")
|
||||
account_id = event["account_id"]
|
||||
if kind in {"message", "edit", "reaction"}:
|
||||
view = await chats_read.get_message(
|
||||
self._pool, account_id, event["chat_id"], event["message_id"]
|
||||
)
|
||||
if view is None:
|
||||
return None
|
||||
return {"type": kind, "message": view.model_dump(mode="json")}
|
||||
if kind == "delete":
|
||||
return {
|
||||
"type": "delete",
|
||||
"chat_id": event.get("chat_id"),
|
||||
"message_ids": event.get("message_ids", []),
|
||||
}
|
||||
if kind == "presence":
|
||||
sample = await presence_read.current_presence(
|
||||
self._pool, account_id, event["chat_id"]
|
||||
)
|
||||
return {
|
||||
"type": "presence",
|
||||
"peer_id": event["chat_id"],
|
||||
"sample": sample.model_dump(mode="json") if sample else None,
|
||||
}
|
||||
if kind == "receipt":
|
||||
return {
|
||||
"type": "receipt",
|
||||
"chat_id": event["chat_id"],
|
||||
"read_up_to": event["message_id"],
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
hub = EventHub()
|
||||
@@ -0,0 +1,30 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query
|
||||
|
||||
from utils.read import analytics
|
||||
from utils.read.models import ResponseStats, VolumeBucket
|
||||
|
||||
router = APIRouter(prefix="/api/analytics", tags=["analytics"], route_class=DishkaRoute)
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
ChatId = Annotated[int, Query()]
|
||||
|
||||
|
||||
@router.get("/volume")
|
||||
async def volume(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
chat_id: ChatId,
|
||||
days: Annotated[int, Query()] = 90,
|
||||
) -> list[VolumeBucket]:
|
||||
return await analytics.message_volume(pool, account_id, chat_id, days=days)
|
||||
|
||||
|
||||
@router.get("/response-time")
|
||||
async def response_time(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: AccountId, chat_id: ChatId
|
||||
) -> ResponseStats:
|
||||
return await analytics.response_stats(pool, account_id, chat_id)
|
||||
@@ -24,6 +24,10 @@ class FetchMediaRequest(BaseModel):
|
||||
message_id: int
|
||||
|
||||
|
||||
class SyncDialogsRequest(BaseModel):
|
||||
account_id: int
|
||||
|
||||
|
||||
class EnqueueResponse(BaseModel):
|
||||
job_id: int
|
||||
|
||||
@@ -78,6 +82,14 @@ async def enqueue_fetch_media(
|
||||
return EnqueueResponse(job_id=job_id)
|
||||
|
||||
|
||||
@router.post("/dialogs/sync", status_code=201)
|
||||
async def enqueue_sync_dialogs(
|
||||
pool: FromDishka[asyncpg.Pool], body: SyncDialogsRequest
|
||||
) -> EnqueueResponse:
|
||||
job_id = await enqueue(pool, body.account_id, "sync_dialogs", {})
|
||||
return EnqueueResponse(job_id=job_id)
|
||||
|
||||
|
||||
@router.get("/jobs")
|
||||
async def list_jobs(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
|
||||
@@ -13,7 +13,9 @@ from utils.read.models import (
|
||||
MessageVersionView,
|
||||
MessageView,
|
||||
Page,
|
||||
PinnedView,
|
||||
)
|
||||
from utils.read.pinned import get_pinned
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["chats"], route_class=DishkaRoute)
|
||||
|
||||
@@ -55,6 +57,13 @@ async def chat_history(
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}/pinned")
|
||||
async def chat_pinned(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||
) -> PinnedView | None:
|
||||
return await get_pinned(pool, account_id, chat_id)
|
||||
|
||||
|
||||
@router.post("/chats/{chat_id}/enrich")
|
||||
async def enrich_chat(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, body: EnrichRequest
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from utils.jobs import enqueue
|
||||
from utils.read.custom_emoji import current_custom_emoji
|
||||
from utils.storage import ContentAddressedStorage
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/api/custom-emoji", tags=["custom-emoji"], route_class=DishkaRoute
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{custom_emoji_id}")
|
||||
async def serve_custom_emoji(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
custom_emoji_id: int,
|
||||
account_id: Annotated[int, Query()],
|
||||
) -> FileResponse:
|
||||
emoji = await current_custom_emoji(pool, custom_emoji_id)
|
||||
if emoji is None or not emoji.downloaded or emoji.storage_key is None:
|
||||
await enqueue(
|
||||
pool, account_id, "fetch_custom_emoji", {"custom_emoji_id": custom_emoji_id}
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=409, detail="custom emoji not downloaded; fetching"
|
||||
)
|
||||
return FileResponse(
|
||||
storage.url(emoji.storage_key),
|
||||
media_type=emoji.mime or "application/octet-stream",
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from api.realtime import Subscriber, hub
|
||||
|
||||
router = APIRouter(prefix="/api/events", tags=["events"])
|
||||
|
||||
HEARTBEAT_SECONDS = 15
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
ChatId = Annotated[int | None, Query()]
|
||||
|
||||
|
||||
async def _stream(sub: Subscriber) -> AsyncGenerator[str]:
|
||||
try:
|
||||
yield ": connected\n\n"
|
||||
while True:
|
||||
try:
|
||||
frame = await asyncio.wait_for(
|
||||
sub.queue.get(), timeout=HEARTBEAT_SECONDS
|
||||
)
|
||||
except TimeoutError:
|
||||
yield ": keepalive\n\n"
|
||||
continue
|
||||
yield f"event: {frame['type']}\ndata: {json.dumps(frame)}\n\n"
|
||||
finally:
|
||||
hub.unsubscribe(sub)
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def events(account_id: AccountId, chat_id: ChatId = None) -> StreamingResponse:
|
||||
sub = hub.subscribe(account_id, chat_id)
|
||||
return StreamingResponse(
|
||||
_stream(sub),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
@@ -36,6 +36,13 @@ async def presence_history(
|
||||
)
|
||||
|
||||
|
||||
@router.get("/current")
|
||||
async def current_presence(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: AccountId, peer_id: PeerId
|
||||
) -> PresenceSample | None:
|
||||
return await presence.current_presence(pool, account_id, peer_id)
|
||||
|
||||
|
||||
@router.get("/hourly")
|
||||
async def presence_hourly(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
|
||||
@@ -3,6 +3,7 @@ from pyrogram.types import Message
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@PyroClient.on_deleted_messages()
|
||||
@@ -21,8 +22,12 @@ async def on_deleted_messages(client: PyroClient, messages: list[Message]) -> No
|
||||
channels.setdefault(chat_id, []).append(message.id)
|
||||
if box:
|
||||
await repository.mark_deleted_box(ctx.pool, ctx.account_id, box)
|
||||
await notify_bg_event(ctx.pool, "delete", ctx.account_id, message_ids=box)
|
||||
for chat_id, ids in channels.items():
|
||||
await repository.mark_deleted_channel(ctx.pool, ctx.account_id, chat_id, ids)
|
||||
await notify_bg_event(
|
||||
ctx.pool, "delete", ctx.account_id, chat_id=chat_id, message_ids=ids
|
||||
)
|
||||
|
||||
|
||||
handlers = on_deleted_messages.handlers
|
||||
|
||||
@@ -5,6 +5,7 @@ from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.chat_meta import meta_from_chat
|
||||
from userbot.modules.capture.message import sender_id
|
||||
from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@PyroClient.on_edited_message()
|
||||
@@ -32,8 +33,13 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
|
||||
has_media=message.media is not None,
|
||||
is_self_destruct=self_destruct_ttl(message) is not None,
|
||||
)
|
||||
if changed and message.media is not None:
|
||||
if not changed:
|
||||
return
|
||||
if message.media is not None:
|
||||
await capture_media(client, message, ctx, chat_id, message.id, toggles)
|
||||
await notify_bg_event(
|
||||
ctx.pool, "edit", ctx.account_id, chat_id=chat_id, message_id=message.id
|
||||
)
|
||||
|
||||
|
||||
handlers = on_edited_message.handlers
|
||||
|
||||
@@ -5,6 +5,7 @@ from userbot.modules.capture import capture_message
|
||||
from userbot.modules.capture.chat_meta import meta_from_chat
|
||||
from userbot.modules.stt import is_transcribable
|
||||
from userbot.modules.stt.gate import safe_transcribe
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@PyroClient.on_message()
|
||||
@@ -18,6 +19,9 @@ async def on_message(client: PyroClient, message: Message) -> None:
|
||||
if not toggles.messages:
|
||||
return
|
||||
await capture_message(client, message, ctx, toggles)
|
||||
await notify_bg_event(
|
||||
ctx.pool, "message", ctx.account_id, chat_id=meta.chat_id, message_id=message.id
|
||||
)
|
||||
if (
|
||||
toggles.stt
|
||||
and is_transcribable(message)
|
||||
|
||||
@@ -4,6 +4,7 @@ from pyrogram.types import User
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.presence import repository
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@PyroClient.on_user_status()
|
||||
@@ -22,6 +23,7 @@ async def on_user_status(client: PyroClient, user: User) -> None:
|
||||
str(user.raw),
|
||||
)
|
||||
await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online")
|
||||
await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id)
|
||||
|
||||
|
||||
handlers = on_user_status.handlers
|
||||
|
||||
@@ -3,6 +3,7 @@ from pyrogram import raw, utils
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.chat_meta import meta_from_peer
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
HANDLES = (raw.types.UpdateMessageReactions,)
|
||||
|
||||
@@ -47,3 +48,10 @@ async def handle(
|
||||
await repository.sync_reactions(
|
||||
ctx.pool, ctx.account_id, meta.chat_id, update.msg_id, current
|
||||
)
|
||||
await notify_bg_event(
|
||||
ctx.pool,
|
||||
"reaction",
|
||||
ctx.account_id,
|
||||
chat_id=meta.chat_id,
|
||||
message_id=update.msg_id,
|
||||
)
|
||||
|
||||
@@ -4,6 +4,7 @@ from pyrogram import raw, utils
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.read_receipts import repository
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
HANDLES = (raw.types.UpdateReadHistoryOutbox,)
|
||||
|
||||
@@ -24,3 +25,6 @@ async def handle(
|
||||
update.max_id,
|
||||
str(update),
|
||||
)
|
||||
await notify_bg_event(
|
||||
ctx.pool, "receipt", ctx.account_id, chat_id=chat_id, message_id=update.max_id
|
||||
)
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import asyncpg
|
||||
|
||||
_GET = """
|
||||
SELECT downloaded FROM custom_emoji WHERE custom_emoji_id = $1
|
||||
"""
|
||||
|
||||
_UPSERT_DOWNLOADED = """
|
||||
INSERT INTO custom_emoji
|
||||
(custom_emoji_id, storage_key, file_size, mime, kind, downloaded, raw)
|
||||
VALUES ($1, $2, $3, $4, $5, true, '{}'::jsonb)
|
||||
ON CONFLICT (custom_emoji_id) DO UPDATE SET
|
||||
storage_key = EXCLUDED.storage_key,
|
||||
file_size = EXCLUDED.file_size,
|
||||
mime = EXCLUDED.mime,
|
||||
kind = EXCLUDED.kind,
|
||||
downloaded = true
|
||||
"""
|
||||
|
||||
|
||||
async def is_downloaded(pool: asyncpg.Pool, custom_emoji_id: int) -> bool:
|
||||
return bool(await pool.fetchval(_GET, custom_emoji_id))
|
||||
|
||||
|
||||
async def upsert_downloaded( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
custom_emoji_id: int,
|
||||
storage_key: str,
|
||||
file_size: int | None,
|
||||
mime: str | None,
|
||||
kind: str,
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
_UPSERT_DOWNLOADED, custom_emoji_id, storage_key, file_size, mime, kind
|
||||
)
|
||||
@@ -2,8 +2,18 @@ from userbot.modules.jobs.handlers import (
|
||||
backfill,
|
||||
enrich_chat,
|
||||
fetch_avatar,
|
||||
fetch_custom_emoji,
|
||||
fetch_media,
|
||||
sync_dialogs,
|
||||
transcribe,
|
||||
)
|
||||
|
||||
__all__ = ["backfill", "enrich_chat", "fetch_avatar", "fetch_media", "transcribe"]
|
||||
__all__ = [
|
||||
"backfill",
|
||||
"enrich_chat",
|
||||
"fetch_avatar",
|
||||
"fetch_custom_emoji",
|
||||
"fetch_media",
|
||||
"sync_dialogs",
|
||||
"transcribe",
|
||||
]
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
from pyrogram.errors import PeerIdInvalid
|
||||
|
||||
from userbot.modules.capture import capture_message
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
@@ -23,11 +25,15 @@ async def backfill(ctx: JobContext) -> None:
|
||||
max_id = (ctx.job.cursor or {}).get("max_id", 0)
|
||||
processed = ctx.job.progress.get("processed", 0)
|
||||
kwargs = {"max_id": max_id} if max_id else {}
|
||||
async for message in client.get_chat_history(chat_id, **kwargs):
|
||||
await capture_message(client, message, capture, toggles)
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
next_max = message.id - 1
|
||||
await ctx.save_cursor({"max_id": next_max})
|
||||
await ctx.report_progress({"processed": processed, "max_id": next_max})
|
||||
try:
|
||||
async for message in client.get_chat_history(chat_id, **kwargs):
|
||||
await capture_message(client, message, capture, toggles)
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
next_max = message.id - 1
|
||||
await ctx.save_cursor({"max_id": next_max})
|
||||
await ctx.report_progress({"processed": processed, "max_id": next_max})
|
||||
except PeerIdInvalid:
|
||||
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
|
||||
return
|
||||
await ctx.report_progress({"processed": processed, "done": True})
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
from io import BytesIO
|
||||
|
||||
from pyrogram.types import Sticker
|
||||
|
||||
from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
|
||||
|
||||
def _kind(sticker: Sticker) -> str:
|
||||
if sticker.is_animated:
|
||||
return "animated"
|
||||
if sticker.is_video:
|
||||
return "video"
|
||||
return "static"
|
||||
|
||||
|
||||
@register("fetch_custom_emoji")
|
||||
async def fetch_custom_emoji(ctx: JobContext) -> None:
|
||||
client = ctx.client
|
||||
if client is None:
|
||||
return
|
||||
capture = getattr(client, "capture", None)
|
||||
if capture is None:
|
||||
return
|
||||
custom_emoji_id = int(ctx.job.params["custom_emoji_id"])
|
||||
if await is_downloaded(ctx.pool, custom_emoji_id):
|
||||
return
|
||||
stickers = await client.get_custom_emoji_stickers([str(custom_emoji_id)])
|
||||
if not stickers:
|
||||
return
|
||||
sticker = stickers[0]
|
||||
buffer = await client.download_media(sticker.file_id, in_memory=True)
|
||||
if not isinstance(buffer, BytesIO):
|
||||
return
|
||||
data = buffer.getvalue()
|
||||
storage_key = capture.storage.put(data)
|
||||
await upsert_downloaded(
|
||||
ctx.pool,
|
||||
custom_emoji_id,
|
||||
storage_key,
|
||||
len(data),
|
||||
sticker.mime_type,
|
||||
_kind(sticker),
|
||||
)
|
||||
@@ -0,0 +1,108 @@
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import BadRequest, Forbidden
|
||||
from pyrogram.types import Chat, User
|
||||
|
||||
from userbot.modules.avatars import note_avatar
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.groups.repository import insert_chat_history
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.profiles.parse import snapshot_from_chat, snapshot_from_high_level
|
||||
from userbot.modules.profiles.repository import write_profile
|
||||
|
||||
SAVE_EVERY = 100
|
||||
USERS_BATCH = 200
|
||||
|
||||
_UPSERT_DIALOG = """
|
||||
INSERT INTO dialogs (account_id, chat_id) VALUES ($1, $2)
|
||||
ON CONFLICT (account_id, chat_id) DO UPDATE SET updated_at = now()
|
||||
"""
|
||||
|
||||
|
||||
async def _save_private(ctx: CaptureContext, chat: Chat, chat_id: int) -> bool:
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat)
|
||||
await write_profile(ctx.pool, ctx.account_id, chat_id, fields, str(chat))
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id
|
||||
)
|
||||
return bool(fields.first_name or fields.last_name or fields.username)
|
||||
|
||||
|
||||
async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) -> None:
|
||||
for start in range(0, len(ids), USERS_BATCH):
|
||||
batch = ids[start : start + USERS_BATCH]
|
||||
try:
|
||||
result = await client.get_users(batch)
|
||||
except (BadRequest, Forbidden):
|
||||
continue
|
||||
users = result if isinstance(result, list) else [result]
|
||||
for user in users:
|
||||
if not isinstance(user, User):
|
||||
continue
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool,
|
||||
ctx.account_id,
|
||||
user.id,
|
||||
"peer",
|
||||
photo_unique_id,
|
||||
photo_file_id,
|
||||
)
|
||||
|
||||
|
||||
async def _save_group(ctx: CaptureContext, chat: Chat, chat_id: int) -> None:
|
||||
photo = chat.photo
|
||||
photo_unique_id = photo.big_photo_unique_id if photo else None
|
||||
photo_file_id = photo.big_file_id if photo else None
|
||||
await insert_chat_history(
|
||||
ctx.pool,
|
||||
ctx.account_id,
|
||||
chat_id,
|
||||
0,
|
||||
"meta",
|
||||
chat.title,
|
||||
photo_unique_id,
|
||||
None,
|
||||
datetime.now(UTC),
|
||||
str(chat),
|
||||
)
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool, ctx.account_id, chat_id, "chat", photo_unique_id, photo_file_id
|
||||
)
|
||||
|
||||
|
||||
@register("sync_dialogs")
|
||||
async def sync_dialogs(ctx: JobContext) -> None:
|
||||
client = ctx.client
|
||||
if client is None:
|
||||
return
|
||||
capture = getattr(client, "capture", None)
|
||||
if capture is None:
|
||||
return
|
||||
processed = ctx.job.progress.get("processed", 0)
|
||||
nameless: list[int] = []
|
||||
async for dialog in client.get_dialogs():
|
||||
chat = dialog.chat
|
||||
if chat is None or chat.id is None:
|
||||
continue
|
||||
chat_id = chat.id
|
||||
try:
|
||||
if chat_id > 0:
|
||||
if not await _save_private(capture, chat, chat_id):
|
||||
nameless.append(chat_id)
|
||||
else:
|
||||
await _save_group(capture, chat, chat_id)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
await ctx.report_progress({"processed": processed})
|
||||
await _enrich_users(client, capture, nameless)
|
||||
await ctx.report_progress({"processed": processed, "done": True})
|
||||
@@ -20,11 +20,20 @@ _MEDIA_ATTRS = (
|
||||
)
|
||||
|
||||
|
||||
_WEB_PAGE_ATTRS = ("photo", "video", "animation", "document", "audio")
|
||||
|
||||
|
||||
def media_object(message: Message) -> tuple[str | None, Any]:
|
||||
for attr in _MEDIA_ATTRS:
|
||||
obj = getattr(message, attr, None)
|
||||
if obj is not None:
|
||||
return attr, obj
|
||||
web_page = getattr(message, "web_page", None)
|
||||
if web_page is not None:
|
||||
for attr in _WEB_PAGE_ATTRS:
|
||||
obj = getattr(web_page, attr, None)
|
||||
if obj is not None:
|
||||
return attr, obj
|
||||
return None, None
|
||||
|
||||
|
||||
@@ -70,7 +79,8 @@ async def capture_media( # noqa: PLR0913
|
||||
file_size = existing["file_size"]
|
||||
downloaded = True
|
||||
else:
|
||||
buffer = await client.download_media(message, in_memory=True)
|
||||
target = message if getattr(message, kind or "", None) is obj else obj
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
if isinstance(buffer, BytesIO):
|
||||
data = buffer.getvalue()
|
||||
storage_key = ctx.storage.put(data)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pyrogram import Client, raw
|
||||
from pyrogram.types import User
|
||||
from pyrogram.types import Chat, User
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -54,3 +54,18 @@ def snapshot_from_high_level(
|
||||
is_deleted_account=bool(user.is_deleted),
|
||||
)
|
||||
return fields, photo_file_id, photo_unique_id
|
||||
|
||||
|
||||
def snapshot_from_chat(chat: Chat) -> tuple[ProfileFields, str | None, str | None]:
|
||||
photo = chat.photo
|
||||
photo_unique_id = photo.big_photo_unique_id if photo else None
|
||||
photo_file_id = photo.big_file_id if photo else None
|
||||
fields = ProfileFields(
|
||||
first_name=chat.first_name,
|
||||
last_name=chat.last_name,
|
||||
username=chat.username,
|
||||
phone=None,
|
||||
photo_unique_id=photo_unique_id,
|
||||
is_deleted_account=False,
|
||||
)
|
||||
return fields, photo_file_id, photo_unique_id
|
||||
|
||||
@@ -13,6 +13,7 @@ from userbot import PyroClient
|
||||
from userbot.modules.capture import CaptureContext, build_capture_context
|
||||
from userbot.modules.jobs import JobConsumer
|
||||
from utils.env import env
|
||||
from utils.jobs import enqueue
|
||||
from utils.logging import logger, setup_logging
|
||||
from utils.read.watches import WATCHES_CHANGED_CHANNEL
|
||||
from utils.storage import ContentAddressedStorage
|
||||
@@ -72,6 +73,17 @@ async def _setup_capture(
|
||||
logger.info("[green]Capture context ready.[/]")
|
||||
|
||||
|
||||
async def _enqueue_sync_dialogs(pool: asyncpg.Pool, account_id: int) -> None:
|
||||
existing = await pool.fetchval(
|
||||
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = 'sync_dialogs' "
|
||||
"AND status IN ('pending', 'running') LIMIT 1",
|
||||
account_id,
|
||||
)
|
||||
if existing is None:
|
||||
await enqueue(pool, account_id, "sync_dialogs", {})
|
||||
logger.info("[green]Queued sync_dialogs.[/]")
|
||||
|
||||
|
||||
async def _listen_changes(
|
||||
clients: list[PyroClient], tasks: set[asyncio.Task]
|
||||
) -> asyncpg.Connection:
|
||||
@@ -134,6 +146,7 @@ async def runner() -> None:
|
||||
await _setup_capture(pool, client, account_id, storage)
|
||||
consumer = JobConsumer(client, pool, account_id)
|
||||
consumer_tasks.append(asyncio.create_task(consumer.run()))
|
||||
await _enqueue_sync_dialogs(pool, account_id)
|
||||
|
||||
if clients:
|
||||
listen_conn = await _listen_changes(clients, reload_tasks)
|
||||
|
||||
@@ -477,3 +477,18 @@ class Alert(SQLModel, table=True):
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class Dialog(SQLModel, table=True):
|
||||
__tablename__ = "dialogs"
|
||||
|
||||
account_id: int = Field(primary_key=True)
|
||||
chat_id: int = Field(primary_key=True)
|
||||
updated_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
import json
|
||||
from typing import Literal
|
||||
|
||||
import asyncpg
|
||||
|
||||
BG_EVENTS_CHANNEL = "bg_events"
|
||||
|
||||
EventKind = Literal["message", "edit", "delete", "reaction", "presence", "receipt"]
|
||||
|
||||
|
||||
async def notify_bg_event( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
kind: EventKind,
|
||||
account_id: int,
|
||||
*,
|
||||
chat_id: int | None = None,
|
||||
message_id: int | None = None,
|
||||
message_ids: list[int] | None = None,
|
||||
) -> None:
|
||||
payload: dict[str, object] = {"kind": kind, "account_id": account_id}
|
||||
if chat_id is not None:
|
||||
payload["chat_id"] = chat_id
|
||||
if message_id is not None:
|
||||
payload["message_id"] = message_id
|
||||
if message_ids is not None:
|
||||
payload["message_ids"] = message_ids
|
||||
await pool.execute(
|
||||
"SELECT pg_notify($1, $2)", BG_EVENTS_CHANNEL, json.dumps(payload)
|
||||
)
|
||||
@@ -3,6 +3,12 @@ import asyncpg
|
||||
from utils.read.models import AccountView
|
||||
|
||||
|
||||
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
|
||||
return await pool.fetchval(
|
||||
"SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id
|
||||
)
|
||||
|
||||
|
||||
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT account_id, label, phone, tg_user_id, is_active FROM accounts "
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.accounts import self_user_id
|
||||
from utils.read.models import ResponseStats, VolumeBucket
|
||||
|
||||
|
||||
async def message_volume(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, *, days: int = 90
|
||||
) -> list[VolumeBucket]:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
rows = await pool.fetch(
|
||||
"SELECT date_trunc('day', date) AS bucket, count(*) AS total, "
|
||||
"count(*) FILTER (WHERE sender_id = $3) AS outgoing "
|
||||
"FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND date >= $4 "
|
||||
"GROUP BY bucket ORDER BY bucket",
|
||||
account_id,
|
||||
chat_id,
|
||||
self_id,
|
||||
datetime.now(UTC) - timedelta(days=days),
|
||||
)
|
||||
return [
|
||||
VolumeBucket(
|
||||
bucket=row["bucket"],
|
||||
total=row["total"],
|
||||
outgoing=row["outgoing"],
|
||||
incoming=row["total"] - row["outgoing"],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
async def response_stats(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> ResponseStats:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
rows = await pool.fetch(
|
||||
"WITH ordered AS ("
|
||||
"SELECT sender_id, date, "
|
||||
"lag(sender_id) OVER w AS prev_sender, "
|
||||
"lag(date) OVER w AS prev_date "
|
||||
"FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
|
||||
"WINDOW w AS (ORDER BY date, message_id)), "
|
||||
"resp AS ("
|
||||
"SELECT (sender_id = $3) AS is_mine, "
|
||||
"EXTRACT(EPOCH FROM (date - prev_date)) AS secs "
|
||||
"FROM ordered "
|
||||
"WHERE prev_sender IS NOT NULL AND prev_sender <> sender_id) "
|
||||
"SELECT is_mine, count(*) AS n, "
|
||||
"percentile_cont(0.5) WITHIN GROUP (ORDER BY secs) AS median_secs "
|
||||
"FROM resp GROUP BY is_mine",
|
||||
account_id,
|
||||
chat_id,
|
||||
self_id,
|
||||
)
|
||||
stats = ResponseStats(
|
||||
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
|
||||
)
|
||||
for row in rows:
|
||||
if row["is_mine"]:
|
||||
stats.mine_median_seconds = row["median_secs"]
|
||||
stats.mine_count = row["n"]
|
||||
else:
|
||||
stats.their_median_seconds = row["median_secs"]
|
||||
stats.their_count = row["n"]
|
||||
return stats
|
||||
@@ -1,5 +1,6 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.accounts import self_user_id
|
||||
from utils.read.message_view import build_message_view, load_raw, media_ref_from
|
||||
from utils.read.models import (
|
||||
ChatListItem,
|
||||
@@ -8,6 +9,7 @@ from utils.read.models import (
|
||||
MessageView,
|
||||
Page,
|
||||
)
|
||||
from utils.read.read_receipts import read_up_to
|
||||
|
||||
_MESSAGE_COLS = (
|
||||
"chat_id, message_id, date, sender_id, text, has_media, is_self_destruct, "
|
||||
@@ -52,35 +54,43 @@ async def list_chats(
|
||||
pool: asyncpg.Pool, account_id: int, page: Page
|
||||
) -> list[ChatListItem]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT m.chat_id, count(*) AS message_count, max(m.date) AS last_date, "
|
||||
"WITH ids AS ("
|
||||
"SELECT DISTINCT chat_id FROM messages WHERE account_id = $1 "
|
||||
"UNION SELECT chat_id FROM dialogs WHERE account_id = $1), "
|
||||
"agg AS (SELECT chat_id, count(*) AS message_count, max(date) AS last_date "
|
||||
"FROM messages WHERE account_id = $1 GROUP BY chat_id) "
|
||||
"SELECT ids.chat_id, COALESCE(agg.message_count, 0) AS message_count, "
|
||||
"agg.last_date AS last_date, "
|
||||
"(SELECT p.first_name FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS first_name, "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS first_name, "
|
||||
"(SELECT p.last_name FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS last_name, "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS last_name, "
|
||||
"(SELECT p.username FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS username, "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS username, "
|
||||
"(SELECT ch.title FROM chat_history ch "
|
||||
"WHERE ch.account_id = $1 AND ch.chat_id = m.chat_id "
|
||||
"WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id "
|
||||
"AND ch.title IS NOT NULL ORDER BY ch.ts DESC LIMIT 1) AS group_title, "
|
||||
"EXISTS (SELECT 1 FROM avatars a "
|
||||
"WHERE a.account_id = $1 AND a.owner_id = m.chat_id) AS has_avatar, "
|
||||
"(SELECT COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool, false) "
|
||||
"FROM peers p WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS is_bot, "
|
||||
"WHERE a.account_id = $1 AND a.owner_id = ids.chat_id) AS has_avatar, "
|
||||
"(SELECT COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool, "
|
||||
"p.raw->>'type' = 'ChatType.BOT', false) "
|
||||
"FROM peers p WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS is_bot, "
|
||||
"(SELECT COALESCE((p.raw->>'is_contact')::bool, (p.raw->>'contact')::bool, "
|
||||
"false) FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS is_contact, "
|
||||
"(SELECT ch.raw->'chat'->>'type' = 'ChatType.CHANNEL' FROM chat_history ch "
|
||||
"WHERE ch.account_id = $1 AND ch.chat_id = m.chat_id "
|
||||
"AND ch.raw->'chat'->>'type' IS NOT NULL "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS is_contact, "
|
||||
"(SELECT COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type') "
|
||||
"= 'ChatType.CHANNEL' FROM chat_history ch "
|
||||
"WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id "
|
||||
"AND COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type') IS NOT NULL "
|
||||
"ORDER BY ch.ts DESC LIMIT 1) AS is_broadcast, "
|
||||
"(SELECT lm.text FROM messages lm "
|
||||
"WHERE lm.account_id = $1 AND lm.chat_id = m.chat_id "
|
||||
"WHERE lm.account_id = $1 AND lm.chat_id = ids.chat_id "
|
||||
"ORDER BY lm.date DESC, lm.message_id DESC LIMIT 1) AS last_text, "
|
||||
"(SELECT lm.sender_id FROM messages lm "
|
||||
"WHERE lm.account_id = $1 AND lm.chat_id = m.chat_id "
|
||||
"WHERE lm.account_id = $1 AND lm.chat_id = ids.chat_id "
|
||||
"ORDER BY lm.date DESC, lm.message_id DESC LIMIT 1) AS last_sender_id "
|
||||
"FROM messages m WHERE m.account_id = $1 "
|
||||
"GROUP BY m.chat_id ORDER BY last_date DESC LIMIT $2 OFFSET $3",
|
||||
"FROM ids LEFT JOIN agg ON agg.chat_id = ids.chat_id "
|
||||
"ORDER BY last_date DESC NULLS LAST, ids.chat_id DESC LIMIT $2 OFFSET $3",
|
||||
account_id,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
@@ -146,9 +156,43 @@ async def get_chat_history(
|
||||
else:
|
||||
views.append(_build_album(members, media_by_key))
|
||||
index = end
|
||||
await _apply_read_status(pool, account_id, chat_id, views)
|
||||
return views
|
||||
|
||||
|
||||
async def get_message(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> MessageView | None:
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT {_MESSAGE_COLS} FROM messages " # noqa: S608
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
media_by_key = await _media_map(pool, account_id, [row])
|
||||
raw = load_raw(row["raw"])
|
||||
view = build_message_view(row, raw, _single_media(row, raw, media_by_key))
|
||||
await _apply_read_status(pool, account_id, chat_id, [view])
|
||||
return view
|
||||
|
||||
|
||||
async def _apply_read_status(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, views: list[MessageView]
|
||||
) -> None:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
if self_id is None:
|
||||
return
|
||||
marker = await read_up_to(pool, account_id, chat_id)
|
||||
if marker is None:
|
||||
return
|
||||
for view in views:
|
||||
if view.sender_id == self_id and view.message_id <= marker:
|
||||
view.read = True
|
||||
|
||||
|
||||
def _build_album(
|
||||
members: list[tuple[asyncpg.Record, dict]],
|
||||
media_by_key: dict[tuple[int, int], asyncpg.Record],
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import CustomEmojiRef
|
||||
|
||||
_GET = """
|
||||
SELECT storage_key, downloaded, mime, kind FROM custom_emoji
|
||||
WHERE custom_emoji_id = $1
|
||||
"""
|
||||
|
||||
|
||||
async def current_custom_emoji(
|
||||
pool: asyncpg.Pool, custom_emoji_id: int
|
||||
) -> CustomEmojiRef | None:
|
||||
row = await pool.fetchrow(_GET, custom_emoji_id)
|
||||
return CustomEmojiRef(**dict(row)) if row else None
|
||||
@@ -1,5 +1,6 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.message_view import load_raw
|
||||
from utils.read.models import MediaVersionView, MediaView
|
||||
|
||||
_MEDIA_COLS = (
|
||||
@@ -9,6 +10,45 @@ _MEDIA_COLS = (
|
||||
|
||||
_VERSION_COLS = "id, kind, storage_key, file_size, mime, observed_at"
|
||||
|
||||
_WEB_PAGE_MEDIA_KINDS = ("photo", "video", "animation", "document", "audio")
|
||||
|
||||
|
||||
async def _web_page_media_stub(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> MediaView | None:
|
||||
row = await pool.fetchrow(
|
||||
"SELECT date, raw FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
web_page = load_raw(row["raw"]).get("web_page")
|
||||
if not isinstance(web_page, dict):
|
||||
return None
|
||||
kind = next(
|
||||
(k for k in _WEB_PAGE_MEDIA_KINDS if isinstance(web_page.get(k), dict)), None
|
||||
)
|
||||
if kind is None:
|
||||
return None
|
||||
obj = web_page[kind]
|
||||
return MediaView(
|
||||
id=0,
|
||||
account_id=account_id,
|
||||
chat_id=chat_id,
|
||||
message_id=message_id,
|
||||
kind=kind,
|
||||
storage_key=None,
|
||||
file_size=obj.get("file_size"),
|
||||
mime=obj.get("mime_type"),
|
||||
ttl_seconds=None,
|
||||
downloaded=False,
|
||||
extracted_text=None,
|
||||
created_at=row["date"],
|
||||
)
|
||||
|
||||
|
||||
async def get_media(pool: asyncpg.Pool, media_id: int) -> MediaView | None:
|
||||
row = await pool.fetchrow(
|
||||
@@ -28,7 +68,9 @@ async def get_message_media(
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return MediaView(**dict(row)) if row else None
|
||||
if row is not None:
|
||||
return MediaView(**dict(row))
|
||||
return await _web_page_media_stub(pool, account_id, chat_id, message_id)
|
||||
|
||||
|
||||
async def get_media_versions(
|
||||
|
||||
@@ -141,6 +141,13 @@ class ServiceView(BaseModel):
|
||||
duration: int | None = None
|
||||
|
||||
|
||||
class PinnedView(BaseModel):
|
||||
message_id: int
|
||||
text: str | None = None
|
||||
media_kind: str | None = None
|
||||
sender_name: str | None = None
|
||||
|
||||
|
||||
class StickerView(BaseModel):
|
||||
emoji: str | None = None
|
||||
set_name: str | None = None
|
||||
@@ -178,6 +185,7 @@ class MessageView(BaseModel):
|
||||
sticker: StickerView | None = None
|
||||
is_sticker: bool = False
|
||||
is_animated_emoji: bool = False
|
||||
read: bool = False
|
||||
|
||||
|
||||
class MessageVersionView(BaseModel):
|
||||
@@ -217,6 +225,13 @@ class AvatarRef(BaseModel):
|
||||
mime: str | None
|
||||
|
||||
|
||||
class CustomEmojiRef(BaseModel):
|
||||
storage_key: str | None
|
||||
downloaded: bool
|
||||
mime: str | None
|
||||
kind: str | None
|
||||
|
||||
|
||||
class CallbackView(BaseModel):
|
||||
position: int
|
||||
label: str | None
|
||||
@@ -256,6 +271,20 @@ class PresenceHourly(BaseModel):
|
||||
last_seen: datetime | None
|
||||
|
||||
|
||||
class VolumeBucket(BaseModel):
|
||||
bucket: datetime
|
||||
total: int
|
||||
outgoing: int
|
||||
incoming: int
|
||||
|
||||
|
||||
class ResponseStats(BaseModel):
|
||||
mine_median_seconds: float | None
|
||||
mine_count: int
|
||||
their_median_seconds: float | None
|
||||
their_count: int
|
||||
|
||||
|
||||
class PeerView(BaseModel):
|
||||
peer_id: int
|
||||
first_name: str | None
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.message_view import _media_kind, _peer_name, load_raw
|
||||
from utils.read.models import PinnedView
|
||||
|
||||
_PINNED_SQL = (
|
||||
"SELECT raw FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 "
|
||||
"AND raw->>'service' LIKE '%PINNED_MESSAGE%' "
|
||||
"ORDER BY date DESC, message_id DESC LIMIT 1"
|
||||
)
|
||||
|
||||
|
||||
async def get_pinned(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> PinnedView | None:
|
||||
row = await pool.fetchrow(_PINNED_SQL, account_id, chat_id)
|
||||
if row is None:
|
||||
return None
|
||||
pinned = load_raw(row["raw"]).get("pinned_message")
|
||||
if not isinstance(pinned, dict):
|
||||
return None
|
||||
message_id = pinned.get("id")
|
||||
if message_id is None:
|
||||
return None
|
||||
sender = pinned.get("from_user")
|
||||
return PinnedView(
|
||||
message_id=message_id,
|
||||
text=pinned.get("text") or pinned.get("caption"),
|
||||
media_kind=_media_kind(pinned),
|
||||
sender_name=_peer_name(sender) if isinstance(sender, dict) else None,
|
||||
)
|
||||
@@ -33,6 +33,19 @@ async def presence_history( # noqa: PLR0913
|
||||
return [PresenceSample(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def current_presence(
|
||||
pool: asyncpg.Pool, account_id: int, peer_id: int
|
||||
) -> PresenceSample | None:
|
||||
row = await pool.fetchrow(
|
||||
"SELECT peer_id, ts, status, last_online_date, next_offline_date "
|
||||
"FROM presence WHERE account_id = $1 AND peer_id = $2 "
|
||||
"ORDER BY ts DESC LIMIT 1",
|
||||
account_id,
|
||||
peer_id,
|
||||
)
|
||||
return PresenceSample(**dict(row)) if row is not None else None
|
||||
|
||||
|
||||
async def presence_hourly(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
import asyncpg
|
||||
|
||||
|
||||
async def read_up_to(pool: asyncpg.Pool, account_id: int, chat_id: int) -> int | None:
|
||||
return await pool.fetchval(
|
||||
"SELECT max(message_id) FROM read_receipts "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND kind = 'read'",
|
||||
account_id,
|
||||
chat_id,
|
||||
)
|
||||
Reference in New Issue
Block a user