chats.py

from pydantic import BaseModel, ConfigDict
from typing import List, Union, Optional

import json
import uuid
import time

from sqlalchemy import Column, String, BigInteger, Boolean
from sqlalchemy.orm import Session

from apps.webui.internal.db import Base


####################
# Chat DB Schema
####################


class Chat(Base):
    __tablename__ = "chat"

    id = Column(String, primary_key=True)
    user_id = Column(String)
    title = Column(String)
    chat = Column(String)  # Save Chat JSON as Text

    created_at = Column(BigInteger)
    updated_at = Column(BigInteger)

    share_id = Column(String, unique=True, nullable=True)
    archived = Column(Boolean, default=False)


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

    id: str
    user_id: str
    title: str
    chat: str

    created_at: int  # timestamp in epoch
    updated_at: int  # timestamp in epoch

    share_id: Optional[str] = None
    archived: bool = False


####################
# Forms
####################


class ChatForm(BaseModel):
    chat: dict


class ChatTitleForm(BaseModel):
    title: str


class ChatResponse(BaseModel):
    id: str
    user_id: str
    title: str
    chat: dict
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch
    share_id: Optional[str] = None  # id of the chat to be shared
    archived: bool


class ChatTitleIdResponse(BaseModel):
    id: str
    title: str
    updated_at: int
    created_at: int


class ChatTable:

    def insert_new_chat(
        self, db: Session, user_id: str, form_data: ChatForm
    ) -> Optional[ChatModel]:
        id = str(uuid.uuid4())
        chat = ChatModel(
            **{
                "id": id,
                "user_id": user_id,
                "title": (
                    form_data.chat["title"] if "title" in form_data.chat else "New Chat"
                ),
                "chat": json.dumps(form_data.chat),
                "created_at": int(time.time()),
                "updated_at": int(time.time()),
            }
        )

        result = Chat(**chat.model_dump())
        db.add(result)
        db.commit()
        db.refresh(result)
        return ChatModel.model_validate(result) if result else None

    def update_chat_by_id(
        self, db: Session, id: str, chat: dict
    ) -> Optional[ChatModel]:
        try:
            db.query(Chat).filter_by(id=id).update(
                {
                    "chat": json.dumps(chat),
                    "title": chat["title"] if "title" in chat else "New Chat",
                    "updated_at": int(time.time()),
                }
            )

            return self.get_chat_by_id(db, id)
        except:
            return None

    def insert_shared_chat_by_chat_id(
        self, db: Session, chat_id: str
    ) -> Optional[ChatModel]:
        # Get the existing chat to share
        chat = db.get(Chat, chat_id)
        # Check if the chat is already shared
        if chat.share_id:
            return self.get_chat_by_id_and_user_id(db, chat.share_id, "shared")
        # Create a new chat with the same data, but with a new ID
        shared_chat = ChatModel(
            **{
                "id": str(uuid.uuid4()),
                "user_id": f"shared-{chat_id}",
                "title": chat.title,
                "chat": chat.chat,
                "created_at": chat.created_at,
                "updated_at": int(time.time()),
            }
        )
        shared_result = Chat(**shared_chat.model_dump())
        db.add(shared_result)
        db.commit()
        db.refresh(shared_result)
        # Update the original chat with the share_id
        result = (
            db.query(Chat).filter_by(id=chat_id).update({"share_id": shared_chat.id})
        )

        return shared_chat if (shared_result and result) else None

    def update_shared_chat_by_chat_id(
        self, db: Session, chat_id: str
    ) -> Optional[ChatModel]:
        try:
            print("update_shared_chat_by_id")
            chat = db.get(Chat, chat_id)
            print(chat)

            db.query(Chat).filter_by(id=chat.share_id).update(
                {"title": chat.title, "chat": chat.chat}
            )

            return self.get_chat_by_id(db, chat.share_id)
        except:
            return None

    def delete_shared_chat_by_chat_id(self, db: Session, chat_id: str) -> bool:
        try:
            db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
            return True
        except:
            return False

    def update_chat_share_id_by_id(
        self, db: Session, id: str, share_id: Optional[str]
    ) -> Optional[ChatModel]:
        try:
            db.query(Chat).filter_by(id=id).update({"share_id": share_id})

            return self.get_chat_by_id(db, id)
        except:
            return None

    def toggle_chat_archive_by_id(self, db: Session, id: str) -> Optional[ChatModel]:
        try:
            chat = self.get_chat_by_id(db, id)
            db.query(Chat).filter_by(id=id).update({"archived": not chat.archived})

            return self.get_chat_by_id(db, id)
        except:
            return None

    def archive_all_chats_by_user_id(self, db: Session, user_id: str) -> bool:
        try:
            db.query(Chat).filter_by(user_id=user_id).update({"archived": True})

            return True
        except:
            return False

    def get_archived_chat_list_by_user_id(
        self, db: Session, user_id: str, skip: int = 0, limit: int = 50
    ) -> List[ChatModel]:
        all_chats = (
            db.query(Chat)
            .filter_by(user_id=user_id, archived=True)
            .order_by(Chat.updated_at.desc())
            # .limit(limit).offset(skip)
            .all()
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def get_chat_list_by_user_id(
        self,
        db: Session,
        user_id: str,
        include_archived: bool = False,
        skip: int = 0,
        limit: int = 50,
    ) -> List[ChatModel]:
        query = db.query(Chat).filter_by(user_id=user_id)
        if not include_archived:
            query = query.filter_by(archived=False)
        all_chats = (
            query.order_by(Chat.updated_at.desc())
            # .limit(limit).offset(skip)
            .all()
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def get_chat_list_by_chat_ids(
        self, db: Session, chat_ids: List[str], skip: int = 0, limit: int = 50
    ) -> List[ChatModel]:
        all_chats = (
            db.query(Chat)
            .filter(Chat.id.in_(chat_ids))
            .filter_by(archived=False)
            .order_by(Chat.updated_at.desc())
            .all()
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def get_chat_by_id(self, db: Session, id: str) -> Optional[ChatModel]:
        try:
            chat = db.get(Chat, id)
            return ChatModel.model_validate(chat)
        except:
            return None

    def get_chat_by_share_id(self, db: Session, id: str) -> Optional[ChatModel]:
        try:
            chat = db.query(Chat).filter_by(share_id=id).first()

            if chat:
                return self.get_chat_by_id(db, id)
            else:
                return None
        except Exception as e:
            return None

    def get_chat_by_id_and_user_id(
        self, db: Session, id: str, user_id: str
    ) -> Optional[ChatModel]:
        try:
            chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
            return ChatModel.model_validate(chat)
        except:
            return None

    def get_chats(self, db: Session, skip: int = 0, limit: int = 50) -> List[ChatModel]:
        all_chats = (
            db.query(Chat)
            # .limit(limit).offset(skip)
            .order_by(Chat.updated_at.desc())
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def get_chats_by_user_id(self, db: Session, user_id: str) -> List[ChatModel]:
        all_chats = (
            db.query(Chat).filter_by(user_id=user_id).order_by(Chat.updated_at.desc())
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def get_archived_chats_by_user_id(
        self, db: Session, user_id: str
    ) -> List[ChatModel]:
        all_chats = (
            db.query(Chat)
            .filter_by(user_id=user_id, archived=True)
            .order_by(Chat.updated_at.desc())
        )
        return [ChatModel.model_validate(chat) for chat in all_chats]

    def delete_chat_by_id(self, db: Session, id: str) -> bool:
        try:
            db.query(Chat).filter_by(id=id).delete()

            return True and self.delete_shared_chat_by_chat_id(db, id)
        except:
            return False

    def delete_chat_by_id_and_user_id(self, db: Session, id: str, user_id: str) -> bool:
        try:
            db.query(Chat).filter_by(id=id, user_id=user_id).delete()

            return True and self.delete_shared_chat_by_chat_id(db, id)
        except:
            return False

    def delete_chats_by_user_id(self, db: Session, user_id: str) -> bool:
        try:

            self.delete_shared_chats_by_user_id(db, user_id)

            db.query(Chat).filter_by(user_id=user_id).delete()
            return True
        except:
            return False

    def delete_shared_chats_by_user_id(self, db: Session, user_id: str) -> bool:
        try:
            chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
            shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]

            db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()

            return True
        except:
            return False


Chats = ChatTable()