add more dataclass (#25039)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

api/core/tools/tool_file_manager.py

@@ -98,6 +98,7 @@ class ToolFileManager:
                 mimetype=mimetype,
                 name=present_filename,
                 size=len(file_binary),
+                original_url=None,
             )
 
             session.add(tool_file)
@@ -131,7 +132,6 @@ class ToolFileManager:
         filename = f"{unique_name}{extension}"
         filepath = f"tools/{tenant_id}/{filename}"
         storage.save(filepath, blob)
-
         with Session(self._engine, expire_on_commit=False) as session:
             tool_file = ToolFile(
                 user_id=user_id,

api/models/tools.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import Any, cast
+from typing import Any, Optional, cast
 from urllib.parse import urlparse
 
 import sqlalchemy as sa
@@ -22,15 +22,15 @@ from .types import StringUUID
 
 
 # system level tool oauth client params (client_id, client_secret, etc.)
-class ToolOAuthSystemClient(Base):
+class ToolOAuthSystemClient(TypeBase):
     __tablename__ = "tool_oauth_system_clients"
     __table_args__ = (
         sa.PrimaryKeyConstraint("id", name="tool_oauth_system_client_pkey"),
         sa.UniqueConstraint("plugin_id", "provider", name="tool_oauth_system_client_plugin_id_provider_idx"),
     )
 
-    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
+    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
-    plugin_id = mapped_column(String(512), nullable=False)
+    plugin_id: Mapped[str] = mapped_column(String(512), nullable=False)
     provider: Mapped[str] = mapped_column(String(255), nullable=False)
     # oauth params of the tool provider
     encrypted_oauth_params: Mapped[str] = mapped_column(sa.Text, nullable=False)
@@ -412,7 +412,7 @@ class ToolConversationVariables(Base):
         return json.loads(self.variables_str)
 
 
-class ToolFile(Base):
+class ToolFile(TypeBase):
     """This table stores file metadata generated in workflows,
     not only files created by agent.
     """
@@ -423,19 +423,19 @@ class ToolFile(Base):
         sa.Index("tool_file_conversation_id_idx", "conversation_id"),
     )
 
-    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
+    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
     # conversation user id
     user_id: Mapped[str] = mapped_column(StringUUID)
     # tenant id
     tenant_id: Mapped[str] = mapped_column(StringUUID)
     # conversation id
-    conversation_id: Mapped[str] = mapped_column(StringUUID, nullable=True)
+    conversation_id: Mapped[Optional[str]] = mapped_column(StringUUID, nullable=True)
     # file key
     file_key: Mapped[str] = mapped_column(String(255), nullable=False)
     # mime type
     mimetype: Mapped[str] = mapped_column(String(255), nullable=False)
     # original url
-    original_url: Mapped[str] = mapped_column(String(2048), nullable=True)
+    original_url: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True, default=None)
     # name
     name: Mapped[str] = mapped_column(default="")
     # size

api/tests/integration_tests/factories/test_storage_key_loader.py

@@ -84,17 +84,17 @@ class TestStorageKeyLoader(unittest.TestCase):
         if tenant_id is None:
             tenant_id = self.tenant_id
 
-        tool_file = ToolFile()
+        tool_file = ToolFile(
+            user_id=self.user_id,
+            tenant_id=tenant_id,
+            conversation_id=self.conversation_id,
+            file_key=file_key,
+            mimetype="text/plain",
+            original_url="http://example.com/file.txt",
+            name="test_tool_file.txt",
+            size=2048,
+        )
         tool_file.id = file_id
-        tool_file.user_id = self.user_id
-        tool_file.tenant_id = tenant_id
-        tool_file.conversation_id = self.conversation_id
-        tool_file.file_key = file_key
-        tool_file.mimetype = "text/plain"
-        tool_file.original_url = "http://example.com/file.txt"
-        tool_file.name = "test_tool_file.txt"
-        tool_file.size = 2048
-
         self.session.add(tool_file)
         self.session.flush()
         self.test_tool_files.append(tool_file)

api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py

@@ -84,16 +84,17 @@ class TestStorageKeyLoader(unittest.TestCase):
         if tenant_id is None:
             tenant_id = self.tenant_id
 
-        tool_file = ToolFile()
+        tool_file = ToolFile(
+            user_id=self.user_id,
+            tenant_id=tenant_id,
+            conversation_id=self.conversation_id,
+            file_key=file_key,
+            mimetype="text/plain",
+            original_url="http://example.com/file.txt",
+            name="test_tool_file.txt",
+            size=2048,
+        )
         tool_file.id = file_id
-        tool_file.user_id = self.user_id
-        tool_file.tenant_id = tenant_id
-        tool_file.conversation_id = self.conversation_id
-        tool_file.file_key = file_key
-        tool_file.mimetype = "text/plain"
-        tool_file.original_url = "http://example.com/file.txt"
-        tool_file.name = "test_tool_file.txt"
-        tool_file.size = 2048
 
         self.session.add(tool_file)
         self.session.flush()

api/tests/unit_tests/core/workflow/nodes/llm/test_file_saver.py

@@ -26,14 +26,13 @@ def _gen_id():
 
 
 class TestFileSaverImpl:
-    def test_save_binary_string(self, monkeypatch):
+    def test_save_binary_string(self, monkeypatch: pytest.MonkeyPatch):
         user_id = _gen_id()
         tenant_id = _gen_id()
         file_type = FileType.IMAGE
         mime_type = "image/png"
         mock_signed_url = "https://example.com/image.png"
         mock_tool_file = ToolFile(
-            id=_gen_id(),
             user_id=user_id,
             tenant_id=tenant_id,
             conversation_id=None,
@@ -43,6 +42,7 @@ class TestFileSaverImpl:
             name=f"{_gen_id()}.png",
             size=len(_PNG_DATA),
         )
+        mock_tool_file.id = _gen_id()
         mocked_tool_file_manager = mock.MagicMock(spec=ToolFileManager)
         mocked_engine = mock.MagicMock(spec=Engine)
 
@@ -80,7 +80,7 @@ class TestFileSaverImpl:
         )
         mocked_sign_file.assert_called_once_with(mock_tool_file.id, ".png")
 
-    def test_save_remote_url_request_failed(self, monkeypatch):
+    def test_save_remote_url_request_failed(self, monkeypatch: pytest.MonkeyPatch):
         _TEST_URL = "https://example.com/image.png"
         mock_request = httpx.Request("GET", _TEST_URL)
         mock_response = httpx.Response(
@@ -99,7 +99,7 @@ class TestFileSaverImpl:
         mock_get.assert_called_once_with(_TEST_URL)
         assert exc.value.response.status_code == 401
 
-    def test_save_remote_url_success(self, monkeypatch):
+    def test_save_remote_url_success(self, monkeypatch: pytest.MonkeyPatch):
         _TEST_URL = "https://example.com/image.png"
         mime_type = "image/png"
         user_id = _gen_id()
@@ -115,7 +115,6 @@ class TestFileSaverImpl:
 
         file_saver = FileSaverImpl(user_id=user_id, tenant_id=tenant_id)
         mock_tool_file = ToolFile(
-            id=_gen_id(),
             user_id=user_id,
             tenant_id=tenant_id,
             conversation_id=None,
@@ -125,6 +124,7 @@ class TestFileSaverImpl:
             name=f"{_gen_id()}.png",
             size=len(_PNG_DATA),
         )
+        mock_tool_file.id = _gen_id()
         mock_get = mock.MagicMock(spec=ssrf_proxy.get, return_value=mock_response)
         monkeypatch.setattr(ssrf_proxy, "get", mock_get)
         mock_save_binary_string = mock.MagicMock(spec=file_saver.save_binary_string, return_value=mock_tool_file)