open-webui/backend/open_webui/models/chat_messages.py

import json
import time
import uuid
from typing import Any, Optional

from sqlalchemy import select, delete, func, cast, Integer, distinct
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from open_webui.utils.response import merge_usage, normalize_usage
from pydantic import BaseModel, ConfigDict
from sqlalchemy import (
    JSON,
    BigInteger,
    Boolean,
    Column,
    ForeignKey,
    Index,
    Integer,
    Text,
    cast,
    delete,
    func,
    select,
)
from sqlalchemy.ext.asyncio import AsyncSession

####################
# Helpers
####################


def _normalize_timestamp(timestamp: int) -> float:
    """Normalize and validate timestamp. Returns current time if invalid."""
    now = time.time()

    # Convert milliseconds to seconds if needed
    if timestamp > 10_000_000_000:
        timestamp = timestamp / 1000

    # Validate: must be after 2020 and not in the future (with 1 day tolerance)
    min_valid = 1577836800  # 2020-01-01 00:00:00 UTC
    max_valid = now + 86400  # 1 day in the future (clock skew tolerance)

    if timestamp < min_valid or timestamp > max_valid:
        return now

    return timestamp


def get_usage(data: dict) -> Optional[dict]:
    """Extract and normalize usage from message data."""
    usage = data.get('usage') or (data.get('info') or {}).get('usage')
    return normalize_usage(usage) if usage else None


def _token_columns(dialect: str):
    """Return (input_tokens, output_tokens) SQL column expressions.

    Falls back to OpenAI-style keys (prompt_tokens / completion_tokens)
    when the normalized keys are absent.
    """
    if dialect == 'sqlite':
        extract = lambda key: cast(func.json_extract(ChatMessage.usage, f'$.{key}'), Integer)
    elif dialect == 'postgresql':
        extract = lambda key: cast(func.json_extract_path_text(ChatMessage.usage, key), Integer)
    else:
        raise NotImplementedError(f'Unsupported dialect: {dialect}')

    return (
        func.coalesce(extract('input_tokens'), extract('prompt_tokens')),
        func.coalesce(extract('output_tokens'), extract('completion_tokens')),
    )


####################
# ChatMessage DB Schema
####################


class ChatMessage(Base):
    __tablename__ = 'chat_message'

    # Identity
    id = Column(Text, primary_key=True)
    chat_id = Column(Text, ForeignKey('chat.id', ondelete='CASCADE'), nullable=False, index=True)
    user_id = Column(Text, index=True)

    # Structure
    role = Column(Text, nullable=False)  # user, assistant, system
    parent_id = Column(Text, nullable=True)

    # Content
    content = Column(JSON, nullable=True)  # Can be str or list of blocks
    output = Column(JSON, nullable=True)

    # Model (for assistant messages)
    model_id = Column(Text, nullable=True, index=True)

    # Attachments
    files = Column(JSON, nullable=True)
    sources = Column(JSON, nullable=True)
    embeds = Column(JSON, nullable=True)

    # Status
    done = Column(Boolean, default=True)
    status_history = Column(JSON, nullable=True)
    error = Column(JSON, nullable=True)

    # Usage (tokens, timing, etc.)
    usage = Column(JSON, nullable=True)

    # Context compaction checkpoint
    context_summary = Column(Text, nullable=True)

    # Timestamps
    created_at = Column(BigInteger, index=True)
    updated_at = Column(BigInteger)

    __table_args__ = (
        Index('chat_message_chat_parent_idx', 'chat_id', 'parent_id'),
        Index('chat_message_model_created_idx', 'model_id', 'created_at'),
        Index('chat_message_user_created_idx', 'user_id', 'created_at'),
    )


####################
# Pydantic Models
####################


class ChatMessageModel(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    chat_id: str
    user_id: str
    role: str
    parent_id: Optional[str] = None
    content: Optional[Any] = None  # str or list of blocks
    output: Optional[list] = None
    model_id: Optional[str] = None
    files: Optional[list] = None
    sources: Optional[list] = None
    embeds: Optional[list] = None
    done: bool = True
    status_history: Optional[list] = None
    error: Optional[dict | str] = None
    usage: Optional[dict] = None
    context_summary: Optional[str] = None
    created_at: int
    updated_at: int


####################
# Table Operations
####################


class ChatMessageTable:
    async def upsert_message(
        self,
        message_id: str,
        chat_id: str,
        user_id: str,
        data: dict,
        db: Optional[AsyncSession] = None,
    ) -> Optional[ChatMessageModel]:
        """Insert or update a chat message."""
        async with get_async_db_context(db) as db:
            now = int(time.time())
            timestamp = data.get('timestamp', now)

            # Use composite ID: {chat_id}-{message_id}
            composite_id = f'{chat_id}-{message_id}'

            existing = await db.get(ChatMessage, composite_id)
            if existing:
                # Update existing
                if 'role' in data:
                    existing.role = data['role']
                if 'parent_id' in data or 'parentId' in data:
                    existing.parent_id = data.get('parent_id') or data.get('parentId')
                if 'content' in data:
                    existing.content = data.get('content')
                if 'output' in data:
                    existing.output = data.get('output')
                if 'model_id' in data or 'model' in data:
                    existing.model_id = data.get('model_id') or data.get('model')
                if 'files' in data:
                    existing.files = data.get('files')
                if 'sources' in data:
                    existing.sources = data.get('sources')
                if 'embeds' in data:
                    existing.embeds = data.get('embeds')
                if 'done' in data:
                    existing.done = data.get('done', True)
                if 'status_history' in data or 'statusHistory' in data:
                    existing.status_history = data.get('status_history') or data.get('statusHistory')
                if 'error' in data:
                    existing.error = data.get('error')
                if 'context_summary' in data or 'contextSummary' in data:
                    existing.context_summary = data.get('context_summary') or data.get('contextSummary')
                # Extract and normalize usage
                usage = get_usage(data)
                if usage:
                    existing_usage = normalize_usage(existing.usage or {}) if existing.usage else {}
                    existing.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage)
                existing.updated_at = now
                await db.commit()
                await db.refresh(existing)
                return ChatMessageModel.model_validate(existing)
            else:
                # Insert new
                # Extract and normalize usage
                usage = get_usage(data)
                message = ChatMessage(
                    id=composite_id,
                    chat_id=chat_id,
                    user_id=user_id,
                    role=data.get('role', 'user'),
                    parent_id=data.get('parent_id') or data.get('parentId'),
                    content=data.get('content'),
                    output=data.get('output'),
                    model_id=data.get('model_id') or data.get('model'),
                    files=data.get('files'),
                    sources=data.get('sources'),
                    embeds=data.get('embeds'),
                    done=data.get('done', True),
                    status_history=data.get('status_history') or data.get('statusHistory'),
                    error=data.get('error'),
                    usage=usage,
                    context_summary=data.get('context_summary') or data.get('contextSummary'),
                    created_at=timestamp,
                    updated_at=now,
                )
                db.add(message)
                await db.commit()
                await db.refresh(message)
                return ChatMessageModel.model_validate(message)

    async def get_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChatMessageModel]:
        async with get_async_db_context(db) as db:
            message = await db.get(ChatMessage, id)
            return ChatMessageModel.model_validate(message) if message else None

    async def get_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> list[ChatMessageModel]:
        async with get_async_db_context(db) as db:
            result = await db.execute(
                select(ChatMessage).filter_by(chat_id=chat_id).order_by(ChatMessage.created_at.asc())
            )
            messages = result.scalars().all()
            return [ChatMessageModel.model_validate(message) for message in messages]

    # DB column names that differ from the JSON message keys.
    DB_TO_JSON_KEY_MAP = {
        'parent_id': 'parentId',
        'model_id': 'model',
        'status_history': 'statusHistory',
        'context_summary': 'contextSummary',
        'created_at': 'timestamp',
    }
    # DB-internal columns excluded from the reconstructed message dict.
    EXCLUDED_COLUMNS = frozenset({'id', 'chat_id', 'user_id', 'updated_at'})

    async def get_messages_map_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> Optional[dict]:
        """Build a {message_id: message_dict} map from chat_message rows.

        Returns the same shape as chat.history.messages so callers
        (get_message_list, middleware) work unchanged.  Returns None if
        no rows exist for the chat (caller should fall back to the
        embedded JSON blob for legacy chats).
        """
        async with get_async_db_context(db) as db:
            result = await db.execute(select(ChatMessage).filter_by(chat_id=chat_id))
            rows = result.scalars().all()

        if not rows:
            return None

        # Strip the composite-id prefix ("{chat_id}-") to recover the
        # original message_id used as map key.
        prefix = f'{chat_id}-'
        prefix_len = len(prefix)
        col_keys = [c.key for c in ChatMessage.__table__.columns]

        messages_map: dict[str, dict] = {}
        for row in rows:
            msg_id = row.id[prefix_len:] if row.id.startswith(prefix) else row.id

            msg: dict = {'id': msg_id}
            for key in col_keys:
                if key in self.EXCLUDED_COLUMNS:
                    continue
                val = getattr(row, key)
                if val is None:
                    continue
                json_key = self.DB_TO_JSON_KEY_MAP.get(key, key)
                msg[json_key] = val

            # Ensure content always has a value
            msg.setdefault('content', '')

            # Mirror usage into info.usage for callers that read it there
            if 'usage' in msg:
                msg['info'] = {'usage': msg['usage']}

            messages_map[msg_id] = msg

        # Reconstruct childrenIds from parentId links so that the map
        # is fully navigable (callers like the frontend rely on this).
        for msg_id, msg in messages_map.items():
            parent_id = msg.get('parentId')
            if parent_id and parent_id in messages_map:
                parent = messages_map[parent_id]
                children = parent.get('childrenIds')
                if children is None:
                    parent['childrenIds'] = [msg_id]
                elif msg_id not in children:
                    children.append(msg_id)

        # Ensure every message has a childrenIds list (leaf nodes get [])
        for msg in messages_map.values():
            if 'childrenIds' not in msg:
                msg['childrenIds'] = []

        return messages_map

    async def get_messages_by_user_id(
        self,
        user_id: str,
        skip: int = 0,
        limit: int = 50,
        db: Optional[AsyncSession] = None,
    ) -> list[ChatMessageModel]:
        async with get_async_db_context(db) as db:
            result = await db.execute(
                select(ChatMessage)
                .filter_by(user_id=user_id)
                .order_by(ChatMessage.created_at.desc())
                .offset(skip)
                .limit(limit)
            )
            messages = result.scalars().all()
            return [ChatMessageModel.model_validate(message) for message in messages]

    async def get_messages_by_model_id(
        self,
        model_id: str,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        skip: int = 0,
        limit: int = 100,
        db: Optional[AsyncSession] = None,
    ) -> list[ChatMessageModel]:
        async with get_async_db_context(db) as db:
            stmt = select(ChatMessage).filter_by(model_id=model_id)
            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            stmt = stmt.order_by(ChatMessage.created_at.desc()).offset(skip).limit(limit)
            result = await db.execute(stmt)
            messages = result.scalars().all()
            return [ChatMessageModel.model_validate(message) for message in messages]

    async def get_chat_ids_by_model_id(
        self,
        model_id: str,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        skip: int = 0,
        limit: int = 50,
        db: Optional[AsyncSession] = None,
    ) -> list[str]:
        """Get distinct chat_ids that used a specific model."""

        async with get_async_db_context(db) as db:
            stmt = select(
                ChatMessage.chat_id,
                func.max(ChatMessage.created_at).label('last_message_at'),
            ).filter(ChatMessage.model_id == model_id)
            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)

            # Group by chat_id and order by most recent message in each chat
            # Secondary sort on chat_id ensures deterministic pagination
            stmt = (
                stmt.group_by(ChatMessage.chat_id)
                .order_by(func.max(ChatMessage.created_at).desc(), ChatMessage.chat_id)
                .offset(skip)
                .limit(limit)
            )
            result = await db.execute(stmt)
            chat_ids = result.all()
            return [chat_id for chat_id, _ in chat_ids]

    async def delete_messages_by_chat_id(self, chat_id: str, db: Optional[AsyncSession] = None) -> bool:
        async with get_async_db_context(db) as db:
            await db.execute(delete(ChatMessage).filter_by(chat_id=chat_id))
            await db.commit()
            return True

    async def delete_message_ids_by_chat_id(
        self,
        chat_id: str,
        message_ids: set[str],
        db: Optional[AsyncSession] = None,
    ) -> bool:
        """Delete specific ``chat_message`` rows by their original message IDs."""
        if not message_ids:
            return True
        async with get_async_db_context(db) as db:
            await db.execute(
                delete(ChatMessage)
                .where(ChatMessage.chat_id == chat_id)
                .where(ChatMessage.id.in_({f'{chat_id}-{mid}' for mid in message_ids}))
            )
            await db.commit()
            return True

    # Analytics methods
    async def get_message_count_by_model(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, int]:
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            stmt = select(ChatMessage.model_id, func.count(ChatMessage.id).label('count')).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.model_id.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.model_id)
            result = await db.execute(stmt)
            return {row.model_id: row.count for row in result.all()}

    async def get_unique_counts_by_model(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, dict]:
        """Count distinct users and chats per model."""
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            stmt = select(
                ChatMessage.model_id,
                func.count(distinct(ChatMessage.user_id)).label('unique_users'),
                func.count(distinct(ChatMessage.chat_id)).label('unique_chats'),
            ).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.model_id.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.model_id)
            result = await db.execute(stmt)
            return {
                row.model_id: {
                    'unique_users': row.unique_users,
                    'unique_chats': row.unique_chats,
                }
                for row in result.all()
            }

    async def get_token_usage_by_model(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, dict]:
        """Aggregate token usage by model using database-level aggregation."""
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            # We need the dialect to determine JSON extraction syntax
            # For async sessions, access via get_bind()
            bind = await db.connection()
            dialect = bind.dialect.name

            input_tokens, output_tokens = _token_columns(dialect)

            stmt = select(
                ChatMessage.model_id,
                func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
                func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
                func.count(ChatMessage.id).label('message_count'),
            ).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.model_id.isnot(None),
                ChatMessage.usage.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.model_id)
            result = await db.execute(stmt)

            return {
                row.model_id: {
                    'input_tokens': row.input_tokens,
                    'output_tokens': row.output_tokens,
                    'total_tokens': row.input_tokens + row.output_tokens,
                    'message_count': row.message_count,
                }
                for row in result.all()
            }

    async def get_token_usage_by_user(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, dict]:
        """Aggregate token usage by user using database-level aggregation."""
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            bind = await db.connection()
            dialect = bind.dialect.name

            input_tokens, output_tokens = _token_columns(dialect)

            stmt = select(
                ChatMessage.user_id,
                func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
                func.coalesce(func.sum(output_tokens), 0).label('output_tokens'),
                func.count(ChatMessage.id).label('message_count'),
            ).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.user_id.isnot(None),
                ChatMessage.usage.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.user_id)
            result = await db.execute(stmt)

            return {
                row.user_id: {
                    'input_tokens': row.input_tokens,
                    'output_tokens': row.output_tokens,
                    'total_tokens': row.input_tokens + row.output_tokens,
                    'message_count': row.message_count,
                }
                for row in result.all()
            }

    async def get_message_count_by_user(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, int]:
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            stmt = select(ChatMessage.user_id, func.count(ChatMessage.id).label('count')).filter(
                ChatMessage.role == 'assistant',
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.user_id)
            result = await db.execute(stmt)
            return {row.user_id: row.count for row in result.all()}

    async def get_message_count_by_chat(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, int]:
        async with get_async_db_context(db) as db:
            from open_webui.models.groups import GroupMember

            stmt = select(ChatMessage.chat_id, func.count(ChatMessage.id).label('count')).filter(
                ChatMessage.role == 'assistant',
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            stmt = stmt.group_by(ChatMessage.chat_id)
            result = await db.execute(stmt)
            return {row.chat_id: row.count for row in result.all()}

    async def get_daily_message_counts_by_model(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        group_id: Optional[str] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, dict[str, int]]:
        """Get message counts grouped by day and model."""
        async with get_async_db_context(db) as db:
            from datetime import datetime, timedelta

            from open_webui.models.groups import GroupMember

            stmt = select(ChatMessage.created_at, ChatMessage.model_id).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.model_id.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)
            if group_id:
                group_users = select(GroupMember.user_id).filter(GroupMember.group_id == group_id).scalar_subquery()
                stmt = stmt.filter(ChatMessage.user_id.in_(group_users))

            result = await db.execute(stmt)
            results = result.all()

            # Group by date -> model -> count
            daily_counts: dict[str, dict[str, int]] = {}
            for timestamp, model_id in results:
                date_str = datetime.fromtimestamp(_normalize_timestamp(timestamp)).strftime('%Y-%m-%d')
                if date_str not in daily_counts:
                    daily_counts[date_str] = {}
                daily_counts[date_str][model_id] = daily_counts[date_str].get(model_id, 0) + 1

            # Fill in missing days
            if start_date and end_date:
                current = datetime.fromtimestamp(_normalize_timestamp(start_date))
                end_dt = datetime.fromtimestamp(_normalize_timestamp(end_date))
                while current <= end_dt:
                    date_str = current.strftime('%Y-%m-%d')
                    if date_str not in daily_counts:
                        daily_counts[date_str] = {}
                    current += timedelta(days=1)

            return daily_counts

    async def get_hourly_message_counts_by_model(
        self,
        start_date: Optional[int] = None,
        end_date: Optional[int] = None,
        db: Optional[AsyncSession] = None,
    ) -> dict[str, dict[str, int]]:
        """Get message counts grouped by hour and model."""
        async with get_async_db_context(db) as db:
            from datetime import datetime, timedelta

            stmt = select(ChatMessage.created_at, ChatMessage.model_id).filter(
                ChatMessage.role == 'assistant',
                ChatMessage.model_id.isnot(None),
            )

            if start_date:
                stmt = stmt.filter(ChatMessage.created_at >= start_date)
            if end_date:
                stmt = stmt.filter(ChatMessage.created_at <= end_date)

            result = await db.execute(stmt)
            results = result.all()

            # Group by hour -> model -> count
            hourly_counts: dict[str, dict[str, int]] = {}
            for timestamp, model_id in results:
                hour_str = datetime.fromtimestamp(_normalize_timestamp(timestamp)).strftime('%Y-%m-%d %H:00')
                if hour_str not in hourly_counts:
                    hourly_counts[hour_str] = {}
                hourly_counts[hour_str][model_id] = hourly_counts[hour_str].get(model_id, 0) + 1

            # Fill in missing hours
            if start_date and end_date:
                current = datetime.fromtimestamp(_normalize_timestamp(start_date)).replace(
                    minute=0, second=0, microsecond=0
                )
                end_dt = datetime.fromtimestamp(_normalize_timestamp(end_date))
                while current <= end_dt:
                    hour_str = current.strftime('%Y-%m-%d %H:00')
                    if hour_str not in hourly_counts:
                        hourly_counts[hour_str] = {}
                    current += timedelta(hours=1)

            return hourly_counts


ChatMessages = ChatMessageTable()