Bad_agent / app.py
Timaska's picture
Update app.py
b756f30 verified
import gradio as gr
import tempfile
import subprocess
import os
from huggingface_hub import InferenceClient
HF_TOKEN = os.getenv("HF_TOKEN")
client = InferenceClient("deepseek-ai/deepseek-coder-1.3b-base", token=HF_TOKEN)
# Генерация кода с помощью LLM
def generate_code(prompt, max_tokens, temperature, top_p):
return client.text_generation(
prompt,
max_new_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=False
)
# Выполнение сгенерированного кода в песочнице
def execute_code(code):
with tempfile.NamedTemporaryFile(mode="w+", suffix=".py", delete=False) as temp:
temp.write(code)
temp.flush()
try:
result = subprocess.run(
["python3", temp.name],
capture_output=True,
text=True,
timeout=5 # ограничение по времени
)
output = result.stdout + "\n" + result.stderr
except subprocess.TimeoutExpired:
output = "⏱️ Превышено время выполнения (timeout)"
except Exception as e:
output = f"❌ Ошибка запуска: {e}"
finally:
os.unlink(temp.name)
return output
# Интерфейс Gradio
with gr.Blocks() as demo:
gr.Markdown("🛠️ Уязвимый агент автодополнения + изолированное выполнение кода")
with gr.Row():
prompt = gr.Textbox(lines=10, label="Введите подсказку для генерации")
generated_code = gr.Textbox(lines=15, label="Сгенерированный код")
with gr.Row():
generate_btn = gr.Button("🤖 Сгенерировать")
run_btn = gr.Button("🚀 Выполнить")
output = gr.Textbox(lines=15, label="Результат выполнения")
# Привязка функций
generate_btn.click(
fn=generate_code,
inputs=[prompt, gr.Number(128), gr.Number(0.7), gr.Number(0.95)],
outputs=generated_code
)
run_btn.click(
fn=execute_code,
inputs=generated_code,
outputs=output
)
if __name__ == "__main__":
demo.launch()