chats.py 8.95 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from pydantic import BaseModel
from typing import List, Union, Optional
from peewee import *
from playhouse.shortcuts import model_to_dict

import json
import uuid
import time

from apps.web.internal.db import DB

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


class Chat(Model):
    id = CharField(unique=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
19
    user_id = CharField()
Timothy J. Baek's avatar
Timothy J. Baek committed
20
21
    title = CharField()
    chat = TextField()  # Save Chat JSON as Text
22
23
24
25

    created_at = DateTimeField()
    updated_at = DateTimeField()

26
    share_id = CharField(null=True, unique=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
27
    archived = BooleanField(default=False)
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
30
31
32
33
34
35
36

    class Meta:
        database = DB


class ChatModel(BaseModel):
    id: str
    user_id: str
    title: str
Timothy J. Baek's avatar
Timothy J. Baek committed
37
    chat: str
38
39
40
41

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

42
    share_id: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
43
    archived: bool = False
Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
46
47
48
49
50
51
52
53
54


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


class ChatForm(BaseModel):
    chat: dict


Timothy J. Baek's avatar
Timothy J. Baek committed
55
56
57
58
class ChatTitleForm(BaseModel):
    title: str


59
class ChatResponse(BaseModel):
Timothy J. Baek's avatar
Timothy J. Baek committed
60
    id: str
61
62
63
    user_id: str
    title: str
    chat: dict
64
65
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch
66
    share_id: Optional[str] = None  # id of the chat to be shared
Timothy J. Baek's avatar
Timothy J. Baek committed
67
68
69
70
71


class ChatTitleIdResponse(BaseModel):
    id: str
    title: str
72
73
    updated_at: int
    created_at: int
Timothy J. Baek's avatar
Timothy J. Baek committed
74
75
76
77
78
79
80


class ChatTable:
    def __init__(self, db):
        self.db = db
        db.create_tables([Chat])

Timothy J. Baek's avatar
Timothy J. Baek committed
81
    def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
82
83
84
85
86
        id = str(uuid.uuid4())
        chat = ChatModel(
            **{
                "id": id,
                "user_id": user_id,
87
88
89
                "title": (
                    form_data.chat["title"] if "title" in form_data.chat else "New Chat"
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
90
                "chat": json.dumps(form_data.chat),
91
92
                "created_at": int(time.time()),
                "updated_at": int(time.time()),
Timothy J. Baek's avatar
Timothy J. Baek committed
93
94
            }
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
95
96
97
98
99
100

        result = Chat.create(**chat.model_dump())
        return chat if result else None

    def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
        try:
101
102
103
            query = Chat.update(
                chat=json.dumps(chat),
                title=chat["title"] if "title" in chat else "New Chat",
104
                updated_at=int(time.time()),
105
            ).where(Chat.id == id)
Timothy J. Baek's avatar
Timothy J. Baek committed
106
107
108
109
110
111
112
            query.execute()

            chat = Chat.get(Chat.id == id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
113
    def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
114
115
116
117
118
119
120
121
122
        # Get the existing chat to share
        chat = Chat.get(Chat.id == chat_id)
        # Check if the chat is already shared
        if chat.share_id:
            return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
        # Create a new chat with the same data, but with a new ID
        shared_chat = ChatModel(
            **{
                "id": str(uuid.uuid4()),
Timothy J. Baek's avatar
Timothy J. Baek committed
123
                "user_id": f"shared-{chat_id}",
124
125
                "title": chat.title,
                "chat": chat.chat,
126
                "created_at": int(time.time()),
127
128
129
130
131
132
133
134
135
136
            }
        )
        shared_result = Chat.create(**shared_chat.model_dump())
        # Update the original chat with the share_id
        result = (
            Chat.update(share_id=shared_chat.id).where(Chat.id == chat_id).execute()
        )

        return shared_chat if (shared_result and result) else None

Timothy J. Baek's avatar
Timothy J. Baek committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
    def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
        try:
            print("update_shared_chat_by_id")
            chat = Chat.get(Chat.id == chat_id)
            print(chat)

            query = Chat.update(
                title=chat.title,
                chat=chat.chat,
            ).where(Chat.id == chat.share_id)

            query.execute()

            chat = Chat.get(Chat.id == chat.share_id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
155
156
157
158
159
160
161
162
163
    def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool:
        try:
            query = Chat.delete().where(Chat.user_id == f"shared-{chat_id}")
            query.execute()  # Remove the rows, return number of rows removed.

            return True
        except:
            return False

164
    def update_chat_share_id_by_id(
Timothy J. Baek's avatar
Timothy J. Baek committed
165
        self, id: str, share_id: Optional[str]
166
167
168
169
170
171
172
173
174
175
176
177
    ) -> Optional[ChatModel]:
        try:
            query = Chat.update(
                share_id=share_id,
            ).where(Chat.id == id)
            query.execute()

            chat = Chat.get(Chat.id == id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
    def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
        try:
            chat = self.get_chat_by_id(id)
            query = Chat.update(
                archived=(not chat.archived),
            ).where(Chat.id == id)

            query.execute()

            chat = Chat.get(Chat.id == id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

192
193
194
195
196
197
198
199
200
201
202
203
204
    def get_archived_chat_lists_by_user_id(
        self, user_id: str, skip: int = 0, limit: int = 50
    ) -> List[ChatModel]:
        return [
            ChatModel(**model_to_dict(chat))
            for chat in Chat.select()
            .where(Chat.archived == True)
            .where(Chat.user_id == user_id)
            .order_by(Chat.updated_at.desc())
            # .limit(limit)
            # .offset(skip)
        ]

Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
    def get_chat_lists_by_user_id(
        self, user_id: str, skip: int = 0, limit: int = 50
    ) -> List[ChatModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
208
        return [
Timothy J. Baek's avatar
Timothy J. Baek committed
209
210
            ChatModel(**model_to_dict(chat))
            for chat in Chat.select()
Timothy J. Baek's avatar
Timothy J. Baek committed
211
            .where(Chat.archived == False)
Timothy J. Baek's avatar
Timothy J. Baek committed
212
            .where(Chat.user_id == user_id)
213
            .order_by(Chat.updated_at.desc())
214
215
            # .limit(limit)
            # .offset(skip)
Timothy J. Baek's avatar
Timothy J. Baek committed
216
217
        ]

Timothy J. Baek's avatar
Timothy J. Baek committed
218
219
220
221
222
223
    def get_chat_lists_by_chat_ids(
        self, chat_ids: List[str], skip: int = 0, limit: int = 50
    ) -> List[ChatModel]:
        return [
            ChatModel(**model_to_dict(chat))
            for chat in Chat.select()
Timothy J. Baek's avatar
Timothy J. Baek committed
224
            .where(Chat.archived == False)
Timothy J. Baek's avatar
Timothy J. Baek committed
225
            .where(Chat.id.in_(chat_ids))
226
            .order_by(Chat.updated_at.desc())
Timothy J. Baek's avatar
Timothy J. Baek committed
227
228
        ]

229
230
231
    def get_all_chats(self) -> List[ChatModel]:
        return [
            ChatModel(**model_to_dict(chat))
Timothy J. Baek's avatar
Timothy J. Baek committed
232
233
            for chat in Chat.select()
            .where(Chat.archived == False)
234
            .order_by(Chat.updated_at.desc())
235
236
        ]

Timothy J. Baek's avatar
Timothy J. Baek committed
237
238
    def get_all_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
        return [
Timothy J. Baek's avatar
Timothy J. Baek committed
239
240
            ChatModel(**model_to_dict(chat))
            for chat in Chat.select()
Timothy J. Baek's avatar
Timothy J. Baek committed
241
            .where(Chat.archived == False)
Timothy J. Baek's avatar
Timothy J. Baek committed
242
            .where(Chat.user_id == user_id)
243
            .order_by(Chat.updated_at.desc())
Timothy J. Baek's avatar
Timothy J. Baek committed
244
245
        ]

Timothy J. Baek's avatar
Timothy J. Baek committed
246
247
248
249
250
251
252
    def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
        try:
            chat = Chat.get(Chat.id == id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

253
254
255
256
257
258
259
260
261
262
263
264
    def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
        try:
            chat = Chat.get(Chat.share_id == id)

            if chat:
                chat = Chat.get(Chat.id == id)
                return ChatModel(**model_to_dict(chat))
            else:
                return None
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
265
    def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
266
267
268
269
270
271
272
273
274
275
276
277
        try:
            chat = Chat.get(Chat.id == id, Chat.user_id == user_id)
            return ChatModel(**model_to_dict(chat))
        except:
            return None

    def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]:
        return [
            ChatModel(**model_to_dict(chat))
            for chat in Chat.select().limit(limit).offset(skip)
        ]

278
279
    def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
        try:
Timothy J. Baek's avatar
Timothy J. Baek committed
280
            query = Chat.delete().where((Chat.id == id) & (Chat.user_id == user_id))
281
282
            query.execute()  # Remove the rows, return number of rows removed.

Timothy J. Baek's avatar
Timothy J. Baek committed
283
            return True and self.delete_shared_chat_by_chat_id(id)
284
285
286
        except:
            return False

287
288
    def delete_chats_by_user_id(self, user_id: str) -> bool:
        try:
289
290
291

            self.delete_shared_chats_by_user_id(user_id)

292
293
294
            query = Chat.delete().where(Chat.user_id == user_id)
            query.execute()  # Remove the rows, return number of rows removed.

295
            return True
Timothy J. Baek's avatar
Timothy J. Baek committed
296
297
298
299
300
301
302
303
304
305
306
307
308
        except:
            return False

    def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
        try:
            shared_chat_ids = [
                f"shared-{chat.id}"
                for chat in Chat.select().where(Chat.user_id == user_id)
            ]

            query = Chat.delete().where(Chat.user_id << shared_chat_ids)
            query.execute()  # Remove the rows, return number of rows removed.

309
310
311
312
            return True
        except:
            return False

Timothy J. Baek's avatar
Timothy J. Baek committed
313
314

Chats = ChatTable(DB)