diff --git a/litellm/tests/test_file_types.py b/litellm/tests/test_file_types.py new file mode 100644 index 00000000000..db83ba0e74b --- /dev/null +++ b/litellm/tests/test_file_types.py @@ -0,0 +1,54 @@ +from litellm.types.files import ( + FILE_EXTENSIONS, + FILE_MIME_TYPES, + FileType, + get_file_extension_from_mime_type, + get_file_type_from_extension, + get_file_extension_for_file_type, + get_file_mime_type_for_file_type, + get_file_mime_type_from_extension, +) +import pytest + + +class TestFileConsts: + def test_all_file_types_have_extensions(self): + for file_type in FileType: + assert file_type in FILE_EXTENSIONS.keys() + + def test_all_file_types_have_mime_types(self): + for file_type in FileType: + assert file_type in FILE_MIME_TYPES.keys() + + def test_get_file_extension_from_mime_type(self): + assert get_file_extension_from_mime_type("audio/aac") == "aac" + assert get_file_extension_from_mime_type("application/pdf") == "pdf" + with pytest.raises(ValueError): + get_file_extension_from_mime_type("application/unknown") + + def test_get_file_type_from_extension(self): + assert get_file_type_from_extension("aac") == FileType.AAC + assert get_file_type_from_extension("pdf") == FileType.PDF + with pytest.raises(ValueError): + get_file_type_from_extension("unknown") + + def test_get_file_extension_for_file_type(self): + assert get_file_extension_for_file_type(FileType.AAC) == "aac" + assert get_file_extension_for_file_type(FileType.PDF) == "pdf" + + def test_get_file_mime_type_for_file_type(self): + assert get_file_mime_type_for_file_type(FileType.AAC) == "audio/aac" + assert get_file_mime_type_for_file_type(FileType.PDF) == "application/pdf" + + def test_get_file_mime_type_from_extension(self): + assert get_file_mime_type_from_extension("aac") == "audio/aac" + assert get_file_mime_type_from_extension("pdf") == "application/pdf" + + def test_uppercase_extensions(self): + # Test that uppercase extensions return the correct file type + assert get_file_type_from_extension("AAC") == FileType.AAC + assert get_file_type_from_extension("PDF") == FileType.PDF + + # Test that uppercase extensions return the correct MIME type + assert get_file_mime_type_from_extension("AAC") == "audio/aac" + assert get_file_mime_type_from_extension("PDF") == "application/pdf" diff --git a/litellm/types/files.py b/litellm/types/files.py index 2da08e76c4d..577b9b55cc1 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -151,23 +151,16 @@ Util Functions """ -def get_file_mime_type_from_extension(extension: str) -> str: - for file_type, extensions in FILE_EXTENSIONS.items(): - if extension in extensions: - return FILE_MIME_TYPES[file_type] - raise ValueError(f"Unknown mime type for extension: {extension}") - - def get_file_extension_from_mime_type(mime_type: str) -> str: for file_type, mime in FILE_MIME_TYPES.items(): - if mime == mime_type: + if mime.lower() == mime_type.lower(): return FILE_EXTENSIONS[file_type][0] raise ValueError(f"Unknown extension for mime type: {mime_type}") def get_file_type_from_extension(extension: str) -> FileType: for file_type, extensions in FILE_EXTENSIONS.items(): - if extension in extensions: + if extension.lower() in extensions: return file_type raise ValueError(f"Unknown file type for extension: {extension}") @@ -181,6 +174,11 @@ def get_file_mime_type_for_file_type(file_type: FileType) -> str: return FILE_MIME_TYPES[file_type] +def get_file_mime_type_from_extension(extension: str) -> str: + file_type = get_file_type_from_extension(extension) + return get_file_mime_type_for_file_type(file_type) + + """ FileType Type Groupings (Videos, Images, etc) """