chats.py 11.7 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
    return Chats.get_chat_list_by_user_id(user.id, skip, limit)


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
############################
# DeleteAllChats
############################


@router.delete("/", response_model=bool)
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,
        )

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


Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
76
77
78
79
80
81
############################
# 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
82
83


84
85
86
87
88
89
############################
# GetArchivedChats
############################


@router.get("/archived", response_model=List[ChatTitleIdResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
90
async def get_archived_session_user_chat_list(
91
92
    user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
Timothy J. Baek's avatar
Timothy J. Baek committed
93
    return Chats.get_archived_chat_list_by_user_id(user.id, skip, limit)
94
95


Timothy J. Baek's avatar
Timothy J. Baek committed
96
############################
Timothy J. Baek's avatar
Timothy J. Baek committed
97
# GetChats
Timothy J. Baek's avatar
Timothy J. Baek committed
98
99
100
101
############################


@router.get("/all", response_model=List[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
102
async def get_user_chats(user=Depends(get_current_user)):
Anuraag Jain's avatar
Anuraag Jain committed
103
    return [
Timothy J. Baek's avatar
Timothy J. Baek committed
104
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
105
        for chat in Chats.get_chats_by_user_id(user.id)
106
    ]
Timothy J. Baek's avatar
Timothy J. Baek committed
107
108


109
110
111
112
113
114
############################
# GetAllChatsInDB
############################


@router.get("/all/db", response_model=List[ChatResponse])
115
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
116
    if not ENABLE_ADMIN_EXPORT:
117
118
119
120
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
121
122
    return [
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
123
        for chat in Chats.get_chats()
124
    ]
125
126


Timothy J. Baek's avatar
Timothy J. Baek committed
127
128
129
130
131
############################
# CreateNewChat
############################


132
@router.post("/new", response_model=Optional[ChatResponse])
133
async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
134
135
136
137
    try:
        chat = Chats.insert_new_chat(user.id, form_data)
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    except Exception as e:
138
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
139
140
141
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
142
143
144
145
146
147
148


############################
# GetChatById
############################


149
@router.get("/{id}", response_model=Optional[ChatResponse])
150
151
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
152

Anuraag Jain's avatar
Anuraag Jain committed
153
    if chat:
Timothy J. Baek's avatar
Timothy J. Baek committed
154
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
155
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
156
157
158
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
159
160
161
162
163
164
165


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


166
@router.post("/{id}", response_model=Optional[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
167
168
169
async def update_chat_by_id(
    id: str, form_data: ChatForm, user=Depends(get_current_user)
):
170
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
Anuraag Jain's avatar
Anuraag Jain committed
171
    if chat:
172
        updated_chat = {**json.loads(chat.chat), **form_data.chat}
Timothy J. Baek's avatar
Timothy J. Baek committed
173

174
        chat = Chats.update_chat_by_id(id, updated_chat)
Timothy J. Baek's avatar
Timothy J. Baek committed
175
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Anuraag Jain's avatar
Anuraag Jain committed
176
    else:
177
178
179
180
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
181
182
183
184
185
186
187
188


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


@router.delete("/{id}", response_model=bool)
Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
async def delete_chat_by_id(request: Request, id: str, user=Depends(get_current_user)):

191
192
193
194
195
196
197
198
199
200
201
202
    if user.role == "admin":
        result = Chats.delete_chat_by_id(id)
        return result
    else:
        if not request.app.state.USER_PERMISSIONS["chat"]["deletion"]:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
            )

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

204

Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
############################
# 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()
        )


222
223
224
225
226
227
228
229
230
231
############################
# 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
232
            shared_chat = Chats.update_shared_chat_by_chat_id(chat.id)
233
234
235
236
            return ChatResponse(
                **{**shared_chat.model_dump(), "chat": json.loads(shared_chat.chat)}
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
237
        shared_chat = Chats.insert_shared_chat_by_chat_id(chat.id)
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
        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
############################


259
260
261
@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)
262
263
264
265
    if chat:
        if not chat.share_id:
            return False

266
267
268
269
        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
270
271
272
273
274
275
276
277
278
279
280
281
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )


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


Timothy J. Baek's avatar
Timothy J. Baek committed
282
283
@router.get("/share/{share_id}", response_model=Optional[ChatResponse])
async def get_shared_chat_by_id(share_id: str, user=Depends(get_current_user)):
284
285
286
287
288
289
290
291
292
    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)
293
294
295
296
297
298
299
300
301

    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
        )


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
############################
# 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:
        log.exception(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )


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


Timothy J. Baek's avatar
Timothy J. Baek committed
324
325
326
327
328
class TagNameForm(BaseModel):
    name: str


@router.post("/tags", response_model=List[ChatTitleIdResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
329
async def get_user_chat_list_by_tag_name(
Timothy J. Baek's avatar
Timothy J. Baek committed
330
331
332
333
    form_data: TagNameForm,
    user=Depends(get_current_user),
    skip: int = 0,
    limit: int = 50,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
334
335
336
):
    chat_ids = [
        chat_id_tag.chat_id
Timothy J. Baek's avatar
Timothy J. Baek committed
337
338
339
        for chat_id_tag in Tags.get_chat_ids_by_tag_name_and_user_id(
            form_data.name, user.id
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
340
341
342
343
344
    ]

    chats = Chats.get_chat_list_by_chat_ids(chat_ids, skip, limit)

    if len(chats) == 0:
Timothy J. Baek's avatar
Timothy J. Baek committed
345
        Tags.delete_tag_by_tag_name_and_user_id(form_data.name, user.id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
346
347
348
349

    return chats


350
351
352
353
354
355
356
357
358
############################
# 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)

359
    if tags != None:
360
361
362
363
364
365
366
367
368
369
370
371
        return tags
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


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


372
@router.post("/{id}/tags", response_model=Optional[ChatIdTagModel])
373
374
375
async def add_chat_tag_by_id(
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
376
377
378
379
    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)
380

381
382
383
384
385
386
387
        if tag:
            return tag
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
388
389
    else:
        raise HTTPException(
390
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
391
392
393
394
395
396
397
398
399
        )


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


@router.delete("/{id}/tags", response_model=Optional[bool])
400
async def delete_chat_tag_by_id(
401
402
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
403
    result = Tags.delete_tag_by_tag_name_and_chat_id_and_user_id(
404
405
406
        form_data.tag_name, id, user.id
    )

407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
    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
426
427
428
429
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )