自己用gemni寫image tool

litellm_endpoint.py

@@ -2,11 +2,15 @@ import os
 import datetime
 import pytz
 import yaml
+from PIL import Image
+from google import genai
+from google.genai import types
+from io import BytesIO
 
 from smolagents import (
     CodeAgent,
     DuckDuckGoSearchTool,
-    load_tool,
+    # load_tool,
     tool,
     LiteLLMModel,
 )
@@ -43,7 +47,7 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 
 # 初始化工具
-search_tool = DuckDuckGoSearchTool()
+web_search = DuckDuckGoSearchTool()
 final_answer = FinalAnswerTool()
 
 # 官方文件寫法
@@ -54,7 +58,31 @@ model = LiteLLMModel(
     api_key=os.getenv("GEMINI_API_KEY"))
 
 # 載入圖像產生工具
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+@tool
+def image_generation_tool(prompt: str) -> Image.Image:
+    """Generates an image based on the provided prompt.
+
+    Args:
+        prompt: The prompt for the image generation.
+    """
+    # Placeholder for actual image generation logic
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-exp-image-generation",
+        contents=prompt,
+        config=types.GenerateContentConfig(
+            response_modalities=["TEXT", "IMAGE"]
+        ),
+    )
+
+    # 處理回應中的圖片
+    for part in response.candidates[0].content.parts:
+        if part.inline_data is not None:
+            image = Image.open(BytesIO(part.inline_data.data))
+            return image
+
 
 # 載入 YAML Prompt 模板
 with open("prompts.yaml", 'r') as stream:
@@ -63,7 +91,7 @@ with open("prompts.yaml", 'r') as stream:
 # 設定代理人
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, search_tool, image_generation_tool, get_current_time_in_timezone],  # Add your tools here (don't remove final answer)
+    tools=[final_answer, web_search, image_generation_tool, get_current_time_in_timezone],  # Add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,

requirements.txt

@@ -4,4 +4,5 @@ requests
 duckduckgo_search
 pandas
 smolagents[openai]
-smolagents[litellm]
+smolagents[litellm]
+google-genai