import requests
import json
import os


url = 'https://internlm-chat.intern-ai.org.cn/puyu/api/v1/chat/completions'
header = {
    'Content-Type': 'application/json',
    "Authorization": f"Bearer {os.environ['INTERNLM_API_KEY']}"
}


def complete(messages, model="internlm3-latest", temperature=None, top_p=None):
    data = {
        "model": model,
        "messages": messages
    }
    if temperature is not None:
        data["temperature"] = temperature
    if top_p is not None:
        data["top_p"] = top_p

    res = requests.post(url, headers=header, data=json.dumps(data))
    if res.status_code == 200:
        return res.json()["choices"][0]['message']["content"]
    else:
        return None


if __name__ == "__main__":
    response = complete(
        messages=[{"role": "user", "content": "你好~"}]
    )
    print(response)