Coverage for open_webui/models/tools.py: 42%
186 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 05:07 +0000
1"""Tool models, forms, and database operations."""
3from __future__ import annotations
5import logging
6import time
8# local imports
9from open_webui.internal.db import Base, JSONField, get_async_db_context
10from open_webui.models.access_grants import AccessGrantModel, AccessGrants
11from open_webui.models.groups import Groups
12from open_webui.models.users import UserResponse, Users
13from open_webui.utils.valves import decrypt_valves, encrypt_valves
14from pydantic import BaseModel, ConfigDict, Field
15from sqlalchemy import BigInteger, Column, String, Text, delete, select, update
16from sqlalchemy.ext.asyncio import AsyncSession
18log = logging.getLogger(__name__)
21class Tool(Base): # database table definition
22 __tablename__ = 'tool'
24 id = Column(String, primary_key=True, unique=True)
25 user_id = Column(String, index=True) # owner user id
26 name = Column(Text) # human-readable label
27 content = Column(Text) # Python source code
28 specs = Column(JSONField) # OpenAPI-style function specs
29 meta = Column(JSONField) # description, manifest, etc.
30 valves = Column(JSONField) # admin-configurable runtime parameters
32 updated_at = Column(BigInteger, nullable=False) # modification timestamp
33 created_at = Column(BigInteger, index=True) # creation timestamp
36class ToolMeta(BaseModel):
37 i18n: dict[str, dict[str, str]] | None = None
38 description: str | None = None
39 manifest: dict | None = {}
40 has_user_valves: bool = False
43class ToolModel(BaseModel):
44 id: str
45 user_id: str | None = None # may be null for legacy/malformed records
46 name: str
47 # None when listed with defer_content=True (source skipped for listings)
48 content: str | None = None
49 specs: list[dict]
50 meta: ToolMeta
51 access_grants: list[AccessGrantModel] = Field(default_factory=list)
53 updated_at: int # timestamp in epoch
54 created_at: int # timestamp in epoch
56 model_config = ConfigDict(from_attributes=True) # enables ORM mapping
59# --- tool request forms ---
60# Forms
61####################
64class ToolUserModel(ToolModel):
65 user: UserResponse | None = None
68class ToolResponse(BaseModel):
69 id: str
70 user_id: str | None = None # may be null for legacy/malformed records
71 name: str
72 meta: ToolMeta
73 access_grants: list[AccessGrantModel] = Field(default_factory=list)
74 updated_at: int # timestamp in epoch
75 created_at: int # timestamp in epoch
78class ToolUserResponse(ToolResponse):
79 user: UserResponse | None = None
81 model_config = ConfigDict(extra='allow')
84class ToolAccessResponse(ToolUserResponse):
85 write_access: bool | None = False
88class ToolForm(BaseModel):
89 id: str
90 name: str
91 content: str
92 meta: ToolMeta
93 access_grants: list[dict] | None = None
96class ToolValves(BaseModel):
97 valves: dict | None = None
100class ToolsTable:
101 async def _get_access_grants(self, tool_id: str, db: AsyncSession | None = None) -> list[AccessGrantModel]:
102 return await AccessGrants.get_grants_by_resource('tool', tool_id, db=db)
104 async def _to_tool_model(
105 self,
106 tool: Tool,
107 access_grants: list[AccessGrantModel] | None = None,
108 db: AsyncSession | None = None,
109 ) -> ToolModel:
110 tool_model = ToolModel.model_validate(tool)
111 tool_model.access_grants = (
112 access_grants if access_grants is not None else await self._get_access_grants(tool_model.id, db=db)
113 )
114 return tool_model
116 async def insert_new_tool(
117 self,
118 user_id: str,
119 form_data: ToolForm,
120 specs: list[dict],
121 db: AsyncSession | None = None,
122 ) -> ToolModel | None:
123 async with get_async_db_context(db) as db:
124 try:
125 result = Tool(
126 **{
127 **form_data.model_dump(exclude={'access_grants'}),
128 'specs': specs,
129 'user_id': user_id,
130 'updated_at': int(time.time()),
131 'created_at': int(time.time()),
132 }
133 )
134 db.add(result)
135 await db.commit()
136 await AccessGrants.set_access_grants('tool', result.id, form_data.access_grants, db=db)
137 if result:
138 return await self._to_tool_model(result, db=db)
139 else:
140 return None
141 except Exception as e:
142 log.exception(f'Error creating a new tool: {e}')
143 return None # creation failed
145 async def get_tool_by_id(
146 self,
147 id: str,
148 db: AsyncSession | None = None,
149 ) -> ToolModel | None:
150 """Fetch a single tool by primary key, including access grants."""
151 try: # single PK lookup + access grants
152 async with get_async_db_context(db) as session:
153 tool = await session.get(Tool, id)
154 if not tool:
155 return None
156 return await self._to_tool_model(tool, db=session)
157 except Exception:
158 return None
160 async def get_tools_by_ids(self, tool_ids: list[str], db: AsyncSession | None = None) -> dict[str, ToolModel]:
161 """Batch-fetch multiple tools by ID, returning a dict keyed by tool ID."""
162 if not tool_ids:
163 return {}
164 async with get_async_db_context(db) as db:
165 result = await db.execute(select(Tool).where(Tool.id.in_(tool_ids)))
166 tools = result.scalars().all()
167 grants_map = await AccessGrants.get_grants_by_resources('tool', [tool.id for tool in tools], db=db)
168 return {
169 tool.id: await self._to_tool_model(tool, access_grants=grants_map.get(tool.id, []), db=db)
170 for tool in tools
171 }
173 async def get_tools(
174 self,
175 defer_content: bool = False,
176 db: AsyncSession | None = None,
177 user_id: str | None = None,
178 user_group_ids: set[str] | None = None,
179 permission: str = 'read',
180 ) -> list[ToolUserModel]:
181 async with get_async_db_context(db) as db:
182 # Skip Tool.content (plugin source, potentially large) via a
183 # column select; Row attributes satisfy from_attributes.
184 stmt = (
185 select(Tool.id, Tool.user_id, Tool.name, Tool.specs, Tool.meta, Tool.updated_at, Tool.created_at)
186 if defer_content
187 else select(Tool)
188 ).order_by(Tool.updated_at.desc())
190 if user_id is not None: 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 if user_group_ids is None:
192 user_group_ids = {group.id for group in await Groups.get_groups_by_member_id(user_id, db=db)}
193 stmt = AccessGrants.has_permission_filter(
194 db=db,
195 query=stmt,
196 DocumentModel=Tool,
197 filter={'user_id': user_id, 'group_ids': user_group_ids},
198 resource_type='tool',
199 permission=permission,
200 )
202 result = await db.execute(stmt)
203 all_tools = result.all() if defer_content else result.scalars().all()
205 user_ids = list(set(tool.user_id for tool in all_tools))
206 tool_ids = [tool.id for tool in all_tools]
208 users = await Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
209 users_dict = {user.id: user for user in users}
210 grants_map = await AccessGrants.get_grants_by_resources('tool', tool_ids, db=db)
212 tools = []
213 for tool in all_tools:
214 user = users_dict.get(tool.user_id)
215 tools.append(
216 ToolUserModel.model_validate(
217 {
218 **(
219 await self._to_tool_model(
220 tool,
221 access_grants=grants_map.get(tool.id, []),
222 db=db,
223 )
224 ).model_dump(),
225 'user': user.model_dump() if user else None,
226 }
227 )
228 )
229 return tools
231 async def get_tools_by_user_id(
232 self,
233 user_id: str,
234 permission: str = 'write',
235 defer_content: bool = False,
236 db: AsyncSession | None = None,
237 ) -> list[ToolUserModel]:
238 user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
239 user_group_ids = {group.id for group in user_groups}
240 return await self.get_tools(
241 defer_content=defer_content,
242 db=db,
243 user_id=user_id,
244 user_group_ids=user_group_ids,
245 permission=permission,
246 )
248 async def get_tool_valves_by_id(self, id: str, db: AsyncSession | None = None) -> dict | None:
249 try:
250 async with get_async_db_context(db) as db:
251 tool = await db.get(Tool, id)
252 return decrypt_valves(tool.valves if tool else None)
253 except Exception:
254 log.exception(f'Error getting tool valves by id {id}')
255 return None
257 async def update_tool_valves_by_id(
258 self, id: str, valves: dict, db: AsyncSession | None = None
259 ) -> ToolValves | None:
260 try:
261 async with get_async_db_context(db) as db:
262 await db.execute(
263 update(Tool).filter_by(id=id).values(valves=encrypt_valves(valves), updated_at=int(time.time()))
264 )
265 await db.commit()
266 return await self.get_tool_by_id(id, db=db)
267 except Exception:
268 return None
270 async def get_user_valves_by_id_and_user_id(
271 self, id: str, user_id: str, db: AsyncSession | None = None
272 ) -> dict | None:
273 try:
274 user = await Users.get_user_by_id(user_id, db=db)
275 user_settings = user.settings.model_dump() if user.settings else {}
277 # Check if user has "tools" and "valves" settings
278 if 'tools' not in user_settings:
279 user_settings['tools'] = {}
280 if 'valves' not in user_settings['tools']:
281 user_settings['tools']['valves'] = {}
283 return decrypt_valves(user_settings['tools']['valves'].get(id))
284 except Exception as e:
285 log.exception(f'Error getting user values by id {id} and user_id {user_id}: {e}')
286 return None
288 async def update_user_valves_by_id_and_user_id(
289 self, id: str, user_id: str, valves: dict, db: AsyncSession | None = None
290 ) -> dict | None:
291 try:
292 user = await Users.get_user_by_id(user_id, db=db)
293 user_settings = user.settings.model_dump() if user.settings else {}
295 # Check if user has "tools" and "valves" settings
296 if 'tools' not in user_settings:
297 user_settings['tools'] = {}
298 if 'valves' not in user_settings['tools']:
299 user_settings['tools']['valves'] = {}
301 user_settings['tools']['valves'][id] = encrypt_valves(valves)
303 # Update the user settings in the database
304 await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db)
306 return valves
307 except Exception as e:
308 log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}')
309 return None
311 async def update_tool_by_id(self, id: str, updated: dict, db: AsyncSession | None = None) -> ToolModel | None:
312 try:
313 async with get_async_db_context(db) as db:
314 access_grants = updated.pop('access_grants', None)
315 await db.execute(update(Tool).filter_by(id=id).values(**updated, updated_at=int(time.time())))
316 await db.commit()
317 if access_grants is not None:
318 await AccessGrants.set_access_grants('tool', id, access_grants, db=db)
320 # populate_existing: the Core update above bypasses any identity-map copy
321 tool = await db.get(Tool, id, populate_existing=True)
322 return await self._to_tool_model(tool, db=db)
323 except Exception:
324 return None
326 async def delete_tool_by_id(self, id: str, db: AsyncSession | None = None) -> bool:
327 try:
328 async with get_async_db_context(db) as db:
329 await AccessGrants.revoke_all_access('tool', id, db=db)
330 await db.execute(delete(Tool).filter_by(id=id))
331 await db.commit()
333 return True
334 except Exception:
335 return False
338Tools = ToolsTable() # singleton tool registry