from dotenv import load_dotenv

load_dotenv()

import chainlit as cl
from openai import OpenAI

client = OpenAI()

import langwatch.evaluation


@cl.on_message
async def main(message: cl.Message):
    msg = cl.Message(
        content="",
    )

    # New API: Use langwatch.evaluation.evaluate() with as_guardrail=True
    jailbreak_guardrail = langwatch.evaluation.evaluate(
        "jailbreak-detection",
        data={"input": message.content},
        as_guardrail=True,
    )
    if not jailbreak_guardrail.passed:
        await msg.stream_token("I'm sorry, I can't help you with that.")
        await msg.update()
        return

    completion = client.chat.completions.create(
        model="gpt-5",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant that only reply in short tweet-like responses, using lots of emojis.",
            },
            {"role": "user", "content": message.content},
        ],
        stream=True,
    )

    for part in completion:
        if token := part.choices[0].delta.content or "":
            await msg.stream_token(token)

    await msg.update()
