Python Use

app.py

@@ -7,6 +7,7 @@ from PIL import Image
 
 from settings_mgr import generate_download_settings_js, generate_upload_settings_js
 from llm import LLM, log_to_console
+from code_exec import eval_restricted_script
 from botocore.config import Config
 
 dump_controls = False
@@ -29,24 +30,39 @@ def process_values_js():
     }
     """
 
-def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temperature, max_tokens, model: str, region):
+def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temperature, max_tokens, model: str, region, python_use):
     try:
         llm = LLM.create_llm(model)
         messages = llm.generate_body(message, history)
-        if system_prompt:
-            sys_prompt = [{"text": system_prompt}]
-        else:
-            sys_prompt = []
+        sys_prompt = [{"text": system_prompt}] if system_prompt else []
 
         config = Config(
             read_timeout = 600,
             connect_timeout = 30,
-            retries = {
-                'max_attempts': 10,
-                'mode': 'adaptive'
-            }
+            retries = {'max_attempts': 10, 'mode': 'adaptive'}
         )
 
+        tool_config = {
+            "tools": [{
+                "toolSpec": {
+                    "name": "eval_python",
+                    "description": "Evaluate RestrictedPython script",
+                    "inputSchema": {
+                        "json": {
+                            "type": "object",
+                            "properties": {
+                                "script": {
+                                    "type": "string",
+                                    "description": "The Python script that will run in a RestrictedPython context"
+                                }
+                            },
+                            "required": ["script"]
+                        }
+                    }
+                }
+            }]
+        } if python_use else None
+
         sess = boto3.Session(
             aws_access_key_id = aws_access,
             aws_secret_access_key = aws_secret,
@@ -54,21 +70,73 @@ def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temp
             region_name = region)
         br = sess.client(service_name="bedrock-runtime", config = config)
 
-        response = br.converse_stream(
-            modelId = model,
-            messages = messages,
-            system = sys_prompt,
-            inferenceConfig = {
-                "temperature": temperature,
-                "maxTokens": max_tokens,
-            }
-        )
-        response_stream = response.get('stream')
-        
-        partial_response = ""
-        for chunk in llm.read_response(response_stream):
-            partial_response += chunk
-            yield partial_response
+        whole_response = ""
+        while True:
+            response = br.converse_stream(
+                modelId = model,
+                messages = messages,
+                system = sys_prompt,
+                inferenceConfig = {
+                    "temperature": temperature,
+                    "maxTokens": max_tokens,
+                },
+                **({'toolConfig': tool_config} if python_use else {})
+            )
+
+            for stop_reason, message in llm.read_response(response.get('stream')):
+                if isinstance(message, str):
+                    whole_response += message
+                    yield whole_response
+
+                if stop_reason:
+                    if stop_reason == "tool_use":
+                        messages.append(message)
+
+                        for content in message['content']:
+                            if 'toolUse' in content:
+                                tool = content['toolUse']
+
+                                if tool['name'] == 'eval_python':
+                                    tool_result = {}
+                                    try:
+                                        tool_script = tool["input"]["script"]
+
+                                        whole_response += f"\n``` script\n{tool_script}\n```\n"
+                                        yield whole_response
+
+                                        tool_result = eval_restricted_script(tool_script)
+                                        tool_result_message = {
+                                            "role": "user",
+                                            "content": [
+                                                {
+                                                    "toolResult": {
+                                                        "toolUseId": tool['toolUseId'],
+                                                        "content": [{"json": tool_result}]
+                                                    }
+                                                }
+                                            ]
+                                        }
+
+                                        whole_response += f"\n``` result\n{tool_result}\n```\n"
+                                        yield whole_response
+                                    except Exception as e:
+                                        tool_result_message = {
+                                            "role": "user",
+                                            "content": [
+                                                {
+                                                    "toolResult": {
+                                                        "content": [{"text":  e.args[0]}],
+                                                        "status": 'error'
+                                                    }
+                                                }
+                                            ]
+                                        }
+                                        whole_response += f"\n``` error\n{e.args[0]}\n```\n"
+                                        yield whole_response
+
+                                    messages.append(tool_result_message)
+                    else:
+                        return
 
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
@@ -136,6 +204,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
                             choices=["eu-central-1", "eu-west-3", "us-east-1", "us-west-1", "us-west-2"])
         temp = gr.Slider(0, 1, label="Temperature", elem_id="temp", value=1)
         max_tokens = gr.Slider(1, 8192, label="Max. Tokens", elem_id="max_tokens", value=4096)
+        python_use = gr.Checkbox(label="Python Use")
         save_button = gr.Button("Save Settings")  
         load_button = gr.Button("Load Settings")  
         dl_settings_button = gr.Button("Download Settings")
@@ -174,7 +243,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
                        ('max_tokens', '#max_tokens input'),
                        ('model', '#model'),
                        ('region', '#region')]
-        controls = [aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region]
+        controls = [aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region, python_use]
 
         dl_settings_button.click(None, controls, js=generate_download_settings_js("amz_chat_settings.bin", control_ids))
         ul_settings_button.click(None, None, None, js=generate_upload_settings_js(control_ids))

code_exec.py

@@ -0,0 +1,109 @@
+from RestrictedPython import compile_restricted
+from RestrictedPython.PrintCollector import PrintCollector
+from RestrictedPython.Guards import safe_globals, safe_builtins, guarded_iter_unpack_sequence
+from RestrictedPython.Eval import default_guarded_getiter
+from RestrictedPython.Utilities import utility_builtins
+from io import StringIO
+
+def eval_restricted_script(script):
+    # Set up print collector and output handling
+    all_prints = StringIO()
+    
+    class CustomPrintCollector:
+        """Collect printed text, accumulating in shared StringIO"""
+        
+        def __init__(self, _getattr_=None):
+            self.txt = []
+            self._getattr_ = _getattr_
+        
+        def write(self, text):
+            all_prints.write(text)
+            self.txt.append(text)
+            
+        def __call__(self):
+            result = ''.join(self.txt)
+            return result
+            
+        def _call_print(self, *objects, **kwargs):
+            if kwargs.get('file', None) is None:
+                kwargs['file'] = self
+            else:
+                self._getattr_(kwargs['file'], 'write')
+            
+            print(*objects, **kwargs)
+
+    # Create the restricted builtins dictionary
+    restricted_builtins = dict(safe_builtins)
+    restricted_builtins.update(utility_builtins)  # Add safe __import__
+    restricted_builtins.update({
+        # Print handling
+        '_print_': CustomPrintCollector,
+        '_getattr_': getattr,
+        '_getiter_': default_guarded_getiter,
+        '_iter_unpack_sequence_': guarded_iter_unpack_sequence,
+        
+        # Define allowed imports
+        '__allowed_modules__': ['math'],
+        '__import__': __import__,
+        
+        # Basic functions
+        'len': len,
+        'range': range,
+        'enumerate': enumerate,
+        'zip': zip,
+        
+        # Math operations
+        'sum': sum,
+        'max': max,
+        'min': min,
+        'abs': abs,
+        'round': round,
+        'pow': pow,
+        
+        # Type conversions
+        'int': int,
+        'float': float,
+        'str': str,
+        'bool': bool,
+        'list': list,
+        'tuple': tuple,
+        'set': set,
+        'dict': dict,
+        'bytes': bytes,
+        'bytearray': bytearray,
+        
+        # Sequence operations
+        'all': all,
+        'any': any,
+        'sorted': sorted,
+        'reversed': reversed,
+        
+        # String operations
+        'chr': chr,
+        'ord': ord,
+        
+        # Other safe operations
+        'isinstance': isinstance,
+        'issubclass': issubclass,
+        'hasattr': hasattr,
+        'callable': callable,
+        'format': format,
+    })
+
+    # Create the restricted globals dictionary
+    restricted_globals = dict(safe_globals)
+    restricted_globals['__builtins__'] = restricted_builtins
+
+    try:
+        byte_code = compile_restricted(script, filename='<inline>', mode='exec')
+        exec(byte_code, restricted_globals)
+        
+        return {
+            'prints': all_prints.getvalue(),
+            'success': True
+        }
+    except Exception as e:
+        return {
+            'error': str(e),
+            'success': False
+        }

llm.py

@@ -232,16 +232,48 @@ class LLM:
         }
 
     def read_response(self, response_stream):
-        for event in response_stream:
-            if 'contentBlockDelta' in event:
-                yield event['contentBlockDelta']['delta']['text']
-            if 'messageStop' in event:
-                if log_to_console:
-                    print(f"\nStop reason: {event['messageStop']['stopReason']}")
-            if 'metadata' in event:
-                metadata = event['metadata']
-                if 'usage' in metadata and log_to_console:
-                    print("\nToken usage:")
-                    print(f"Input tokens: {metadata['usage']['inputTokens']}")
-                    print(f"Output tokens: {metadata['usage']['outputTokens']}")
-                    print(f"Total tokens: {metadata['usage']['totalTokens']}")
+        """
+        Handles response stream that may contain both regular text and tool use requests.
+        Yields tuples of (text, tool_request, stop_reason) where:
+        - text: accumulated text response
+        - tool_request: dict with tool use details if present, None otherwise
+        - stop_reason: string indicating why stream stopped, None while streaming
+        """
+        message = {}
+        content = []
+        message['content'] = content
+        tool_use = {}
+        text = ''
+        stop_reason = None
+
+        for chunk in response_stream:
+            if 'messageStart' in chunk:
+                message['role'] = chunk['messageStart']['role']
+            elif 'contentBlockStart' in chunk:
+                tool = chunk['contentBlockStart']['start']['toolUse']
+                tool_use['toolUseId'] = tool['toolUseId']
+                tool_use['name'] = tool['name']
+            elif 'contentBlockDelta' in chunk:
+                delta = chunk['contentBlockDelta']['delta']
+                if 'toolUse' in delta:
+                    if 'input' not in tool_use:
+                        tool_use['input'] = ''
+                    tool_use['input'] += delta['toolUse']['input']
+                elif 'text' in delta:
+                    text += delta['text']
+                    yield None, delta['text']
+            elif 'contentBlockStop' in chunk:
+                if 'input' in tool_use:
+                    tool_use['input'] = json.loads(tool_use['input'])
+                    content.append({'toolUse': tool_use})
+                    tool_use = {}
+                else:
+                    content.append({'text': text})
+            elif 'messageStop' in chunk:
+                stop_reason = chunk['messageStop']['stopReason']
+                yield stop_reason, message
+            elif 'metadata' in chunk and 'usage' in chunk['metadata'] and log_to_console:
+                print("\nToken usage:")
+                print(f"Input tokens: {metadata['usage']['inputTokens']}")
+                print(f"Output tokens: {metadata['usage']['outputTokens']}")
+                print(f"Total tokens: {metadata['usage']['totalTokens']}")

requirements.txt

@@ -2,4 +2,5 @@ gradio == 5.1
 langchain
 boto3>1.34.54
 lxml
-PyMuPDF
+PyMuPDF
+RestrictedPython