fix: Fix test managers state transition tests (#3551)
This commit is contained in:
@@ -151,7 +151,8 @@ class FileManager:
|
||||
Enforces state transition rules (when enforce_state_transitions=True):
|
||||
- PENDING -> PARSING -> EMBEDDING -> COMPLETED (normal flow)
|
||||
- Any non-terminal state -> ERROR
|
||||
- ERROR and COMPLETED are terminal (no transitions allowed)
|
||||
- Same-state transitions are allowed (e.g., EMBEDDING -> EMBEDDING)
|
||||
- ERROR and COMPLETED are terminal (no status transitions allowed, metadata updates blocked)
|
||||
|
||||
Args:
|
||||
file_id: ID of the file to update
|
||||
@@ -196,28 +197,31 @@ class FileManager:
|
||||
]
|
||||
|
||||
# only add state transition validation if enforce_state_transitions is True
|
||||
if enforce_state_transitions:
|
||||
# prevent updates to terminal states (ERROR, COMPLETED)
|
||||
if enforce_state_transitions and processing_status is not None:
|
||||
# enforce specific transitions based on target status
|
||||
if processing_status == FileProcessingStatus.PARSING:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.PENDING, FileProcessingStatus.PARSING])
|
||||
)
|
||||
elif processing_status == FileProcessingStatus.EMBEDDING:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.PARSING, FileProcessingStatus.EMBEDDING])
|
||||
)
|
||||
elif processing_status == FileProcessingStatus.COMPLETED:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.EMBEDDING, FileProcessingStatus.COMPLETED])
|
||||
)
|
||||
elif processing_status == FileProcessingStatus.ERROR:
|
||||
# ERROR can be set from any non-terminal state
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED])
|
||||
)
|
||||
elif enforce_state_transitions and processing_status is None:
|
||||
# If only updating metadata fields (not status), prevent updates to terminal states
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED])
|
||||
)
|
||||
|
||||
if processing_status is not None:
|
||||
# enforce specific transitions based on target status
|
||||
if processing_status == FileProcessingStatus.PARSING:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.PENDING, FileProcessingStatus.PARSING])
|
||||
)
|
||||
elif processing_status == FileProcessingStatus.EMBEDDING:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.PARSING, FileProcessingStatus.EMBEDDING])
|
||||
)
|
||||
elif processing_status == FileProcessingStatus.COMPLETED:
|
||||
where_conditions.append(
|
||||
FileMetadataModel.processing_status.in_([FileProcessingStatus.EMBEDDING, FileProcessingStatus.COMPLETED])
|
||||
)
|
||||
# ERROR can be set from any non-terminal state (already handled by terminal check above)
|
||||
|
||||
# fast in-place update with state validation
|
||||
stmt = (
|
||||
update(FileMetadataModel)
|
||||
|
||||
@@ -55,6 +55,7 @@ from letta.schemas.block import BlockUpdate, CreateBlock
|
||||
from letta.schemas.embedding_config import EmbeddingConfig
|
||||
from letta.schemas.enums import ActorType, AgentStepStatus, FileProcessingStatus, JobStatus, JobType, MessageRole, ProviderType
|
||||
from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate, SandboxEnvironmentVariableUpdate
|
||||
from letta.schemas.file import FileMetadata
|
||||
from letta.schemas.file import FileMetadata as PydanticFileMetadata
|
||||
from letta.schemas.identity import IdentityCreate, IdentityProperty, IdentityPropertyType, IdentityType, IdentityUpdate, IdentityUpsert
|
||||
from letta.schemas.job import BatchJob
|
||||
|
||||
Reference in New Issue
Block a user