tools.py 11.1 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
from datetime import datetime, timedelta
from typing import List, Union, Optional

from fastapi import APIRouter
from pydantic import BaseModel
import json

Timothy J. Baek's avatar
Timothy J. Baek committed
9
from apps.webui.models.users import Users
Timothy J. Baek's avatar
Timothy J. Baek committed
10
from apps.webui.models.tools import Tools, ToolForm, ToolModel, ToolResponse
Timothy J. Baek's avatar
Timothy J. Baek committed
11
from apps.webui.utils import load_toolkit_module_by_id
Timothy J. Baek's avatar
Timothy J. Baek committed
12

Timothy J. Baek's avatar
Timothy J. Baek committed
13
from utils.utils import get_admin_user, get_verified_user
Timothy J. Baek's avatar
Timothy J. Baek committed
14
15
16
17
18
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
19
from pathlib import Path
Timothy J. Baek's avatar
Timothy J. Baek committed
20

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
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])
36
async def get_toolkits(user=Depends(get_verified_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
37
38
39
40
41
42
43
44
45
46
    toolkits = [toolkit for toolkit in Tools.get_tools()]
    return toolkits


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


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


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


@router.post("/create", response_model=Optional[ToolResponse])
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
58
async def create_new_toolkit(
59
60
61
    request: Request,
    form_data: ToolForm,
    user=Depends(get_admin_user),
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
62
):
Timothy J. Baek's avatar
Timothy J. Baek committed
63
64
65
66
67
68
    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
69
70
    form_data.id = form_data.id.lower()

71
    toolkit = Tools.get_tool_by_id(form_data.id)
Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
76
77
    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)

78
79
            toolkit_module, frontmatter = load_toolkit_module_by_id(form_data.id)
            form_data.meta.manifest = frontmatter
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
80
81

            TOOLS = request.app.state.TOOLS
Timothy J. Baek's avatar
Timothy J. Baek committed
82
83
84
            TOOLS[form_data.id] = toolkit_module

            specs = get_tools_specs(TOOLS[form_data.id])
85
            toolkit = Tools.insert_new_tool(user.id, form_data, specs)
Timothy J. Baek's avatar
Timothy J. Baek committed
86

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


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


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

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


128
129
130
131
132
133
134
############################
# UpdateToolkitById
############################


@router.post("/id/{id}/update", response_model=Optional[ToolModel])
async def update_toolkit_by_id(
135
136
137
138
    request: Request,
    id: str,
    form_data: ToolForm,
    user=Depends(get_admin_user),
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
):
    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)

        toolkit_module, frontmatter = load_toolkit_module_by_id(id)
        form_data.meta.manifest = frontmatter

        TOOLS = request.app.state.TOOLS
        TOOLS[id] = toolkit_module

        specs = get_tools_specs(TOOLS[id])

        updated = {
            **form_data.model_dump(exclude={"id"}),
            "specs": specs,
        }

        print(updated)
160
        toolkit = Tools.update_tool_by_id(id, updated)
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

        if toolkit:
            return toolkit
        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)
183
async def delete_toolkit_by_id(request: Request, id: str, user=Depends(get_admin_user)):
184
    result = Tools.delete_tool_by_id(id)
185
186
187
188
189
190
191
192
193
194
195
196
197

    if result:
        TOOLS = request.app.state.TOOLS
        if id in TOOLS:
            del TOOLS[id]

        # delete the toolkit file
        toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
        os.remove(toolkit_path)

    return result


198
199
200
201
202
203
204
205
206
207
############################
# GetToolValves
############################


@router.get("/id/{id}/valves", response_model=Optional[dict])
async def get_toolkit_valves_by_id(id: str, user=Depends(get_admin_user)):
    toolkit = Tools.get_tool_by_id(id)
    if toolkit:
        try:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
208
209
            valves = Tools.get_tool_valves_by_id(id)
            return valves
210
211
212
213
214
215
216
217
218
219
220
221
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
############################
# GetToolValvesSpec
############################


@router.get("/id/{id}/valves/spec", response_model=Optional[dict])
async def get_toolkit_valves_spec_by_id(
    request: Request, id: str, user=Depends(get_admin_user)
):
    toolkit = Tools.get_tool_by_id(id)
    if toolkit:
        if id in request.app.state.TOOLS:
            toolkit_module = request.app.state.TOOLS[id]
        else:
236
            toolkit_module, frontmatter = load_toolkit_module_by_id(id)
Timothy J. Baek's avatar
Timothy J. Baek committed
237
238
            request.app.state.TOOLS[id] = toolkit_module

Timothy J. Baek's avatar
Timothy J. Baek committed
239
240
241
        if hasattr(toolkit_module, "Valves"):
            Valves = toolkit_module.Valves
            return Valves.schema()
Timothy J. Baek's avatar
Timothy J. Baek committed
242
243
244
245
246
247
248
249
        return None
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


250
251
252
253
254
255
256
############################
# UpdateToolValves
############################


@router.post("/id/{id}/valves/update", response_model=Optional[dict])
async def update_toolkit_valves_by_id(
Timothy J. Baek's avatar
Timothy J. Baek committed
257
    request: Request, id: str, form_data: dict, user=Depends(get_admin_user)
258
259
260
):
    toolkit = Tools.get_tool_by_id(id)
    if toolkit:
Timothy J. Baek's avatar
Timothy J. Baek committed
261
262
263
        if id in request.app.state.TOOLS:
            toolkit_module = request.app.state.TOOLS[id]
        else:
264
            toolkit_module, frontmatter = load_toolkit_module_by_id(id)
Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
267
268
269
270
            request.app.state.TOOLS[id] = toolkit_module

        if hasattr(toolkit_module, "Valves"):
            Valves = toolkit_module.Valves

            try:
Timothy J. Baek's avatar
Timothy J. Baek committed
271
                form_data = {k: v for k, v in form_data.items() if v is not None}
Timothy J. Baek's avatar
Timothy J. Baek committed
272
273
274
275
276
277
278
279
280
281
                valves = Valves(**form_data)
                Tools.update_tool_valves_by_id(id, valves.model_dump())
                return valves.model_dump()
            except Exception as e:
                print(e)
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=ERROR_MESSAGES.DEFAULT(e),
                )
        else:
282
            raise HTTPException(
Timothy J. Baek's avatar
Timothy J. Baek committed
283
284
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
285
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
286

287
288
289
290
291
292
293
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
296
297
298
299
300
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
############################
# ToolUserValves
############################


@router.get("/id/{id}/valves/user", response_model=Optional[dict])
async def get_toolkit_user_valves_by_id(id: str, user=Depends(get_verified_user)):
    toolkit = Tools.get_tool_by_id(id)
    if toolkit:
        try:
            user_valves = Tools.get_user_valves_by_id_and_user_id(id, user.id)
            return user_valves
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


@router.get("/id/{id}/valves/user/spec", response_model=Optional[dict])
async def get_toolkit_user_valves_spec_by_id(
    request: Request, id: str, user=Depends(get_verified_user)
):
    toolkit = Tools.get_tool_by_id(id)
    if toolkit:
        if id in request.app.state.TOOLS:
            toolkit_module = request.app.state.TOOLS[id]
        else:
327
            toolkit_module, frontmatter = load_toolkit_module_by_id(id)
Timothy J. Baek's avatar
Timothy J. Baek committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
            request.app.state.TOOLS[id] = toolkit_module

        if hasattr(toolkit_module, "UserValves"):
            UserValves = toolkit_module.UserValves
            return UserValves.schema()
        return None
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


@router.post("/id/{id}/valves/user/update", response_model=Optional[dict])
async def update_toolkit_user_valves_by_id(
    request: Request, id: str, form_data: dict, user=Depends(get_verified_user)
):
    toolkit = Tools.get_tool_by_id(id)

    if toolkit:
        if id in request.app.state.TOOLS:
            toolkit_module = request.app.state.TOOLS[id]
        else:
351
            toolkit_module, frontmatter = load_toolkit_module_by_id(id)
Timothy J. Baek's avatar
Timothy J. Baek committed
352
353
354
355
356
357
            request.app.state.TOOLS[id] = toolkit_module

        if hasattr(toolkit_module, "UserValves"):
            UserValves = toolkit_module.UserValves

            try:
Timothy J. Baek's avatar
Timothy J. Baek committed
358
                form_data = {k: v for k, v in form_data.items() if v is not None}
Timothy J. Baek's avatar
Timothy J. Baek committed
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
                user_valves = UserValves(**form_data)
                Tools.update_user_valves_by_id_and_user_id(
                    id, user.id, user_valves.model_dump()
                )
                return user_valves.model_dump()
            except Exception as e:
                print(e)
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=ERROR_MESSAGES.DEFAULT(e),
                )
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )