refac

backend/open_webui/models/chats.py

@@ -14,7 +14,7 @@ from open_webui.models.chat_messages import ChatMessage, ChatMessages
 from open_webui.models.folders import Folders
 from open_webui.models.tags import Tag, TagModel, Tags
 from open_webui.utils.misc import sanitize_data_for_db, sanitize_text_for_db
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, field_validator
 from sqlalchemy import (
     JSON,
     BigInteger,
@@ -58,6 +58,7 @@ class Chat(Base):  # database table mapping for chat entity
     pinned = Column(Boolean, default=False, nullable=True)
 
     meta = Column(JSON, server_default='{}')
+    variables = Column(JSON, nullable=True)
     folder_id = Column(Text, nullable=True)
 
     tasks = Column(JSON, nullable=True)
@@ -95,6 +96,7 @@ class ChatModel(BaseModel):
     pinned: bool | None = False
 
     meta: dict = {}
+    variables: dict = {}
     folder_id: str | None = None
 
     tasks: list | None = None
@@ -103,6 +105,11 @@ class ChatModel(BaseModel):
 
     last_read_at: int | None = None
 
+    @field_validator('variables', mode='before')
+    @classmethod
+    def normalize_variables(cls, value):
+        return value if isinstance(value, dict) else {}
+
 
 class ChatFile(Base):
     __tablename__ = 'chat_file'
@@ -141,6 +148,7 @@ class ChatFileModel(BaseModel):
 
 class ChatForm(BaseModel):
     chat: dict
+    variables: dict | None = None
     folder_id: str | None = None
 
 
@@ -176,6 +184,7 @@ class ChatResponse(BaseModel):
     archived: bool
     pinned: bool | None = False
     meta: dict = {}
+    variables: dict = {}
     folder_id: str | None = None
 
     tasks: list | None = None
@@ -183,6 +192,11 @@ class ChatResponse(BaseModel):
     current_message_id: str | None = None
     context_usage: dict | None = None
 
+    @field_validator('variables', mode='before')
+    @classmethod
+    def normalize_variables(cls, value):
+        return value if isinstance(value, dict) else {}
+
 
 class ChatTitleIdResponse(BaseModel):
     id: str
@@ -394,6 +408,7 @@ class ChatTable:
                     'chat': self._clean_null_bytes(form_data.chat),
                     'folder_id': form_data.folder_id,
                     'meta': internal_meta or {},
+                    'variables': form_data.variables or {},
                     'current_message_id': self.get_current_message_id(form_data.chat),
                     'created_at': int(time.time()),
                     'updated_at': int(time.time()),
@@ -481,6 +496,7 @@ class ChatTable:
                 'title': self._clean_null_bytes(form_data.chat['title'] if 'title' in form_data.chat else 'New Chat'),
                 'chat': self._clean_null_bytes(form_data.chat),
                 'meta': form_data.meta,
+                'variables': form_data.variables or {},
                 'pinned': form_data.pinned,
                 'folder_id': form_data.folder_id,
                 'current_message_id': form_data.current_message_id or self.get_current_message_id(form_data.chat),
@@ -574,6 +590,29 @@ class ChatTable:
         except Exception:
             return
 
+    async def update_chat_variables_by_id(
+        self,
+        id: str,
+        variables: dict | None,
+        db: AsyncSession | None = None,
+        *,
+        touch: bool = True,
+    ) -> ChatModel | None:
+        try:
+            async with get_async_db_context(db) as session:
+                chat_item = await session.get(Chat, id)
+                if chat_item is None:
+                    return None
+
+                chat_item.variables = variables if isinstance(variables, dict) else {}
+                if touch:
+                    chat_item.updated_at = int(time.time())
+
+                await session.commit()
+                return ChatModel.model_validate(chat_item)
+        except Exception:
+            return None
+
     async def update_chat_last_read_at_by_id(self, id: str, user_id: str, db: AsyncSession | None = None) -> bool:
         try:
             async with get_async_db_context(db) as session: