diff --git a/backend/spielberg/agents/upload.py b/backend/spielberg/agents/upload.py index 26295c2..1313caf 100644 --- a/backend/spielberg/agents/upload.py +++ b/backend/spielberg/agents/upload.py @@ -15,6 +15,10 @@ "type": "string", "description": "URL to upload the content", }, + "name": { + "type": "string", + "description": "Name of the content to upload", + }, "media_type": { "type": "string", "enum": ["video", "audio", "image"], @@ -36,7 +40,7 @@ def __init__(self, session: Session, **kwargs): self.parameters = UPLOAD_AGENT_PARAMETERS super().__init__(session=session, **kwargs) - def _upload(self, url: str, media_type: str): + def _upload(self, url: str, media_type: str, name: str): """Upload the media with the given URL.""" try: if media_type == "video": @@ -51,7 +55,7 @@ def _upload(self, url: str, media_type: str): content.status_message = f"Uploading {media_type}..." self.output_message.push_update() - upload_data = self.videodb_tool.upload(url, media_type) + upload_data = self.videodb_tool.upload(url, media_type, name=name) content.status_message = f"{upload_data['name']} uploaded successfully" if media_type == "video": @@ -108,7 +112,13 @@ def _upload_yt_playlist(self, playlist_info: dict, media_type): ) def __call__( - self, url: str, media_type="video", collection_id: str = None, *args, **kwargs + self, + url: str, + media_type="video", + collection_id: str = None, + name: str = None, + *args, + **kwargs, ) -> AgentResponse: """ Upload the media with the given URL. @@ -134,7 +144,7 @@ def __call__( return self._upload_yt_playlist(playlist_info, media_type) # upload the media - upload_data = self._upload(url, media_type) + upload_data = self._upload(url, media_type, name) self.output_message.publish() return AgentResponse( diff --git a/backend/spielberg/tools/videodb_tool.py b/backend/spielberg/tools/videodb_tool.py index 556fab8..4b54559 100644 --- a/backend/spielberg/tools/videodb_tool.py +++ b/backend/spielberg/tools/videodb_tool.py @@ -64,8 +64,12 @@ def get_videos(self): for video in videos ] - def upload(self, url, media_type): - media = self.conn.upload(url=url, media_type=media_type) + 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 + else: + media = self.conn.upload(url=url, media_type=media_type, name=name) if media_type == "video": return { @@ -73,7 +77,7 @@ def upload(self, url, media_type): "collection_id": media.collection_id, "stream_url": media.stream_url, "player_url": media.player_url, - "name": media.name, + "name": name, "description": media.description, "thumbnail_url": media.thumbnail_url, "length": media.length,