### Import Section ###
"""
IMPORTS HERE
"""
# Example Imports (adjust based on actual needs)
import chainlit as cl
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import Runnable
from langchain.schema.runnable.config import RunnableConfig
from typing import cast

### Global Section ###
"""
GLOBAL CODE HERE
"""
# Initialize a language model or chain globally
llm = ChatOpenAI(temperature=0.9)
conversation_chain = ConversationChain(llm=llm)

# Any global variables like API keys, configurations, etc.
# API_KEY = "your_api_key_here"


### On Chat Start (Session Start) Section ###
@cl.on_chat_start
async def on_chat_start():
    """ SESSION SPECIFIC CODE HERE """
    await cl.Message(content="Welcome! How can I assist you today?").send()

### Rename Chains ###
@cl.author_rename
def rename(orig_author: str):
    if orig_author == "user":
        return "You"
    elif orig_author == "system":
        return "Assistant"
    return orig_author

### On Message Section ###
@cl.on_message
async def on_message(message: cl.Message):
    runnable = cast(Runnable, cl.user_session.get("runnable"))  

    msg = cl.Message(content="")

    async for chunk in runnable.astream(
        {"question": message.content},
        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
    ):
        await msg.stream_token(chunk)

    await msg.send()