Merge pull request #13 from video-db/ankit/refactor

feat: add search data and shot data
video-db · Oct 24, 2024 · f14e400 · f14e400
2 parents bf56081 + da3c1fb
commit f14e400
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
     """