import streamlit as st
from langchain.callbacks.tracers.stdout import FunctionCallbackHandler


class ModelTraceCallback(FunctionCallbackHandler):
    """
    This class implements a callback to write model traces as they are produced. A model trace
    contains logs produced by an LLM chain execution. This callback is mostly intended for
    testing purposes.
    """

    name: str = "model_trace_callback_handler"

    def __init__(self):
        """
        Creates the callback handler.
        """
        super().__init__(self._callback)

        # The message placeholder must be overriden by the calling app at each turn. It represents
        # an empty slot that will be filled with logs produced by the an LLM on-the-fly.
        self._message_placeholder: st.container = None
        self._full_text = ""

    def reset(self, message_placeholder: st.container):
        """
        Clears previously saved tokens.

        :param message_placeholder: clean, empty container that will receive the new logs.
        """

        self._message_placeholder = message_placeholder
        self._full_text = ""

    def _callback(self, trace: str) -> None:
        """
        Appends a new trace to the current partially produced log message.

        :param trace: new trace.
        """

        self._full_text += trace
        if self._message_placeholder:
            self._message_placeholder.write(self._full_text)

    def on_chat_model_start(*args, **kwargs):
        """
        Empty method. This method is inherited by the super class but never implemented. No
        implementation is needed for our use cases. Copied here just to avoid warnings.
        """
        return super().on_chat_model_start(*args, **kwargs)
