Merge pull request #47 from mobiusml/add_media_id_model

Media ID and Quesition Models

Author: movchan74

aana/configs/db.py

@@ -5,12 +5,32 @@
 
 from alembic import command
 from alembic.config import Config
-from sqlalchemy import String, create_engine
+from sqlalchemy import String, TypeDecorator, create_engine
 
-# These are here so we can change types in a single place.
+from aana.models.pydantic.media_id import MediaId
 
-media_id_type: TypeAlias = str
-MediaIdSqlType: TypeAlias = String
+
+class MediaIdType(TypeDecorator):
+    """Custom type for handling MediaId objects with SQLAlchemy."""
+
+    impl = String
+
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect):
+        """Convert a MediaId instance to a string value for storage."""
+        if value is None:
+            return value
+        return str(value)
+
+    def process_result_value(self, value, dialect):
+        """Convert a string value from the database back into a MediaId instance."""
+        if value is None:
+            return value
+        return MediaId(value)
+
+
+MediaIdSqlType: TypeAlias = MediaIdType
 
 
 class SQLiteConfig(TypedDict):

aana/configs/pipeline.py

@@ -14,7 +14,9 @@
 from aana.models.pydantic.captions import CaptionsList, VideoCaptionsList
 from aana.models.pydantic.chat_message import ChatDialog
 from aana.models.pydantic.image_input import ImageInputList
+from aana.models.pydantic.media_id import MediaId
 from aana.models.pydantic.prompt import Prompt
+from aana.models.pydantic.question import Question
 from aana.models.pydantic.sampling_params import SamplingParams
 from aana.models.pydantic.video_input import VideoInput, VideoInputList
 from aana.models.pydantic.video_metadata import VideoMetadata
@@ -525,6 +527,7 @@
                 "name": "media_id",
                 "key": "media_id",
                 "path": "media_id",
+                "data_model": MediaId,
             }
         ],
     },
@@ -537,6 +540,7 @@
                 "name": "question",
                 "key": "question",
                 "path": "question",
+                "data_model": Question,
             }
         ],
     },

aana/exceptions/database.py

@@ -1,17 +1,17 @@
 from mobius_pipeline.exceptions import BaseException
 
-from aana.configs.db import media_id_type
+from aana.models.pydantic.media_id import MediaId
 
 
 class NotFoundException(BaseException):
     """Raised when an item searched by id is not found."""
 
-    def __init__(self, table_name: str, id: int | media_id_type):  # noqa: A002
+    def __init__(self, table_name: str, id: int | MediaId):  # noqa: A002
         """Constructor.
 
         Args:
             table_name (str): the name of the table being queried.
-            id (media_id_type): the id of the item to be retrieved.
+            id (int | MediaId): the id of the item to be retrieved.
         """
         super().__init__(table=table_name, id=id)
         self.table_name = table_name
@@ -26,12 +26,12 @@ def __reduce__(self):
 class MediaIdAlreadyExistsException(BaseException):
     """Raised when a media_id already exists."""
 
-    def __init__(self, table_name: str, media_id: media_id_type):
+    def __init__(self, table_name: str, media_id: MediaId):
         """Constructor.
 
         Args:
             table_name (str): the name of the table being queried.
-            media_id (media_id_type): the id of the item to be retrieved.
+            media_id (MediaId): the id of the item to be retrieved.
         """
         super().__init__(table=table_name, id=media_id)
         self.table_name = table_name

aana/exceptions/general.py

@@ -153,28 +153,6 @@ def __reduce__(self):
         return (self.__class__, (self.prompt_len, self.max_len))
 
 
-class MediaIdNotFoundException(BaseException):
-    """Exception raised when a media ID is not found.
-
-    Attributes:
-        media_id (str): the media ID
-    """
-
-    def __init__(self, media_id: str):
-        """Initialize the exception.
-
-        Args:
-            media_id (str): the media ID
-        """
-        super().__init__(media_id=media_id)
-        self.media_id = media_id
-        self.http_status_code = 404
-
-    def __reduce__(self):
-        """Used for pickling."""
-        return (self.__class__, (self.media_id,))
-
-
 class EndpointNotFoundException(BaseException):
     """Exception raised when an endpoint is not found.
 

aana/models/core/image.py

@@ -130,7 +130,7 @@ class Image(Media):
         url (str): The URL of the image.
         content (bytes): The content of the image in bytes (image file as bytes).
         numpy (np.ndarray): The image as a numpy array.
-        media_id (str): The ID of the image, generated automatically if not provided.
+        media_id (MediaId): The ID of the image, generated automatically if not provided.
     """
 
     media_dir: Path | None = settings.image_dir

aana/models/core/media.py

@@ -1,8 +1,8 @@
 import hashlib
-import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
 
+from aana.models.pydantic.media_id import MediaId
 from aana.utils.general import download_file
 
 
@@ -19,13 +19,13 @@ class Media:
         path (Path): the path to the media file
         url (str): the URL of the media
         content (bytes): the content of the media in bytes
-        media_id (str): the ID of the media. If not provided, it will be generated automatically.
+        media_id (MediaId): the ID of the media. If not provided, it will be generated automatically.
     """
 
     path: Path | None = None
     url: str | None = None
     content: bytes | None = None
-    media_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    media_id: MediaId = field(default_factory=lambda: MediaId.random())
     save_on_disk: bool = True
     is_saved: bool = False
     media_dir: Path | None = None

aana/models/core/video.py

@@ -20,7 +20,7 @@ class Video(Media):
         path (Path): the path to the video file
         url (str): the URL of the video
         content (bytes): the content of the video in bytes
-        media_id (str): the ID of the video. If not provided, it will be generated automatically.
+        media_id (MediaId): the ID of the video. If not provided, it will be generated automatically.
         title (str): the title of the video
         description (str): the description of the video
         media_dir (Path): the directory to save the video in

aana/models/db/caption.py

@@ -5,11 +5,12 @@
 from sqlalchemy import CheckConstraint, Column, Float, ForeignKey, Integer, String
 from sqlalchemy.orm import relationship
 
-from aana.configs.db import MediaIdSqlType, media_id_type
+from aana.configs.db import MediaIdSqlType
 from aana.models.db.base import BaseEntity, TimeStampEntity
 
 if typing.TYPE_CHECKING:
     from aana.models.pydantic.captions import Caption
+    from aana.models.pydantic.media_id import MediaId
 
 
 class CaptionEntity(BaseEntity, TimeStampEntity):
@@ -52,7 +53,7 @@ class CaptionEntity(BaseEntity, TimeStampEntity):
     def from_caption_output(
         cls,
         model_name: str,
-        media_id: media_id_type,
+        media_id: MediaId,
         video_id: int,
         frame_id: int,
         frame_timestamp: float,

aana/models/db/timeline.py.bak

@@ -5,7 +5,7 @@
 from sqlalchemy import JSON, CheckConstraint, Column, Float, ForeignKey, Integer, String
 from sqlalchemy.orm import relationship
 
-from aana.configs.db import MediaIdSqlType, media_id_type
+from aana.configs.db import MediaIdSqlType
 from aana.models.db.base import BaseEntity, TimeStampEntity
 
 if TYPE_CHECKING:
@@ -14,6 +14,7 @@
         AsrTranscription,
         AsrTranscriptionInfo,
     )
+    from aana.models.pydantic.media_id import MediaId
 
 
 class TranscriptEntity(BaseEntity, TimeStampEntity):
@@ -54,7 +55,7 @@ class TranscriptEntity(BaseEntity, TimeStampEntity):
     def from_asr_output(
         cls,
         model_name: str,
-        media_id: media_id_type,
+        media_id: MediaId,
         video_id: int,
         info: AsrTranscriptionInfo,
         transcription: AsrTranscription,

aana/models/db/transcript.py

@@ -1,4 +1,5 @@
 from pydantic import BaseModel
+from pyparsing import Any
 
 
 class BaseListModel(BaseModel):
@@ -34,3 +35,75 @@ def __contains__(self, item):
     def __add__(self, other):
         """Add two models."""
         return self.__class__(__root__=self.__root__ + other.__root__)
+
+
+class BaseStringModel(BaseModel):
+    """The base model for pydantic models that are just strings."""
+
+    __root__: str
+
+    def __init__(self, __root__value: Any = None, **data):
+        """Initialize the model."""
+        if __root__value is not None:
+            super().__init__(__root__=__root__value, **data)
+        else:
+            super().__init__(**data)
+
+    def __str__(self) -> str:
+        """Convert to a string."""
+        return self.__root__
+
+    def __repr__(self) -> str:
+        """Convert to a string representation."""
+        return f"{self.__class__.__name__}({self.__root__!r})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check if two models are equal."""
+        if isinstance(other, self.__class__):
+            return self.__root__ == other.__root__
+        if isinstance(other, str):
+            return self.__root__ == other
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        """Get hash of model."""
+        return hash(self.__root__)
+
+    def __getitem__(self, key):
+        """Get item at key of model."""
+        return self.__root__[key]
+
+    def __len__(self) -> int:
+        """Get length of model."""
+        return len(self.__root__)
+
+    def __iter__(self):
+        """Get iterator for model."""
+        return iter(self.__root__)
+
+    def __contains__(self, item):
+        """Check if modle contains item."""
+        return item in self.__root__
+
+    def __add__(self, other):
+        """Add two models or a model and a string."""
+        if isinstance(other, self.__class__):
+            return self.__class__(__root__=self.__root__ + other.__root__)
+        if isinstance(other, str):
+            return str(self.__root__) + other
+        return NotImplemented
+
+    def __getattr__(self, item):
+        """Automatically delegate method calls to self.__root__ if they are not found in the model.
+
+        Check if the attribute is a callable (method) of __root__ and return a wrapped call if it is.
+        This will handle methods like startswith, endswith, and split.
+        """
+        attr = getattr(self.__root__, item)
+        if callable(attr):
+
+            def wrapper(*args, **kwargs):
+                return attr(*args, **kwargs)
+
+            return wrapper
+        return attr

aana/models/pydantic/base.py

@@ -1,5 +1,4 @@
 import io
-import uuid
 from pathlib import Path
 from types import MappingProxyType
 
@@ -9,6 +8,7 @@
 
 from aana.models.core.image import Image
 from aana.models.pydantic.base import BaseListModel
+from aana.models.pydantic.media_id import MediaId
 
 
 class ImageInput(BaseModel):
@@ -45,8 +45,8 @@ class ImageInput(BaseModel):
             "Set this field to 'file' to upload files to the endpoint."
         ),
     )
-    media_id: str = Field(
-        default_factory=lambda: str(uuid.uuid4()),
+    media_id: MediaId = Field(
+        default_factory=lambda: MediaId.random(),
         description="The ID of the image. If not provided, it will be generated automatically.",
     )
 
@@ -55,7 +55,7 @@ def media_id_must_not_be_empty(cls, media_id):
         """Validates that the media_id is not an empty string.
 
         Args:
-            media_id (str): The value of the media_id field.
+            media_id (MediaId): The value of the media_id field.
 
         Raises:
             ValueError: If the media_id is an empty string.