spb/uqo-chat

1"""execute_python — run Python in the sandbox-runner, collect outputs and figures."""23from __future__ import annotations45import base646import re7from typing import Any89from pydantic import BaseModel, Field1011from app.core.ratelimit import MSG_SANDBOX, limiter12from app.llm.schemas import Artifact, ToolResult13from app.sandbox.client import SandboxClient14from app.services import files as file_service15from app.tools.registry import ToolContext, registry161718class ExecArgs(BaseModel):19    code: str = Field(..., description="Python source code")20    description: str = Field("", max_length=300)21    heavy: bool = False222324IMAGE_EXT = {".png", ".jpg", ".jpeg", ".svg"}252627async def run(args: dict[str, Any], ctx: ToolContext) -> ToolResult:28    if ctx.role == "student":29        limiter.check(f"sandbox:{ctx.user_id}", ctx.settings.RATE_SANDBOX_PER_HOUR, 3600,30                      MSG_SANDBOX)31    await ctx.report("running", "Exécution du code Python…")32    code = args["code"].strip()33    m = re.match(r"^```(?:python|py)?\s*\n(.*?)\n```\s*$", code, re.S)34    if m:35        code = m.group(1)36    args["code"] = code37    client = SandboxClient(ctx.settings)38    files_in = await file_service.sandbox_inputs(ctx.user_id, ctx.file_ids)39    timeout = ctx.settings.SANDBOX_HEAVY_TIMEOUT_S if args.get("heavy") else None40    res = await client.run(args["code"], files_in, timeout)41    if res.error and res.exit_code == 127:42        return ToolResult(content=f"Le sandbox n'a pas pu exécuter le code : {res.error}. "43                          "Explique le calcul à la main ou réessaie plus tard.", error=True,44                          payload={"code": args["code"], "stderr": res.error})4546    artifacts: list[Artifact] = []47    file_lines: list[str] = []48    for f in res.files_out:49        name = f.get("name", "sortie")50        data = base64.b64decode(f.get("content_b64", "")) if f.get("content_b64") else b""51        if not data:52            continue53        ext = "." + name.rsplit(".", 1)[-1].lower() if "." in name else ""54        ftype = "image" if ext in IMAGE_EXT else ("xlsx" if ext == ".xlsx" else55                                                  "csv" if ext == ".csv" else "file")56        stored = await file_service.store_artifact(57            user_id=ctx.user_id, conversation_id=ctx.conversation_id, filename=name,58            data=data, ftype=ftype)59        artifacts.append(Artifact(type=ftype, file_id=stored.id, filename=name,60                                  url=f"/api/v1/files/{stored.id}"))61        file_lines.append(f"- {name} ({len(data)} octets)")6263    status = "ok" if res.exit_code == 0 else f"exit {res.exit_code}"64    content = (f"Exécution terminée ({status}, {res.duration_ms} ms).\n"65               f"STDOUT:\n{res.stdout or '(vide)'}\n")66    if res.stderr:67        content += f"STDERR:\n{res.stderr}\n"68    if file_lines:69        content += "Fichiers produits (déjà affichés à l'étudiant) :\n" + "\n".join(file_lines)70    if res.truncated:71        content += "\n(sortie tronquée)"72    return ToolResult(73        content=content, artifacts=artifacts, error=res.exit_code != 0,74        payload={"code": args["code"], "description": args.get("description", ""),75                 "stdout": res.stdout, "stderr": res.stderr, "exit_code": res.exit_code,76                 "duration_ms": res.duration_ms, "artifacts": [a.to_dict() for a in artifacts]},77        meta={"summary": args.get("description") or "Exécution Python"},78    )798081registry.register("execute_python", run, ExecArgs, heavy=True)82