Spaces:
Paused
Paused
| import logging | |
| import time | |
| from typing import Optional | |
| from open_webui.internal.db import Base, JSONField, get_db | |
| from open_webui.models.users import Users, UserResponse | |
| from open_webui.env import SRC_LOG_LEVELS | |
| from pydantic import BaseModel, ConfigDict | |
| from sqlalchemy import BigInteger, Column, String, Text, JSON | |
| from open_webui.utils.access_control import has_access | |
| log = logging.getLogger(__name__) | |
| log.setLevel(SRC_LOG_LEVELS["MODELS"]) | |
| #################### | |
| # Tools DB Schema | |
| #################### | |
| class Tool(Base): | |
| __tablename__ = "tool" | |
| id = Column(String, primary_key=True) | |
| user_id = Column(String) | |
| name = Column(Text) | |
| content = Column(Text) | |
| specs = Column(JSONField) | |
| meta = Column(JSONField) | |
| valves = Column(JSONField) | |
| access_control = Column(JSON, nullable=True) # Controls data access levels. | |
| # Defines access control rules for this entry. | |
| # - `None`: Public access, available to all users with the "user" role. | |
| # - `{}`: Private access, restricted exclusively to the owner. | |
| # - Custom permissions: Specific access control for reading and writing; | |
| # Can specify group or user-level restrictions: | |
| # { | |
| # "read": { | |
| # "group_ids": ["group_id1", "group_id2"], | |
| # "user_ids": ["user_id1", "user_id2"] | |
| # }, | |
| # "write": { | |
| # "group_ids": ["group_id1", "group_id2"], | |
| # "user_ids": ["user_id1", "user_id2"] | |
| # } | |
| # } | |
| updated_at = Column(BigInteger) | |
| created_at = Column(BigInteger) | |
| class ToolMeta(BaseModel): | |
| description: Optional[str] = None | |
| manifest: Optional[dict] = {} | |
| class ToolModel(BaseModel): | |
| id: str | |
| user_id: str | |
| name: str | |
| content: str | |
| specs: list[dict] | |
| meta: ToolMeta | |
| access_control: Optional[dict] = None | |
| updated_at: int # timestamp in epoch | |
| created_at: int # timestamp in epoch | |
| model_config = ConfigDict(from_attributes=True) | |
| #################### | |
| # Forms | |
| #################### | |
| class ToolUserModel(ToolModel): | |
| user: Optional[UserResponse] = None | |
| class ToolResponse(BaseModel): | |
| id: str | |
| user_id: str | |
| name: str | |
| meta: ToolMeta | |
| access_control: Optional[dict] = None | |
| updated_at: int # timestamp in epoch | |
| created_at: int # timestamp in epoch | |
| class ToolUserResponse(ToolResponse): | |
| user: Optional[UserResponse] = None | |
| class ToolForm(BaseModel): | |
| id: str | |
| name: str | |
| content: str | |
| meta: ToolMeta | |
| access_control: Optional[dict] = None | |
| class ToolValves(BaseModel): | |
| valves: Optional[dict] = None | |
| class ToolsTable: | |
| def insert_new_tool( | |
| self, user_id: str, form_data: ToolForm, specs: list[dict] | |
| ) -> Optional[ToolModel]: | |
| with get_db() as db: | |
| 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(**tool.model_dump()) | |
| db.add(result) | |
| db.commit() | |
| db.refresh(result) | |
| if result: | |
| return ToolModel.model_validate(result) | |
| else: | |
| return None | |
| except Exception as e: | |
| print(f"Error creating tool: {e}") | |
| return None | |
| def get_tool_by_id(self, id: str) -> Optional[ToolModel]: | |
| try: | |
| with get_db() as db: | |
| tool = db.get(Tool, id) | |
| return ToolModel.model_validate(tool) | |
| except Exception: | |
| return None | |
| def get_tools(self) -> list[ToolUserModel]: | |
| with get_db() as db: | |
| tools = [] | |
| for tool in db.query(Tool).order_by(Tool.updated_at.desc()).all(): | |
| user = Users.get_user_by_id(tool.user_id) | |
| tools.append( | |
| ToolUserModel.model_validate( | |
| { | |
| **ToolModel.model_validate(tool).model_dump(), | |
| "user": user.model_dump() if user else None, | |
| } | |
| ) | |
| ) | |
| return tools | |
| def get_tools_by_user_id( | |
| self, user_id: str, permission: str = "write" | |
| ) -> list[ToolUserModel]: | |
| tools = self.get_tools() | |
| return [ | |
| tool | |
| for tool in tools | |
| if tool.user_id == user_id | |
| or has_access(user_id, permission, tool.access_control) | |
| ] | |
| def get_tool_valves_by_id(self, id: str) -> Optional[dict]: | |
| try: | |
| with get_db() as db: | |
| tool = db.get(Tool, id) | |
| return tool.valves if tool.valves else {} | |
| except Exception as e: | |
| print(f"An error occurred: {e}") | |
| return None | |
| def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]: | |
| try: | |
| with get_db() as db: | |
| db.query(Tool).filter_by(id=id).update( | |
| {"valves": valves, "updated_at": int(time.time())} | |
| ) | |
| db.commit() | |
| return self.get_tool_by_id(id) | |
| except Exception: | |
| return None | |
| def get_user_valves_by_id_and_user_id( | |
| self, id: str, user_id: str | |
| ) -> Optional[dict]: | |
| try: | |
| user = Users.get_user_by_id(user_id) | |
| user_settings = user.settings.model_dump() if user.settings else {} | |
| # 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"] = {} | |
| return user_settings["tools"]["valves"].get(id, {}) | |
| except Exception as e: | |
| print(f"An error occurred: {e}") | |
| return None | |
| def update_user_valves_by_id_and_user_id( | |
| self, id: str, user_id: str, valves: dict | |
| ) -> Optional[dict]: | |
| try: | |
| user = Users.get_user_by_id(user_id) | |
| user_settings = user.settings.model_dump() if user.settings else {} | |
| # 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"] = {} | |
| user_settings["tools"]["valves"][id] = valves | |
| # Update the user settings in the database | |
| Users.update_user_by_id(user_id, {"settings": user_settings}) | |
| return user_settings["tools"]["valves"][id] | |
| except Exception as e: | |
| print(f"An error occurred: {e}") | |
| return None | |
| def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]: | |
| try: | |
| with get_db() as db: | |
| db.query(Tool).filter_by(id=id).update( | |
| {**updated, "updated_at": int(time.time())} | |
| ) | |
| db.commit() | |
| tool = db.query(Tool).get(id) | |
| db.refresh(tool) | |
| return ToolModel.model_validate(tool) | |
| except Exception: | |
| return None | |
| def delete_tool_by_id(self, id: str) -> bool: | |
| try: | |
| with get_db() as db: | |
| db.query(Tool).filter_by(id=id).delete() | |
| db.commit() | |
| return True | |
| except Exception: | |
| return False | |
| Tools = ToolsTable() | |