tools.py 5.07 KB
Newer Older
Timothy J. Baek's avatar
refac  
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
8
9
from datetime import datetime, timedelta
from typing import List, Union, Optional

from fastapi import APIRouter
from pydantic import BaseModel
import json

from apps.webui.models.tools import Tools, ToolForm, ToolModel, ToolResponse
Timothy J. Baek's avatar
Timothy J. Baek committed
10
from apps.webui.utils import load_toolkit_module_by_id
Timothy J. Baek's avatar
Timothy J. Baek committed
11
12
13
14
15
16
17

from utils.utils import get_current_user, get_admin_user
from utils.tools import get_tools_specs
from constants import ERROR_MESSAGES

from importlib import util
import os
Timothy J. Baek's avatar
Timothy J. Baek committed
18
from pathlib import Path
Timothy J. Baek's avatar
Timothy J. Baek committed
19

Timothy J. Baek's avatar
Timothy J. Baek committed
20
from config import DATA_DIR, CACHE_DIR
Timothy J. Baek's avatar
Timothy J. Baek committed
21

Timothy J. Baek's avatar
Timothy J. Baek committed
22

Timothy J. Baek's avatar
Timothy J. Baek committed
23
24
25
26
27
28
29
30
31
32
33
34
35
TOOLS_DIR = f"{DATA_DIR}/tools"
os.makedirs(TOOLS_DIR, exist_ok=True)


router = APIRouter()

############################
# GetToolkits
############################


@router.get("/", response_model=List[ToolResponse])
async def get_toolkits(user=Depends(get_current_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
36
37
38
39
40
41
42
43
44
45
    toolkits = [toolkit for toolkit in Tools.get_tools()]
    return toolkits


############################
# ExportToolKits
############################


@router.get("/export", response_model=List[ToolModel])
Timothy J. Baek's avatar
Timothy J. Baek committed
46
async def get_toolkits(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
47
    toolkits = [toolkit for toolkit in Tools.get_tools()]
Timothy J. Baek's avatar
Timothy J. Baek committed
48
49
50
51
52
53
54
55
56
    return toolkits


############################
# CreateNewToolKit
############################


@router.post("/create", response_model=Optional[ToolResponse])
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
57
58
59
async def create_new_toolkit(
    request: Request, form_data: ToolForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
Timothy J. Baek committed
60
61
62
63
64
65
    if not form_data.id.isidentifier():
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Only alphanumeric characters and underscores are allowed in the id",
        )

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
66
67
    form_data.id = form_data.id.lower()

Timothy J. Baek's avatar
Timothy J. Baek committed
68
69
70
71
72
73
74
    toolkit = Tools.get_tool_by_id(form_data.id)
    if toolkit == None:
        toolkit_path = os.path.join(TOOLS_DIR, f"{form_data.id}.py")
        try:
            with open(toolkit_path, "w") as tool_file:
                tool_file.write(form_data.content)

Timothy J. Baek's avatar
Timothy J. Baek committed
75
            toolkit_module = load_toolkit_module_by_id(form_data.id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
76
77

            TOOLS = request.app.state.TOOLS
Timothy J. Baek's avatar
Timothy J. Baek committed
78
79
80
81
82
            TOOLS[form_data.id] = toolkit_module

            specs = get_tools_specs(TOOLS[form_data.id])
            toolkit = Tools.insert_new_tool(user.id, form_data, specs)

Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
85
            tool_cache_dir = Path(CACHE_DIR) / "tools" / form_data.id
            tool_cache_dir.mkdir(parents=True, exist_ok=True)

Timothy J. Baek's avatar
Timothy J. Baek committed
86
            if toolkit:
Timothy J. Baek's avatar
Timothy J. Baek committed
87
                return toolkit
Timothy J. Baek's avatar
Timothy J. Baek committed
88
89
90
            else:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
91
                    detail=ERROR_MESSAGES.DEFAULT("Error creating toolkit"),
Timothy J. Baek's avatar
Timothy J. Baek committed
92
93
                )
        except Exception as e:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
94
            print(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
95
96
97
98
99
100
101
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
Timothy J. Baek's avatar
Timothy J. Baek committed
102
            detail=ERROR_MESSAGES.ID_TAKEN,
Timothy J. Baek's avatar
Timothy J. Baek committed
103
104
105
106
107
108
109
110
        )


############################
# GetToolkitById
############################


Timothy J. Baek's avatar
Timothy J. Baek committed
111
@router.get("/id/{id}", response_model=Optional[ToolModel])
Timothy J. Baek's avatar
Timothy J. Baek committed
112
113
114
115
async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)):
    toolkit = Tools.get_tool_by_id(id)

    if toolkit:
Timothy J. Baek's avatar
Timothy J. Baek committed
116
        return toolkit
Timothy J. Baek's avatar
Timothy J. Baek committed
117
118
119
120
121
122
123
124
125
126
127
128
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# UpdateToolkitById
############################


Timothy J. Baek's avatar
Timothy J. Baek committed
129
@router.post("/id/{id}/update", response_model=Optional[ToolModel])
Timothy J. Baek's avatar
Timothy J. Baek committed
130
async def update_toolkit_by_id(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
131
    request: Request, id: str, form_data: ToolForm, user=Depends(get_admin_user)
Timothy J. Baek's avatar
Timothy J. Baek committed
132
133
134
135
136
137
138
):
    toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")

    try:
        with open(toolkit_path, "w") as tool_file:
            tool_file.write(form_data.content)

Timothy J. Baek's avatar
Timothy J. Baek committed
139
        toolkit_module = load_toolkit_module_by_id(id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
140
141

        TOOLS = request.app.state.TOOLS
Timothy J. Baek's avatar
Timothy J. Baek committed
142
143
144
        TOOLS[id] = toolkit_module

        specs = get_tools_specs(TOOLS[id])
Timothy J. Baek's avatar
Timothy J. Baek committed
145
146

        updated = {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
147
            **form_data.model_dump(exclude={"id"}),
Timothy J. Baek's avatar
Timothy J. Baek committed
148
149
150
151
152
            "specs": specs,
        }

        print(updated)
        toolkit = Tools.update_tool_by_id(id, updated)
Timothy J. Baek's avatar
Timothy J. Baek committed
153
154

        if toolkit:
Timothy J. Baek's avatar
Timothy J. Baek committed
155
            return toolkit
Timothy J. Baek's avatar
Timothy J. Baek committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT("Error updating toolkit"),
            )

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


############################
# DeleteToolkitById
############################


@router.delete("/id/{id}/delete", response_model=bool)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
175
async def delete_toolkit_by_id(request: Request, id: str, user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
176
    result = Tools.delete_tool_by_id(id)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
177
178
179

    if result:
        TOOLS = request.app.state.TOOLS
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
180
181
        if id in TOOLS:
            del TOOLS[id]
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
182

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
183
184
185
186
        # delete the toolkit file
        toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
        os.remove(toolkit_path)

Timothy J. Baek's avatar
Timothy J. Baek committed
187
    return result