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
    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
191
192
193
194
195
196
197
198
199
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,
        )

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

203

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


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

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


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

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


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


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

    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
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
############################
# 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
############################


@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)
    ]

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

    if len(chats) == 0:
        Tags.delete_tag_by_tag_name_and_user_id(tag_name, user.id)

    return chats


340
341
342
343
344
345
346
347
348
############################
# 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)

349
    if tags != None:
350
351
352
353
354
355
356
357
358
359
360
361
        return tags
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


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


362
@router.post("/{id}/tags", response_model=Optional[ChatIdTagModel])
363
364
365
async def add_chat_tag_by_id(
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
366
367
368
369
    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)
370

371
372
373
374
375
376
377
        if tag:
            return tag
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
378
379
    else:
        raise HTTPException(
380
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
381
382
383
384
385
386
387
388
389
        )


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


@router.delete("/{id}/tags", response_model=Optional[bool])
390
async def delete_chat_tag_by_id(
391
392
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
393
    result = Tags.delete_tag_by_tag_name_and_chat_id_and_user_id(
394
395
396
        form_data.tag_name, id, user.id
    )

397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
    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
416
417
418
419
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )