import os

from openai import AzureOpenAI

# Set up credentials
# NOTE: When running locally, these have to be set in the environment
client = AzureOpenAI(
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_KEY"),
    api_version="2024-02-01",
)

deployment_name = "sqlai"


def send_message(message: str) -> str:
    """Send a message to the openai chat completion API and return the response.

    Parameters
    ----------
    message : str
        The user's message to be sent to the chat completion API.

    Returns
    -------
    str
        The content of the assistant's response message.
    """
    response = client.chat.completions.create(
        model=deployment_name,
        messages=[
            {"role": "system", "content": "Du bist ein hilfreicher Assistent."},
            {"role": "user", "content": message},
        ],
    )
    return response.choices[0].message.content