Implement retry

src/app.py

@@ -97,26 +97,35 @@ 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",
+    "question_rephrase": ":mag: Collecting data",
+    "function_agent": ":bulb: Preparing response",
+    "output_node": ":bulb: Done!",
 }
 
 def main():
     st.title("LLM-ADE 9B Demo")
 
     input_text = st.text_area("Enter your text here:", value="", height=200)
+    
+    def get_response(input_text: str) -> str:
+        try:
+            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}")
+            return value["generation"]
+        except Exception as e:
+            logger.error(e)
+            logger.info("Retrying..")
+            get_response(input_text)
 
     if st.button("Generate"):
         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=":question: Breaking down question")
+                response = get_response(input_text)
+                st.write(response)
                 config.status.update(label="Finished!", state="complete", expanded=True)
         else:
             st.warning("Please enter some text to generate a response.")

tests/qa_questions.json

@@ -1,6 +1,6 @@
 [
     {
-        "question": "What do you think about MSTR?",
+        "question": "Should I buy bitcoin or MSTR?",
         "expecteds": 
             {
                 "includes": ["microstrategy", "bitcoin"],
@@ -39,52 +39,51 @@
         "question": "Should I buy BTC or ETH?",
         "expecteds": 
             {
-                "includes": [],
+                "includes": ["BTC", "ETH"],
                 "excludes": [],
                 "functions_used": ["get_analysis"]
             }
     },
     {
-        "question": "Is Bitcoin cheap?",
+        "question": "Why is Bitcoin down today?",
         "expecteds": 
             {
                 "includes": ["BTC"],
                 "excludes": [],
-                "functions_used": ["get_current_stock_price"]
+                "functions_used": ["get_analysis", "get_current_stock_price"]
             }
-    },
+    }
     {
-        "question": "Why is Bitcoin down today?",
+        "question": "Why is Bitcoin up today?",
         "expecteds": 
             {
                 "includes": ["BTC"],
                 "excludes": [],
-                "functions_used": ["get_current_stock_price", "get_analysis"]
+                "functions_used": ["get_analysis", "get_current_stock_price"]
             }
-    },
+    }
     {
-        "question": "Why is Bitcoin up today??",
+        "question": "What is the latest news on Bitcoin?",
         "expecteds": 
             {
                 "includes": ["BTC"],
                 "excludes": [],
-                "functions_used": ["get_current_stock_price", "get_analysis"]
+                "functions_used": ["get_analysis"]
             }
-    },
     {
-        "question": "What is the latest news on Bitcoin?",
+        "question": "What does bitcoin halving mean?",
         "expecteds": 
             {
-                "includes": ["BTC"],
+                "includes": ["BTC", "mining"],
                 "excludes": [],
                 "functions_used": ["get_analysis"]
             }
-    },
+    }
     {
-        "question": "What happened during the halvening?",
+        "question": "What bitcoin ETF should I buy?",
         "expecteds": 
             {
-                "includes": ["mining", "BTC"],
+                "includes": ["BTC", "ETF"],
                 "excludes": [],
                 "functions_used": ["get_analysis"]
             }