Spaces:
Sleeping
Sleeping
File size: 1,180 Bytes
06520ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
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'])
|