From da3c1fbaf434f0792f19bc06cd5daf3b24dcff28 Mon Sep 17 00:00:00 2001
From: Ankit raj <113342181+ankit-v2-3@users.noreply.github.com>
Date: Thu, 24 Oct 2024 14:00:20 +0530
Subject: [PATCH] feat: add search data and shot data

- Add SearchData and ShotData for SearchResultsContent
- Fix upload agent
---
 backend/spielberg/agents/search.py | 17 +++++++++++++++--
 backend/spielberg/agents/upload.py |  5 ++---
 backend/spielberg/core/session.py  | 23 +++++++++++++++++++++--
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/backend/spielberg/agents/search.py b/backend/spielberg/agents/search.py
index ebc20c9..b5b3e1d 100644
--- a/backend/spielberg/agents/search.py
+++ b/backend/spielberg/agents/search.py
@@ -7,6 +7,8 @@
     MsgStatus,
     TextContent,
     SearchResultsContent,
+    SearchData,
+    ShotData,
     VideoContent,
     VideoData,
     ContextMessage,
@@ -110,10 +112,21 @@ def run(
                             }
                         ],
                     }
-            search_result_content.search_results = list(search_result_videos.values())
+            search_result_content.search_results = [
+                SearchData(
+                    video_id=sr["video_id"],
+                    video_title=sr["video_title"],
+                    stream_url=sr["stream_url"],
+                    duration=sr["duration"],
+                    shots=[ShotData(**shot) for shot in sr["shots"]],
+                )
+                for sr in search_result_videos.values()
+            ]
             search_result_content.status = MsgStatus.success
             search_result_content.status_message = "Search done."
-            self.output_message.actions.append("Generating search result compilation clip..")
+            self.output_message.actions.append(
+                "Generating search result compilation clip.."
+            )
             self.output_message.push_update()
             compilation_stream_url = search_results.compile()
             compilation_content.video = VideoData(stream_url=compilation_stream_url)
diff --git a/backend/spielberg/agents/upload.py b/backend/spielberg/agents/upload.py
index 2e88b25..227cbe5 100644
--- a/backend/spielberg/agents/upload.py
+++ b/backend/spielberg/agents/upload.py
@@ -23,7 +23,7 @@
         },
         "name": {
             "type": "string",
-            "description": "Name of the content to upload",
+            "description": "Name of the content to upload, optional parameter",
         },
         "media_type": {
             "type": "string",
@@ -155,5 +155,4 @@ def run(
             return self._upload_yt_playlist(playlist_info, media_type)
 
         # upload the media
-        return self._upload(url, media_type)
-
+        return self._upload(url, media_type, name)
diff --git a/backend/spielberg/core/session.py b/backend/spielberg/core/session.py
index af60f43..c9029b0 100644
--- a/backend/spielberg/core/session.py
+++ b/backend/spielberg/core/session.py
@@ -91,8 +91,27 @@ class ImageContent(BaseContent):
     type: ContentType = ContentType.image
 
 
+class ShotData(BaseModel):
+    """Shot data model class for search results content."""
+
+    search_score: Union[int, float]
+    start: Union[int, float]
+    end: Union[int, float]
+    text: str
+
+
+class SearchData(BaseModel):
+    """Search data model class for search results content."""
+
+    video_id: str
+    video_title: str
+    stream_url: str
+    duration: Union[int, float]
+    shots: List[ShotData]
+
+
 class SearchResultsContent(BaseContent):
-    search_results: dict = {}
+    search_results: Optional[List[SearchData]] = None
     type: ContentType = ContentType.search_results
 
 
@@ -127,7 +146,7 @@ class BaseMessage(BaseModel):
 
 class InputMessage(BaseMessage):
     """Input message to the agent
-    
+
     :param BaseDB db: Database instance
     :param MsgType msg_type: :class:`MsgType` of the message
     """