adds ui config for videos content

Laksh-star · Jan 3, 2025 · b0f62cd · b0f62cd
1 parent dbba8c6
commit b0f62cd
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/backend/director/agents/comparison.py b/backend/director/agents/comparison.py
@@ -2,7 +2,13 @@
 import asyncio
 
 from director.agents.base import BaseAgent, AgentResponse, AgentStatus
-from director.core.session import Session, VideosContent, VideoData, MsgStatus
+from director.core.session import (
+    Session,
+    VideosContent,
+    VideoData,
+    MsgStatus,
+    VideosContentUIConfig,
+)
 from director.agents.video_generation import VideoGenerationAgent
 from director.agents.video_generation import VIDEO_GENERATION_AGENT_PARAMETERS
 
@@ -94,6 +100,7 @@ def run(
             if job_type == "video_generation_comparison":
                 self.videos_content = VideosContent(
                     agent_name=self.agent_name,
+                    ui_config=VideosContentUIConfig(columns=3),
                     status=MsgStatus.progress,
                     status_message="Generating videos (Usually takes 3-7 mins)",
                     videos=[],

diff --git a/backend/director/core/session.py b/backend/director/core/session.py
@@ -89,10 +89,15 @@ class VideoContent(BaseContent):
     type: ContentType = ContentType.video
 
 
+class VideosContentUIConfig(BaseModel):
+    columns: Optional[int] = 4
+
+
 class VideosContent(BaseContent):
     """Videos content model class for videos content."""
 
     videos: Optional[List[VideoData]] = None
+    ui_config: VideosContentUIConfig = VideosContentUIConfig()
     type: ContentType = ContentType.videos
 
 
@@ -152,7 +157,14 @@ class BaseMessage(BaseModel):
     actions: List[str] = []
     agents: List[str] = []
     content: List[
-        Union[dict, TextContent, ImageContent, VideoContent, VideosContent, SearchResultsContent]
+        Union[
+            dict,
+            TextContent,
+            ImageContent,
+            VideoContent,
+            VideosContent,
+            SearchResultsContent,
+        ]
     ] = []
     status: MsgStatus = MsgStatus.success
     msg_id: str = Field(