feat: improve multipart file extension detection

Refactor _get_multipart_ext method to use known suffixes list for more reliable file extension detection. The new implementation handles compound file extensions like '.metadata.json.bak' and '.safetensors' by checking against predefined suffixes in order of length. Falls back to existing logic for unknown file types. This improves accuracy when working with model files that have complex naming conventions.
This commit is contained in:
Will Miao
2025-10-27 11:15:16 +08:00
parent f4dcd89835
commit 49bdf77040
2 changed files with 99 additions and 6 deletions

View File

@@ -236,10 +236,20 @@ class ModelLifecycleService:
def _get_multipart_ext(filename: str) -> str:
"""Return the extension for files with compound suffixes."""
parts = filename.split(".")
if len(parts) == 3:
return "." + ".".join(parts[-2:])
if len(parts) >= 4:
return "." + ".".join(parts[-3:])
return os.path.splitext(filename)[1]
known_suffixes = [
".metadata.json.bak",
".metadata.json",
".safetensors",
*PREVIEW_EXTENSIONS,
]
for suffix in sorted(known_suffixes, key=len, reverse=True):
if filename.endswith(suffix):
return suffix
basename = os.path.basename(filename)
dot_index = basename.find(".")
if dot_index != -1:
return basename[dot_index:]
return os.path.splitext(basename)[1]