Skip to content

Commit f949d17

Browse files
tjbcktjbck
committed
refac
1 parent 6d7744c commit f949d17

1 file changed

Lines changed: 99 additions & 52 deletions

File tree

‎backend/open_webui/routers/tools.py‎

Lines changed: 99 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -762,21 +762,37 @@ async def update_tools_valves_by_id(
762762
@router.get('/id/{id}/valves/user', response_model=Optional[dict])
763763
async def get_tools_user_valves_by_id(id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)):
764764
tools = Tools.get_tool_by_id(id, db=db)
765-
if tools:
766-
try:
767-
user_valves = Tools.get_user_valves_by_id_and_user_id(id, user.id, db=db)
768-
return user_valves
769-
except Exception as e:
770-
raise HTTPException(
771-
status_code=status.HTTP_400_BAD_REQUEST,
772-
detail=ERROR_MESSAGES.DEFAULT(str(e)),
773-
)
774-
else:
765+
if not tools:
775766
raise HTTPException(
776-
status_code=status.HTTP_401_UNAUTHORIZED,
767+
status_code=status.HTTP_404_NOT_FOUND,
777768
detail=ERROR_MESSAGES.NOT_FOUND,
778769
)
779770

771+
if (
772+
tools.user_id != user.id
773+
and not AccessGrants.has_access(
774+
user_id=user.id,
775+
resource_type='tool',
776+
resource_id=tools.id,
777+
permission='read',
778+
db=db,
779+
)
780+
and user.role != 'admin'
781+
):
782+
raise HTTPException(
783+
status_code=status.HTTP_401_UNAUTHORIZED,
784+
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
785+
)
786+
787+
try:
788+
user_valves = Tools.get_user_valves_by_id_and_user_id(id, user.id, db=db)
789+
return user_valves
790+
except Exception as e:
791+
raise HTTPException(
792+
status_code=status.HTTP_400_BAD_REQUEST,
793+
detail=ERROR_MESSAGES.DEFAULT(str(e)),
794+
)
795+
780796

781797
@router.get('/id/{id}/valves/user/spec', response_model=Optional[dict])
782798
async def get_tools_user_valves_spec_by_id(
@@ -786,26 +802,42 @@ async def get_tools_user_valves_spec_by_id(
786802
db: Session = Depends(get_session),
787803
):
788804
tools = Tools.get_tool_by_id(id, db=db)
789-
if tools:
790-
if id in request.app.state.TOOLS:
791-
tools_module = request.app.state.TOOLS[id]
792-
else:
793-
tools_module, _ = load_tool_module_by_id(id)
794-
request.app.state.TOOLS[id] = tools_module
795-
796-
if hasattr(tools_module, 'UserValves'):
797-
UserValves = tools_module.UserValves
798-
schema = UserValves.schema()
799-
# Resolve dynamic options for select dropdowns
800-
schema = resolve_valves_schema_options(UserValves, schema, user)
801-
return schema
802-
return None
803-
else:
805+
if not tools:
804806
raise HTTPException(
805-
status_code=status.HTTP_401_UNAUTHORIZED,
807+
status_code=status.HTTP_404_NOT_FOUND,
806808
detail=ERROR_MESSAGES.NOT_FOUND,
807809
)
808810

811+
if (
812+
tools.user_id != user.id
813+
and not AccessGrants.has_access(
814+
user_id=user.id,
815+
resource_type='tool',
816+
resource_id=tools.id,
817+
permission='read',
818+
db=db,
819+
)
820+
and user.role != 'admin'
821+
):
822+
raise HTTPException(
823+
status_code=status.HTTP_401_UNAUTHORIZED,
824+
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
825+
)
826+
827+
if id in request.app.state.TOOLS:
828+
tools_module = request.app.state.TOOLS[id]
829+
else:
830+
tools_module, _ = load_tool_module_by_id(id)
831+
request.app.state.TOOLS[id] = tools_module
832+
833+
if hasattr(tools_module, 'UserValves'):
834+
UserValves = tools_module.UserValves
835+
schema = UserValves.schema()
836+
# Resolve dynamic options for select dropdowns
837+
schema = resolve_valves_schema_options(UserValves, schema, user)
838+
return schema
839+
return None
840+
809841

810842
@router.post('/id/{id}/valves/user/update', response_model=Optional[dict])
811843
async def update_tools_user_valves_by_id(
@@ -816,33 +848,48 @@ async def update_tools_user_valves_by_id(
816848
db: Session = Depends(get_session),
817849
):
818850
tools = Tools.get_tool_by_id(id, db=db)
851+
if not tools:
852+
raise HTTPException(
853+
status_code=status.HTTP_404_NOT_FOUND,
854+
detail=ERROR_MESSAGES.NOT_FOUND,
855+
)
819856

820-
if tools:
821-
if id in request.app.state.TOOLS:
822-
tools_module = request.app.state.TOOLS[id]
823-
else:
824-
tools_module, _ = load_tool_module_by_id(id)
825-
request.app.state.TOOLS[id] = tools_module
826-
827-
if hasattr(tools_module, 'UserValves'):
828-
UserValves = tools_module.UserValves
829-
830-
try:
831-
form_data = {k: v for k, v in form_data.items() if v is not None}
832-
user_valves = UserValves(**form_data)
833-
user_valves_dict = user_valves.model_dump(exclude_unset=True)
834-
Tools.update_user_valves_by_id_and_user_id(id, user.id, user_valves_dict, db=db)
835-
return user_valves_dict
836-
except Exception as e:
837-
log.exception(f'Failed to update user valves by id {id}: {e}')
838-
raise HTTPException(
839-
status_code=status.HTTP_400_BAD_REQUEST,
840-
detail=ERROR_MESSAGES.DEFAULT(str(e)),
841-
)
842-
else:
857+
if (
858+
tools.user_id != user.id
859+
and not AccessGrants.has_access(
860+
user_id=user.id,
861+
resource_type='tool',
862+
resource_id=tools.id,
863+
permission='read',
864+
db=db,
865+
)
866+
and user.role != 'admin'
867+
):
868+
raise HTTPException(
869+
status_code=status.HTTP_401_UNAUTHORIZED,
870+
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
871+
)
872+
873+
if id in request.app.state.TOOLS:
874+
tools_module = request.app.state.TOOLS[id]
875+
else:
876+
tools_module, _ = load_tool_module_by_id(id)
877+
request.app.state.TOOLS[id] = tools_module
878+
879+
if hasattr(tools_module, 'UserValves'):
880+
UserValves = tools_module.UserValves
881+
882+
try:
883+
form_data = {k: v for k, v in form_data.items() if v is not None}
884+
user_valves = UserValves(**form_data)
885+
user_valves_dict = user_valves.model_dump(exclude_unset=True)
886+
Tools.update_user_valves_by_id_and_user_id(id, user.id, user_valves_dict, db=db)
887+
return user_valves_dict
888+
except Exception as e:
889+
log.exception(f'Failed to update user valves by id {id}: {e}')
843890
raise HTTPException(
844-
status_code=status.HTTP_401_UNAUTHORIZED,
845-
detail=ERROR_MESSAGES.NOT_FOUND,
891+
status_code=status.HTTP_400_BAD_REQUEST,
892+
detail=ERROR_MESSAGES.DEFAULT(str(e)),
846893
)
847894
else:
848895
raise HTTPException(

0 commit comments

Comments
 (0)