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
fix  
Timothy J. Baek committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
############################
# GetSharedChatById
############################


@router.get("/share/{share_id}", response_model=Optional[ChatResponse])
async def get_shared_chat_by_id(share_id: str, user=Depends(get_current_user)):
    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)

    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
Timothy J. Baek committed
121
############################
Timothy J. Baek's avatar
Timothy J. Baek committed
122
# GetChats
Timothy J. Baek's avatar
Timothy J. Baek committed
123
124
125
126
############################


@router.get("/all", response_model=List[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
127
async def get_user_chats(user=Depends(get_current_user)):
Anuraag Jain's avatar
Anuraag Jain committed
128
    return [
Timothy J. Baek's avatar
Timothy J. Baek committed
129
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
130
        for chat in Chats.get_chats_by_user_id(user.id)
131
    ]
Timothy J. Baek's avatar
Timothy J. Baek committed
132
133


134
135
136
137
138
139
############################
# GetAllChatsInDB
############################


@router.get("/all/db", response_model=List[ChatResponse])
140
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
141
    if not ENABLE_ADMIN_EXPORT:
142
143
144
145
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
146
147
    return [
        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
148
        for chat in Chats.get_chats()
149
    ]
150
151


Timothy J. Baek's avatar
Timothy J. Baek committed
152
153
154
155
156
############################
# CreateNewChat
############################


157
@router.post("/new", response_model=Optional[ChatResponse])
158
async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
159
160
161
162
    try:
        chat = Chats.insert_new_chat(user.id, form_data)
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
    except Exception as e:
163
        log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
164
165
166
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
167
168


Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
############################
# GetChatsByTags
############################


class TagNameForm(BaseModel):
    name: str
    skip: Optional[int] = 0
    limit: Optional[int] = 50


@router.post("/tags", response_model=List[ChatTitleIdResponse])
async def get_user_chat_list_by_tag_name(
    form_data: TagNameForm, user=Depends(get_current_user)
):

    print(form_data)
    chat_ids = [
        chat_id_tag.chat_id
        for chat_id_tag in Tags.get_chat_ids_by_tag_name_and_user_id(
            form_data.name, user.id
        )
    ]

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

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

    return chats


############################
# 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()
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
218
219
220
221
222
############################
# GetChatById
############################


223
@router.get("/{id}", response_model=Optional[ChatResponse])
224
225
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
226

Anuraag Jain's avatar
Anuraag Jain committed
227
    if chat:
Timothy J. Baek's avatar
Timothy J. Baek committed
228
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Timothy J. Baek's avatar
Timothy J. Baek committed
229
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
230
231
232
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
233
234
235
236
237
238
239


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


240
@router.post("/{id}", response_model=Optional[ChatResponse])
Timothy J. Baek's avatar
Timothy J. Baek committed
241
242
243
async def update_chat_by_id(
    id: str, form_data: ChatForm, user=Depends(get_current_user)
):
244
    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
Anuraag Jain's avatar
Anuraag Jain committed
245
    if chat:
246
        updated_chat = {**json.loads(chat.chat), **form_data.chat}
Timothy J. Baek's avatar
Timothy J. Baek committed
247

248
        chat = Chats.update_chat_by_id(id, updated_chat)
Timothy J. Baek's avatar
Timothy J. Baek committed
249
        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
Anuraag Jain's avatar
Anuraag Jain committed
250
    else:
251
252
253
254
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
255
256
257
258
259
260
261
262


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


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

265
266
267
268
269
270
271
272
273
274
275
276
    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
277

278

Timothy J. Baek's avatar
Timothy J. Baek committed
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
############################
# 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()
        )


296
297
298
299
300
301
302
303
304
305
############################
# 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
306
            shared_chat = Chats.update_shared_chat_by_chat_id(chat.id)
307
308
309
310
            return ChatResponse(
                **{**shared_chat.model_dump(), "chat": json.loads(shared_chat.chat)}
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
311
        shared_chat = Chats.insert_shared_chat_by_chat_id(chat.id)
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        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
############################


333
334
335
@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)
336
337
338
339
    if chat:
        if not chat.share_id:
            return False

340
341
342
343
        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
344
345
346
347
348
349
350
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )


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

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


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


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

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


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


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

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