Fal tool

cookbook/agents/43_generate_replicate_video.py

@@ -1,16 +1,16 @@
 from phi.agent import Agent
 from phi.model.openai import OpenAIChat
-from phi.tools.replicate import ReplicateToolKit
+from phi.tools.replicate import ReplicateToolkit
 
 """Create an agent specialized for Replicate AI content generation"""
 
 video_agent = Agent(
     name="Video Generator Agent",
     model=OpenAIChat(id="gpt-4o"),
     tools=[
-        ReplicateToolKit(model="tencent/hunyuan-video:847dfa8b01e739637fc76f480ede0c1d76408e1d694b830b5dfb8e547bf98405")
+        ReplicateToolkit(model="tencent/hunyuan-video:847dfa8b01e739637fc76f480ede0c1d76408e1d694b830b5dfb8e547bf98405")
     ],
-    description="You are an AI agent that can generate images and videos using the Replicate API.",
+    description="You are an AI agent that can generate videos using the Replicate API.",
     instructions=[
         "When the user asks you to create a video, use the `generate_media` tool to create the video.",
         "Return the URL as raw to the user.",

cookbook/agents/44_generate_replicate_image.py

@@ -0,0 +1,24 @@
+from phi.agent import Agent
+from phi.model.openai import OpenAIChat
+from phi.tools.replicate import ReplicateToolkit
+
+"""Create an agent specialized for Replicate AI content generation"""
+
+video_agent = Agent(
+    name="Image Generator Agent",
+    model=OpenAIChat(id="gpt-4o"),
+    tools=[
+        ReplicateToolkit(model="luma/photon-flash")
+    ],
+    description="You are an AI agent that can generate images using the Replicate API.",
+    instructions=[
+        "When the user asks you to create an image, use the `generate_media` tool to create the image.",
+        "Return the URL as raw to the user.",
+        "Don't convert image URL to markdown or anything else.",
+    ],
+    markdown=True,
+    debug_mode=True,
+    show_tool_calls=True,
+)
+
+video_agent.print_response("Generate an image of a horse in the dessert.")

phi/tools/replicate.py

@@ -1,8 +1,10 @@
+import os
 from os import getenv
+from urllib.parse import urlparse
 from uuid import uuid4
 
 from phi.agent import Agent
-from phi.model.content import Video
+from phi.model.content import Video, Image
 from phi.tools import Toolkit
 from phi.utils.log import logger
 
@@ -13,19 +15,16 @@
     raise ImportError("`replicate` not installed. Please install using `pip install replicate`.")
 
 
-class ReplicateToolKit(Toolkit):
+class ReplicateToolkit(Toolkit):
     def __init__(
         self,
         model: str = "minimax/video-01",
-        wait_for_completion: bool = True,
-        max_wait_time: int = 300,  # 5 minutes
     ):
         super().__init__(name="replicate_toolkit")
         self.api_key = getenv("REPLICATE_API_TOKEN")
         if not self.api_key:
             logger.error("REPLICATE_API_TOKEN not set. Please set the REPLICATE_API_TOKEN environment variable.")
         self.model = model
-
         self.register(self.generate_media)
 
     def generate_media(self, agent: Agent, prompt: str) -> str:
@@ -38,12 +37,36 @@ def generate_media(self, agent: Agent, prompt: str) -> str:
         """
         output: FileOutput = replicate.run(ref=self.model, input={"prompt": prompt})
 
-        # Update the run response with the video URLs
-        agent.add_video(
-            Video(
-                id=str(uuid4()),
-                url=output.url,
+        # Parse the URL to extract the file extension
+        parsed_url = urlparse(output.url)
+        path = parsed_url.path
+        _, ext = os.path.splitext(path)
+        ext = ext.lower()
+
+        # Define supported extensions
+        image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
+        video_extensions = {'.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm'}
+
+        media_id = str(uuid4())
+
+        if ext in image_extensions:
+            agent.add_image(
+                Image(
+                    id=media_id,
+                    url=output.url,
+                )
+            )
+            media_type = "image"
+        elif ext in video_extensions:
+            agent.add_video(
+                Video(
+                    id=media_id,
+                    url=output.url,
+                )
             )
-        )
+            media_type = "video"
+        else:
+            logger.error(f"Unsupported media type with extension '{ext}' for URL: {output.url}")
+            return f"Unsupported media type with extension '{ext}'."
 
-        return f"Media generated successfully at {output.url}"
+        return f"{media_type.capitalize()} generated successfully at {output.url}"