refac
This commit is contained in:
@@ -2,7 +2,7 @@ import logging
|
|||||||
import time
|
import time
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session, defer
|
||||||
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
|
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
|
||||||
from open_webui.models.users import Users, UserResponse
|
from open_webui.models.users import Users, UserResponse
|
||||||
from open_webui.models.groups import Groups
|
from open_webui.models.groups import Groups
|
||||||
@@ -156,9 +156,14 @@ class ToolsTable:
|
|||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def get_tools(self, db: Optional[Session] = None) -> list[ToolUserModel]:
|
def get_tools(
|
||||||
|
self, defer_content: bool = False, db: Optional[Session] = None
|
||||||
|
) -> list[ToolUserModel]:
|
||||||
with get_db_context(db) as db:
|
with get_db_context(db) as db:
|
||||||
all_tools = db.query(Tool).order_by(Tool.updated_at.desc()).all()
|
query = db.query(Tool).order_by(Tool.updated_at.desc())
|
||||||
|
if defer_content:
|
||||||
|
query = query.options(defer(Tool.content), defer(Tool.specs))
|
||||||
|
all_tools = query.all()
|
||||||
|
|
||||||
user_ids = list(set(tool.user_id for tool in all_tools))
|
user_ids = list(set(tool.user_id for tool in all_tools))
|
||||||
tool_ids = [tool.id for tool in all_tools]
|
tool_ids = [tool.id for tool in all_tools]
|
||||||
@@ -185,9 +190,9 @@ class ToolsTable:
|
|||||||
return tools
|
return tools
|
||||||
|
|
||||||
def get_tools_by_user_id(
|
def get_tools_by_user_id(
|
||||||
self, user_id: str, permission: str = "write", db: Optional[Session] = None
|
self, user_id: str, permission: str = "write", defer_content: bool = False, db: Optional[Session] = None
|
||||||
) -> list[ToolUserModel]:
|
) -> list[ToolUserModel]:
|
||||||
tools = self.get_tools(db=db)
|
tools = self.get_tools(defer_content=defer_content, db=db)
|
||||||
user_group_ids = {
|
user_group_ids = {
|
||||||
group.id for group in Groups.get_groups_by_member_id(user_id, db=db)
|
group.id for group in Groups.get_groups_by_member_id(user_id, db=db)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -64,13 +64,13 @@ async def get_tools(
|
|||||||
tools = []
|
tools = []
|
||||||
|
|
||||||
# Local Tools
|
# Local Tools
|
||||||
for tool in Tools.get_tools(db=db):
|
for tool in Tools.get_tools(defer_content=True, db=db):
|
||||||
tool_module = get_tool_module(request, tool.id)
|
tool_module = request.app.state.TOOLS.get(tool.id) if hasattr(request.app.state, 'TOOLS') else None
|
||||||
tools.append(
|
tools.append(
|
||||||
ToolUserResponse(
|
ToolUserResponse(
|
||||||
**{
|
**{
|
||||||
**tool.model_dump(),
|
**tool.model_dump(),
|
||||||
"has_user_valves": hasattr(tool_module, "UserValves"),
|
"has_user_valves": hasattr(tool_module, "UserValves") if tool_module else False,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -196,27 +196,35 @@ async def get_tool_list(
|
|||||||
user=Depends(get_verified_user), db: Session = Depends(get_session)
|
user=Depends(get_verified_user), db: Session = Depends(get_session)
|
||||||
):
|
):
|
||||||
if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
|
if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
|
||||||
tools = Tools.get_tools(db=db)
|
tools = Tools.get_tools(defer_content=True, db=db)
|
||||||
else:
|
else:
|
||||||
tools = Tools.get_tools_by_user_id(user.id, "read", db=db)
|
tools = Tools.get_tools_by_user_id(user.id, "read", defer_content=True, db=db)
|
||||||
|
|
||||||
return [
|
user_group_ids = {
|
||||||
ToolAccessResponse(
|
group.id for group in Groups.get_groups_by_member_id(user.id, db=db)
|
||||||
**tool.model_dump(),
|
}
|
||||||
write_access=(
|
|
||||||
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
result = []
|
||||||
or user.id == tool.user_id
|
for tool in tools:
|
||||||
or AccessGrants.has_access(
|
has_write = (
|
||||||
user_id=user.id,
|
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||||
resource_type="tool",
|
or user.id == tool.user_id
|
||||||
resource_id=tool.id,
|
or any(
|
||||||
permission="write",
|
g.permission == "write"
|
||||||
db=db,
|
and (
|
||||||
|
(g.principal_type == "user" and (g.principal_id == user.id or g.principal_id == "*"))
|
||||||
|
or (g.principal_type == "group" and g.principal_id in user_group_ids)
|
||||||
)
|
)
|
||||||
),
|
for g in tool.access_grants
|
||||||
|
)
|
||||||
)
|
)
|
||||||
for tool in tools
|
result.append(
|
||||||
]
|
ToolAccessResponse(
|
||||||
|
**tool.model_dump(),
|
||||||
|
write_access=has_write,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
############################
|
############################
|
||||||
|
|||||||
Reference in New Issue
Block a user