Add local upload support and improve playlist exception

backend/director/agents/upload.py

@@ -17,9 +17,14 @@
 UPLOAD_AGENT_PARAMETERS = {
     "type": "object",
     "properties": {
-        "url": {
+        "source": {
             "type": "string",
-            "description": "URL to upload the content",
+            "description": "URL or local path to upload the content",
+        },
+        "source_type": {
+            "type": "string",
+            "description": "Type of given source.",
+            "enum": ["url", "local_file"],
         },
         "name": {
             "type": "string",
@@ -42,11 +47,16 @@
 class UploadAgent(BaseAgent):
     def __init__(self, session: Session, **kwargs):
         self.agent_name = "upload"
-        self.description = "Uploads media content to the VideoDB. This agent takes a URL of the media content and uploads it to the VideoDB. The media content can be a video, audio, or image file. Youtube playlist and links are also supported."
+        self.description = (
+            "This agent uploads the media content to VideoDB. "
+            "This agent takes a source which can be a URL or local path of the media content. "
+            "The media content can be a video, audio, or image file. "
+            "Youtube playlist and links are also supported. "
+        )
         self.parameters = UPLOAD_AGENT_PARAMETERS
         super().__init__(session=session, **kwargs)
 
-    def _upload(self, url: str, media_type: str, name: str):
+    def _upload(self, source: str, source_type: str, media_type: str, name: str = None):
         """Upload the media with the given URL."""
         try:
             if media_type == "video":
@@ -61,14 +71,16 @@ def _upload(self, url: str, media_type: str, name: str):
             content.status_message = f"Uploading {media_type}..."
             self.output_message.push_update()
 
-            upload_data = self.videodb_tool.upload(url, media_type, name=name)
+            upload_data = self.videodb_tool.upload(
+                source, source_type, media_type, name=name
+            )
 
             content.status_message = f"{upload_data['name']} uploaded successfully"
             if media_type == "video":
                 content.video = VideoData(**upload_data)
             else:
                 content.text = (
-                    f"\n ID: {upload_data['id']}, TITLE: {upload_data['name']}"
+                    f"\n ID: {upload_data['id']}, Title: {upload_data['name']}"
                 )
             content.status = MsgStatus.success
             self.output_message.publish()
@@ -107,52 +119,57 @@ def _get_yt_playlist_videos(self, playlist_url: str):
 
     def _upload_yt_playlist(self, playlist_info: dict, media_type):
         """Upload the videos in a youtube playlist."""
-        try:
-            for media in playlist_info:
+        for media in playlist_info:
+            try:
                 self.output_message.actions.append(
                     f"Uploading video: {media['title']} as {media_type}"
                 )
-                self._upload(media["url"], media_type)
-        except Exception as e:
-            logger.exception(f"Error in uploading playlist: {e}")
-            return AgentResponse(status=AgentStatus.ERROR, message=str(e))
-
+                self._upload(media["url"], "url", media_type)
+            except Exception as e:
+                self.output_message.actions.append(
+                    f"Upload failed for {media['title']}"
+                )
+                logger.exception(f"Error in uploading {media['title']}: {e}")
         return AgentResponse(
             status=AgentStatus.SUCCESS,
             message="All the videos in the playlist uploaded successfully as {media_type}",
         )
 
     def run(
         self,
-        url: str,
+        collection_id: str,
+        source: str,
+        source_type: str,
         media_type="video",
-        collection_id: str = None,
         name: str = None,
         *args,
         **kwargs,
     ) -> AgentResponse:
         """
-        Upload the media with the given URL.
-
-        :param url: The URL of the media to upload.
-        :type url: str
-        :param media_type: The type of media to upload, defaults to "video"
-        :type media_type: str, optional
-        :return: The response containing information about the upload operation.
-        :rtype: AgentResponse
+        Upload the media with the given source.
+
+        :param collection_id: str - collection_id in which the upload is required.
+        :param source: str - The URL or local path of the media to upload.
+        :param source_type: str - The type indicating the source of the media.
+        :param media_type: str, optional - The type of media to upload, defaults to "video".
+        :param name: str, optional - Name required for uploaded file.
+        :return: AgentResponse - The response containing information about the upload operation.
         """
 
-        if collection_id is None:
-            self.videodb_tool = VideoDBTool()
+        self.videodb_tool = VideoDBTool(collection_id=collection_id)
+
+        if source_type == "local_file":
+            return self._upload(source, source_type, media_type, name)
+        elif source_type == "url":
+            playlist_info = self._get_yt_playlist_videos(source)
+            if playlist_info:
+                self.output_message.actions.append("YouTube Playlist detected")
+                self.output_message.push_update()
+                return self._upload_yt_playlist(playlist_info, media_type)
+            return self._upload(source, source_type, media_type, name)
         else:
-            self.videodb_tool = VideoDBTool(collection_id=collection_id)
-
-        # check if the url is a youtube playlist
-        playlist_info = self._get_yt_playlist_videos(url)
-        if playlist_info:
-            self.output_message.actions.append("Youtube playlist detected")
-            self.output_message.push_update()
-            return self._upload_yt_playlist(playlist_info, media_type)
-
-        # upload the media
-        return self._upload(url, media_type, name)
+            error_message = f"Invalid source type {source_type}"
+            logger.error(error_message)
+            return AgentResponse(
+                status=AgentStatus.ERROR, message=error_message, data={}
+            )

backend/director/tools/videodb_tool.py

@@ -64,13 +64,16 @@ def get_videos(self):
             for video in videos
         ]
 
-    def upload(self, url, media_type, name=None):
-        if name is None:
-            media = self.conn.upload(url=url, media_type=media_type)
-            name = media.name
+    def upload(self, source, source_type="url", media_type="video", name=None):
+        upload_args = {"media_type": media_type}
+        if name:
+            upload_args["name"] = name
+        if source_type == "url":
+            upload_args["url"] = source
         else:
-            media = self.conn.upload(url=url, media_type=media_type, name=name)
-
+            upload_args["file_path"] = source
+        media = self.conn.upload(**upload_args)
+        name = media.name
         if media_type == "video":
             return {
                 "id": media.id,