extend code eval

app.py

@@ -46,14 +46,19 @@ def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temp
             "tools": [{
                 "toolSpec": {
                     "name": "eval_python",
-                    "description": "Evaluate RestrictedPython script",
+                    "description": "Evaluate a simple script written in a conservative, restricted subset of Python."
+                                "Note: Augmented assignments, in-place operations (e.g., +=, -=), lambdas (e.g. list comprehensions) are not supported. "
+                                "Use regular assignments and operations instead. Only 'import math' is allowed. "
+                                "Returns: unquoted results without HTML encoding.",
                     "inputSchema": {
                         "json": {
                             "type": "object",
                             "properties": {
                                 "script": {
                                     "type": "string",
-                                    "description": "The Python script that will run in a RestrictedPython context"
+                                    "description": "The Python script that will run in a RestrictedPython context. "
+                                                "Avoid using augmented assignments or in-place operations (+=, -=, etc.), as well as lambdas (e.g. list comprehensions). "
+                                                "Use regular assignments and operations instead. Only 'import math' is allowed."
                                 }
                             },
                             "required": ["script"]
@@ -111,13 +116,13 @@ def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temp
                                                 {
                                                     "toolResult": {
                                                         "toolUseId": tool['toolUseId'],
-                                                        "content": [{"json": tool_result}]
+                                                        "content": [{"json": tool_result }]
                                                     }
                                                 }
                                             ]
                                         }
 
-                                        whole_response += f"\n``` result\n{tool_result}\n```\n"
+                                        whole_response += f"\n``` result\n{tool_result if not tool_result['success'] else tool_result['prints']}\n```\n"
                                         yield whole_response
                                     except Exception as e:
                                         tool_result_message = {

code_exec.py

@@ -1,8 +1,10 @@
 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.Eval import default_guarded_getiter, default_guarded_getitem
 from RestrictedPython.Utilities import utility_builtins
+from RestrictedPython.Guards import guarded_unpack_sequence
+from RestrictedPython.Limits import limited_range, limited_list, limited_tuple
 from io import StringIO
 
 def eval_restricted_script(script):
@@ -39,8 +41,13 @@ def eval_restricted_script(script):
         # Print handling
         '_print_': CustomPrintCollector,
         '_getattr_': getattr,
+        '_getitem_': default_guarded_getitem,
         '_getiter_': default_guarded_getiter,
         '_iter_unpack_sequence_': guarded_iter_unpack_sequence,
+        '_unpack_sequence_': guarded_unpack_sequence,
+        '_inplacevar_': protected_inplacevar,
+        '_apply_': _apply,
+        '_write_': _default_write_,
         
         # Define allowed imports
         '__allowed_modules__': ['math'],
@@ -48,7 +55,7 @@ def eval_restricted_script(script):
         
         # Basic functions
         'len': len,
-        'range': range,
+        'range': limited_range,
         'enumerate': enumerate,
         'zip': zip,
         
@@ -65,8 +72,8 @@ def eval_restricted_script(script):
         'float': float,
         'str': str,
         'bool': bool,
-        'list': list,
-        'tuple': tuple,
+        'list': limited_list,
+        'tuple': limited_tuple,
         'set': set,
         'dict': dict,
         'bytes': bytes,
@@ -107,3 +114,129 @@ def eval_restricted_script(script):
             'error': str(e),
             'success': False
         }
+
+def _default_write_(obj):
+    if isinstance(obj, types.ModuleType):
+        raise ValueError("Modules are not allowed in to be written to.")
+    
+    return obj
+
+""" 
+Borrowed implementation of _inplacevar_ from the Zope Foundations's AccessControl module
+https://github.com/zopefoundation/AccessControl/blob/f9ae58816f0712eb6ea97459b4ccafbf4662d9db/src/AccessControl/ZopeGuards.py#L530
+"""
+
+valid_inplace_types = (list, set)
+
+
+inplace_slots = {
+    '+=': '__iadd__',
+    '-=': '__isub__',
+    '*=': '__imul__',
+    '/=': (1 / 2 == 0) and '__idiv__' or '__itruediv__',
+    '//=': '__ifloordiv__',
+    '%=': '__imod__',
+    '**=': '__ipow__',
+    '<<=': '__ilshift__',
+    '>>=': '__irshift__',
+    '&=': '__iand__',
+    '^=': '__ixor__',
+    '|=': '__ior__',
+}
+
+
+def __iadd__(x, y):
+    x += y
+    return x
+
+
+def __isub__(x, y):
+    x -= y
+    return x
+
+
+def __imul__(x, y):
+    x *= y
+    return x
+
+
+def __idiv__(x, y):
+    x /= y
+    return x
+
+
+def __ifloordiv__(x, y):
+    x //= y
+    return x
+
+
+def __imod__(x, y):
+    x %= y
+    return x
+
+
+def __ipow__(x, y):
+    x **= y
+    return x
+
+
+def __ilshift__(x, y):
+    x <<= y
+    return x
+
+
+def __irshift__(x, y):
+    x >>= y
+    return x
+
+
+def __iand__(x, y):
+    x &= y
+    return x
+
+
+def __ixor__(x, y):
+    x ^= y
+    return x
+
+
+def __ior__(x, y):
+    x |= y
+    return x
+
+
+inplace_ops = {
+    '+=': __iadd__,
+    '-=': __isub__,
+    '*=': __imul__,
+    '/=': __idiv__,
+    '//=': __ifloordiv__,
+    '%=': __imod__,
+    '**=': __ipow__,
+    '<<=': __ilshift__,
+    '>>=': __irshift__,
+    '&=': __iand__,
+    '^=': __ixor__,
+    '|=': __ior__,
+}
+
+
+def protected_inplacevar(op, var, expr):
+    """Do an inplace operation
+
+    If the var has an inplace slot, then disallow the operation
+    unless the var an instance of ``valid_inplace_types``.
+    """
+    if hasattr(var, inplace_slots[op]) and \
+    not isinstance(var, valid_inplace_types):
+        try:
+            cls = var.__class__
+        except AttributeError:
+            cls = type(var)
+        raise TypeError(
+            "Augmented assignment to %s objects is not allowed"
+            " in untrusted code" % cls.__name__)
+    return inplace_ops[op](var, expr)
+
+def _apply(f, *a, **kw):
+    return f(*a, **kw)