chats.py 11.5 KB
Newer Older
Anuraag Jain's avatar
Anuraag Jain committed
1
from fastapi import Depends, Request, HTTPException, status
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
from datetime import datetime, timedelta
from typing import List, Union, Optional
4
from utils.utils import get_current_user, get_admin_user
Timothy J. Baek's avatar
Timothy J. Baek committed
5
6
from fastapi import APIRouter
from pydantic import BaseModel
7
import json
8
import logging
Timothy J. Baek's avatar
Timothy J. Baek committed
9
10
11
12

from apps.web.models.users import Users
from apps.web.models.chats import (
    ChatModel,
13
    ChatResponse,
Timothy J. Baek's avatar
Timothy J. Baek committed
14
    ChatTitleForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
15
16
17
18
19
    ChatForm,
    ChatTitleIdResponse,
    Chats,
)

20
21
22

from apps.web.models.tags import (
    TagModel,
23
    ChatIdTagModel,
24
25
26
27
28
    ChatIdTagForm,
    ChatTagsResponse,
    Tags,
)

Timothy J. Baek's avatar
Timothy J. Baek committed
29
30
from constants import ERROR_MESSAGES

31
from config import SRC_LOG_LEVELS, ENABLE_ADMIN_EXPORT
Timothy J. Baek's avatar
Timothy J. Baek committed
32

33
34
35
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])

Timothy J. Baek's avatar
Timothy J. Baek committed
36
37
38
router = APIRouter()

############################
Timothy J. Baek's avatar
Timothy J. Baek committed
39
# GetChatList
Timothy J. Baek's avatar
Timothy J. Baek committed
40
41
42
43
############################


@router.get("/", response_model=List[ChatTitleIdResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
@router.get("/list", response_model=List[ChatTitleIdResponse])
async def get_session_user_chat_list(
Timothy J. Baek's avatar
Timothy J. Baek committed
46
47
    user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
Timothy J. Baek's avatar
Timothy J. Baek committed
48
49
50
51
52
53
54
55
56
57
58
59
60
    return Chats.get_chat_list_by_user_id(user.id, skip, limit)


############################
# GetUserChatList
############################


@router.get("/list/user/{user_id}", response_model=List[ChatTitleIdResponse])
async def get_user_chat_list_by_user_id(
    user_id: str, user=Depends(get_admin_user), skip: int = 0, limit: int = 50
):
    return Chats.get_chat_list_by_user_id(user_id, skip, limit)
Timothy J. Baek's avatar
Timothy J. Baek committed
61
62


63
64
65
66
67
68
############################
# GetArchivedChats
############################


@router.get("/archived", response_model=List[ChatTitleIdResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
69
async def get_archived_session_user_chat_list(
70
71
    user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
Timothy J. Baek's avatar
Timothy J. Baek committed
72
    return Chats.get_archived_chat_list_by_user_id(user.id, skip, limit)
73
74


Timothy J. Baek's avatar
Timothy J. Baek committed
75
############################
Timothy J. Baek's avatar
Timothy J. Baek committed
76
# GetChats
Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
79
80
############################


@router.get("/all", response_model=List[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
81
async def get_user_chats(user=Depends(get_current_user)):
Anuraag Jain's avatar
Anuraag Jain committed
82
    return [
Timothy J. Baek's avatar
Timothy J. Baek committed
83
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
84
        for chat in Chats.get_chats_by_user_id(user.id)
85
    ]
Timothy J. Baek's avatar
Timothy J. Baek committed
86
87


88
89
90
91
92
93
############################
# GetAllChatsInDB
############################


@router.get("/all/db", response_model=List[ChatResponse])
94
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
95
    if not ENABLE_ADMIN_EXPORT:
96
97
98
99
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
100
101
    return [
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
102
        for chat in Chats.get_chats()
103
    ]
104
105


Timothy J. Baek's avatar
Timothy J. Baek committed
106
107
108
109
110
############################
# CreateNewChat
############################


111
@router.post("/new", response_model=Optional[ChatResponse])
112
async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
113
114
115
116
    try:
        chat = Chats.insert_new_chat(user.id, form_data)
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    except Exception as e:
117
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
118
119
120
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
121
122


Timothy J. Baek's avatar
Timothy J. Baek committed
123
124
125
126
127
128
129
130
131
132
133
############################
# GetAllTags
############################


@router.get("/tags/all", response_model=List[TagModel])
async def get_all_tags(user=Depends(get_current_user)):
    try:
        tags = Tags.get_tags_by_user_id(user.id)
        return tags
    except Exception as e:
134
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )


############################
# GetChatsByTags
############################


@router.get("/tags/tag/{tag_name}", response_model=List[ChatTitleIdResponse])
async def get_user_chats_by_tag_name(
    tag_name: str, user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
    chat_ids = [
        chat_id_tag.chat_id
        for chat_id_tag in Tags.get_chat_ids_by_tag_name_and_user_id(tag_name, user.id)
    ]

Timothy J. Baek's avatar
Timothy J. Baek committed
154
    chats = Chats.get_chat_list_by_chat_ids(chat_ids, skip, limit)
Timothy J. Baek's avatar
Timothy J. Baek committed
155

Timothy J. Baek's avatar
Timothy J. Baek committed
156
157
158
159
    if len(chats) == 0:
        Tags.delete_tag_by_tag_name_and_user_id(tag_name, user.id)

    return chats
Timothy J. Baek's avatar
Timothy J. Baek committed
160
161


Timothy J. Baek's avatar
Timothy J. Baek committed
162
163
164
165
166
############################
# GetChatById
############################


167
@router.get("/{id}", response_model=Optional[ChatResponse])
168
169
async def get_chat_by_id(id: str, user=Depends(get_current_user)):
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
Timothy J. Baek's avatar
Timothy J. Baek committed
170

Anuraag Jain's avatar
Anuraag Jain committed
171
    if chat:
Timothy J. Baek's avatar
Timothy J. Baek committed
172
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
173
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
174
175
176
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
177
178
179
180
181
182
183


############################
# UpdateChatById
############################


184
@router.post("/{id}", response_model=Optional[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
185
186
187
async def update_chat_by_id(
    id: str, form_data: ChatForm, user=Depends(get_current_user)
):
188
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
Anuraag Jain's avatar
Anuraag Jain committed
189
    if chat:
190
        updated_chat = {**json.loads(chat.chat), **form_data.chat}
Timothy J. Baek's avatar
Timothy J. Baek committed
191

192
        chat = Chats.update_chat_by_id(id, updated_chat)
Timothy J. Baek's avatar
Timothy J. Baek committed
193
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Anuraag Jain's avatar
Anuraag Jain committed
194
    else:
195
196
197
198
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
199
200
201
202
203
204
205
206


############################
# DeleteChatById
############################


@router.delete("/{id}", response_model=bool)
Timothy J. Baek's avatar
Timothy J. Baek committed
207
208
209
210
211
212
213
214
215
216
217
async def delete_chat_by_id(request: Request, id: str, user=Depends(get_current_user)):

    if (
        user.role == "user"
        and not request.app.state.USER_PERMISSIONS["chat"]["deletion"]
    ):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )

218
219
    result = Chats.delete_chat_by_id_and_user_id(id, user.id)
    return result
Timothy J. Baek's avatar
Timothy J. Baek committed
220

221

Timothy J. Baek's avatar
Timothy J. Baek committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
############################
# ArchiveChat
############################


@router.get("/{id}/archive", response_model=Optional[ChatResponse])
async def archive_chat_by_id(id: str, user=Depends(get_current_user)):
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
    if chat:
        chat = Chats.toggle_chat_archive_by_id(id)
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
        )


239
240
241
242
243
244
245
246
247
248
############################
# ShareChatById
############################


@router.post("/{id}/share", response_model=Optional[ChatResponse])
async def share_chat_by_id(id: str, user=Depends(get_current_user)):
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
    if chat:
        if chat.share_id:
Timothy J. Baek's avatar
Timothy J. Baek committed
249
            shared_chat = Chats.update_shared_chat_by_chat_id(chat.id)
250
251
252
253
            return ChatResponse(
                **{**shared_chat.model_dump(), "chat": json.loads(shared_chat.chat)}
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
254
        shared_chat = Chats.insert_shared_chat_by_chat_id(chat.id)
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        if not shared_chat:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )

        return ChatResponse(
            **{**shared_chat.model_dump(), "chat": json.loads(shared_chat.chat)}
        )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )


############################
# DeletedSharedChatById
############################


276
277
278
@router.delete("/{id}/share", response_model=Optional[bool])
async def delete_shared_chat_by_id(id: str, user=Depends(get_current_user)):
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
279
280
281
282
    if chat:
        if not chat.share_id:
            return False

283
284
285
286
        result = Chats.delete_shared_chat_by_chat_id(id)
        update_result = Chats.update_chat_share_id_by_id(id, None)

        return result and update_result != None
287
288
289
290
291
292
293
294
295
296
297
298
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )


############################
# GetSharedChatById
############################


Timothy J. Baek's avatar
Timothy J. Baek committed
299
300
@router.get("/share/{share_id}", response_model=Optional[ChatResponse])
async def get_shared_chat_by_id(share_id: str, user=Depends(get_current_user)):
301
302
303
304
305
306
307
308
309
    if user.role == "pending":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )

    if user.role == "user":
        chat = Chats.get_chat_by_share_id(share_id)
    elif user.role == "admin":
        chat = Chats.get_chat_by_id(share_id)
310
311
312
313
314
315
316
317
318

    if chat:
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


319
320
321
322
323
324
325
326
327
############################
# GetChatTagsById
############################


@router.get("/{id}/tags", response_model=List[TagModel])
async def get_chat_tags_by_id(id: str, user=Depends(get_current_user)):
    tags = Tags.get_tags_by_chat_id_and_user_id(id, user.id)

328
    if tags != None:
329
330
331
332
333
334
335
336
337
338
339
340
        return tags
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


############################
# AddChatTagById
############################


341
@router.post("/{id}/tags", response_model=Optional[ChatIdTagModel])
342
343
344
async def add_chat_tag_by_id(
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
345
346
347
348
    tags = Tags.get_tags_by_chat_id_and_user_id(id, user.id)

    if form_data.tag_name not in tags:
        tag = Tags.add_tag_to_chat(user.id, form_data)
349

350
351
352
353
354
355
356
        if tag:
            return tag
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
357
358
    else:
        raise HTTPException(
359
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
360
361
362
363
364
365
366
367
368
        )


############################
# DeleteChatTagById
############################


@router.delete("/{id}/tags", response_model=Optional[bool])
369
async def delete_chat_tag_by_id(
370
371
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
372
    result = Tags.delete_tag_by_tag_name_and_chat_id_and_user_id(
373
374
375
        form_data.tag_name, id, user.id
    )

376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
    if result:
        return result
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


############################
# DeleteAllChatTagsById
############################


@router.delete("/{id}/tags/all", response_model=Optional[bool])
async def delete_all_chat_tags_by_id(id: str, user=Depends(get_current_user)):
    result = Tags.delete_tags_by_chat_id_and_user_id(id, user.id)

    if result:
        return result
395
396
397
398
399
400
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
401
402
403
404
405
406
############################
# DeleteAllChats
############################


@router.delete("/", response_model=bool)
407
408
409
410
411
412
413
414
415
416
417
async def delete_all_user_chats(request: Request, user=Depends(get_current_user)):

    if (
        user.role == "user"
        and not request.app.state.USER_PERMISSIONS["chat"]["deletion"]
    ):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )

418
419
    result = Chats.delete_chats_by_user_id(user.id)
    return result