feat: refactor to oop

src/streamlit_app.py

@@ -6,101 +6,121 @@ import pandas as pd
 import streamlit as st
 from openai import OpenAI
 import pickle
-# import logging
-# logging.basicConfig(filename="/app/logs/app.log", level=logging.INFO)
-
-client = OpenAI()
-
-df = pickle.load(open("fraudTrainData.pkl", "rb"))
-
-st.title("Financial Agent")
-
-st.session_state["openai_model"] = "gpt-5-mini-2025-08-07"
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-def streamAnswer(instructions, input):
-    chunk = client.responses.create(
-        model=st.session_state["openai_model"],
-        instructions=instructions,
-        input=input,
-        stream=True,
-                # if contextType == "both" || contextType == "pdf"add vector store  tools file search
-        tools=[{
-            "type": "file_search",
-            "vector_store_ids": ['vs_68bf713eea2c81919ac08298a05d6704']
-        }]
-    )
-
-    for chunk in chunk:
-        if chunk.type == 'response.output_text.delta':
-            yield chunk.delta
-
-if prompt := st.chat_input("What is up?"):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    response = client.responses.parse(
-        model=st.session_state["openai_model"],
-        instructions=REFINERY_PROMPT.format(
-            df_head=df.head().to_markdown(),
-            df_columns=df.columns.tolist(),
-            df_sample=df.sample(5).to_markdown()
-        ),
-        input=[
-            {"role": m["role"], "content": m["content"]}
-            for m in st.session_state.messages
-        ],
-        stream=False,
-        text_format=ResponseState
-    )
-
-    context_prompt = ""
-    # logging.info(response)
-    responseState : ResponseState = response.output_parsed
-    if responseState.isNeedContext:
-        contextType = responseState.contextType
-        st.session_state.messages.append({"role": "assistant", "content": responseState.response})
-        # with st.chat_message("assistant"):
-        #     st.markdown(responseState.response)
-        if contextType == "data" or contextType == "both":
-            local_scope = {
-                "df": df,
-                "np": np,
-                "pd": pd,
-                "plt": plt
-            }
-            exec(responseState.code, {}, local_scope)
-            fig = plt.gcf()  # get current active figure
-            if fig.get_axes():  # check if plot was created
+
+
+class FinancialAgentApp:
+    def __init__(self):
+        self.client = OpenAI()
+        self.df = pickle.load(open("fraudTrainData.pkl", "rb"))
+        self.model_name = "gpt-5-mini-2025-08-07"
+
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+        st.session_state["openai_model"] = self.model_name
+
+    def render_header(self):
+        st.title("Financial Agent")
+
+    def render_messages(self):
+        """Render previous chat messages."""
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+
+    def stream_answer(self, instructions, input_messages):
+        """Stream OpenAI response as a generator."""
+        response_stream = self.client.responses.create(
+            model=self.model_name,
+            instructions=instructions,
+            input=input_messages,
+            stream=True,
+            tools=[{
+                "type": "file_search",
+                "vector_store_ids": ['vs_68bf713eea2c81919ac08298a05d6704']
+            }]
+        )
+        for chunk in response_stream:
+            if chunk.type == 'response.output_text.delta':
+                yield chunk.delta
+
+    def process_prompt(self, prompt):
+        """Main pipeline for processing a new user input."""
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        # Step 1: Run refinery prompt
+        response = self.client.responses.parse(
+            model=self.model_name,
+            instructions=REFINERY_PROMPT.format(
+                df_head=self.df.head().to_markdown(),
+                df_columns=self.df.columns.tolist(),
+                df_sample=self.df.sample(5).to_markdown()
+            ),
+            input=[{"role": m["role"], "content": m["content"]} for m in st.session_state.messages],
+            stream=False,
+            text_format=ResponseState
+        )
+
+        response_state: ResponseState = response.output_parsed
+
+        # Step 2: Check if context is needed
+        if response_state.isNeedContext:
+            context_prompt = self.handle_context(response_state)
+            self.generate_final_answer(context_prompt)
+        else:
+            self.display_final_answer(response_state.response)
+
+    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}
+            exec(response_state.code, {}, local_scope)
+
+            fig = plt.gcf()
+            if fig.get_axes():  # if a chart was generated
                 with st.chat_message("assistant"):
                     st.pyplot(fig)
-                plt.close(fig)  # prevent duplicate rendering later
+                plt.close(fig)
+
             context_prompt = "## CONTEXT DATAFRAME.\n"
-            context_prompt += local_scope.get("result", "")
-            # logging.info("context from data: " + context_prompt)
-        # elif format == "pdf" or format == "both":
-        #     context_prompt = "Provide the relevant information from the PDF documents to answer the user's question."
-        # st.session_state.messages.append({"role": "user", "content": context_prompt})
-        # with st.chat_message("user"):
-        #     st.markdown(context_prompt)
-        answer = ""
+            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."
+
+        return context_prompt
+
+    def generate_final_answer(self, context_prompt: str):
+        """Generate and stream the final answer with context."""
         with st.chat_message("assistant"):
-            answer = st.write_stream(streamAnswer(
-                instructions=FINAL_PROMPT,
-                input=[
-                    {"role": m["role"], "content": m["content"]}
-                    for m in st.session_state.messages
-                ] + [{"role": "user", "content": context_prompt}]
-            ))
+            answer = st.write_stream(
+                self.stream_answer(
+                    instructions=FINAL_PROMPT,
+                    input=[
+                        {"role": m["role"], "content": m["content"]}
+                        for m in st.session_state.messages
+                    ] + [{"role": "user", "content": context_prompt}]
+                )
+            )
+        st.session_state.messages.append({"role": "assistant", "content": answer})
+
+    def display_final_answer(self, answer: str):
+        """Display a non-streamed assistant answer."""
         st.session_state.messages.append({"role": "assistant", "content": answer})
-    else: #only write the response
-        st.session_state.messages.append({"role": "assistant", "content": responseState.response})
         with st.chat_message("assistant"):
-            st.markdown(responseState.response)
+            st.markdown(answer)
+
+    def run(self):
+        """Run the app."""
+        self.render_header()
+        self.render_messages()
+
+        if prompt := st.chat_input("What is up?"):
+            self.process_prompt(prompt)
+
+if __name__ == "__main__":
+    app = FinancialAgentApp()
+    app.run()