# You can find this code for Chainlit python streaming here (https://docs.chainlit.io/concepts/streaming/python)

# OpenAI Chat completion
import os
from openai import AsyncOpenAI  # importing openai for API usage
import chainlit as cl  # importing chainlit for our app
from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
from dotenv import load_dotenv
load_dotenv()

# ChatOpenAI Templates
system_template = """\
###Instruction###
You are an expert assistant answering technical questions on machine learning and deep learning subject.
Ensure that your response is unbiased and generic, you will be 'AWARDED' for giving really good clarity and correct answers.
##EXAMPLES##
If users asks 'explain neural networks', your response should be with an overview of neural networks, discussing how they are computational models inspired by the human brain that are used to recognize patterns and solve complex problems in machine learning.
If users ask 'code convolutional neural network', your response should contain example of the code necessary to create a CNN.
If users ask 'resource gradient descent', your response should offer links to tutorials, video lectures, or articles that explain gradient descent, which is an optimization algorithm used to minimize a function by iteratively moving in the direction of steepest descent.
If users ask 'project sentiment analysis', your response should discuss the steps involved in creating a sentiment analysis model, such as data collection, preprocessing, model selection, training, and evaluation, and potentially offer advice on best practices or methodologies to consider.
"""

user_template = """{input} \n + 
Think and give only explanation or code or links for resources or steps for project, for the questions asked along with response.
"""


@cl.on_chat_start  # marks a function that will be executed at the start of a user session
async def start_chat():
    settings = {
        "model": "gpt-3.5-turbo",
        "temperature": 0,
        "max_tokens": 500,
        "top_p": 1,
        "frequency_penalty": 0,
        "presence_penalty": 0,
    }

    cl.user_session.set("settings", settings)

@cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user
async def main(message: cl.Message):
    settings = cl.user_session.get("settings")

    client = AsyncOpenAI()

    #print(message.content)

    prompt = Prompt(
        provider=ChatOpenAI.id,
        messages=[
            PromptMessage(
                role="system",
                template=system_template,
                formatted=system_template,
            ),
            PromptMessage(
                role="user",
                template=user_template,
                formatted=user_template.format(input=message.content),
            ),
        ],
        inputs={"input": message.content},
        settings=settings,
    )

    #print([m.to_openai() for m in prompt.messages])

    msg = cl.Message(content="")
    await msg.send()
    

    # Call OpenAI
    async for stream_resp in await client.chat.completions.create(
        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
    ):
        token = stream_resp.choices[0].delta.content
        if not token:
            token = ""
        await msg.stream_token(token)

    # Update the prompt object with the completion
    prompt.completion = msg.content
    msg.prompt = prompt

    # Send and close the message stream
    await msg.send()