Update app.py

app.py

@@ -8,56 +8,36 @@ import tempfile
 from tools.final_answer import FinalAnswerTool # Assuming this exists in ./tools/final_answer.py
 from tools.visit_webpage import VisitWebpageTool # Assuming this exists in ./tools/visit_webpage.py
 from tools.web_search import DuckDuckGoSearchTool
-from smolagents import GradioUI
+
 import gradio as gr
 import json # Not used directly in the provided snippet, but can be common
-import os # os is imported twice, which is fine.
+
 from typing import Dict, List, Optional, Union, Any
 import re # Added for regex in Gradio UI if not already globally available
 
-'''
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
-@tool
-def my_custom_tool(x:str, y:int)-> int: #it's import to specify the return type
-    #Keep this format for the description / args / args description but feel free to modify the tool
-    """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:
-        # Create timezone object
-        tz = pytz.timezone(timezone)
-        # Get current time in that 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)}"
-'''
+from smolagents.tools import Tool as SmolToolBase # Renamed to avoid conflict with @tool decorator
+from smolagents.agent_types import AgentAudio, AgentImage, AgentText # For GradioUI output handling
+
+from Gradio_UI import GradioUI # This will import the class from your modified file
+
+
 @tool
 def create_document(text: str, format: str = "docx") -> str:
     """Creates a document with the provided text and allows download.
-    
     Args:
         text: The text content to write to the document.
-        format: The output format, either 'docx', 'pdf', or 'txt'.
+        format: The output format, either 'docx', 'pdf', or 'txt'. Default is 'docx'.
     """
     try:
-        # Create a temp directory to store files
         temp_dir = tempfile.mkdtemp()
+        file_name = "generated_document" # Consistent filename
         
         if format.lower() == "txt":
-            txt_path = os.path.join(temp_dir, "generated_document.txt")
-            with open(txt_path, "w", encoding="utf-8") as f:
+            path = os.path.join(temp_dir, f"{file_name}.txt")
+            with open(path, "w", encoding="utf-8") as f:
                 f.write(text)
-            return txt_path
+            print(f"Document created (txt): {path}")
+            return path
 
         elif format.lower() in ["docx", "pdf"]:
             try:
@@ -68,7 +48,6 @@ def create_document(text: str, format: str = "docx") -> str:
                         f"Please install it (e.g., 'pip install python-docx'). "
                         f"You can try creating a 'txt' file instead by specifying format='txt'.")
 
-            # Create a new document
             doc = docx.Document()
             doc.add_heading('Generated Document', 0)
             style = doc.styles['Normal']
@@ -77,25 +56,30 @@ def create_document(text: str, format: str = "docx") -> str:
             font.size = Pt(11)
             
             for paragraph in text.split('\n'):
-                if paragraph.strip():  # Skip empty paragraphs
+                if paragraph.strip():
                     doc.add_paragraph(paragraph)
             
-            docx_path = os.path.join(temp_dir, "generated_document.docx")
+            docx_path = os.path.join(temp_dir, f"{file_name}.docx")
             doc.save(docx_path)
-            
+            print(f"Document created (docx): {docx_path}")
+
             if format.lower() == "pdf":
                 try:
                     from docx2pdf import convert
-                    pdf_path = os.path.join(temp_dir, "generated_document.pdf")
+                    pdf_path = os.path.join(temp_dir, f"{file_name}.pdf")
                     convert(docx_path, pdf_path)
-                    # Optional: os.remove(docx_path) # If PDF is main target and DOCX is intermediate
+                    print(f"Document converted to PDF: {pdf_path}")
                     return pdf_path
                 except ImportError:
-                    return (f"ERROR: PDF conversion requires the 'docx2pdf' package. "
-                            f"Please install it (e.g., 'pip install docx2pdf'). "
-                            f"Document saved as DOCX instead at: {docx_path}")
+                    err_msg = (f"ERROR: PDF conversion requires the 'docx2pdf' package. "
+                               f"Please install it (e.g., 'pip install docx2pdf'). "
+                               f"Document saved as DOCX instead at: {docx_path}")
+                    print(err_msg)
+                    return err_msg # Or return docx_path with a note
                 except Exception as e_pdf:
-                    return f"Error converting DOCX to PDF: {str(e_pdf)}. Document saved as DOCX at: {docx_path}"
+                    err_msg = f"Error converting DOCX to PDF: {str(e_pdf)}. Document saved as DOCX at: {docx_path}"
+                    print(err_msg)
+                    return err_msg # Or return docx_path
             
             return docx_path
         
@@ -103,68 +87,94 @@ def create_document(text: str, format: str = "docx") -> str:
             return f"Error: Unsupported format '{format}'. Supported formats are 'docx', 'pdf', 'txt'."
             
     except Exception as e:
-        # Catch-all for other unexpected errors during file operations, etc.
+        print(f"General error in create_document: {str(e)}")
         return f"Error creating document: {str(e)}"
 
-# Custom file download tool to help with file handling
 @tool
 def get_file_download_link(file_path: str) -> str:
-    """Creates a download link for a file.
-    
+    """Informs that a file is ready for download and its type. (Used by agent to tell user).
     Args:
-        file_path: Path to the file that should be made available for download
+        file_path: Path to the file that should be made available for download.
     """
     if not os.path.exists(file_path):
+        print(f"get_file_download_link: File not found at {file_path}")
         return f"Error: File not found at {file_path}"
     
-    # Get file extension and set up appropriate mime type
     _, file_extension = os.path.splitext(file_path)
     mime_types = {
         '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
         '.pdf': 'application/pdf',
-        '.txt': 'text/plain',  # Added for .txt files
+        '.txt': 'text/plain',
     }
     mime_type = mime_types.get(file_extension.lower(), 'application/octet-stream')
-    
-    # Return information that can be used by the agent to instruct the user
-    return f"File ready for download: {os.path.basename(file_path)} ({mime_type})"
+    msg = f"File '{os.path.basename(file_path)}' is ready for download (type: {mime_type})."
+    print(f"get_file_download_link: {msg}")
+    return msg
         
-
 final_answer = FinalAnswerTool()
-web_search=DuckDuckGoSearchTool() # This was commented out in original, keeping it so
+web_search=DuckDuckGoSearchTool()
 visit_webpage=VisitWebpageTool()
 
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
+# Load image generation tool from Hub
+try:
+    image_generation_tool = load_tool("stabilityai/stable-diffusion-xl-base-1.0", trust_remote_code=True)
+    # Common alternative: "agents-course/text-to-image" (uses a smaller SD model)
+    # image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+    print(f"Successfully loaded image generation tool. Name: {getattr(image_generation_tool, 'name', 'N/A')}")
+except Exception as e:
+    print(f"Error loading image generation tool: {e}")
+    print("Image generation will not be available.")
+    image_generation_tool = None # Set to None if loading fails
+
+# --- Configure LLM Model ---
+# Ensure GEMINI_KEY is set in your environment variables
+if not os.getenv("GEMINI_KEY"):
+    print("CRITICAL: GEMINI_KEY environment variable not set. The LLM will not work.")
+    # exit(1) # Or handle gracefully
 
 model = LiteLLMModel(
-    model_id="gemini/gemini-2.0-flash", # As per original, though gemini-1.5-flash-latest might be more common
+    model_id="gemini/gemini-1.5-flash-latest", # Using a common and capable Gemini model
     api_key=os.getenv("GEMINI_KEY"),
-    max_tokens=8192,
-    temperature=0.7
+    max_tokens=4096, # Max output tokens for the LLM
+    temperature=0.6 # Slightly lower temperature for more factual tool use
 )
 
+# --- Load Prompts ---
+prompts_file = "prompts.yaml"
+if not os.path.exists(prompts_file):
+    print(f"Warning: '{prompts_file}' not found. Using default agent prompts.")
+    prompt_templates = None # Agent will use its defaults
+else:
+    with open(prompts_file, 'r') as stream:
+        prompt_templates = yaml.safe_load(stream)
 
-# Import tool from Hub
-# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True) # This tool is loaded but not added to the agent
-
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
+# --- Initialize Agent ---
+agent_tools = [
+    final_answer_tool,
+    web_search_tool,
+    visit_webpage_tool,
+    create_document, # This is a @tool decorated function, so it's directly usable
+    get_file_download_link, # Also a @tool decorated function
+]
 
-#,web_search
+if image_generation_tool:
+    agent_tools.append(image_generation_tool)
+else:
+    print("Note: Image generation tool was not loaded, so it's not added to the agent.")
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer,visit_webpage,create_document,get_file_download_link,web_search], ## add your tools here (don't remove final answer)
-    max_steps=10,
-    verbosity_level=1,
+    tools=agent_tools,
+    max_steps=10, # Max iterations for the agent
+    verbosity_level=1, # Set to 1 or higher to populate interaction_logs for GradioUI
     prompt_templates=prompt_templates
 )
 
-
-
-# Ensure this conditional execution if the script can also be imported
+# --- Main Application Entry Point ---
 if __name__ == "__main__":
-    # Note: The original script had GradioUI(agent).launch() directly.
-    # Using the custom UI:
-    GradioUI(agent).launch()
+    print("Starting Gradio UI...")
+    # Initialize GradioUI with the agent
+    # The GradioUI class should be imported from your modified Gradio_UI.py
+    # Ensure Gradio_UI.py is in the same directory or python path
+    ui = GradioUI(agent, file_upload_folder="uploaded_files") # Enable file uploads
+    ui.launch() # share=False is default and safer for Spaces