feat: user valves integration

This commit is contained in:
Timothy J. Baek
2024-06-22 12:08:32 -07:00
parent 15fc23df87
commit d362fd027e
6 changed files with 375 additions and 37 deletions

View File

@@ -8,6 +8,8 @@ from apps.webui.internal.db import DB, JSONField
from apps.webui.models.users import Users
import json
import copy
from config import SRC_LOG_LEVELS
@@ -121,14 +123,15 @@ class FunctionsTable:
) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump()
# Check if user has "functions" and "valves" settings
if "functions" not in user.settings:
user.settings["functions"] = {}
if "valves" not in user.settings["functions"]:
user.settings["functions"]["valves"] = {}
if "functions" not in user_settings:
user_settings["functions"] = {}
if "valves" not in user_settings["functions"]:
user_settings["functions"]["valves"] = {}
return user.settings["functions"]["valves"].get(id, {})
return user_settings["functions"]["valves"].get(id, {})
except Exception as e:
print(f"An error occurred: {e}")
return None
@@ -138,20 +141,21 @@ class FunctionsTable:
) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump()
# Check if user has "functions" and "valves" settings
if "functions" not in user.settings:
user.settings["functions"] = {}
if "valves" not in user.settings["functions"]:
user.settings["functions"]["valves"] = {}
if "functions" not in user_settings:
user_settings["functions"] = {}
if "valves" not in user_settings["functions"]:
user_settings["functions"]["valves"] = {}
user.settings["functions"]["valves"][id] = valves
user_settings["functions"]["valves"][id] = valves
# Update the user settings in the database
query = Users.update_user_by_id(user_id, {"settings": user.settings})
query = Users.update_user_by_id(user_id, {"settings": user_settings})
query.execute()
return user.settings["functions"]["valves"][id]
return user_settings["functions"]["valves"][id]
except Exception as e:
print(f"An error occurred: {e}")
return None

View File

@@ -8,6 +8,8 @@ from apps.webui.internal.db import DB, JSONField
from apps.webui.models.users import Users
import json
import copy
from config import SRC_LOG_LEVELS
@@ -112,14 +114,15 @@ class ToolsTable:
) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump()
# Check if user has "tools" and "valves" settings
if "tools" not in user.settings:
user.settings["tools"] = {}
if "valves" not in user.settings["tools"]:
user.settings["tools"]["valves"] = {}
if "tools" not in user_settings:
user_settings["tools"] = {}
if "valves" not in user_settings["tools"]:
user_settings["tools"]["valves"] = {}
return user.settings["tools"]["valves"].get(id, {})
return user_settings["tools"]["valves"].get(id, {})
except Exception as e:
print(f"An error occurred: {e}")
return None
@@ -129,20 +132,21 @@ class ToolsTable:
) -> Optional[dict]:
try:
user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump()
# Check if user has "tools" and "valves" settings
if "tools" not in user.settings:
user.settings["tools"] = {}
if "valves" not in user.settings["tools"]:
user.settings["tools"]["valves"] = {}
if "tools" not in user_settings:
user_settings["tools"] = {}
if "valves" not in user_settings["tools"]:
user_settings["tools"]["valves"] = {}
user.settings["tools"]["valves"][id] = valves
user_settings["tools"]["valves"][id] = valves
# Update the user settings in the database
query = Users.update_user_by_id(user_id, {"settings": user.settings})
query = Users.update_user_by_id(user_id, {"settings": user_settings})
query.execute()
return user.settings["tools"]["valves"][id]
return user_settings["tools"]["valves"][id]
except Exception as e:
print(f"An error occurred: {e}")
return None

View File

@@ -145,7 +145,7 @@ async def get_function_user_valves_by_id(id: str, user=Depends(get_verified_user
async def get_function_user_valves_spec_by_id(
request: Request, id: str, user=Depends(get_verified_user)
):
function = Functions.get_tool_by_id(id)
function = Functions.get_function_by_id(id)
if function:
if id in request.app.state.FUNCTIONS:
function_module = request.app.state.FUNCTIONS[id]
@@ -168,8 +168,7 @@ async def get_function_user_valves_spec_by_id(
async def update_function_user_valves_by_id(
request: Request, id: str, form_data: dict, user=Depends(get_verified_user)
):
function = Functions.get_tool_by_id(id)
function = Functions.get_function_by_id(id)
if function:
if id in request.app.state.FUNCTIONS:
@@ -211,7 +210,7 @@ async def update_function_user_valves_by_id(
@router.post("/id/{id}/update", response_model=Optional[FunctionModel])
async def update_toolkit_by_id(
async def update_function_by_id(
request: Request, id: str, form_data: FunctionForm, user=Depends(get_admin_user)
):
function_path = os.path.join(FUNCTIONS_DIR, f"{id}.py")