import os
from fastapi.responses import JSONResponse

from chainlit.auth import create_jwt
from chainlit.server import app
import chainlit as cl


import chainlit as cl
import langroid as lr
from langroid.agent.callbacks.chainlit import add_instructions


@cl.on_chat_start
async def on_chat_start():
    config = lr.ChatAgentConfig(
        name="Demo",
        system_message="You are a helpful assistant. Be concise in your answers.",
    )
    agent = lr.ChatAgent(config)

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

    await add_instructions(
        title="Instructions",
        content="Interact with a **Langroid ChatAgent**",
    )


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    # important: only apply callbacks after getting first msg.
    lr.ChainlitAgentCallbacks(agent, message)
    await agent.llm_response_async(message.content)