tinyagent-py 0.0.9__py3-none-any.whl → 0.0.12__py3-none-any.whl

tinyagent/code_agent/tiny_code_agent.py

@@ -62,6 +62,7 @@ class TinyCodeAgent:
         self.user_variables = user_variables or {}
         self.pip_packages = pip_packages or []
         self.local_execution = local_execution
+        self.provider = provider  # Store provider type for reuse
         
         # Create the code execution provider
         self.code_provider = self._create_provider(provider, self.provider_config)
@@ -96,8 +97,13 @@ class TinyCodeAgent:
             config_pip_packages = config.get("pip_packages", [])
             final_pip_packages = list(set(self.pip_packages + config_pip_packages))
             
+            # Merge authorized_imports from both sources (direct parameter and provider_config)
+            config_authorized_imports = config.get("authorized_imports", [])
+            final_authorized_imports = list(set(self.authorized_imports + config_authorized_imports))
+            
             final_config = config.copy()
             final_config["pip_packages"] = final_pip_packages
+            final_config["authorized_imports"] = final_authorized_imports
             
             return ModalProvider(
                 log_manager=self.log_manager,
@@ -258,6 +264,14 @@ class TinyCodeAgent:
                 result = await self.code_provider.execute_python(code_lines, timeout)
                 return str(result)
             except Exception as e:
+                print("!"*100)
+                COLOR = {
+                        "RED": "\033[91m",
+                        "ENDC": "\033[0m",
+                    }
+                print(f"{COLOR['RED']}{str(e)}{COLOR['ENDC']}")
+                print(f"{COLOR['RED']}{traceback.format_exc()}{COLOR['ENDC']}")
+                print("!"*100)
                 return f"Error executing code: {str(e)}"
         
         self.agent.add_tool(run_python)
@@ -441,6 +455,69 @@ class TinyCodeAgent:
         """
         return self.pip_packages.copy()
     
+    def add_authorized_imports(self, imports: List[str]):
+        """
+        Add additional authorized imports to the execution environment.
+        
+        Args:
+            imports: List of import names to authorize
+        """
+        self.authorized_imports.extend(imports)
+        self.authorized_imports = list(set(self.authorized_imports))  # Remove duplicates
+        
+        # Update the provider with the new authorized imports
+        # This requires recreating the provider
+        print("⚠️  Warning: Adding authorized imports after initialization requires recreating the Modal environment.")
+        print("   For better performance, set authorized_imports during TinyCodeAgent initialization.")
+        
+        # Recreate the provider with new authorized imports
+        self.code_provider = self._create_provider(self.provider, self.provider_config)
+        
+        # Re-set user variables if they exist
+        if self.user_variables:
+            self.code_provider.set_user_variables(self.user_variables)
+        
+        # Rebuild system prompt to include new authorized imports
+        self.system_prompt = self._build_system_prompt()
+        # Update the agent's system prompt
+        self.agent.system_prompt = self.system_prompt
+    
+    def get_authorized_imports(self) -> List[str]:
+        """
+        Get a copy of current authorized imports.
+        
+        Returns:
+            List of authorized imports
+        """
+        return self.authorized_imports.copy()
+    
+    def remove_authorized_import(self, import_name: str):
+        """
+        Remove an authorized import.
+        
+        Args:
+            import_name: Import name to remove
+        """
+        if import_name in self.authorized_imports:
+            self.authorized_imports.remove(import_name)
+            
+            # Update the provider with the new authorized imports
+            # This requires recreating the provider
+            print("⚠️  Warning: Removing authorized imports after initialization requires recreating the Modal environment.")
+            print("   For better performance, set authorized_imports during TinyCodeAgent initialization.")
+            
+            # Recreate the provider with updated authorized imports
+            self.code_provider = self._create_provider(self.provider, self.provider_config)
+            
+            # Re-set user variables if they exist
+            if self.user_variables:
+                self.code_provider.set_user_variables(self.user_variables)
+            
+            # Rebuild system prompt to reflect updated authorized imports
+            self.system_prompt = self._build_system_prompt()
+            # Update the agent's system prompt
+            self.agent.system_prompt = self.system_prompt
+    
     async def close(self):
         """Clean up resources."""
         await self.code_provider.cleanup()
@@ -498,6 +575,7 @@ async def run_example():
         user_variables={
             "sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
         },
+        authorized_imports=["tinyagent", "gradio", "requests", "numpy", "pandas"],  # Explicitly specify authorized imports
         local_execution=False  # Remote execution via Modal (default)
     )
     
@@ -524,6 +602,7 @@ async def run_example():
         user_variables={
             "sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
         },
+        authorized_imports=["tinyagent", "gradio", "requests"],  # More restricted imports for local execution
         local_execution=True  # Local execution
     )
     
@@ -550,6 +629,18 @@ async def run_example():
     agent_remote.add_code_tool(validator)
     agent_local.add_code_tool(validator)
     
+    # Demonstrate adding authorized imports dynamically
+    print("\n" + "="*80)
+    print("🔧 Testing with dynamically added authorized imports")
+    agent_remote.add_authorized_imports(["matplotlib", "seaborn"])
+    
+    # Test with visualization libraries
+    viz_prompt = "Create a simple plot of the sample_data and save it as a base64 encoded image string."
+    
+    response_viz = await agent_remote.run(viz_prompt)
+    print("Remote Agent Visualization Response:")
+    print(response_viz)
+    
     print("\n" + "="*80)
     print("🔧 Testing with dynamically added tools")
     

tinyagent/code_agent/utils.py

@@ -1,6 +1,7 @@
 import sys
 import cloudpickle
-from typing import Dict, Any
+from typing import Dict, Any, List
+from .safety import validate_code_safety, function_safety_context
 
 
 def clean_response(resp: Dict[str, Any]) -> Dict[str, Any]:
@@ -40,7 +41,14 @@ def make_session_blob(ns: dict) -> bytes:
     return cloudpickle.dumps(clean)
 
 
-def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dict[str, Any] = None):
+def _run_python(
+    code: str,
+    globals_dict: Dict[str, Any] | None = None,
+    locals_dict: Dict[str, Any] | None = None,
+    authorized_imports: List[str] | None = None,
+    authorized_functions: List[str] | None = None,
+    trusted_code: bool = False,
+):
     """
     Execute Python code in a controlled environment with proper error handling.
     
@@ -48,6 +56,9 @@ def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dic
         code: Python code to execute
         globals_dict: Global variables dictionary
         locals_dict: Local variables dictionary
+        authorized_imports: List of authorized imports that user code may access. Wildcards (e.g. "numpy.*") are supported. A value of None disables the allow-list and only blocks dangerous modules.
+        authorized_functions: List of authorized dangerous functions that user code may access. A value of None disables the allow-list and blocks all dangerous functions.
+        trusted_code: If True, skip security checks. Should only be used for framework code, tools, or default executed code.
         
     Returns:
         Dictionary containing execution results
@@ -56,16 +67,27 @@ def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dic
     import traceback
     import io
     import ast
-    
+    import builtins  # Needed for import hook
+    import sys
+
+    # ------------------------------------------------------------------
+    # 1. Static safety analysis – refuse code containing dangerous imports or functions
+    # ------------------------------------------------------------------
+    validate_code_safety(code, authorized_imports=authorized_imports, 
+                        authorized_functions=authorized_functions, trusted_code=trusted_code)
+
     # Make copies to avoid mutating the original parameters
     globals_dict = globals_dict or {}
     locals_dict = locals_dict or {}
     updated_globals = globals_dict.copy()
     updated_locals = locals_dict.copy()
     
-    # Pre-import essential modules into the global namespace
-    # This ensures they're available for imports inside functions
-    essential_modules = ['requests', 'json', 'os', 'sys', 'time', 'datetime', 're', 'random', 'math']
+    # Only pre-import a **minimal** set of safe modules so that common helper
+    # functions work out of the box without giving user code access to the
+    # full standard library.  Anything outside this list must be imported
+    # explicitly by the user – and will be blocked by the safety layer above
+    # if considered dangerous.
+    essential_modules = ['requests', 'json', 'time', 'datetime', 're', 'random', 'math','cloudpickle']
     
     for module_name in essential_modules:
         try:
@@ -75,12 +97,30 @@ def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dic
         except ImportError:
             print(f"⚠️  Warning: {module_name} module not available")
     
+    # Variable to store print output
+    output_buffer = []
+    
+    # Create a custom print function that captures output
+    def custom_print(*args, **kwargs):
+        # Get the sep and end kwargs, defaulting to ' ' and '\n'
+        sep = kwargs.get('sep', ' ')
+        end = kwargs.get('end', '\n')
+        
+        # Convert all arguments to strings and join them
+        output = sep.join(str(arg) for arg in args) + end
+        
+        # Store the output
+        output_buffer.append(output)
+    
+    # Add the custom print function to the globals
+    #updated_globals['print'] = custom_print
+    
+    # Parse the code
     tree = ast.parse(code, mode="exec")
     compiled = compile(tree, filename="<ast>", mode="exec")
     stdout_buf = io.StringIO()
-    stderr_buf = io.StringIO()
-
-    # Execute with stdout+stderr capture and exception handling
+    stderr_buf = io.StringIO()   
+    # Execute with exception handling
     error_traceback = None
     output = None
 
@@ -92,8 +132,15 @@ def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dic
             merged_globals = updated_globals.copy()
             merged_globals.update(updated_locals)
             
+            # Add 'exec' to authorized_functions for internal use
+            internal_authorized_functions = ['exec','eval']
+            if authorized_functions is not None and not isinstance(authorized_functions, bool):
+                internal_authorized_functions.extend(authorized_functions)
+            
             # Execute with only globals - this fixes generator expression scoping issues
-            output = exec(code, merged_globals)
+            # Use the function_safety_context to block dangerous functions during execution
+            with function_safety_context(authorized_functions=internal_authorized_functions, trusted_code=trusted_code):
+                output = exec(compiled, merged_globals)
             
             # Update both dictionaries with any new variables created during execution
             for key, value in merged_globals.items():
@@ -106,6 +153,8 @@ def _run_python(code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dic
             # Capture the full traceback as a string
             error_traceback = traceback.format_exc()
 
+    # Join all captured output
+    #printed_output = ''.join(output_buffer)  
     printed_output = stdout_buf.getvalue()
     stderr_output = stderr_buf.getvalue()
     error_traceback_output = error_traceback

tinyagent/hooks/gradio_callback.py

@@ -5,6 +5,7 @@ import os
 import re
 import shutil
 import time
+import io
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Union
 
@@ -36,6 +37,7 @@ class GradioCallback:
         show_thinking: bool = True,
         show_tool_calls: bool = True,
         logger: Optional[logging.Logger] = None,
+        log_manager: Optional[Any] = None,
     ):
         """
         Initialize the Gradio callback.
@@ -46,6 +48,7 @@ class GradioCallback:
             show_thinking: Whether to show the thinking process
             show_tool_calls: Whether to show tool calls
             logger: Optional logger to use
+            log_manager: Optional LoggingManager instance to capture logs from
         """
         self.logger = logger or logging.getLogger(__name__)
         self.show_thinking = show_thinking
@@ -81,6 +84,37 @@ class GradioCallback:
         # References to Gradio UI components (will be set in create_app)
         self._chatbot_component = None
         self._token_usage_component = None
+        
+        # Log stream for displaying logs in the UI
+        self.log_stream = io.StringIO()
+        self._log_component = None
+        
+        # Setup logging
+        self.log_manager = log_manager
+        if log_manager:
+            # Create a handler that writes to our StringIO stream
+            self.log_handler = logging.StreamHandler(self.log_stream)
+            self.log_handler.setFormatter(
+                logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+            )
+            self.log_handler.setLevel(logging.DEBUG)
+            
+            # Add the handler to the LoggingManager
+            log_manager.configure_handler(
+                self.log_handler,
+                format_string='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+                level=logging.DEBUG
+            )
+            self.logger.debug("Added log handler to LoggingManager")
+        elif logger:
+            # Fall back to single logger if no LoggingManager is provided
+            self.log_handler = logging.StreamHandler(self.log_stream)
+            self.log_handler.setFormatter(
+                logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+            )
+            self.log_handler.setLevel(logging.DEBUG)
+            logger.addHandler(self.log_handler)
+            self.logger.debug("Added log handler to logger")
 
         self.logger.debug("GradioCallback initialized")
     
@@ -525,7 +559,7 @@ class GradioCallback:
         typing_message_index = len(chatbot_history) - 1
  
         # Initial yield to show user message and typing indicator
-        yield chatbot_history, self._get_token_usage_text()
+        yield chatbot_history, self._get_token_usage_text(), self.log_stream.getvalue() if self._log_component else None
  
         # Kick off the agent in the background
         loop = asyncio.get_event_loop()
@@ -632,9 +666,10 @@ class GradioCallback:
                             del in_progress_tool_calls[tid]
                             self.logger.debug(f"Updated tool call to completed: {tname}")
  
-                # yield updated history + token usage
+                # yield updated history + token usage + logs
                 token_text = self._get_token_usage_text()
-                yield chatbot_history, token_text
+                logs = self.log_stream.getvalue() if self._log_component else None
+                yield chatbot_history, token_text, logs
                 self.last_update_yield_time = now
  
             await asyncio.sleep(update_interval)
@@ -657,8 +692,9 @@ class GradioCallback:
             )
             self.logger.debug(f"Added final result: {final_text[:50]}...")
  
-        # final token usage
-        yield chatbot_history, self._get_token_usage_text()
+        # final token usage and logs
+        logs = self.log_stream.getvalue() if self._log_component else None
+        yield chatbot_history, self._get_token_usage_text(), logs
 
     def _format_response(self, response_text):
         """
@@ -839,6 +875,22 @@ class GradioCallback:
                     # Clear button
                     clear_btn = gr.Button("Clear Conversation")
                     
+                    # Log accordion - similar to the example provided
+                    with gr.Accordion("Agent Logs", open=False) as log_accordion:
+                        self._log_component = gr.Code(
+                            label="Live Logs",
+                            lines=15,
+                            interactive=False,
+                            value=self.log_stream.getvalue()
+                        )
+                        refresh_logs_btn = gr.Button("🔄 Refresh Logs")
+                        refresh_logs_btn.click(
+                            fn=lambda: self.log_stream.getvalue(),
+                            inputs=None,
+                            outputs=[self._log_component],
+                            queue=False
+                        )
+                    
                     # Store processed input temporarily between steps
                     processed_input_state = gr.State("")
 
@@ -859,7 +911,7 @@ class GradioCallback:
                         # 3. Run the main interaction loop (this yields updates)
                         fn=self.interact_with_agent,
                         inputs=[processed_input_state, self._chatbot_component],
-                        outputs=[self._chatbot_component, self._token_usage_component], # Update chat and tokens
+                        outputs=[self._chatbot_component, self._token_usage_component, self._log_component], # Update chat, tokens, and logs
                         queue=True # Explicitly enable queue for this async generator
                     ).then(
                         # 4. Re-enable the button after interaction finishes
@@ -885,7 +937,7 @@ class GradioCallback:
                         # 3. Run the main interaction loop (this yields updates)
                         fn=self.interact_with_agent,
                         inputs=[processed_input_state, self._chatbot_component],
-                        outputs=[self._chatbot_component, self._token_usage_component], # Update chat and tokens
+                        outputs=[self._chatbot_component, self._token_usage_component, self._log_component], # Update chat, tokens, and logs
                         queue=True # Explicitly enable queue for this async generator
                     ).then(
                         # 4. Re-enable the button after interaction finishes
@@ -899,8 +951,8 @@ class GradioCallback:
                     clear_btn.click(
                         fn=self.clear_conversation,
                         inputs=None, # No inputs needed
-                        # Outputs: Clear chatbot and reset token text
-                        outputs=[self._chatbot_component, self._token_usage_component],
+                        # Outputs: Clear chatbot, reset token text, and update logs
+                        outputs=[self._chatbot_component, self._token_usage_component, self._log_component],
                         queue=False # Run quickly
                     )
 
@@ -917,6 +969,12 @@ class GradioCallback:
         self.assistant_text_responses = []
         self.token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
         self.is_running = False
+        
+        # Clear log stream
+        if hasattr(self, 'log_stream'):
+            self.log_stream.seek(0)
+            self.log_stream.truncate(0)
+            self.logger.info("Log stream cleared")
 
         # Completely reset the agent state with a new session
         try:
@@ -965,8 +1023,9 @@ class GradioCallback:
         except Exception as e:
             self.logger.error(f"Failed to reset TinyAgent completely: {e}")
 
-        # Return cleared UI components: empty chat + fresh token usage
-        return [], self._get_token_usage_text()
+        # Return cleared UI components: empty chat + fresh token usage + empty logs
+        logs = self.log_stream.getvalue() if hasattr(self, 'log_stream') else ""
+        return [], self._get_token_usage_text(), logs
 
     def launch(self, agent, title="TinyAgent Chat", description=None, share=False, **kwargs):
         """
@@ -1028,21 +1087,31 @@ async def run_example():
     from tinyagent import TinyAgent  # Assuming TinyAgent is importable
     from tinyagent.hooks.logging_manager import LoggingManager  # Assuming LoggingManager exists
 
-    # --- Logging Setup (Simplified) ---
+    # --- Logging Setup (Similar to the example provided) ---
     log_manager = LoggingManager(default_level=logging.INFO)
     log_manager.set_levels({
         'tinyagent.hooks.gradio_callback': logging.DEBUG,
         'tinyagent.tiny_agent': logging.DEBUG,
         'tinyagent.mcp_client': logging.DEBUG,
+        'tinyagent.code_agent': logging.DEBUG,
     })
+    
+    # Console handler for terminal output
     console_handler = logging.StreamHandler(sys.stdout)
     log_manager.configure_handler(
         console_handler,
         format_string='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
         level=logging.DEBUG
     )
+    
+    # The Gradio UI will automatically set up its own log handler
+    # through the LoggingManager when we pass it to GradioCallback
+    
+    # Get loggers for different components
     ui_logger = log_manager.get_logger('tinyagent.hooks.gradio_callback')
     agent_logger = log_manager.get_logger('tinyagent.tiny_agent')
+    mcp_logger = log_manager.get_logger('tinyagent.mcp_client')
+    
     ui_logger.info("--- Starting GradioCallback Example ---")
     # --- End Logging Setup ---
 
@@ -1064,12 +1133,13 @@ async def run_example():
 
     agent.add_tool(get_weather)
 
-    # Create the Gradio callback
+    # Create the Gradio callback with LoggingManager integration
     gradio_ui = GradioCallback(
         file_upload_folder=upload_folder,
         show_thinking=True,
         show_tool_calls=True,
-        logger=ui_logger # Pass the specific logger
+        logger=ui_logger,
+        log_manager=log_manager  # Pass the LoggingManager for comprehensive logging
     )
     agent.add_callback(gradio_ui)
 
@@ -1084,25 +1154,9 @@ async def run_example():
         ui_logger.error(f"Failed to connect to MCP servers: {e}", exc_info=True)
         # Continue without servers - we still have the local get_weather tool
 
-    # Create the Gradio app but don't launch it yet
-    #app = gradio_ui.create_app(
-    #    agent,
-    #    title="TinyAgent Chat Interface",
-    #    description="Chat with TinyAgent. Try asking: 'Plan a trip to Toronto for 7 days in the next month.'",
-    #)
-    
-    # Configure the queue without extra parameters
-    #app.queue()
-    
-    # Launch the app in a way that doesn't block our event loop
+    # Launch the Gradio interface
     ui_logger.info("Launching Gradio interface...")
     try:
-        # Launch without blocking
-        #app.launch(
-        #    share=False,
-        #    prevent_thread_lock=True,  # Critical to not block our event loop
-        #    show_error=True
-        #)
         gradio_ui.launch(
             agent,
             title="TinyAgent Chat Interface",
@@ -1113,9 +1167,19 @@ async def run_example():
         )
         ui_logger.info("Gradio interface launched (non-blocking).")
         
+        # Generate some log messages to demonstrate the log panel
+        # These will appear in both the terminal and the Gradio UI log panel
+        ui_logger.info("UI component initialized successfully")
+        agent_logger.debug("Agent ready to process requests")
+        mcp_logger.info("MCP connection established")
+        
+        for i in range(3):
+            ui_logger.info(f"Example log message {i+1} from UI logger")
+            agent_logger.debug(f"Example debug message {i+1} from agent logger")
+            mcp_logger.warning(f"Example warning {i+1} from MCP logger")
+            await asyncio.sleep(1)
+        
         # Keep the main event loop running to handle both Gradio and MCP operations
-        # This is the key part - we need to keep our main event loop running
-        # but also allow it to process both Gradio and MCP client operations
         while True:
             await asyncio.sleep(1)  # More efficient than an Event().wait()