Fix threading issue

src/app.py

@@ -35,7 +35,7 @@ def question_node(state: GraphState) -> Dict[str, str]:
     Generate a question for the function agent.
     """
     logger.info("Generating question for function agent")
-    config.status.update(label=":question: Breaking down question")
+    # config.status.update(label=":question: Breaking down question")
     question = state["question"]
     logger.info(f"Original question: {question}")
     rephrased_question = agents["rephrase_agent"].invoke({"question": question})
@@ -49,7 +49,7 @@ def function_agent_node(state: GraphState) -> Literal["finished"]:
     logger.info("Calling function agent")
     question = state["rephrased_question"]
     response = agents["function_agent"].invoke({"input": question, "tools": tools_dict}).get("output")
-    config.status.update(label=":brain: Analysing data..")
+    # config.status.update(label=":brain: Analysing data..")
     logger.info(f"Function agent output: {response}")
     return {"function_agent_output": response}
 
@@ -58,7 +58,7 @@ def output_node(state: GraphState) -> Dict[str, str]:
     Generate the final output
     """
     logger.info("Generating output")
-    config.status.update(label=":bulb: Preparing response..")
+    # config.status.update(label=":bulb: Preparing response..")
     generation = agents["output_agent"].invoke({"context": state["function_agent_output"],
                                                 "question": state["rephrased_question"]})
     return {"generation": generation}
@@ -70,7 +70,7 @@ def route_question(state: GraphState) -> Literal["vectorstore", "websearch"]:
     Route quesition to web search or RAG
     """
     logger.info("Routing question")
-    config.state.update(label=":chart_with_upwards_trend: Routing question")
+    # config.state.update(label=":chart_with_upwards_trend: Routing question")
     question = state["question"]
     logger.info(f"Question: {question}")
     source = agents["router_agent"].invoke({"question": question})
@@ -96,6 +96,12 @@ workflow.set_finish_point("output_node")
 
 flow = workflow.compile()
 
+progress_map = {
+    "question_rephrase": ":question: Breaking down question",
+    "function_agent": ":mag: Collecting data",
+    "output_node": ":bulb: Preparing response",
+}
+
 def main():
     st.title("LLM-ADE 9B Demo")
 
@@ -105,8 +111,10 @@ def main():
         if input_text:
             with st.status("Generating response...") as status:
                 config.status = status
+                config.status.update(label=":brain: Thinking")
                 for output in flow.stream({"question": input_text}):
                     for key, value in output.items():
+                        config.status.update(label=progress_map[key])
                         pprint(f"Finished running: {key}")
                 st.write(value["generation"])
                 config.status.update(label="Finished!", state="complete", expanded=True)

src/functioncall.py

@@ -54,7 +54,7 @@ class ModelInference:
             raise ValueError("Assistant message is None")
         
     def execute_function_call(self, tool_call):
-        config.status.update(label=":mag: Gathering information..")
+        # config.status.update(label=":mag: Gathering information..")
         function_name = tool_call.get("name")
         function_to_call = getattr(functions, function_name, None)
         function_args = tool_call.get("arguments", {})
@@ -81,7 +81,7 @@ class ModelInference:
             chat = [{"role": "user", "content": user_message}]
             tools = functions.get_openai_tools()
             prompt = self.prompter.generate_prompt(chat, tools, num_fewshot)
-            config.status.update(label=":brain: Thinking..")
+            # config.status.update(label=":brain: Thinking..")
             completion = self.run_inference(prompt)
 
             def recursive_loop(prompt, completion, depth):
@@ -115,7 +115,7 @@ class ModelInference:
                         completion = self.run_inference(prompt)
                         return completion
 
-                    config.status.update(label=":brain: Analysing information..")
+                    # config.status.update(label=":brain: Analysing information..")
                     completion = self.run_inference(prompt)
                     return recursive_loop(prompt, completion, depth)
                 elif error_message:

src/functions.py

@@ -113,7 +113,7 @@ def get_current_stock_price(symbol: str) -> float:
     float: The current stock price, or None if an error occurs.
   """
   try:
-    config.status.update(label=":chart_with_upwards_trend: Getting price")
+    # config.status.update(label=":chart_with_upwards_trend: Getting price")
     stock = yf.Ticker(symbol)
     # Use "regularMarketPrice" for regular market hours, or "currentPrice" for pre/post market
     current_price = stock.info.get("regularMarketPrice", stock.info.get("currentPrice"))
@@ -199,7 +199,7 @@ def get_key_financial_ratios(symbol: str) -> dict:
     dict: Dictionary containing key financial ratios.
     """
     try:
-        config.status.update(label=":chart_with_upwards_trend: Gathering financial data")
+        # config.status.update(label=":chart_with_upwards_trend: Gathering financial data")
         stock = yf.Ticker(symbol)
         key_ratios = stock.info
         return key_ratios
@@ -238,7 +238,7 @@ def get_company_news(symbol: str) -> pd.DataFrame:
     Returns:
     pd.DataFrame: DataFrame containing company news and press releases.
     """
-    config.status.update(label=":newspaper: Getting news")
+    # config.status.update(label=":newspaper: Getting news")
     try:
         news = yf.Ticker(symbol).news
         return news

src/prompts/output_agent_template.yaml

@@ -22,6 +22,7 @@ Do not be formal, answer as you would in a conversation with a friend.
 Keep your response short unless the question requires a longer response.
 #Audience:
 You are talking to a high net worth individual with high risk tolerance and lots of investing experience - they do not need traditional, generic financial advice.
+Address the user directly.
 \n ------- \n
 Your thoughts are as follows:
 {context}