enabled showing TA plots

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 
 from src.stock_analysis_agent import StockAnalyst
+from src.technical_analysis import TechnicalAnalysis
 
 
 def initialize_agent():
@@ -14,12 +15,22 @@ def initialize_agent():
         message = f"Error initializing Stock-Analysis Agent:\n{e}"
     return message, stock_analyst
 
-def interact_with_agent(agent_instance, ticker):
+def ask_stock_agent(agent_instance, ticker):
     if agent_instance is None:
         return 'Stock-Analysis Agent is not initialized. Please initialize first.'
     response = agent_instance.get_formatted_stock_summary(ticker)
     return response
 
+def plot_stock(ticker):
+    # run the technical analysis
+    _, fig = TechnicalAnalysis(
+        ticker=ticker,
+        fetchperiodinweeks=12,
+        plot_ta=True,
+        savefig=False,
+        debug=False).run()
+    return fig
+
 with gr.Blocks() as demo:
     gr.Markdown("# Stock Analysis Agent")
     # State to store the RAG instance
@@ -35,12 +46,19 @@ with gr.Blocks() as demo:
     # gr.Text(label='Status', value=message)
     gr.Markdown("Enter a stock symbol ('ticker') to be analyzed (GOOG, MSFT, etc.)")
     ticker = gr.Textbox(label="ticker")
+    # plot_button = gr.Button('Plot')
     analyze_button = gr.Button('Analyze Stock')
-    output = gr.Markdown(label="Output Box")
+    plot_output = gr.Plot(label=ticker, format="png")
+    analyze_button.click(
+            fn=plot_stock,
+            inputs=[ticker],
+            outputs=plot_output, 
+            api_name="analyze_stock")
+    md_output = gr.Markdown(label="Output Box")
     analyze_button.click(
-            fn=interact_with_agent,
+            fn=ask_stock_agent,
             inputs=[agent_instance, ticker],
-            outputs=output, 
+            outputs=md_output, 
             api_name="analyze_stock")
 
 

src/stock_analysis_agent.py

@@ -25,10 +25,11 @@ def get_stock_prices(
     ticker: str
         The stock ticker symbol to fetch data for.
     """
-    df = TechnicalAnalysis(
+    df, _ = TechnicalAnalysis(
         ticker=ticker,
         fetchperiodinweeks=12,
         plot_ta=False,
+        savefig=False,
         debug=False).run()
     if df.shape[0] > 0:
         df['Date'] = df.index.astype(str)

src/technical_analysis.py

@@ -18,6 +18,7 @@ class TechnicalAnalysis():
             ticker:str,
             fetchperiodinweeks:int=12,
             plot_ta:bool=True,
+            savefig:bool=False,
             debug=False):
         # input arguments
         """
@@ -44,6 +45,7 @@ class TechnicalAnalysis():
         self.ticker = ticker
         self.fetchperiodinweeks = fetchperiodinweeks
         self.plot_ta = plot_ta
+        self.savefig = savefig
         # done initializing
         self.logger.info(f'Initialized TechnicalAnalysis object for ticker: {ticker}')
 
@@ -65,15 +67,21 @@ class TechnicalAnalysis():
             # plot the results
             if self.plot_ta:
                 os.makedirs('plots', exist_ok=True)
-                self.plot_stock_metrics(
+                fig = self.plot_stock_metrics(
                     self.df,
                     datasets={
                             'Volume': ['Volume'],
                             'Prices': ['Close', 'VWAP'],  # 'High','Low', 
                             'Indices': ['RSI', 'StochOsc'],
-                            'Trend': ['MACD', 'MACDsig', 'MACDdif']}
-                            )
-        return self.df
+                            'Trend': ['MACD', 'MACDsig', 'MACDdif']},
+                    savefig=self.savefig
+                    )
+            else:
+                fig = None
+        else:
+            fig = None
+
+        return self.df, fig
 
     def fetch_data(
             self
@@ -183,8 +191,9 @@ class TechnicalAnalysis():
             datasets={
                 'Volume': ['Volume'],
                 'Price': ['Close']  # 'High','Low'
-                }
-                ) -> None:
+                },
+            savefig=False
+            ) -> None:
         """
         Plots the given stock metrics datasets as subplots.
         This method takes in a DataFrame and a dictionary of datasets, where
@@ -199,6 +208,8 @@ class TechnicalAnalysis():
         datasets (dict)
             A dictionary of datasets, where each key is a dataset name and
             the value is a list of column names to be plotted
+        savefig (bool)
+            Whether to save the figure to a file
         """
         numax = len(datasets)
         fig, axes = plt.subplots(
@@ -214,8 +225,11 @@ class TechnicalAnalysis():
                 df,
                 colstoplot)
         plt.tight_layout()
-        plt.savefig(os.path.join('plots',  f'{self.ticker}.png'))
-        plt.close()
+        if savefig:
+            plt.savefig(os.path.join('plots',  f'{self.ticker}.png'))
+            plt.close()
+            fig = None
+        return fig
 
     def plot_stock_metrics_ax(
             self,
@@ -259,7 +273,7 @@ class TechnicalAnalysis():
         ax.xaxis.set_minor_locator(mdates.DayLocator())
         
         ax.set_title(dataset)
-        ax.set_xlabel('Date')
+        # ax.set_xlabel('Date')
         ax.set_ylabel(dataset)
         if len(colstoplot) > 1:
             ax.legend()