Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add local upload support and improve playlist exception #46

Merged
merged 3 commits into from
Oct 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 53 additions & 36 deletions backend/director/agents/upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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":
Expand All @@ -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()
Expand Down Expand Up @@ -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={}
)
15 changes: 9 additions & 6 deletions backend/director/tools/videodb_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down