chats.py 11.1 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
39
40
41
42
43
router = APIRouter()

############################
# GetChats
############################


@router.get("/", response_model=List[ChatTitleIdResponse])
44
async def get_user_chats(
Timothy J. Baek's avatar
Timothy J. Baek committed
45
46
    user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
47
    return Chats.get_chat_lists_by_user_id(user.id, skip, limit)
Timothy J. Baek's avatar
Timothy J. Baek committed
48
49


50
51
52
53
54
55
56
57
58
59
60
61
############################
# GetArchivedChats
############################


@router.get("/archived", response_model=List[ChatTitleIdResponse])
async def get_archived_user_chats(
    user=Depends(get_current_user), skip: int = 0, limit: int = 50
):
    return Chats.get_archived_chat_lists_by_user_id(user.id, skip, limit)


Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
65
66
67
############################
# GetAllChats
############################


@router.get("/all", response_model=List[ChatResponse])
68
async def get_all_user_chats(user=Depends(get_current_user)):
Anuraag Jain's avatar
Anuraag Jain committed
69
    return [
Timothy J. Baek's avatar
Timothy J. Baek committed
70
71
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
        for chat in Chats.get_all_chats_by_user_id(user.id)
72
    ]
Timothy J. Baek's avatar
Timothy J. Baek committed
73
74


75
76
77
78
79
80
############################
# GetAllChatsInDB
############################


@router.get("/all/db", response_model=List[ChatResponse])
81
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
82
    if not ENABLE_ADMIN_EXPORT:
83
84
85
86
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
87
88
89
90
    return [
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
        for chat in Chats.get_all_chats()
    ]
91
92


Timothy J. Baek's avatar
Timothy J. Baek committed
93
94
95
96
97
############################
# CreateNewChat
############################


98
@router.post("/new", response_model=Optional[ChatResponse])
99
async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102
103
    try:
        chat = Chats.insert_new_chat(user.id, form_data)
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    except Exception as e:
104
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
107
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
108
109


Timothy J. Baek's avatar
Timothy J. Baek committed
110
111
112
113
114
115
116
117
118
119
120
############################
# 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:
121
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
        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
141
    chats = Chats.get_chat_lists_by_chat_ids(chat_ids, skip, limit)
Timothy J. Baek's avatar
Timothy J. Baek committed
142

Timothy J. Baek's avatar
Timothy J. Baek committed
143
144
145
146
    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
147
148


Timothy J. Baek's avatar
Timothy J. Baek committed
149
150
151
152
153
############################
# GetChatById
############################


154
@router.get("/{id}", response_model=Optional[ChatResponse])
155
156
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
157

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


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


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

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


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


@router.delete("/{id}", response_model=bool)
Timothy J. Baek's avatar
Timothy J. Baek committed
194
195
196
197
198
199
200
201
202
203
204
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,
        )

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

208

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


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

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


263
264
265
@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)
266
267
268
269
    if chat:
        if not chat.share_id:
            return False

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


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


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

    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
        )


306
307
308
309
310
311
312
313
314
############################
# 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)

315
    if tags != None:
316
317
318
319
320
321
322
323
324
325
326
327
        return tags
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


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


328
@router.post("/{id}/tags", response_model=Optional[ChatIdTagModel])
329
330
331
async def add_chat_tag_by_id(
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
332
333
334
335
    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)
336

337
338
339
340
341
342
343
        if tag:
            return tag
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
344
345
    else:
        raise HTTPException(
346
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()
347
348
349
350
351
352
353
354
355
        )


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


@router.delete("/{id}/tags", response_model=Optional[bool])
356
async def delete_chat_tag_by_id(
357
358
    id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
):
359
    result = Tags.delete_tag_by_tag_name_and_chat_id_and_user_id(
360
361
362
        form_data.tag_name, id, user.id
    )

363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
    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
382
383
384
385
386
387
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
388
389
390
391
392
393
############################
# DeleteAllChats
############################


@router.delete("/", response_model=bool)
394
395
396
397
398
399
400
401
402
403
404
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,
        )

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