shell-agent.py

#!/usr/bin/env python
import argparse
import json
import re
import signal
import subprocess
import sys
from urllib.request import Request, urlopen

MODEL_NAME = "qwen2.5-coder:7b"
CHAT_COMPLETIONS_URL = "http://<vanished>:11434/v1/chat/completions"
SYSTEM_PROMPT = (
    "You are a helpful assistant. Return only raw bash code without explanation."
)


def extract_code(s: str) -> str:
    if match := re.search(r"```(?:bash|sh)?\n(.*?)```", s, re.DOTALL):
        s = match.group(1)
    return s.strip()


def generate_code(prompt: str) -> list[str]:
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {
                "role": "system",
                "content": SYSTEM_PROMPT,
            },
            {"role": "user", "content": prompt},
        ],
        "stream": False,
        "options": {"temperature": 0.0},
    }
    req = Request(
        CHAT_COMPLETIONS_URL,
        method="POST",
        data=json.dumps(payload).encode(),
        headers={"content-type": "application/json"},
    )
    with urlopen(req, timeout=90) as res:
        data = json.loads(res.read())

    return [extract_code(choice["message"]["content"]) for choice in data["choices"]]


def signal_handler(sig, frame):
    print("\nBye!")
    sys.exit(0)


def confirm_and_run(cmd: str) -> int:
    print(cmd)
    print()
    if input("Confirm [Y/n]: ").strip().lower() not in ("", "y"):
        return 0
    ret = subprocess.run(["bash"], input=cmd, text=True)
    return ret.returncode


def main(argv: list[str]) -> None | int:
    parser = argparse.ArgumentParser()
    parser.add_argument("prompt", type=str.strip, help="Prompt")
    args = parser.parse_args(argv)
    if not args.prompt:
        parser.error("Empty prompt")

    results = generate_code(args.prompt)
    if not results:
        print("No variants returned from API.", file=sys.stderr)
        return 1

    if len(results) > 1:
        while True:
            for n, result in enumerate(results, 1):
                print(f"Variant #{n}:\n\n{result}\n\n" + "=" * 40 + "\n")
            try:
                i = int(input(f"Choose [1-{len(results)}] (0 to cancel): "))
                if i == 0:
                    return 0
                if 1 <= i <= len(results):
                    cmd = results[i - 1]
                    break
            except ValueError:
                pass
            print("Invalid choice. Try again...\n")

    else:
        cmd = results[0]

    return confirm_and_run(cmd)


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    sys.exit(main(sys.argv[1:]))