Skip to content

Commit

Permalink
fix: add UUID validation for tool file ID extraction (#12011)
Browse files Browse the repository at this point in the history
Signed-off-by: -LAN- <laipz8200@outlook.com>
  • Loading branch information
laipz8200 authored Dec 23, 2024
1 parent 75bce28 commit d0dd8b7
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions api/core/workflow/nodes/tool/tool_node.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from collections.abc import Mapping, Sequence
from typing import Any
from uuid import UUID

from sqlalchemy import select
from sqlalchemy.orm import Session
Expand Down Expand Up @@ -231,6 +232,10 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
url = str(response.message)
transfer_method = FileTransferMethod.TOOL_FILE
tool_file_id = url.split("/")[-1].split(".")[0]
try:
UUID(tool_file_id)
except ValueError:
raise ToolFileError(f"cannot extract tool file id from url {url}")
with Session(db.engine) as session:
stmt = select(ToolFile).where(ToolFile.id == tool_file_id)
tool_file = session.scalar(stmt)
Expand Down

0 comments on commit d0dd8b7

Please sign in to comment.