mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-16 03:47:49 +01:00
feat: tools backend
This commit is contained in:
Timothy J. Baek
61
backend/apps/webui/internal/migrations/012_add_tools.py
Normal file
61
backend/apps/webui/internal/migrations/012_add_tools.py
Normal file
@@ -0,0 +1,61 @@
|
||||
"""Peewee migrations -- 009_add_models.py.
|
||||
|
||||
Some examples (model - class or model name)::
|
||||
|
||||
> Model = migrator.orm['table_name'] # Return model in current state by name
|
||||
> Model = migrator.ModelClass # Return model in current state by name
|
||||
|
||||
> migrator.sql(sql) # Run custom SQL
|
||||
> migrator.run(func, *args, **kwargs) # Run python function with the given args
|
||||
> migrator.create_model(Model) # Create a model (could be used as decorator)
|
||||
> migrator.remove_model(model, cascade=True) # Remove a model
|
||||
> migrator.add_fields(model, **fields) # Add fields to a model
|
||||
> migrator.change_fields(model, **fields) # Change fields
|
||||
> migrator.remove_fields(model, *field_names, cascade=True)
|
||||
> migrator.rename_field(model, old_field_name, new_field_name)
|
||||
> migrator.rename_table(model, new_table_name)
|
||||
> migrator.add_index(model, *col_names, unique=False)
|
||||
> migrator.add_not_null(model, *field_names)
|
||||
> migrator.add_default(model, field_name, default)
|
||||
> migrator.add_constraint(model, name, sql)
|
||||
> migrator.drop_index(model, *col_names)
|
||||
> migrator.drop_not_null(model, *field_names)
|
||||
> migrator.drop_constraints(model, *constraints)
|
||||
|
||||
"""
|
||||
|
||||
from contextlib import suppress
|
||||
|
||||
import peewee as pw
|
||||
from peewee_migrate import Migrator
|
||||
|
||||
|
||||
with suppress(ImportError):
|
||||
import playhouse.postgres_ext as pw_pext
|
||||
|
||||
|
||||
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your migrations here."""
|
||||
|
||||
@migrator.create_model
|
||||
class Tool(pw.Model):
|
||||
id = pw.TextField(unique=True)
|
||||
user_id = pw.TextField()
|
||||
|
||||
name = pw.TextField()
|
||||
content = pw.TextField()
|
||||
specs = pw.TextField()
|
||||
|
||||
meta = pw.TextField()
|
||||
|
||||
created_at = pw.BigIntegerField(null=False)
|
||||
updated_at = pw.BigIntegerField(null=False)
|
||||
|
||||
class Meta:
|
||||
table_name = "tool"
|
||||
|
||||
|
||||
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your rollback migrations here."""
|
||||
|
||||
migrator.remove_model("tool")
|
||||
@@ -6,6 +6,7 @@ from apps.webui.routers import (
|
||||
users,
|
||||
chats,
|
||||
documents,
|
||||
tools,
|
||||
models,
|
||||
prompts,
|
||||
configs,
|
||||
@@ -26,8 +27,8 @@ from config import (
|
||||
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
|
||||
JWT_EXPIRES_IN,
|
||||
WEBUI_BANNERS,
|
||||
AppConfig,
|
||||
ENABLE_COMMUNITY_SHARING,
|
||||
AppConfig,
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
@@ -70,6 +71,7 @@ app.include_router(users.router, prefix="/users", tags=["users"])
|
||||
app.include_router(chats.router, prefix="/chats", tags=["chats"])
|
||||
|
||||
app.include_router(documents.router, prefix="/documents", tags=["documents"])
|
||||
app.include_router(tools.router, prefix="/tools", tags=["tools"])
|
||||
app.include_router(models.router, prefix="/models", tags=["models"])
|
||||
app.include_router(prompts.router, prefix="/prompts", tags=["prompts"])
|
||||
app.include_router(memories.router, prefix="/memories", tags=["memories"])
|
||||
|
||||
131
backend/apps/webui/models/tools.py
Normal file
131
backend/apps/webui/models/tools.py
Normal file
@@ -0,0 +1,131 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
import logging
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
|
||||
import json
|
||||
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
|
||||
####################
|
||||
# Tools DB Schema
|
||||
####################
|
||||
|
||||
|
||||
class Tool(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
name = TextField()
|
||||
content = TextField()
|
||||
specs = JSONField()
|
||||
meta = JSONField()
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
|
||||
|
||||
class ToolMeta(BaseModel):
|
||||
description: Optional[str] = None
|
||||
|
||||
|
||||
class ToolModel(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
name: str
|
||||
content: str
|
||||
specs: dict
|
||||
meta: ToolMeta
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
####################
|
||||
|
||||
|
||||
class ToolResponse(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
name: str
|
||||
meta: ToolMeta
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
|
||||
class ToolForm(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
content: str
|
||||
meta: ToolMeta
|
||||
|
||||
|
||||
class ToolsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Tool])
|
||||
|
||||
def insert_new_tool(
|
||||
self, user_id: str, form_data: ToolForm, specs: dict
|
||||
) -> Optional[ToolModel]:
|
||||
tool = ToolModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"specs": specs,
|
||||
"user_id": user_id,
|
||||
"updated_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Tool.create(**tool.model_dump())
|
||||
if result:
|
||||
return tool
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
|
||||
try:
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_tools(self) -> List[ToolModel]:
|
||||
return [ToolModel(**model_to_dict(tool)) for tool in Tool.select()]
|
||||
|
||||
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
|
||||
try:
|
||||
query = Tool.update(
|
||||
**updated,
|
||||
updated_at=int(time.time()),
|
||||
).where(Tool.id == id)
|
||||
query.execute()
|
||||
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_tool_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Tool.delete().where((Tool.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Tools = ToolsTable(DB)
|
||||
162
backend/apps/webui/routers/tools.py
Normal file
162
backend/apps/webui/routers/tools.py
Normal file
@@ -0,0 +1,162 @@
|
||||
from fastapi import Depends, FastAPI, HTTPException, status
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
from config import DATA_DIR
|
||||
|
||||
TOOLS_DIR = f"{DATA_DIR}/tools"
|
||||
os.makedirs(TOOLS_DIR, exist_ok=True)
|
||||
|
||||
TOOLS = {}
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def load_toolkit_module_from_path(tools_id, tools_path):
|
||||
spec = util.spec_from_file_location(tools_id, tools_path)
|
||||
module = util.module_from_spec(spec)
|
||||
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
print(f"Loaded module: {module.__name__}")
|
||||
if hasattr(module, "Tools"):
|
||||
return module.Tools()
|
||||
else:
|
||||
raise Exception("No Tools class found")
|
||||
except Exception as e:
|
||||
print(f"Error loading module: {tools_id}")
|
||||
|
||||
# Move the file to the error folder
|
||||
os.rename(tools_path, f"{tools_path}.error")
|
||||
raise e
|
||||
|
||||
|
||||
############################
|
||||
# GetToolkits
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ToolResponse])
|
||||
async def get_toolkits(user=Depends(get_current_user)):
|
||||
toolkits = [ToolResponse(**toolkit) for toolkit in Tools.get_tools()]
|
||||
return toolkits
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewToolKit
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[ToolResponse])
|
||||
async def create_new_toolkit(form_data: ToolForm, user=Depends(get_admin_user)):
|
||||
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)
|
||||
|
||||
toolkit_module = load_toolkit_module_from_path(form_data.id, toolkit_path)
|
||||
TOOLS[form_data.id] = toolkit_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[form_data.id])
|
||||
toolkit = Tools.insert_new_tool(user.id, form_data, specs)
|
||||
|
||||
if toolkit:
|
||||
return ToolResponse(**toolkit)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.FILE_EXISTS,
|
||||
)
|
||||
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_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NAME_TAG_TAKEN,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetToolkitById
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/id/{id}", response_model=Optional[ToolResponse])
|
||||
async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)):
|
||||
toolkit = Tools.get_tool_by_id(id)
|
||||
|
||||
if toolkit:
|
||||
return ToolResponse(**toolkit)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateToolkitById
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[ToolResponse])
|
||||
async def update_toolkit_by_id(
|
||||
id: str, form_data: ToolForm, user=Depends(get_admin_user)
|
||||
):
|
||||
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 = load_toolkit_module_from_path(id, toolkit_path)
|
||||
TOOLS[id] = toolkit_module
|
||||
|
||||
specs = get_tools_specs(TOOLS[id])
|
||||
toolkit = Tools.update_tool_by_id(
|
||||
id, {**form_data.model_dump(), "specs": specs}
|
||||
)
|
||||
|
||||
if toolkit:
|
||||
return ToolResponse(**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)
|
||||
async def delete_toolkit_by_id(id: str, user=Depends(get_admin_user)):
|
||||
result = Tools.delete_tool_by_id(id)
|
||||
return result
|
||||
Reference in New Issue
Block a user