feat: improve file type check (#563)

Author: micpst

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Switch from imghdr to filetype for image file type check (#563)
 - Remove prompt lab (#549)
 - Add batched() helper method to utils (#555)
 - Rename DocumentMeta create_text_document_from_literal to from_literal (#561)

packages/ragbits-core/pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
     "tomli>=2.0.2,<3.0.0",
     "litellm>=1.55.0,<2.0.0",
     "aiohttp>=3.10.8,<4.0.0",
-    "standard-imghdr>=3.10.14,<3.14"
+    "filetype>=1.2.0,<2.0.0",
 ]
 
 [project.urls]

packages/ragbits-core/src/ragbits/core/prompt/prompt.py

@@ -1,11 +1,11 @@
 import asyncio
 import base64
-import imghdr
 import textwrap
 from abc import ABCMeta
 from collections.abc import Awaitable, Callable
 from typing import Any, Generic, cast, get_args, get_origin, overload
 
+import filetype
 from jinja2 import Environment, Template, meta
 from pydantic import BaseModel
 from typing_extensions import TypeVar, get_original_bases
@@ -288,10 +288,10 @@ def list_images(self) -> list[str]:
     @staticmethod
     def _create_message_with_image(image: str | bytes) -> dict:
         if isinstance(image, bytes):
-            image_type = imghdr.what(None, image)
-            if not image_type:
+            detected_type = filetype.guess(image)
+            if not detected_type or not detected_type.mime.startswith("image/"):
                 raise PromptWithImagesOfInvalidFormat()
-            image_url = f"data:image/{image_type};base64,{base64.b64encode(image).decode('utf-8')}"
+            image_url = f"data:{detected_type.mime};base64,{base64.b64encode(image).decode('utf-8')}"
         else:
             image_url = image
         return {

packages/ragbits-document-search/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Improve document file type check (#563)
 - Fix reranker options typing (#562)
 - Add query rephraser options (#560)
 - Rename DocumentMeta create_text_document_from_literal to from_literal (#561)

packages/ragbits-document-search/pyproject.toml

@@ -31,7 +31,13 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ["unstructured>=0.16.9", "unstructured-client>=0.26.0", "rerankers>=0.6.1", "ragbits-core==0.17.1"]
+dependencies = [
+    "unstructured>=0.16.9",
+    "unstructured-client>=0.26.0",
+    "rerankers>=0.6.1",
+    "filetype>=1.2.0,<2.0.0",
+    "ragbits-core==0.17.1",
+]
 
 [project.urls]
 "Homepage" = "https://github.com/deepsense-ai/ragbits"

packages/ragbits-document-search/src/ragbits/document_search/documents/document.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Annotated, Any
 
+import filetype
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
@@ -121,7 +122,7 @@ def from_local_path(cls, local_path: Path) -> "DocumentMeta":
             The document metadata.
         """
         return cls(
-            document_type=DocumentType(local_path.suffix[1:]),
+            document_type=cls._infer_document_type(local_path),
             source=LocalFileSource(path=local_path),
         )
 
@@ -139,10 +140,25 @@ async def from_source(cls, source: Source) -> "DocumentMeta":
         path = await source.fetch()
 
         return cls(
-            document_type=DocumentType(path.suffix[1:]),
+            document_type=cls._infer_document_type(path),
             source=source,
         )
 
+    @staticmethod
+    def _infer_document_type(path: Path) -> DocumentType:
+        """
+        Infer the document type by checking the file signature. Use the file extension as a fallback.
+
+        Args:
+            path: The path to the file.
+
+        Returns:
+            The inferred document type.
+        """
+        if kind := filetype.guess(path):
+            return DocumentType(kind.extension)
+        return DocumentType(path.suffix[1:])
+
 
 class Document(BaseModel):
     """