Update app.py

app.py

@@ -2,40 +2,72 @@ import os
 import importlib.util
 import sys
 import tempfile
-import traceback
-from pathlib import Path
-
+import re
+import ast
 import gradio as gr
-import openai
-from dotenv import load_dotenv
+from openai import OpenAI
 
-from utils import sanitize_code, extract_code_blocks, validate_gradio_code
+# Extract code blocks from text
+def extract_code_blocks(text):
+    pattern = r'```(?:python)?\s*([\s\S]*?)```'
+    matches = re.findall(pattern, text)
+    
+    if not matches and text.strip():
+        if re.search(r'import\s+\w+|def\s+\w+\(|class\s+\w+:|if\s+__name__\s*==\s*[\'"]__main__[\'"]:', text):
+            return [text.strip()]
+    
+    return [match.strip() for match in matches]
 
-# Load environment variables (for local development)
-load_dotenv()
+# Validate that the code is safe to run
+def validate_gradio_code(code):
+    try:
+        tree = ast.parse(code)
+        
+        # Check imports
+        allowed_modules = ['gradio', 'numpy', 'pandas', 'matplotlib', 'PIL', 'os', 'io', 'base64', 
+                          'time', 'datetime', 'json', 'random', 'math', 'sys', 're', 'pathlib',
+                          'collections', 'typing', 'warnings']
+        
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                for name in node.names:
+                    if name.name not in allowed_modules:
+                        return False, f"Unauthorized import: {name.name}"
+            
+            elif isinstance(node, ast.ImportFrom):
+                if node.module not in allowed_modules and node.module is not None:
+                    return False, f"Unauthorized import from: {node.module}"
+        
+        # Check for harmful operations
+        code_str = code.lower()
+        harmful_ops = ['subprocess', 'system(', 'popen(', 'execve(', 'chmod(', 'rmdir(', 
+                      'remove(', 'unlink(', 'rmtree(', 'socket', 'eval(', 'exec(']
+        
+        for op in harmful_ops:
+            if op in code_str:
+                return False, f"Potentially harmful operation detected: {op}"
+        
+        return True, None
+        
+    except SyntaxError as e:
+        return False, f"Syntax error in the code: {str(e)}"
+    except Exception as e:
+        return False, f"Error validating code: {str(e)}"
 
-# Global variables
+# Global variables for app storage
 generated_app = None
 current_code = ""
-user_api_key = ""
 
+# Generate app using OpenAI
 def generate_gradio_app(api_key, prompt):
-    """Generate Gradio app code using OpenAI API"""
-    global user_api_key
-    
-    # Validate API key
     if not api_key or len(api_key) < 20:
         return None, "Please provide a valid OpenAI API key"
     
-    # Store API key for this session
-    user_api_key = api_key
-    
     try:
-        # Configure client with user's API key
-        client = openai.OpenAI(api_key=api_key)
+        client = OpenAI(api_key=api_key)
         
         response = client.chat.completions.create(
-            model="gpt-4o",  # Using gpt-4o for best code generation
+            model="gpt-4o",
             messages=[
                 {"role": "system", "content": """You are an expert Gradio developer. 
                 Create a standalone Gradio application based on the user's prompt.
@@ -43,13 +75,10 @@ def generate_gradio_app(api_key, prompt):
                 The code must:
                 1. Import all necessary libraries
                 2. Define a complete, functional Gradio interface
-                3. Launch the interface with share=False and show_api=False
-                4. Use gr.Blocks() for complex interfaces
-                5. Handle errors gracefully
-                6. Use relative paths for any file operations
-                7. NOT use external APIs or services unless specifically requested
-                8. Be completely self-contained in a single script
-                9. End with a simple if __name__ == "__main__": block that launches the app
+                3. Launch the interface with share=False
+                4. Handle errors gracefully
+                5. Be completely self-contained in a single script
+                6. End with a simple if __name__ == "__main__": block that launches the app
                 """
                 },
                 {"role": "user", "content": prompt}
@@ -69,127 +98,118 @@ def generate_gradio_app(api_key, prompt):
     except Exception as e:
         return None, str(e)
 
+# Load and run generated app
 def load_and_run_gradio_app(code):
-    """Load and run the generated Gradio app code"""
     global generated_app, current_code
     
-    # Check if code is safe to execute
+    # Validate code
     is_valid, error_msg = validate_gradio_code(code)
     if not is_valid:
         return None, error_msg
     
-    # Clean up previous app if it exists
+    # Clean up previous app
     if generated_app:
         try:
             generated_app.close()
         except:
             pass
     
-    # Save code to a temporary file
+    # Save code to temp file
     with tempfile.NamedTemporaryFile(suffix='.py', delete=False) as f:
         f.write(code.encode('utf-8'))
         temp_file = f.name
     
     try:
-        # Import the module
+        # Import module
         module_name = os.path.basename(temp_file).replace('.py', '')
         spec = importlib.util.spec_from_file_location(module_name, temp_file)
         module = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = module
         spec.loader.exec_module(module)
         
-        # Find the Gradio interface
+        # Find Gradio interface
         for attr_name in dir(module):
             attr = getattr(module, attr_name)
-            if isinstance(attr, gr.Blocks) or isinstance(attr, gr.Interface):
-                # Save a reference to the app
+            if isinstance(attr, (gr.Blocks, gr.Interface)):
                 generated_app = attr
                 current_code = code
-                
-                # Return the app
                 return attr, None
         
         return None, "No Gradio interface found in the generated code"
         
     except Exception as e:
-        error_details = traceback.format_exc()
-        return None, f"Error executing the generated code: {str(e)}\n{error_details}"
+        return None, f"Error executing the generated code: {str(e)}"
     finally:
-        # Clean up the temporary file
         try:
             os.unlink(temp_file)
         except:
             pass
 
-def create_ui():
-    """Create the main Gradio interface"""
-    with gr.Blocks(title="Dynamic Gradio App Generator") as interface:
-        gr.Markdown("# 🤖 Dynamic Gradio App Generator")
-        gr.Markdown("Describe the Gradio app you want to create, and the AI will generate and run it for you.")
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                api_key = gr.Textbox(
-                    label="OpenAI API Key", 
-                    placeholder="sk-...",
-                    type="password",
-                    info="Your key is used for this session only and not stored"
-                )
-                
-                prompt = gr.Textbox(
-                    label="App Description", 
-                    placeholder="Describe the Gradio app you want to create...",
-                    lines=5
-                )
-                
-                with gr.Row():
-                    submit_btn = gr.Button("Generate & Run App", variant="primary")
-                    clear_btn = gr.Button("Clear", variant="secondary")
-                
-                with gr.Accordion("Generated Code", open=False):
-                    code_output = gr.Code(language="python", label="Generated Code")
+# Create the main UI
+with gr.Blocks(title="AI Gradio App Generator") as demo:
+    gr.Markdown("# 🤖 AI Gradio App Generator")
+    gr.Markdown("Describe the app you want, and I'll generate it for you using OpenAI's API.")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            api_key = gr.Textbox(
+                label="OpenAI API Key", 
+                placeholder="sk-...",
+                type="password",
+                info="Your key is used for this session only and not stored"
+            )
             
-            with gr.Column(scale=3):
-                output = gr.HTML("<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>")
-                error_output = gr.Markdown(visible=False)
-        
-        def on_submit(api_key_input, prompt_text):
-            # Generate the Gradio app code
-            code, error = generate_gradio_app(api_key_input, prompt_text)
-            if error:
-                return None, f"⚠️ **Error generating code**: {error}", gr.update(visible=True), gr.update(value="<div style='text-align: center; padding: 50px;'><h3>Error generating app</h3></div>")
+            prompt = gr.Textbox(
+                label="App Description", 
+                placeholder="Describe the Gradio app you want to create...",
+                lines=5
+            )
+            
+            with gr.Row():
+                submit_btn = gr.Button("Generate App", variant="primary")
+                clear_btn = gr.Button("Clear", variant="secondary")
             
-            # Load and run the generated app
-            app, run_error = load_and_run_gradio_app(code)
-            if run_error:
-                return code, f"⚠️ **Error running the generated app**: {run_error}", gr.update(visible=True), gr.update(value="<div style='text-align: center; padding: 50px;'><h3>Error running app</h3></div>")
+            with gr.Accordion("Generated Code", open=False):
+                code_output = gr.Code(language="python", label="Generated Code")
             
-            # Create an iframe to display the app
-            iframe_html = f'''
-            <div style="border: 1px solid #ddd; border-radius: 8px; padding: 0; overflow: hidden;">
-                <iframe id="appFrame" src="/generated_app" width="100%" height="800px" frameborder="0"></iframe>
-            </div>
-            '''
-            return code, "", gr.update(visible=False), iframe_html
+            status_output = gr.Markdown("")
         
-        def on_clear():
-            return gr.update(value=""), "", gr.update(visible=False), "<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>"
+        with gr.Column(scale=3):
+            app_container = gr.HTML("<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>")
+    
+    def on_submit(api_key_input, prompt_text):
+        # Generate code
+        code, error = generate_gradio_app(api_key_input, prompt_text)
+        if error:
+            return None, f"⚠️ Error generating code: {error}", "<div style='text-align: center; padding: 50px;'><h3>Error generating app</h3></div>"
         
-        submit_btn.click(
-            on_submit, 
-            inputs=[api_key, prompt], 
-            outputs=[code_output, error_output, error_output, output]
-        )
+        # Load and run app
+        app, run_error = load_and_run_gradio_app(code)
+        if run_error:
+            return code, f"⚠️ Error running the generated app: {run_error}", "<div style='text-align: center; padding: 50px;'><h3>Error running app</h3></div>"
         
-        clear_btn.click(
-            on_clear,
-            inputs=[],
-            outputs=[prompt, code_output, error_output, output]
-        )
+        # Create iframe to display app
+        iframe_html = f'''
+        <div style="border: 1px solid #ddd; border-radius: 8px; padding: 0; overflow: hidden;">
+            <iframe id="appFrame" src="/generated_app" width="100%" height="800px" frameborder="0"></iframe>
+        </div>
+        '''
+        return code, "✅ App generated successfully!", iframe_html
+    
+    def on_clear():
+        return "", "", "<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>"
+    
+    submit_btn.click(
+        on_submit,
+        inputs=[api_key, prompt],
+        outputs=[code_output, status_output, app_container]
+    )
     
-    return interface
+    clear_btn.click(
+        on_clear,
+        inputs=[],
+        outputs=[prompt, status_output, app_container]
+    )
 
-# Create and launch the Gradio interface
 if __name__ == "__main__":
-    demo = create_ui()
     demo.launch(server_name="0.0.0.0", server_port=7860)