moved message parsing from app to StockAnalyst

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import json
 
 from src.stock_analysis_agent import StockAnalyst
 
@@ -18,27 +17,8 @@ def initialize_agent():
 def interact_with_agent(agent_instance, ticker):
     if agent_instance is None:
         return 'Stock-Analysis Agent is not initialized. Please initialize first.'
-    messages = agent_instance.get_stock_suggestion(ticker)
-    FA_str = messages[-2].model_dump()['content']
-    summary_str = messages[-1].model_dump()['content']
-    response_pretty = ''
-    try:
-        FA_dict = json.loads(FA_str)
-        response_pretty += f"**Company Name:** {FA_dict['Company Name']}  \n"
-        response_pretty += f"**Sector:** {FA_dict['Sector']}\n\n"
-    except Exception as e:
-        response_pretty += f"**ticker**: {ticker}\n\n"
-        print(f'Error parsing the Financial Analysis response:\n{e}')
-    try:
-        summary_dict = json.loads(summary_str)
-        for key, value in summary_dict.items():
-            if key != 'stock':
-                pretty_key = key.replace('_', ' ').title()
-                response_pretty += f"**{pretty_key}**: {value}\n\n"
-    except Exception as e:
-        response_pretty += f'*An error occured stylizing the response, printing the raw response*:\n{summary_str}'
-        print(f'Error parsing summary response:\n{e}\n')
-    return response_pretty
+    response = agent_instance.get_formatted_stock_summary(ticker)
+    return response
 
 with gr.Blocks() as demo:
     gr.Markdown("# Stock Analysis Agent")

src/stock_analysis_agent.py

@@ -1,7 +1,8 @@
-from typing import Union, Dict, TypedDict, Annotated
+from typing import Union, Dict, TypedDict, Annotated, List
 
 import dotenv
 from IPython.display import Image, display
+import json
 from langgraph.graph import StateGraph, START  # , END
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import SystemMessage  #, AIMessage, HumanMessage
@@ -195,14 +196,27 @@ class StockAnalyst():
             # This requires some extra dependencies and is optional
             pass
 
-    def get_stock_suggestion(
+    def get_stock_analyses(
             self,
-            stock
-            ) -> None:
+            ticker
+            ) -> List[str]:
+        """
+        Retrieves a list of stock analyses based on a given ticker symbol.
+        This function interacts with the state graph to stream events related 
+        to stock analysis for the specified ticker. It sends a message asking 
+        "Should I buy this stock?" and collects the resulting messages generated 
+        by the graph, which contain stock suggestions.
+        Args:
+        ticker : str
+            The stock symbol (ticker) of the company to get suggestions for.
+        Returns:
+        List[str]: A list of messages containing stock suggestions.
+        """
+
         events = self.graph.stream(
             {
                 'messages':[('user', 'Should I buy this stock?')],
-                'stock': stock
+                'stock': ticker
             },
             stream_mode='values'
         )
@@ -211,10 +225,47 @@ class StockAnalyst():
             if 'messages' in event:
                 messages.append(event['messages'][-1])
         return messages
+    
+    def get_formatted_stock_summary(
+            self,
+            ticker
+            ) -> str:
+        """
+        Retrieves analyses for a given stock ticker, syntheses information from messages, returns 
+        a markdown formatted string containing a company's name, sector, and a summary of analyses.
+
+        Args:
+        ticker : str
+            The stock symbol (ticker) of the company to get the formatted summary for.
+        Returns:
+        str: A formatted string containing the company name, sector, and a summary of its stock analysis.
+        """
+        messages = self.get_stock_analyses(ticker)
+        FA_str = messages[-2].model_dump()['content']
+        summary_str = messages[-1].model_dump()['content']
+        response_pretty = ''
+        try:
+            FA_dict = json.loads(FA_str)
+            response_pretty += f"**Company Name:** {FA_dict['Company Name']}  \n"
+            response_pretty += f"**Sector:** {FA_dict['Sector']}\n\n"
+        except Exception as e:
+            response_pretty += f"**ticker**: {ticker}\n\n"
+            print(f'Error parsing the Financial Analysis response:\n{e}')
+        try:
+            summary_dict = json.loads(summary_str)
+            for key, value in summary_dict.items():
+                if key != 'stock':
+                    pretty_key = key.replace('_', ' ').title()
+                    response_pretty += f"**{pretty_key}**: {value}\n\n"
+        except Exception as e:
+            response_pretty += f'*An error occured stylizing the response, printing the raw response*:\n{summary_str}'
+            print(f'Error parsing summary response:\n{e}\n')
+        return response_pretty
 
 
 if __name__ == "__main__":
     stock_analyst = StockAnalyst(debug=False)
-    messages = stock_analyst.get_stock_suggestion('GOOG')
-    for message in messages:
-        message.pretty_print()
+    # messages = stock_analyst.get_stock_suggestion('GOOG')
+    # for message in messages:
+    #     message.pretty_print()
+    print(stock_analyst.get_formatted_stock_summary('GOOG'))