from langchain_core.callbacks import BaseCallbackHandler
import streamlit as st

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."""
        st.session_state.messages.append({"role": "assistant", "content": outputs['input']})
        st.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."""
        st.session_state.messages.append({"role": "assistant", "content": inputs['input']})
        st.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."""
        st.session_state.messages.append({"role": self.agent_name, "content": outputs['output']})
        st.chat_message(self.agent_name).write(outputs['output'])