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

Ankit/refactor
video-db · Oct 23, 2024 · bf56081 · bf56081
2 parents d977dc1 + f93fec7
commit bf56081
Show file tree

Hide file tree

Showing 23 changed files with 213 additions and 127 deletions.
diff --git a/backend/.env.sample b/backend/.env.sample
@@ -1,9 +1,16 @@
+# VideoDB Integration
+VIDEO_DB_API_KEY=
 
-# OpenAI config
+# LLM Integrations
 OPENAI_API_KEY=
 
-# VideoDB config
-VIDEO_DB_API_KEY=
-# only for dev
-VIDEO_DB_BASE_URL=
+# Tools
+REPLICATE_API_TOKEN=
+
+# Brandkit Agent
+INTRO_VIDEO_ID=
+OUTRO_VIDEO_ID=
+BRAND_IMAGE_ID=
 
+# Profanity Remover Agent
+BEEP_AUDIO_ID=
diff --git a/backend/spielberg/agents/base.py b/backend/spielberg/agents/base.py
@@ -10,13 +10,13 @@
 logger = logging.getLogger(__name__)
 
 
-class AgentResult:
+class AgentStatus:
     SUCCESS = "success"
     ERROR = "error"
 
 
 class AgentResponse(BaseModel):
-    result: str = AgentResult.SUCCESS
+    status: str = AgentStatus.SUCCESS
     message: str = ""
     data: dict = {}
 
@@ -30,7 +30,7 @@ def __init__(self, session: Session, **kwargs):
 
     def get_parameters(self):
         function_inferrer = FunctionInferrer.infer_from_function_reference(
-            self.__call__
+            self.run
         )
         function_json = function_inferrer.to_json_schema()
         parameters = function_json.get("parameters")
@@ -58,12 +58,12 @@ def agent_description(self):
 
     def safe_call(self, *args, **kwargs):
         try:
-            return self.__call__(*args, **kwargs)
+            return self.run(*args, **kwargs)
 
         except Exception as e:
             logger.exception(f"error in {self.agent_name} agent: {e}")
-            return AgentResponse(result=AgentResult.ERROR, message=str(e))
+            return AgentResponse(status=AgentStatus.ERROR, message=str(e))
 
     @abstractmethod
-    def __call__(*args, **kwargs) -> AgentResponse:
+    def run(*args, **kwargs) -> AgentResponse:
         pass
diff --git a/backend/spielberg/agents/brandkit.py b/backend/spielberg/agents/brandkit.py
@@ -1,8 +1,8 @@
 import os
 import logging
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
-from spielberg.core.session import Session, MsgStatus, VideoContent
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
+from spielberg.core.session import Session, MsgStatus, VideoContent, VideoData
 from spielberg.tools.videodb_tool import VideoDBTool
 
 logger = logging.getLogger(__name__)
@@ -22,7 +22,7 @@ def __init__(self, session: Session, **kwargs):
         self.parameters = self.get_parameters()
         super().__init__(session=session, **kwargs)
 
-    def __call__(
+    def run(
         self,
         collection_id: str,
         video_id: str,
@@ -55,7 +55,7 @@ def __call__(
                     "Branding elementes not provided, either you can provide provide IDs for intro video, outro video and branding image"
                     " or you can set INTRO_VIDEO_ID, OUTRO_VIDEO_ID and BRAND_IMAGE_ID in .env of backend directory."
                 )
-                return AgentResponse(result=AgentResult.ERROR, message=message)
+                return AgentResponse(status=AgentStatus.ERROR, message=message)
             video_content = VideoContent(
                 agent_name=self.agent_name,
                 status=MsgStatus.progress,
@@ -67,7 +67,7 @@ def __call__(
             brandkit_stream = videodb_tool.add_brandkit(
                 video_id, intro_video_id, outro_video_id, brand_image_id
             )
-            video_content.video = {"stream_url": brandkit_stream}
+            video_content.video = VideoData(stream_url=brandkit_stream)
             video_content.status = MsgStatus.success
             video_content.status_message = "Here is your brandkit stream"
             self.output_message.publish()
@@ -77,9 +77,9 @@ def __call__(
             error_message = "Error in adding branding."
             video_content.status_message = error_message
             self.output_message.publish()
-            return AgentResponse(result=AgentResult.ERROR, message=error_message)
+            return AgentResponse(status=AgentStatus.ERROR, message=error_message)
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message=f"Agent {self.name} completed successfully.",
             data={"stream_url": brandkit_stream},
         )
diff --git a/backend/spielberg/agents/download.py b/backend/spielberg/agents/download.py
@@ -1,6 +1,6 @@
 import logging
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
 from spielberg.core.session import Session
 from spielberg.tools.videodb_tool import VideoDBTool
 
@@ -14,7 +14,7 @@ def __init__(self, session: Session, **kwargs):
         self.parameters = self.get_parameters()
         super().__init__(session=session, **kwargs)
 
-    def __call__(
+    def run(
         self, stream_link: str, name: str = None, *args, **kwargs
     ) -> AgentResponse:
         """
@@ -34,9 +34,9 @@ def __call__(
             download_response = videodb_tool.download(stream_link, name)
         except Exception as e:
             logger.exception(f"error in {self.agent_name} agent: {e}")
-            return AgentResponse(result=AgentResult.ERROR, message=str(e))
+            return AgentResponse(status=AgentStatus.ERROR, message=str(e))
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message="Download successful",
             data=download_response,
         )
diff --git a/backend/spielberg/agents/image_generation.py b/backend/spielberg/agents/image_generation.py
@@ -1,7 +1,7 @@
 import logging
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
-from spielberg.core.session import Session, MsgStatus, ImageContent
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
+from spielberg.core.session import Session, MsgStatus, ImageContent, ImageData
 from spielberg.tools.replicate import flux_dev
 
 logger = logging.getLogger(__name__)
@@ -14,7 +14,7 @@ def __init__(self, session: Session, **kwargs):
         self.parameters = self.get_parameters()
         super().__init__(session=session, **kwargs)
 
-    def __call__(self, prompt: str, *args, **kwargs) -> AgentResponse:
+    def run(self, prompt: str, *args, **kwargs) -> AgentResponse:
         """
         Process the prompt to generate the image.
 
@@ -39,9 +39,9 @@ def __call__(self, prompt: str, *args, **kwargs) -> AgentResponse:
                 image_content.status_message = "Error in generating image."
                 self.output_message.publish()
                 error_message = "Agent failed with error in replicate."
-                return AgentResponse(result=AgentResult.ERROR, message=error_message)
+                return AgentResponse(status=AgentStatus.ERROR, message=error_message)
             image_url = flux_output[0].url
-            image_content.image = {"url": image_url}
+            image_content.image = ImageData(url=image_url)
             image_content.status = MsgStatus.success
             image_content.status_message = "Here is your generated image"
             self.output_message.publish()
@@ -51,9 +51,9 @@ def __call__(self, prompt: str, *args, **kwargs) -> AgentResponse:
             image_content.status_message = "Error in generating image."
             self.output_message.publish()
             error_message = f"Agent failed with error {e}"
-            return AgentResponse(result=AgentResult.ERROR, message=error_message)
+            return AgentResponse(status=AgentStatus.ERROR, message=error_message)
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message=f"Agent {self.name} completed successfully.",
-            data={"image_url": image_url},
+            data={},
         )
diff --git a/backend/spielberg/agents/index.py b/backend/spielberg/agents/index.py
@@ -1,6 +1,6 @@
 import logging
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
 
 from spielberg.core.session import Session
 from spielberg.tools.videodb_tool import VideoDBTool
@@ -34,7 +34,7 @@ def __init__(self, session: Session, **kwargs):
         self.parameters = INDEX_AGENT_PARAMETERS
         super().__init__(session=session, **kwargs)
 
-    def __call__(
+    def run(
         self, video_id: str, index_type: str, collection_id=None, *args, **kwargs
     ) -> AgentResponse:
         """
@@ -65,10 +65,10 @@ def __call__(
 
         except Exception as e:
             logger.exception(f"error in {self.agent_name} agent: {e}")
-            return AgentResponse(result=AgentResult.ERROR, message=str(e))
+            return AgentResponse(status=AgentStatus.ERROR, message=str(e))
 
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message=f"{index_type} indexing successful",
             data=scene_data,
         )
diff --git a/backend/spielberg/agents/pricing.py b/backend/spielberg/agents/pricing.py
@@ -1,6 +1,6 @@
 import logging
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
 from spielberg.core.session import (
     Session,
     MsgStatus,
@@ -97,7 +97,7 @@ def __init__(self, session: Session, **kwargs):
         self.llm = OpenAI()
         super().__init__(session=session, **kwargs)
 
-    def __call__(self, query: str, *args, **kwargs) -> AgentResponse:
+    def run(self, query: str, *args, **kwargs) -> AgentResponse:
         """
         Get the answer to the query of agent
 
@@ -125,7 +125,7 @@ def __call__(self, query: str, *args, **kwargs) -> AgentResponse:
                 text_content.status_message = "Failed to generate the response."
                 self.output_message.publish()
                 return AgentResponse(
-                    result=AgentResult.ERROR,
+                    status=AgentStatus.ERROR,
                     message="Pricing failed due to LLM error.",
                 )
             text_content.text = llm_response.content
@@ -138,10 +138,10 @@ def __call__(self, query: str, *args, **kwargs) -> AgentResponse:
             text_content.status = MsgStatus.error
             text_content.status_message = error_message
             self.output_message.publish()
-            return AgentResponse(result=AgentResult.ERROR, message=error_message)
+            return AgentResponse(status=AgentStatus.ERROR, message=error_message)
 
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message="Fetch successful and output displayed above.",
             data={},
         )
diff --git a/backend/spielberg/agents/profanity_remover.py b/backend/spielberg/agents/profanity_remover.py
@@ -4,11 +4,12 @@
 
 from videodb.asset import VideoAsset, AudioAsset
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
 from spielberg.core.session import (
     Session,
     MsgStatus,
     VideoContent,
+    VideoData,
     ContextMessage,
     RoleTypes,
 )
@@ -45,7 +46,7 @@ def add_beep(self, videodb_tool, video_id, beep_audio_id, timestamps):
         stream_url = timeline.generate_stream()
         return stream_url
 
-    def __call__(
+    def run(
         self,
         collection_id: str,
         video_id: str,
@@ -68,7 +69,7 @@ def __call__(
             beep_audio_id = beep_audio_id or BEEP_AUDIO_ID
             if not beep_audio_id:
                 return AgentResponse(
-                    result=AgentResult.failed,
+                    status=AgentStatus.failed,
                     message="Please provide the beep_audio_id or setup BEEP_AUDIO_ID in .env of backend directory.",
                 )
             self.output_message.actions.append("Started process to remove profanity..")
@@ -100,7 +101,7 @@ def __call__(
             clean_stream = self.add_beep(
                 videodb_tool, video_id, beep_audio_id, profanity_timeline
             )
-            video_content.video = {"stream_url": clean_stream}
+            video_content.video = VideoData(stream_url=clean_stream)
             video_content.status = MsgStatus.success
             video_content.status_message = "Here is the clean stream"
             self.output_message.content.append(video_content)
@@ -111,9 +112,9 @@ def __call__(
             video_content.status_message = "Failed to generate clean stream"
             self.output_message.publish()
             error_message = f"Error in generating the clean stream due to {e}."
-            return AgentResponse(result=AgentResult.ERROR, message=error_message)
+            return AgentResponse(status=AgentStatus.ERROR, message=error_message)
         return AgentResponse(
-            result=AgentResult.SUCCESS,
+            status=AgentStatus.SUCCESS,
             message=f"Agent {self.name} completed successfully.",
             data={"stream_url": clean_stream},
         )
diff --git a/backend/spielberg/agents/prompt_clip.py b/backend/spielberg/agents/prompt_clip.py
@@ -2,13 +2,14 @@
 import json
 import concurrent.futures
 
-from spielberg.agents.base import BaseAgent, AgentResponse, AgentResult
+from spielberg.agents.base import BaseAgent, AgentResponse, AgentStatus
 from spielberg.core.session import (
     Session,
     ContextMessage,
     RoleTypes,
     MsgStatus,
     VideoContent,
+    VideoData,
 )
 from spielberg.tools.videodb_tool import VideoDBTool
 from spielberg.llm.openai import OpenAI
@@ -120,7 +121,7 @@ def _text_prompter(self, transcript_text, prompt):
                     continue
         return matches
 
-    def __call__(
+    def run(
         self, prompt: str, video_id: str, collection_id: str, *args, **kwargs
     ) -> AgentResponse:
         try:
@@ -178,27 +179,25 @@ def __call__(
                         video_id=video_id, timeline=timeline
                     )
                     video_content.status_message = "Clip generated successfully."
-                    video_content.video = {
-                        "stream_url": stream_url,
-                    }
+                    video_content.video = VideoData(stream_url=stream_url)
                     video_content.status = MsgStatus.success
                     self.output_message.publish()
 
                 except Exception as e:
                     logger.exception(f"Error in creating video content: {e}")
-                    return AgentResponse(result=AgentResult.ERROR, message=str(e))
+                    return AgentResponse(status=AgentStatus.ERROR, message=str(e))
+
+                return AgentResponse(
+                    status=AgentStatus.SUCCESS,
+                    message=f"Agent {self.name} completed successfully.",
+                    data={"stream_url": stream_url},
+                )
             else:
                 return AgentResponse(
-                    result=AgentResult.ERROR,
+                    status=AgentStatus.ERROR,
                     message="No relevant moments found.",
                 )
 
         except Exception as e:
             logger.exception(f"error in {self.agent_name}")
-            return AgentResponse(result=AgentResult.ERROR, message=str(e))
-
-        return AgentResponse(
-            result=AgentResult.SUCCESS,
-            message=f"Agent {self.name} completed successfully.",
-            data={},
-        )
+            return AgentResponse(status=AgentStatus.ERROR, message=str(e))