"""Call a configured AI Experiments gateway endpoint."""

import os

from openai import OpenAI

from get_token import get_token


def main() -> None:
    tracking_uri = os.environ["MLFLOW_TRACKING_URI"].rstrip("/")
    endpoint_name = os.environ["MLFLOW_GATEWAY_ENDPOINT"]
    client = OpenAI(
        base_url=f"{tracking_uri}/gateway/mlflow/v1",
        api_key=get_token(),
    )
    response = client.chat.completions.create(
        model=endpoint_name,
        messages=[{"role": "user", "content": "Hello!"}],
    )
    print(response.choices[0].message.content)


if __name__ == "__main__":
    main()
