"""Print greedy answers to fixed prompts from an OpenAI-compatible server, as JSON.

usage: python3 scripts/greedy.py http://127.0.0.1:8000 MODEL > answers.json
"""

import json
import sys
import urllib.request

PROMPTS = [
    "Explain what a mixture-of-experts language model is in two sentences.",
    "Write a Python function that returns the n-th Fibonacci number.",
    "Name three prime numbers greater than 100.",
]


def complete(base_url, model, prompt):
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
                       "temperature": 0, "max_tokens": 64}).encode()
    request = urllib.request.Request(f"{base_url}/v1/chat/completions", body,
                                     {"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=1800) as response:
        return json.load(response)["choices"][0]["message"]["content"]


if __name__ == "__main__":
    base_url, model = sys.argv[1], sys.argv[2]
    print(json.dumps([complete(base_url, model, p) for p in PROMPTS], indent=1))