ChatGeneration

import os

from chainlit.playground.providers import ChatOpenAI
import chainlit as cl

# If no OPENAI_API_KEY is available, the ChatOpenAI provider won't be available in the prompt playground
os.environ["OPENAI_API_KEY"] = "sk-..."

template = """Hello, this is a template.
This is a variable1 {variable1}
And this is variable2 {variable2}
"""

variables = {
    "variable1": "variable1 value",
    "variable2": "variable2 value",
}

settings = {
    "model": "gpt-3.5-turbo",
    "temperature": 0,
    # ... more settings
}


@cl.step(type="llm")
async def fake_llm():
    # Pretend we're calling the LLM
    await cl.sleep(2)

    fake_completion = "The openai completion"

    prompt = template.format(**variables)

    # Add the generation to the current step
    cl.context.current_step.generation = cl.ChatGeneration(
        model=settings["model"],
        provider=ChatOpenAI.id,
        message_completion={
          "content": fake_completion, "role": "assistant"
        },
        variables=variables,
        settings=settings,
        messages=[
            {
                "content": "You are a helpful assistant.", "role": "system"
            },
            {
                "content": prompt, "role": "user"
            },
        ],
    )

    return fake_completion


@cl.on_chat_start
async def start():
    await fake_llm()