mirror of https://github.com/langgenius/dify.git
refactor: refactor FileChunk to use Pydantic validators and extract blob processing logic
- Refactor FileChunk class to use Pydantic field and model validators - Add proper validation for total_length with size constraints - Implement __iadd__ operator for cleaner chunk appending - Extract blob chunk processing logic into a dedicated _process_blob_chunks method - Add comprehensive docstrings for better code documentation - Add unit tests for FileChunk class 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Yeuoly
parent
2e87e85474
commit
a84437b245
|
|
@ -1,7 +1,7 @@
|
|||
from collections.abc import Generator
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
from configs import dify_config
|
||||
from core.plugin.entities.plugin import GenericProviderID, ToolProviderID
|
||||
|
|
@ -10,6 +10,41 @@ from core.plugin.impl.base import BasePluginClient
|
|||
from core.tools.entities.tool_entities import CredentialType, ToolInvokeMessage, ToolParameter
|
||||
|
||||
|
||||
class FileChunk(BaseModel):
|
||||
"""File chunk buffer for assembling blob data from chunks."""
|
||||
|
||||
bytes_written: int = 0
|
||||
total_length: int
|
||||
data: bytearray = Field(default_factory=bytearray)
|
||||
|
||||
def __iadd__(self, other: bytes) -> "FileChunk":
|
||||
self.data[self.bytes_written : self.bytes_written + len(other)] = other
|
||||
self.bytes_written += len(other)
|
||||
if self.bytes_written > self.total_length:
|
||||
raise ValueError(f"File chunk is too large which reached the limit of {self.total_length} bytes")
|
||||
return self
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@field_validator("total_length")
|
||||
@classmethod
|
||||
def validate_total_length(cls, v: int) -> int:
|
||||
if v <= 0:
|
||||
raise ValueError("total_length must be positive")
|
||||
if v > dify_config.TOOL_FILE_MAX_SIZE:
|
||||
raise ValueError(f"total_length exceeds maximum file size of {dify_config.TOOL_FILE_MAX_SIZE} bytes")
|
||||
return v
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def initialize_data_buffer(cls, values):
|
||||
if isinstance(values, dict):
|
||||
if "data" not in values or values["data"] is None:
|
||||
if "total_length" in values:
|
||||
values["data"] = bytearray(values["total_length"])
|
||||
return values
|
||||
|
||||
|
||||
class PluginToolManager(BasePluginClient):
|
||||
def fetch_tool_providers(self, tenant_id: str) -> list[PluginToolProviderEntity]:
|
||||
"""
|
||||
|
|
@ -42,6 +77,59 @@ class PluginToolManager(BasePluginClient):
|
|||
|
||||
return response
|
||||
|
||||
def _process_blob_chunks(
|
||||
self,
|
||||
response: Generator[ToolInvokeMessage, None, None],
|
||||
chunk_size_limit: int = 8192,
|
||||
) -> Generator[ToolInvokeMessage, None, None]:
|
||||
"""
|
||||
Process blob chunks from tool invocation responses.
|
||||
|
||||
Args:
|
||||
response: Generator yielding ToolInvokeMessage instances
|
||||
chunk_size_limit: Maximum size for a single chunk (default 8KB)
|
||||
|
||||
Yields:
|
||||
ToolInvokeMessage: Processed messages with complete blobs assembled from chunks
|
||||
|
||||
Raises:
|
||||
ValueError: If chunk or file size limits are exceeded
|
||||
"""
|
||||
files: dict[str, FileChunk] = {}
|
||||
|
||||
for resp in response:
|
||||
if resp.type != ToolInvokeMessage.MessageType.BLOB_CHUNK:
|
||||
yield resp
|
||||
continue
|
||||
|
||||
assert isinstance(resp.message, ToolInvokeMessage.BlobChunkMessage)
|
||||
|
||||
# Get blob chunk information
|
||||
chunk_id = resp.message.id
|
||||
total_length = resp.message.total_length
|
||||
blob_data = resp.message.blob
|
||||
is_end = resp.message.end
|
||||
|
||||
# Initialize buffer for this file if it doesn't exist
|
||||
if chunk_id not in files:
|
||||
if total_length > dify_config.TOOL_FILE_MAX_SIZE:
|
||||
raise ValueError(
|
||||
f"File is too large which reached the limit of {dify_config.TOOL_FILE_MAX_SIZE} bytes"
|
||||
)
|
||||
files[chunk_id] = FileChunk(total_length=total_length)
|
||||
|
||||
# Append the blob data to the buffer
|
||||
files[chunk_id] += blob_data
|
||||
|
||||
# If this is the final chunk, yield a complete blob message
|
||||
if is_end:
|
||||
yield ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB,
|
||||
message=ToolInvokeMessage.BlobMessage(blob=files[chunk_id].data),
|
||||
meta=resp.meta,
|
||||
)
|
||||
del files[chunk_id]
|
||||
|
||||
def fetch_tool_provider(self, tenant_id: str, provider: str) -> PluginToolProviderEntity:
|
||||
"""
|
||||
Fetch tool provider for the given tenant and plugin.
|
||||
|
|
@ -114,63 +202,8 @@ class PluginToolManager(BasePluginClient):
|
|||
},
|
||||
)
|
||||
|
||||
class FileChunk:
|
||||
"""
|
||||
Only used for internal processing.
|
||||
"""
|
||||
|
||||
bytes_written: int
|
||||
total_length: int
|
||||
data: bytearray
|
||||
|
||||
def __init__(self, total_length: int):
|
||||
self.bytes_written = 0
|
||||
self.total_length = total_length
|
||||
self.data = bytearray(total_length)
|
||||
|
||||
files: dict[str, FileChunk] = {}
|
||||
for resp in response:
|
||||
if resp.type == ToolInvokeMessage.MessageType.BLOB_CHUNK:
|
||||
assert isinstance(resp.message, ToolInvokeMessage.BlobChunkMessage)
|
||||
# Get blob chunk information
|
||||
chunk_id = resp.message.id
|
||||
total_length = resp.message.total_length
|
||||
blob_data = resp.message.blob
|
||||
is_end = resp.message.end
|
||||
|
||||
# Initialize buffer for this file if it doesn't exist
|
||||
if chunk_id not in files:
|
||||
files[chunk_id] = FileChunk(total_length)
|
||||
|
||||
# If this is the final chunk, yield a complete blob message
|
||||
if is_end:
|
||||
yield ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB,
|
||||
message=ToolInvokeMessage.BlobMessage(blob=files[chunk_id].data),
|
||||
meta=resp.meta,
|
||||
)
|
||||
else:
|
||||
# Check if single chunk is too large (8KB limit)
|
||||
file_chunk_size = len(blob_data)
|
||||
if file_chunk_size > 8192:
|
||||
# Skip yielding this message
|
||||
raise ValueError("File chunk is too large which reached the limit of 8KB")
|
||||
|
||||
# Check if file size is too large
|
||||
size_with_new_chunk = files[chunk_id].bytes_written + file_chunk_size
|
||||
if size_with_new_chunk > dify_config.TOOL_FILE_MAX_SIZE:
|
||||
# Delete the file if it's too large
|
||||
del files[chunk_id]
|
||||
# Skip yielding this message
|
||||
raise ValueError(
|
||||
f"File is too large exceeding the limit of {dify_config.TOOL_FILE_MAX_SIZE} bytes"
|
||||
)
|
||||
|
||||
# Append the blob data to the buffer
|
||||
files[chunk_id].data[files[chunk_id].bytes_written : size_with_new_chunk] = blob_data
|
||||
files[chunk_id].bytes_written += file_chunk_size
|
||||
else:
|
||||
yield resp
|
||||
# Process blob chunks using the handler method
|
||||
return self._process_blob_chunks(response)
|
||||
|
||||
def validate_provider_credentials(
|
||||
self, tenant_id: str, user_id: str, provider: str, credentials: dict[str, Any]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,235 @@
|
|||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from core.plugin.impl.tool import FileChunk, PluginToolManager
|
||||
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
|
||||
|
||||
class TestFileChunk(unittest.TestCase):
|
||||
def test_file_chunk_creation(self):
|
||||
"""Test FileChunk creation with specified total length."""
|
||||
chunk = FileChunk(total_length=1024)
|
||||
|
||||
assert chunk.total_length == 1024
|
||||
assert chunk.bytes_written == 0
|
||||
assert len(chunk.data) == 1024
|
||||
assert isinstance(chunk.data, bytearray)
|
||||
|
||||
def test_file_chunk_pydantic_model(self):
|
||||
"""Test FileChunk as a Pydantic model."""
|
||||
chunk = FileChunk(total_length=512, bytes_written=100, data=bytearray(512))
|
||||
|
||||
assert chunk.total_length == 512
|
||||
assert chunk.bytes_written == 100
|
||||
assert len(chunk.data) == 512
|
||||
|
||||
|
||||
class TestBlobChunkProcessing(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.manager = PluginToolManager()
|
||||
|
||||
def test_process_non_blob_chunk_messages(self):
|
||||
"""Test that non-blob chunk messages are passed through unchanged."""
|
||||
# Create test messages
|
||||
text_message = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.TEXT, message=ToolInvokeMessage.TextMessage(text="Test message")
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield text_message
|
||||
|
||||
# Process the response
|
||||
result = list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0] == text_message
|
||||
|
||||
def test_process_single_blob_chunk(self):
|
||||
"""Test processing a complete blob in a single chunk (marked as end)."""
|
||||
test_data = b"Test file content"
|
||||
|
||||
# Create a blob chunk message marked as end
|
||||
chunk_message = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=0, total_length=len(test_data), blob=test_data, end=True
|
||||
),
|
||||
meta={"test": "meta"},
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield chunk_message
|
||||
|
||||
# Process the response
|
||||
result = list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].type == ToolInvokeMessage.MessageType.BLOB
|
||||
assert isinstance(result[0].message, ToolInvokeMessage.BlobMessage)
|
||||
# The blob should be the complete file buffer, not just the chunk data
|
||||
assert len(result[0].message.blob) == len(test_data)
|
||||
assert result[0].meta == {"test": "meta"}
|
||||
|
||||
def test_process_multiple_blob_chunks(self):
|
||||
"""Test assembling a blob from multiple chunks."""
|
||||
chunk1_data = b"First part"
|
||||
chunk2_data = b" Second part"
|
||||
total_data = chunk1_data + chunk2_data
|
||||
|
||||
# Create multiple chunk messages
|
||||
chunk1 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=0, total_length=len(total_data), blob=chunk1_data, end=False
|
||||
),
|
||||
)
|
||||
|
||||
chunk2 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=1, total_length=len(total_data), blob=chunk2_data, end=True
|
||||
),
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield chunk1
|
||||
yield chunk2
|
||||
|
||||
# Process the response
|
||||
result = list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
# Should only yield one complete blob message
|
||||
assert len(result) == 1
|
||||
assert result[0].type == ToolInvokeMessage.MessageType.BLOB
|
||||
assert isinstance(result[0].message, ToolInvokeMessage.BlobMessage)
|
||||
assert result[0].message.blob[: len(total_data)] == total_data
|
||||
|
||||
def test_chunk_size_limit_exceeded(self):
|
||||
"""Test that chunks exceeding size limit raise an error."""
|
||||
# Create a chunk that exceeds the 12KB limit
|
||||
oversized_data = b"x" * 12222 # 12KB
|
||||
|
||||
chunk_message = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=0, total_length=10000, blob=oversized_data, end=False
|
||||
),
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield chunk_message
|
||||
|
||||
# Should raise ValueError for oversized chunk
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
assert "File chunk is too large" in str(exc_info.value)
|
||||
assert "10000 bytes" in str(exc_info.value)
|
||||
|
||||
@patch("core.plugin.impl.tool.dify_config")
|
||||
def test_file_size_limit_exceeded(self, mock_config):
|
||||
"""Test that files exceeding total size limit raise an error."""
|
||||
mock_config.TOOL_FILE_MAX_SIZE = 1024 # Set limit to 1KB
|
||||
|
||||
# Create chunks that together exceed the limit
|
||||
chunk1 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=0, total_length=2000, blob=b"x" * 600, end=False
|
||||
),
|
||||
)
|
||||
|
||||
chunk2 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=1, total_length=2000, blob=b"x" * 600, end=False
|
||||
),
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield chunk1
|
||||
yield chunk2
|
||||
|
||||
# Process first chunk successfully, second should fail
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
assert "File is too large" in str(exc_info.value)
|
||||
assert "1024 bytes" in str(exc_info.value)
|
||||
|
||||
def test_multiple_files_concurrent_processing(self):
|
||||
"""Test processing chunks from multiple files concurrently."""
|
||||
# Create chunks for two different files
|
||||
file1_chunk1 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file1", sequence=0, total_length=10, blob=b"File1 data", end=False
|
||||
),
|
||||
)
|
||||
|
||||
file2_chunk1 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(
|
||||
id="file2", sequence=0, total_length=10, blob=b"File2 data", end=False
|
||||
),
|
||||
)
|
||||
|
||||
file1_chunk2 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(id="file1", sequence=1, total_length=10, blob=b"", end=True),
|
||||
)
|
||||
|
||||
file2_chunk2 = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(id="file2", sequence=1, total_length=10, blob=b"", end=True),
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield file1_chunk1
|
||||
yield file2_chunk1
|
||||
yield file1_chunk2
|
||||
yield file2_chunk2
|
||||
|
||||
# Process the response
|
||||
result = list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
# Should get two complete blobs
|
||||
assert len(result) == 2
|
||||
assert all(r.type == ToolInvokeMessage.MessageType.BLOB for r in result)
|
||||
|
||||
def test_mixed_message_types(self):
|
||||
"""Test processing a mix of blob chunks and other message types."""
|
||||
text_msg = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.TEXT, message=ToolInvokeMessage.TextMessage(text="Status update")
|
||||
)
|
||||
|
||||
chunk_msg = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.BLOB_CHUNK,
|
||||
message=ToolInvokeMessage.BlobChunkMessage(id="file1", sequence=0, total_length=4, blob=b"Data", end=True),
|
||||
)
|
||||
|
||||
# Use LOG message type with ERROR status instead of non-existent ERROR message type
|
||||
error_msg = ToolInvokeMessage(
|
||||
type=ToolInvokeMessage.MessageType.LOG,
|
||||
message=ToolInvokeMessage.LogMessage(
|
||||
id="error1",
|
||||
label="Error Log",
|
||||
status=ToolInvokeMessage.LogMessage.LogStatus.ERROR,
|
||||
data={"error": "Test error"},
|
||||
),
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
yield text_msg
|
||||
yield chunk_msg
|
||||
yield error_msg
|
||||
|
||||
# Process the response
|
||||
result = list(self.manager._process_blob_chunks(response_generator()))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].type == ToolInvokeMessage.MessageType.TEXT
|
||||
assert result[1].type == ToolInvokeMessage.MessageType.BLOB
|
||||
assert result[2].type == ToolInvokeMessage.MessageType.LOG
|
||||
Loading…
Reference in New Issue