feat: backup, save image in state

src/FinancialAgentApp.py

@@ -26,7 +26,9 @@ class FinancialAgentApp (ABC):
     def render_messages(self):
         """Render previous chat messages."""
         for message in self.st.session_state.messages:
-            with self.st.chat_message(message["role"]):
+            if message.get("type") == "plot":
+                self.st.image(message["content"])
+            else:
                 self.st.markdown(message["content"])
 
     @abstractmethod
@@ -71,7 +73,26 @@ class FinancialAgentApp (ABC):
     @abstractmethod
     def __handle_context__(self, response_state: ResponseState) -> str:
         """Handle context if need to add context from data/pdf"""
-        pass
+        context_prompt = ""
+        if response_state.contextType in ("data", "both"):
+            local_scope = {"df": self.df, "np": np, "pd": pd, "plt": plt, "savefig": self.__safe_savefig__}
+            exec(response_state.code, {}, local_scope)
+
+            fig = plt.gcf()
+            if fig.get_axes():  # if a chart was generated
+                with self.st.chat_message("assistant"):
+                    self.st.pyplot(fig)
+                plt.close(fig)
+                self.st.session_state.messages.append({
+                    "role": "assistant",
+                    "type": "plot",
+                    "content": self.__safe_savefig__()
+                })
+
+            context_prompt = "## CONTEXT DATAFRAME.\n"
+            context_prompt += str(local_scope.get("result", ""))
+
+        return context_prompt
 
     def generate_final_answer(self, context_prompt: str):
         """Generate and stream the final answer with context."""
@@ -116,20 +137,7 @@ class HFFinancialRAG(FinancialAgentApp):
 
     def __handle_context__(self, response_state: ResponseState) -> str:
         """Handle additional context (data, PDF, etc.)."""
-        context_prompt = ""
-        if response_state.contextType in ("data", "both"):
-            local_scope = {"df": self.df, "np": np, "pd": pd, "plt": plt, "savefig": self.__safe_savefig__}
-            exec(response_state.code, {}, local_scope)
-
-            fig = plt.gcf()
-            if fig.get_axes():  # if a chart was generated
-                with self.st.chat_message("assistant"):
-                    self.st.pyplot(fig)
-                plt.close(fig)
-
-            context_prompt = "## CONTEXT DATAFRAME.\n"
-            context_prompt += str(local_scope.get("result", ""))
-
+        context_prompt = super().__handle_context(response_state)
         if response_state.contextType in ("pdf", "both"):
             context_prompt += "## CONTEXT PDF.\n"
             results = self.vector_db.similarity_search(response_state.retriverKey, k=5)
@@ -171,22 +179,5 @@ class OpenAIFinancialRAG(FinancialAgentApp):
 
     def __handle_context__(self, response_state: ResponseState):
         """Handle additional context (data, PDF, etc.)."""
-        context_prompt = ""
-        if response_state.contextType in ("data", "both"):
-            local_scope = {"df": self.df, "np": np, "pd": pd, "plt": plt, "savefig": self.__safe_savefig__}
-            exec(response_state.code, {}, local_scope)
-
-            fig = plt.gcf()
-            if fig.get_axes():  # if a chart was generated
-                with self.st.chat_message("assistant"):
-                    self.st.pyplot(fig)
-                plt.close(fig)
-
-            context_prompt = "## CONTEXT DATAFRAME.\n"
-            context_prompt += str(local_scope.get("result", ""))
-
-        # Placeholder for PDF or other context handling
-        # elif response_state.contextType in ("pdf", "both"):
-        #     context_prompt = "Provide the relevant information from the PDF documents."
-
+        context_prompt = super().__handle_context(response_state)
         return context_prompt