# /// script
# requires-python = ">=3.11"
# dependencies = ["exedev>=0.4,<0.5", "openai>=2,<3"]
# ///
"""Answer one question with code: a sandbox VM for the question, deleted after.

    EXE_TOKEN=exe1.... uv run ask.py "What is the 1000th prime?"

Needs sandbox.py beside it. On an exe.dev VM the model is reached through the
LLM gateway with no key; elsewhere, use OpenAI() with your OPENAI_API_KEY.
https://exe.dev/deploy/code-interpreter/ask.py
"""

import json
import sys

import exedev
from openai import OpenAI

from sandbox import Sandbox

exe = exedev.Client()  # reads EXE_TOKEN
llm = OpenAI(base_url="https://llm.int.exe.xyz/v1", api_key="unused")

TOOLS = [{
    "type": "function",
    "function": {
        "name": "run_python",
        "description": "Run a cell in a Jupyter kernel and return its output. State persists between calls.",
        "parameters": {"type": "object", "properties": {"code": {"type": "string"}}, "required": ["code"]},
    },
}]


def ask(question: str) -> str:
    messages = [
        {"role": "system", "content": "Answer by running Python with run_python; don't guess."},
        {"role": "user", "content": question},
    ]
    with Sandbox.create(exe, comment="ask.py") as sb:  # a new VM, deleted at the end of the block
        for _ in range(20):
            msg = llm.chat.completions.create(model="gpt-5.5", messages=messages, tools=TOOLS).choices[0].message
            messages.append(msg)
            if not msg.tool_calls:
                return msg.content
            for call in msg.tool_calls:
                code = json.loads(call.function.arguments)["code"]
                print(f"--- running:\n{code}\n", file=sys.stderr)
                result = sb.run_code(code, timeout=120).for_model()
                messages.append({"role": "tool", "tool_call_id": call.id, "content": result})
    return "Gave up after 20 rounds of code."


if __name__ == "__main__":
    print(ask(" ".join(sys.argv[1:])))