refac

backend/open_webui/env.py

@@ -683,6 +683,8 @@ WEBUI_SECRET_KEY = os.getenv(
     os.getenv('WEBUI_JWT_SECRET_KEY', ''),
 )
 
+ENABLE_VALVE_ENCRYPTION = os.getenv('ENABLE_VALVE_ENCRYPTION', 'False').lower() == 'true'
+
 WEBUI_SESSION_COOKIE_SAME_SITE = os.getenv('WEBUI_SESSION_COOKIE_SAME_SITE', 'lax')
 WEBUI_SESSION_COOKIE_SECURE = os.getenv('WEBUI_SESSION_COOKIE_SECURE', 'false').lower() == 'true'
 WEBUI_AUTH_COOKIE_SAME_SITE = os.getenv('WEBUI_AUTH_COOKIE_SAME_SITE', WEBUI_SESSION_COOKIE_SAME_SITE)

backend/open_webui/models/functions.py

@@ -7,7 +7,8 @@ import time
 
 # local imports
 from open_webui.internal.db import Base, JSONField, get_async_db_context
-from open_webui.models.users import UserModel, UserResponse, Users
+from open_webui.models.users import UserResponse, Users
+from open_webui.utils.valves import decrypt_valves, encrypt_valves
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy import BigInteger, Boolean, Column, Index, String, Text, delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -143,7 +144,8 @@ class FunctionsTable:
         functions: list[FunctionWithValvesModel],
         db: AsyncSession | None = None,
     ) -> list[FunctionWithValvesModel]:
-        # Synchronize functions for a user by updating existing ones, inserting new ones, and removing those that are no longer present.
+        # Synchronize functions by updating existing ones, inserting new ones,
+        # and removing those that are no longer present.
         try:
             async with get_async_db_context(db) as db:
                 # Get existing functions
@@ -156,24 +158,19 @@ class FunctionsTable:
 
                 # Update or insert functions
                 for func in functions:
+                    func_data = func.model_dump()
+                    func_data['valves'] = encrypt_valves(func_data['valves']) if func_data.get('valves') else None
+                    func_data['user_id'] = user_id
+                    func_data['updated_at'] = int(time.time())
+
                     if func.id in existing_ids:
                         await db.execute(
                             update(Function)
                             .filter_by(id=func.id)
-                            .values(
-                                **func.model_dump(),
-                                user_id=user_id,
-                                updated_at=int(time.time()),
-                            )
+                            .values(**func_data)
                         )
                     else:
-                        new_func = Function(
-                            **{
-                                **func.model_dump(),
-                                'user_id': user_id,
-                                'updated_at': int(time.time()),
-                            }
-                        )
+                        new_func = Function(**func_data)
                         db.add(new_func)
 
                 # Remove functions that are no longer present
@@ -227,7 +224,15 @@ class FunctionsTable:
             functions = result.scalars().all()
 
             if include_valves:
-                return [FunctionWithValvesModel.model_validate(function) for function in functions]
+                return [
+                    FunctionWithValvesModel.model_validate(
+                        {
+                            **FunctionModel.model_validate(function).model_dump(),
+                            'valves': decrypt_valves(function.valves),
+                        }
+                    )
+                    for function in functions
+                ]
             else:
                 return [FunctionModel.model_validate(function) for function in functions]
 
@@ -283,7 +288,7 @@ class FunctionsTable:
         async with get_async_db_context(db) as db:
             try:
                 function = await db.get(Function, id)
-                return function.valves if function.valves else {}
+                return decrypt_valves(function.valves if function else None)
             except Exception as e:
                 log.exception(f'Error getting function valves by id {id}: {e}')
                 return None
@@ -300,7 +305,7 @@ class FunctionsTable:
             async with get_async_db_context(db) as db:
                 result = await db.execute(select(Function.id, Function.valves).filter(Function.id.in_(ids)))
                 functions = result.all()
-                return {f.id: (f.valves if f.valves else {}) for f in functions}
+                return {f.id: decrypt_valves(f.valves) for f in functions}
         except Exception as e:
             log.exception(f'Error batch-fetching function valves: {e}')
             return {}
@@ -311,7 +316,7 @@ class FunctionsTable:
         async with get_async_db_context(db) as db:
             try:
                 function = await db.get(Function, id)
-                function.valves = valves
+                function.valves = encrypt_valves(valves)
                 function.updated_at = int(time.time())
                 await db.commit()
                 await db.refresh(function)
@@ -355,8 +360,8 @@ class FunctionsTable:
             if 'valves' not in user_settings['functions']:
                 user_settings['functions']['valves'] = {}
 
-            return user_settings['functions']['valves'].get(id, {})
-        except Exception as e:
+            return decrypt_valves(user_settings['functions']['valves'].get(id))
+        except Exception:
             log.exception(f'Error getting user values by id {id} and user id {user_id}')
             return None
 
@@ -373,12 +378,12 @@ class FunctionsTable:
             if 'valves' not in user_settings['functions']:
                 user_settings['functions']['valves'] = {}
 
-            user_settings['functions']['valves'][id] = valves
+            user_settings['functions']['valves'][id] = encrypt_valves(valves)
 
             # Update the user settings in the database
             await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db)
 
-            return user_settings['functions']['valves'][id]
+            return valves
         except Exception as e:
             log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}')
             return None

backend/open_webui/models/tools.py

@@ -10,6 +10,7 @@ from open_webui.internal.db import Base, JSONField, get_async_db_context
 from open_webui.models.access_grants import AccessGrantModel, AccessGrants
 from open_webui.models.groups import Groups
 from open_webui.models.users import UserResponse, Users
+from open_webui.utils.valves import decrypt_valves, encrypt_valves
 from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy import BigInteger, Column, String, Text, delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -232,8 +233,8 @@ class ToolsTable:
         try:
             async with get_async_db_context(db) as db:
                 tool = await db.get(Tool, id)
-                return tool.valves if tool.valves else {}
-        except Exception as e:
+                return decrypt_valves(tool.valves if tool else None)
+        except Exception:
             log.exception(f'Error getting tool valves by id {id}')
             return None
 
@@ -242,7 +243,11 @@ class ToolsTable:
     ) -> ToolValves | None:
         try:
             async with get_async_db_context(db) as db:
-                await db.execute(update(Tool).filter_by(id=id).values(valves=valves, updated_at=int(time.time())))
+                await db.execute(
+                    update(Tool)
+                    .filter_by(id=id)
+                    .values(valves=encrypt_valves(valves), updated_at=int(time.time()))
+                )
                 await db.commit()
                 return await self.get_tool_by_id(id, db=db)
         except Exception:
@@ -261,7 +266,7 @@ class ToolsTable:
             if 'valves' not in user_settings['tools']:
                 user_settings['tools']['valves'] = {}
 
-            return user_settings['tools']['valves'].get(id, {})
+            return decrypt_valves(user_settings['tools']['valves'].get(id))
         except Exception as e:
             log.exception(f'Error getting user values by id {id} and user_id {user_id}: {e}')
             return None
@@ -279,12 +284,12 @@ class ToolsTable:
             if 'valves' not in user_settings['tools']:
                 user_settings['tools']['valves'] = {}
 
-            user_settings['tools']['valves'][id] = valves
+            user_settings['tools']['valves'][id] = encrypt_valves(valves)
 
             # Update the user settings in the database
             await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db)
 
-            return user_settings['tools']['valves'][id]
+            return valves
         except Exception as e:
             log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}')
             return None

backend/open_webui/utils/valves.py

@@ -0,0 +1,41 @@
+import base64
+import hashlib
+import json
+import logging
+from functools import lru_cache
+
+from cryptography.fernet import Fernet, InvalidToken
+from open_webui.env import ENABLE_VALVE_ENCRYPTION, WEBUI_SECRET_KEY
+
+log = logging.getLogger(__name__)
+
+
+@lru_cache(maxsize=1)
+def _fernet() -> Fernet:
+    key = WEBUI_SECRET_KEY.encode()
+    if len(WEBUI_SECRET_KEY) != 44:
+        key = base64.urlsafe_b64encode(hashlib.sha256(key).digest())
+    return Fernet(key)
+
+
+def encrypt_valves(valves: dict) -> dict | str:
+    if not ENABLE_VALVE_ENCRYPTION:
+        return valves
+    return _fernet().encrypt(json.dumps(valves).encode()).decode()
+
+
+def decrypt_valves(valves) -> dict:
+    if not valves:
+        return {}
+    if isinstance(valves, dict):
+        return valves
+    if not isinstance(valves, str):
+        return {}
+
+    try:
+        decrypted = json.loads(_fernet().decrypt(valves.encode()).decode())
+    except (InvalidToken, json.JSONDecodeError) as e:
+        log.warning('Failed to decrypt valves: %s', type(e).__name__)
+        return {}
+
+    return decrypted if isinstance(decrypted, dict) else {}