from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import Runnable
from langchain.schema.runnable.config import RunnableConfig

import chainlit as cl


@cl.on_chat_start
async def on_chat_start():
    model = ChatOpenAI(openai_api_base="http://localhost:8888/v1",streaming=True)
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You're a very knowledgeable historian who provides accurate and eloquent answers to historical questions.",
            ),
            ("human", "{question}"),
        ]
    )
    runnable = prompt | model | StrOutputParser()
    cl.user_session.set("runnable", runnable)


@cl.on_message
async def on_message(message: cl.Message):
    runnable = cl.user_session.get("runnable")  # type: Runnable

    msg = cl.Message(content="")

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

    await msg.send()

langchain.prompts.ChatPromptTemplate.from_messages langchain.schema.StrOutputParser langchain.chat_models.ChatOpenAI
[((328, 398), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_base': '"""http://localhost:8888/v1"""', 'streaming': '(True)'}), "(openai_api_base='http://localhost:8888/v1', streaming=True)\n", (338, 398), False, 'from langchain.chat_models import ChatOpenAI\n'), ((411, 600), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (['[(\'system\',\n    "You\'re a very knowledgeable historian who provides accurate and eloquent answers to historical questions."\n    ), (\'human\', \'{question}\')]'], {}), '([(\'system\',\n    "You\'re a very knowledgeable historian who provides accurate and eloquent answers to historical questions."\n    ), (\'human\', \'{question}\')])\n', (443, 600), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((742, 783), 'chainlit.user_session.set', 'cl.user_session.set', (['"""runnable"""', 'runnable'], {}), "('runnable', runnable)\n", (761, 783), True, 'import chainlit as cl\n'), ((859, 890), 'chainlit.user_session.get', 'cl.user_session.get', (['"""runnable"""'], {}), "('runnable')\n", (878, 890), True, 'import chainlit as cl\n'), ((920, 942), 'chainlit.Message', 'cl.Message', ([], {'content': '""""""'}), "(content='')\n", (930, 942), True, 'import chainlit as cl\n'), ((720, 737), 'langchain.schema.StrOutputParser', 'StrOutputParser', ([], {}), '()\n', (735, 737), False, 'from langchain.schema import StrOutputParser\n'), ((967, 997), 'chainlit.make_async', 'cl.make_async', (['runnable.stream'], {}), '(runnable.stream)\n', (980, 997), True, 'import chainlit as cl\n'), ((1079, 1108), 'chainlit.LangchainCallbackHandler', 'cl.LangchainCallbackHandler', ([], {}), '()\n', (1106, 1108), True, 'import chainlit as cl\n')]