Update app.py

app.py

@@ -1,44 +1,69 @@
 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)
 
-# Клиент Hugging Face для выбранной модели
-client = InferenceClient(
-    model="deepseek-ai/deepseek-coder-1.3b-base",  # Можешь заменить на starcoder2 или codegemma
-    token=HF_TOKEN
-)
-
-# Функция генерации без stream (deepseek не поддерживает stream=True)
-def complete_code(prompt, max_tokens, temperature, top_p):
-    response = client.text_generation(
+# Генерация кода с помощью 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  # ⚠️ обязательно False для этой модели
+        stream=False
     )
-    return response
+
+# Выполнение сгенерированного кода в песочнице
+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
-demo = gr.Interface(
-    fn=complete_code,
-    inputs=[
-        gr.Textbox(lines=10, label="Введите код или комментарий"),
-        gr.Slider(minimum=1, maximum=1024, value=128, step=1, label="Максимум новых токенов"),
-        gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
-    ],
-    outputs=gr.Textbox(lines=15, label="Дополненный код"),
-    title="🛠️ Уязвимый агент автодополнения кода",
-    description=(
-        "Этот агент демонстрирует поведение LLM без фильтрации: он уязвим к prompt injection, "
-        "может генерировать вредоносные команды и не имеет встроенных ограничений. "
-        "Использовать только в исследовательских целях!"
-    ),
-)
+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()