add Python-Use

app.py

@@ -9,6 +9,7 @@ import io
 from settings_mgr import generate_download_settings_js, generate_upload_settings_js
 
 from doc2json import process_docx
+from code_exec import eval_restricted_script
 
 dump_controls = False
 log_to_console = False
@@ -143,7 +144,7 @@ def process_values_js():
     }
     """
 
-def bot(message, history, oai_key, system_prompt, seed, temperature, max_tokens, model):
+def bot(message, history, oai_key, system_prompt, seed, temperature, max_tokens, model, python_use):
     try:
         client = OpenAI(
             api_key=oai_key
@@ -192,6 +193,31 @@ def bot(message, history, oai_key, system_prompt, seed, temperature, max_tokens,
             if seed:
                 seed_i = int(seed)
 
+            tools = None if not python_use else [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "eval_python",
+                        "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.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "python_source_code": {
+                                    "type": "string",
+                                    "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. Results need to be reported through print()."
+                                }
+                            },
+                            "required": ["python_source_code"]
+                        }
+                    }
+                }
+            ]
+
             if log_to_console:
                 print(f"bot history: {str(history)}")
 
@@ -242,30 +268,107 @@ def bot(message, history, oai_key, system_prompt, seed, temperature, max_tokens,
                 if log_to_console:
                         print(f"usage: {response.usage}")
             else:
-                response = client.chat.completions.create(
-                    model=model,
-                    messages= history_openai_format,
-                    temperature=temperature,
-                    seed=seed_i,
-                    max_tokens=max_tokens,
-                    stream=True,
-                    stream_options={"include_usage": True}
-                )
-
-                partial_response=""
-                for chunk in response:
-                    if chunk.choices:
-                        txt = ""
-                        for choice in chunk.choices:
-                            cont = choice.delta.content
-                            if cont:
-                                txt += cont
-
-                        partial_response += txt
-                        yield partial_response
-
-                    if chunk.usage and log_to_console:
-                        print(f"usage: {chunk.usage}")
+                whole_response = ""
+                while True:
+                    response = client.chat.completions.create(
+                        model=model,
+                        messages= history_openai_format,
+                        temperature=temperature,
+                        seed=seed_i,
+                        max_tokens=max_tokens,
+                        stream=True,
+                        stream_options={"include_usage": True},
+                        tools = tools,
+                        tool_choice = "auto" if python_use else None
+                    )
+
+                    # Accumulators for partial model responses
+                    tool_name_accum = None
+                    tool_args_accum = ""
+                    tool_call_id = None
+                    # process
+                    for chunk in response:
+                        if chunk.choices:
+                            txt = ""
+                            for choice in chunk.choices:
+                                delta = choice.delta
+                                if not delta:
+                                    continue
+
+                                cont = delta.content
+                                if cont:
+                                    txt += cont
+                                
+                                if delta.tool_calls:
+                                    for tc in delta.tool_calls:
+                                        if tc.function.name:
+                                            tool_name_accum = tc.function.name
+                                        if tc.function.arguments:
+                                            tool_args_accum += tc.function.arguments
+                                        if tc.id:
+                                            tool_call_id = tc.id
+
+                            finish_reason = choice.finish_reason
+                            if finish_reason:
+                                if finish_reason == "tool_calls":
+                                    try:
+                                        parsed_args = json.loads(tool_args_accum)
+                                        tool_script = parsed_args.get("python_source_code", "")
+
+                                        whole_response += f"\n``` script\n{tool_script}\n```\n"
+                                        yield whole_response
+
+                                        tool_result = eval_restricted_script(tool_script)
+
+                                        whole_response += f"\n``` result\n{tool_result if not tool_result['success'] else tool_result['prints']}\n```\n"
+                                        yield whole_response
+
+                                        history_openai_format.extend([
+                                            {
+                                                "role": "assistant",
+                                                "content": txt,
+                                                "tool_calls": [
+                                                    {
+                                                        "id": tool_call_id,
+                                                        "type": "function",
+                                                        "function": {
+                                                            "name": tool_name_accum,
+                                                            "arguments": json.dumps(parsed_args)
+                                                        }
+                                                    }
+                                                ]
+                                            },
+                                            {
+                                                "role": "tool",
+                                                "tool_call_id": tool_call_id,
+                                                "name": tool_name_accum,
+                                                "content": json.dumps(tool_result)
+                                            }
+                                        ])
+
+                                    except Exception as e:
+                                        history_openai_format.extend([{
+                                            "role": "tool",
+                                            "tool_call_id": tool_call_id,
+                                            "name": tool_name_accum,
+                                            "content": [
+                                                {
+                                                    "toolResult": {
+                                                        "content": [{"text":  e.args[0]}],
+                                                        "status": 'error'
+                                                    }
+                                                }
+                                            ]
+                                        }])
+                                        whole_response += f"\n``` error\n{e.args[0]}\n```\n"
+                                        yield whole_response
+                                else:
+                                    return
+                            else:
+                                whole_response += txt
+                                yield whole_response
+                        if chunk.usage and log_to_console:
+                            print(f"usage: {chunk.usage}")
 
         if log_to_console:
             print(f"br_result: {str(history)}")
@@ -311,6 +414,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
         seed = gr.Textbox(label="Seed", elem_id="seed")
         temp = gr.Slider(0, 2, label="Temperature", elem_id="temp", value=1)
         max_tokens = gr.Slider(1, 16384, label="Max. Tokens", elem_id="max_tokens", value=800)
+        python_use = gr.Checkbox(label="Python Use", value=False)
         save_button = gr.Button("Save Settings")  
         load_button = gr.Button("Load Settings")  
         dl_settings_button = gr.Button("Download Settings")
@@ -345,7 +449,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
                        ('temp', '#temp input'),
                        ('max_tokens', '#max_tokens input'),
                        ('model', '#model')]
-        controls = [oai_key, system_prompt, seed, temp, max_tokens, model]
+        controls = [oai_key, system_prompt, seed, temp, max_tokens, model, python_use]
 
         dl_settings_button.click(None, controls, js=generate_download_settings_js("oai_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,243 @@
+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, 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
+import types
+
+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,
+        '_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', 'datetime'],
+        '__import__': __import__,
+        
+        # Basic functions
+        'len': len,
+        'range': limited_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': limited_list,
+        'tuple': limited_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
+        }
+
+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)

requirements.txt

@@ -1,4 +1,5 @@
 gradio == 5.1
 openai >= 1.0.0
 lxml
-PyMuPDF
+PyMuPDF
+RestrictedPython