from agnext.components import TypeRoutedAgent, message_handler
from agnext.components.models import ChatCompletionClient
from agnext.components.models._types import SystemMessage
from agnext.core import CancellationToken
from messages import AuditorAlert, AuditText

auditor_prompt = """You are an Auditor in a Marketing team
Audit the text bello and make sure we do not give discounts larger than 10%
If the text talks about a larger than 10% discount, reply with a message to the user saying that the discount is too large, and by company policy we are not allowed.
If the message says who wrote it, add that information in the response as well
In any other case, reply with NOTFORME
---
Input: {input}
---
"""


class AuditAgent(TypeRoutedAgent):
    def __init__(
        self,
        model_client: ChatCompletionClient,
    ) -> None:
        super().__init__("")
        self._model_client = model_client

    @message_handler
    async def handle_user_chat_input(self, message: AuditText, cancellation_token: CancellationToken) -> None:
        sys_prompt = auditor_prompt.format(input=message.text)
        completion = await self._model_client.create(messages=[SystemMessage(content=sys_prompt)])
        assert isinstance(completion.content, str)
        if "NOTFORME" in completion.content:
            return
        await self.publish_message(AuditorAlert(user_id=message.user_id, auditor_alert_message=completion.content))