from langchain_core.callbacks import BaseCallbackHandler

class CustomHandler(BaseCallbackHandler):
    """A custom handler for logging interactions within the process chain."""

    def __init__(self, agent_name: str) -> None:
        super().__init__()
        self.agent_name = agent_name

    def on_chain_start(self, serialized, outputs, **kwargs) -> None:
        """Log the start of a chain with user input."""
        from streamlit import session_state, chat_message
        session_state.messages.append({"role": "assistant", "content": outputs['input']})
        chat_message("assistant").write(outputs['input'])

    def on_agent_action(self, serialized, inputs, **kwargs) -> None:
        """Log the action taken by an agent during a chain run."""
        from streamlit import session_state, chat_message
        session_state.messages.append({"role": "assistant", "content": inputs['input']})
        chat_message("assistant").write(inputs['input'])

    def on_chain_end(self, outputs, **kwargs) -> None:
        """Log the end of a chain with the output generated by an agent."""
        import streamlit as st
        from streamlit import session_state, chat_message
        session_state.messages.append({"role": self.agent_name, "content": outputs['output']})
        output = outputs['output']
        st.write("**********")
        st.write(self.agent_name)
        st.write("**********")
        if self.agent_name == 'Visual Content Creator':
            chat_message(self.agent_name).image(f"{self.agent_name } : {output}")
        else : 
            chat_message(self.agent_name).markdown(f"{self.agent_name } : {output}")