Update app.py

app.py

@@ -7,6 +7,9 @@ from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 #from tools.web_shearch import DuckDuckGoSearchTool
 from smolagents import GradioUI
+import tempfile
+import gradio as gr
+import os
 
 '''
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
@@ -36,6 +39,80 @@ def get_current_time_in_timezone(timezone: str) -> str:
         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'
+    """
+    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":
+            try:
+                from docx2pdf import convert
+                pdf_path = os.path.join(temp_dir, "generated_document.pdf")
+                convert(docx_path, pdf_path)
+                return pdf_path
+            except ImportError:
+                return f"PDF conversion requires docx2pdf package. Document saved as DOCX instead at: {docx_path}"
+        
+        return docx_path
+    
+    except Exception as 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.
+    
+    Args:
+        file_path: Path to the file that should be made available for download
+    """
+    if not os.path.exists(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',
+    }
+    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})"
+        
 
 final_answer = FinalAnswerTool()
 #web_search=DuckDuckGoSearchTool()
@@ -63,7 +140,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer,visit_webpage], ## add your tools here (don't remove final answer)
+    tools=[final_answer,visit_webpage,create_document,get_file_download_link], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -73,5 +150,58 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
+# Custom Gradio UI with file download capability
+class CustomGradioUI(GradioUI):
+    def build_interface(self):
+        with gr.Blocks() as interface:
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("# AI Assistant")
+                
+            chatbot = gr.Chatbot(height=600)
+            msg = gr.Textbox(
+                placeholder="Ask me anything...",
+                container=False,
+                scale=7,
+            )
+            
+            # Add a file download component
+            download_btn = gr.Button("Download File", visible=False)
+            file_output = gr.File(label="Generated Document", visible=False)
+            
+            # Store the latest file path
+            self._latest_file_path = None
+            
+            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)
+                if paths:
+                    self._latest_file_path = paths[0]
+                    show_download = True
+                
+                return chat_history, gr.Button.update(visible=show_download), gr.File.update(visible=False)
+            
+            def prepare_download():
+                if self._latest_file_path:
+                    return gr.File.update(value=self._latest_file_path, visible=True)
+                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")
+            
+        return interface
 
 GradioUI(agent).launch()