import os

from scale_gp import SGPClient
from scale_gp.types.completion_create_params import ModelParameters
from scale_gp.types.chat_completion_create_params import (
    ModelParameters as ChatModelParameters,
    MessageEgpAPIBackendServerAPIModelsEgpModelsUserMessage,
)

client = SGPClient(api_key=api_key)

model_parameters = ModelParameters(max_tokens=200, temperature=0.5, top_k=1, top_p=1)

completion = client.completions.create(
    model="gpt-4",
    prompt="Why is the sky blue?",
    account_id=account_id,
    model_parameters=model_parameters,
)

print(completion)

chat_model_parameters = ChatModelParameters(max_tokens=200, temperature=0.5, top_k=1, top_p=1)
message = MessageEgpAPIBackendServerAPIModelsEgpModelsUserMessage(role="user", content="What is the capital of Canada?")

chat_completion = client.chat_completions.create(
    model="gpt-4",
    messages=[
        message
    ],  # messages is a list of historical messages in the conversation, with roles usually alternating between user, assistant and system
    account_id=account_id,
    model_parameters=model_parameters,
    instructions="Answer the question like an elementary school teacher.",
)

print(chat_completion)

CompletionsResponse(
    completion=Completion(text="The sky appears blue because of the way Earth's atmosphere scatters sunlight in all directions and blue light is scattered more than other colors because it travels in smaller, shorter waves. This is known as Rayleigh scattering.", finish_reason='stop', response_metadata=None),
    token_usage=TokenUsage(total=56, completion=43, prompt=13)
)