models.py 3.08 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
from fastapi import Depends, FastAPI, HTTPException, status, Request
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
4
5
6
7
from datetime import datetime, timedelta
from typing import List, Union, Optional

from fastapi import APIRouter
from pydantic import BaseModel
import json
8
9

from apps.webui.internal.db import get_db
10
from apps.webui.models.models import Models, ModelModel, ModelForm, ModelResponse
Timothy J. Baek's avatar
Timothy J. Baek committed
11
12
13
14
15
16
17

from utils.utils import get_verified_user, get_admin_user
from constants import ERROR_MESSAGES

router = APIRouter()

###########################
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
18
# getModels
Timothy J. Baek's avatar
Timothy J. Baek committed
19
20
21
22
###########################


@router.get("/", response_model=List[ModelResponse])
23
24
async def get_models(user=Depends(get_verified_user), db=Depends(get_db)):
    return Models.get_all_models(db)
Timothy J. Baek's avatar
Timothy J. Baek committed
25
26
27
28
29
30
31
32


############################
# AddNewModel
############################


@router.post("/add", response_model=Optional[ModelModel])
Timothy J. Baek's avatar
Timothy J. Baek committed
33
async def add_new_model(
34
35
36
37
    request: Request,
    form_data: ModelForm,
    user=Depends(get_admin_user),
    db=Depends(get_db),
Timothy J. Baek's avatar
Timothy J. Baek committed
38
39
):
    if form_data.id in request.app.state.MODELS:
Timothy J. Baek's avatar
Timothy J. Baek committed
40
41
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
Timothy J. Baek's avatar
Timothy J. Baek committed
42
            detail=ERROR_MESSAGES.MODEL_ID_TAKEN,
Timothy J. Baek's avatar
Timothy J. Baek committed
43
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
44
    else:
45
        model = Models.insert_new_model(db, form_data, user.id)
Timothy J. Baek's avatar
Timothy J. Baek committed
46
47
48
49
50
51
52
53

        if model:
            return model
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
54
55
56
57
58
59
60


############################
# GetModelById
############################


61
62
63
@router.get("/{id}", response_model=Optional[ModelModel])
async def get_model_by_id(id: str, user=Depends(get_verified_user), db=Depends(get_db)):
    model = Models.get_model_by_id(db, id)
Timothy J. Baek's avatar
Timothy J. Baek committed
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

    if model:
        return model
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# UpdateModelById
############################


Timothy J. Baek's avatar
Timothy J. Baek committed
79
@router.post("/update", response_model=Optional[ModelModel])
Timothy J. Baek's avatar
Timothy J. Baek committed
80
async def update_model_by_id(
81
82
83
84
85
    request: Request,
    id: str,
    form_data: ModelForm,
    user=Depends(get_admin_user),
    db=Depends(get_db),
Timothy J. Baek's avatar
Timothy J. Baek committed
86
):
87
    model = Models.get_model_by_id(db, id)
Timothy J. Baek's avatar
Timothy J. Baek committed
88
    if model:
89
        model = Models.update_model_by_id(db, id, form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
90
91
        return model
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
92
        if form_data.id in request.app.state.MODELS:
93
            model = Models.insert_new_model(db, form_data, user.id)
Timothy J. Baek's avatar
Timothy J. Baek committed
94
95
96
97
98
99
100
101
102
103
104
105
            if model:
                return model
            else:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=ERROR_MESSAGES.DEFAULT(),
                )
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
106
107
108
109
110
111
112


############################
# DeleteModelById
############################


Timothy J. Baek's avatar
Timothy J. Baek committed
113
@router.delete("/delete", response_model=bool)
114
115
async def delete_model_by_id(id: str, user=Depends(get_admin_user), db=Depends(get_db)):
    result = Models.delete_model_by_id(db, id)
Timothy J. Baek's avatar
Timothy J. Baek committed
116
    return result