from autogen import UserProxyAgent
import streamlit as st
from autogen.agentchat.contrib.qdrant_retrieve_user_proxy_agent import QdrantRetrieveUserProxyAgent


class TrackableQdrantRetrieveUserProxyAgent(QdrantRetrieveUserProxyAgent):
    def _process_received_message(self, message, sender, silent):
        with st.chat_message("ai"):
            st.write(message["content"])
        return super()._process_received_message(message, sender, silent)