Update app.py

app.py

@@ -1,17 +1,18 @@
 from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool, LiteLLMModel
 import datetime
-import requests
-import pytz
+import requests # Not used in the provided snippet directly, but often a common import
+import pytz # Used in the commented-out example tool
 import yaml
 import os
 import tempfile
-from tools.final_answer import FinalAnswerTool
-from tools.visit_webpage import VisitWebpageTool
+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 smolagents import GradioUI
 import gradio as gr
-import json
-import os
+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 !
@@ -24,7 +25,6 @@ def my_custom_tool(x:str, y:int)-> int: #it's import to specify the return type
         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.
@@ -39,58 +39,70 @@ def get_current_time_in_timezone(timezone: str) -> str:
         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 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' or 'pdf'
+        text: The text content to write to the document.
+        format: The output format, either 'docx', 'pdf', or 'txt'.
     """
     try:
-        import docx
-        from docx.shared import Pt
-        
         # Create a temp directory to store files
         temp_dir = tempfile.mkdtemp()
         
-        # Create a new document
-        doc = docx.Document()
-        
-        # Add a heading
-        doc.add_heading('Generated Document', 0)
-        
-        # Set font style for regular text
-        style = doc.styles['Normal']
-        font = style.font
-        font.name = 'Calibri'
-        font.size = Pt(11)
-        
-        # Add paragraphs from the input text
-        # Split by newlines to maintain paragraph structure
-        for paragraph in text.split('\n'):
-            if paragraph.strip():  # Skip empty paragraphs
-                doc.add_paragraph(paragraph)
-        
-        # Save the document
-        docx_path = os.path.join(temp_dir, "generated_document.docx")
-        doc.save(docx_path)
-        
-        # Convert to PDF if requested
-        if format.lower() == "pdf":
+        if format.lower() == "txt":
+            txt_path = os.path.join(temp_dir, "generated_document.txt")
+            with open(txt_path, "w", encoding="utf-8") as f:
+                f.write(text)
+            return txt_path
+
+        elif format.lower() in ["docx", "pdf"]:
             try:
-                from docx2pdf import convert
-                pdf_path = os.path.join(temp_dir, "generated_document.pdf")
-                convert(docx_path, pdf_path)
-                return pdf_path
+                import docx
+                from docx.shared import Pt
             except ImportError:
-                return f"PDF conversion requires docx2pdf package. Document saved as DOCX instead at: {docx_path}"
+                return (f"ERROR: To create DOCX or PDF files, the 'python-docx' package is required. "
+                        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']
+            font = style.font
+            font.name = 'Calibri'
+            font.size = Pt(11)
+            
+            for paragraph in text.split('\n'):
+                if paragraph.strip():  # Skip empty paragraphs
+                    doc.add_paragraph(paragraph)
+            
+            docx_path = os.path.join(temp_dir, "generated_document.docx")
+            doc.save(docx_path)
+            
+            if format.lower() == "pdf":
+                try:
+                    from docx2pdf import convert
+                    pdf_path = os.path.join(temp_dir, "generated_document.pdf")
+                    convert(docx_path, pdf_path)
+                    # Optional: os.remove(docx_path) # If PDF is main target and DOCX is intermediate
+                    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}")
+                except Exception as e_pdf:
+                    return f"Error converting DOCX to PDF: {str(e_pdf)}. Document saved as DOCX at: {docx_path}"
+            
+            return docx_path
         
-        return docx_path
-    
+        else:
+            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.
         return f"Error creating document: {str(e)}"
 
 # Custom file download tool to help with file handling
@@ -109,6 +121,7 @@ def get_file_download_link(file_path: str) -> str:
     mime_types = {
         '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
         '.pdf': 'application/pdf',
+        '.txt': 'text/plain',  # Added for .txt files
     }
     mime_type = mime_types.get(file_extension.lower(), 'application/octet-stream')
     
@@ -117,14 +130,14 @@ def get_file_download_link(file_path: str) -> str:
         
 
 final_answer = FinalAnswerTool()
-#web_search=DuckDuckGoSearchTool()
+#web_search=DuckDuckGoSearchTool() # This was commented out in original, keeping it so
 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' 
 
 model = LiteLLMModel(
-    model_id="gemini/gemini-2.0-flash",
+    model_id="gemini/gemini-2.0-flash", # As per original, though gemini-1.5-flash-latest might be more common
     api_key=os.getenv("GEMINI_KEY"),
     max_tokens=8192,
     temperature=0.7
@@ -132,7 +145,7 @@ model = LiteLLMModel(
 
 
 # Import tool from Hub
-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) # This tool is loaded but not added to the agent
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
@@ -167,38 +180,45 @@ class CustomGradioUI(GradioUI):
             file_output = gr.File(label="Generated Document", visible=False)
             
             # Store the latest file path
-            self._latest_file_path = None
+            self._latest_file_path = None # Instance variable to store path between calls
             
             def respond(message, chat_history):
                 agent_response = self.agent.run(message)
                 chat_history.append((message, agent_response))
                 
-                # Check if response contains a file path
-                import re
-                file_paths = re.findall(r'File ready for download: .+ \((application/[\w.+-]+)\)', agent_response)
-                
                 show_download = False
-                self._latest_file_path = None
-                
-                # Look for generated file paths in the response
-                paths = re.findall(r'/tmp/\w+/generated_document\.(docx|pdf)', agent_response)
+                # Reset latest_file_path to ensure it's fresh for each response
+                # self._latest_file_path = None # Actually, should be set if file is found
+
+                # Look for generated file paths in the agent's response.
+                # tempfile.mkdtemp() usually creates paths like /tmp/some_random_string/
+                # The tool consistently names files "generated_document.ext"
+                paths = re.findall(r'/tmp/[^/]+/generated_document\.(docx|pdf|txt)', agent_response)
                 if paths:
-                    self._latest_file_path = paths[0]
+                    self._latest_file_path = paths[0] # Store the full path
                     show_download = True
+                else:
+                    self._latest_file_path = None # Clear if no path found in this response
                 
-                return chat_history, gr.Button.update(visible=show_download), gr.File.update(visible=False)
+                return chat_history, gr.Button.update(visible=show_download), gr.File.update(visible=False) # Keep file output hidden until download clicked
             
             def prepare_download():
-                if self._latest_file_path:
+                if self._latest_file_path and os.path.exists(self._latest_file_path):
                     return gr.File.update(value=self._latest_file_path, visible=True)
+                # Handle case where file might have been cleaned up or path is stale
+                gr.Warning("File not available for download. It might have been temporary or an error occurred.")
                 return gr.File.update(visible=False)
             
             # Connect the components
             msg.submit(respond, [msg, chatbot], [chatbot, download_btn, file_output])
             download_btn.click(prepare_download, [], [file_output])
             
-            gr.Markdown("Powered by smolagents and Qwen")
+            gr.Markdown("Powered by smolagents and Qwen") # Qwen is mentioned, perhaps the HF endpoint uses it.
             
         return interface
 
-GradioUI(agent).launch()
+# Ensure this conditional execution if the script can also be imported
+if __name__ == "__main__":
+    # Note: The original script had GradioUI(agent).launch() directly.
+    # Using the custom UI:
+    CustomGradioUI(agent).launch()