Spaces:
Sleeping
Sleeping
File size: 1,657 Bytes
3524557 |
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 26 27 28 29 30 31 32 33 34 35 |
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}")
|