import builtins
from functools import partial
from viper.src.model_processes import forward
from viper.src.utils import find_forbidden_imports, restricted_import
from viper.src.api import *
from viper.configs import config
ALLOWED_IMPORTS = {"math"}
# Assumes no batching; if enable batching within model, can just pass list of prompts
def generate_code(prompt, action):
"""
Generates the code to be executed based on the provided prompts and actions.
"""
code = []
for p, a in zip(prompt, action):
response = forward(config.codex.name, p, a)
if response is None:
raise KeyError("No code generated. Did you set the OPENAI_API_KEY environment variable?")
if isinstance(response, list):
code.extend(response)
else:
code.append(response)
return code
def execute_code(code, image):
"""
Executes the generated code on the provided images.
"""
results = []
for c, img in zip(code, image):
# 1) static scan
problems = find_forbidden_imports(c, ALLOWED_IMPORTS)
if problems:
msgs = ", ".join([f"{mod} (line {ln})" for ln, mod in problems])
raise ImportError(f"Forbidden imports detected: {msgs}")
# 2) prune builtins and intercept imports
allowed_builtin_names = {
"abs","all","any","bool","dict","enumerate","float","int","len","list",
"max","min","print","range","repr","set","str","sum","zip"
}
safe_builtins = {k: getattr(builtins, k) for k in allowed_builtin_names}
safe_builtins["__import__"] = partial(restricted_import, allowed=ALLOWED_IMPORTS)
g = {"__builtins__": safe_builtins}
output = {}
exec(compile(c, 'Codex', 'exec'), globals() | g, output)
res = output['execute_command'](img) # Generated by the compilation
results.append(res)
return results
