TobiasAlke's picture
Update app.py
0bd130c verified
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import tempfile
import os
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
"""A tool that does nothing yet
Args:
arg1: the first argument
arg2: the second argument
"""
return "What magic will you build ?"
@tool
def get_current_time_in_timezone(timezone: str) -> str:
"""A tool that fetches the current local time in a specified timezone.
Args:
timezone: A string representing a valid timezone (e.g., 'America/New_York').
"""
try:
tz = pytz.timezone(timezone)
local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
return f"The current local time in {timezone} is: {local_time}"
except Exception as e:
return f"Error fetching time for timezone '{timezone}': {str(e)}"
@tool
def generate_and_save_image(prompt: str, filename: str = "generated_image.png") -> str:
"""Generate an image and save it to file for better display
Args:
prompt: Description of the image to generate
filename: Name of the file to save (optional)
"""
try:
# Verwende das geladene image_generation_tool
image = image_generation_tool(prompt=prompt)
# Speichere in temporärem Verzeichnis für bessere Organisation
temp_dir = tempfile.gettempdir()
full_path = os.path.join(temp_dir, filename)
# Speichere das Bild
image.save(full_path)
# Erstelle auch eine Kopie im aktuellen Verzeichnis für einfachen Zugriff
image.save(filename)
return f"Image generated and saved as {filename} (also available at {full_path}). You can download it from the current directory."
except Exception as e:
return f"Error generating image: {str(e)}"
# Load tools
final_answer = FinalAnswerTool()
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
# Debug: Prüfe den Namen des geladenen Tools
print(f"Loaded image tool name: {image_generation_tool.name}")
# Erstelle einen Enhanced Image Generator mit Download-Funktionalität
@tool
def enhanced_image_generator(prompt: str, filename: str = None) -> str:
"""Enhanced image generator that creates images and provides download information
Args:
prompt: Text description of the image to generate
filename: Optional custom filename (without extension)
"""
try:
image = image_generation_tool(prompt=prompt)
# Generiere Dateinamen falls nicht angegeben
if filename is None:
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
filename = f"generated_image_{timestamp}"
# Stelle sicher dass Dateiname .png Extension hat
if not filename.endswith('.png'):
filename += '.png'
# Speichere das Bild für Download
image.save(filename)
# Gib eine klare Download-Anweisung zurück
return f"✅ Image generated successfully!\n📁 File saved as: {filename}\n💡 You can find and download the image file '{filename}' from your current directory.\n🎨 Prompt used: '{prompt}'"
except Exception as e:
return f"❌ Error generating image: {str(e)}"
@tool
def quick_image_generator(prompt: str) -> str:
"""Quick image generator with automatic filename
Args:
prompt: Text description of the image to generate
"""
try:
image = image_generation_tool(prompt=prompt)
# Automatischer Dateiname basierend auf Prompt (erste 3 Wörter)
words = prompt.split()[:3]
safe_name = "_".join(word.lower().replace(",", "").replace(".", "") for word in words)
timestamp = datetime.datetime.now().strftime('%H%M%S')
filename = f"{safe_name}_{timestamp}.png"
# Speichere das Bild
image.save(filename)
return f"🎨 Image created and saved as '{filename}'! You can download it from your file directory. Generated from: '{prompt}'"
except Exception as e:
return f"❌ Error generating image: {str(e)}"
# Test das Tool beim Start (optional - kann entfernt werden)
print("Testing image generation tool...")
try:
test_result = image_generation_tool(prompt="test")
print(f"Image tool loaded successfully: {type(test_result)}")
except Exception as e:
print(f"Warning: Image tool test failed: {e}")
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct', # it is possible that this model may be overloaded
custom_role_conversions=None,
)
# Load system prompt from prompt.yaml file
with open("prompts.yaml", 'r') as stream:
prompt_templates = yaml.safe_load(stream)
agent = CodeAgent(
model=model,
tools=[
final_answer,
image_generation_tool, # Das Original-Tool
enhanced_image_generator, # Tool mit besseren Download-Infos
quick_image_generator, # Schnelle Bildgenerierung mit automatischen Namen
get_current_time_in_timezone,
my_custom_tool
],
max_steps=6,
verbosity_level=1,
grammar=None,
planning_interval=None,
name=None,
description=None,
prompt_templates=prompt_templates
)
# Debugging Info
print("Agent initialized with tools:")
for tool in agent.tools:
print(f"- {tool.name if hasattr(tool, 'name') else str(tool)}")
print("\nStarting Gradio UI...")
GradioUI(agent).launch()
# Beispiel für direkten Test (kann auskommentiert werden):
"""
# Direkter Test ohne UI:
if __name__ == "__main__":
print("Testing agent directly...")
response = agent.run("Generiere ein Bild eines Fisches")
print(f"Agent response: {response}")
"""