diff --git a/backend/director/agents/upload.py b/backend/director/agents/upload.py index 19f71042..522cc562 100644 --- a/backend/director/agents/upload.py +++ b/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 = None): + 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 = None): 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,16 +119,17 @@ 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}", @@ -124,35 +137,39 @@ def _upload_yt_playlist(self, playlist_info: dict, 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={} + ) diff --git a/backend/director/tools/videodb_tool.py b/backend/director/tools/videodb_tool.py index 68e3636e..b2a68dab 100644 --- a/backend/director/tools/videodb_tool.py +++ b/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,