streaming outputs

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import time 
 import json
 from cerebras.cloud.sdk import Cerebras
-from typing import List, Dict, Tuple, Any
+from typing import List, Dict, Tuple, Any, Generator
 from tenacity import retry, stop_after_attempt, wait_fixed
 
 def make_api_call(api_key: str, messages: List[Dict[str, str]], max_tokens: int, is_final_answer: bool = False) -> Dict[str, Any]:
@@ -26,9 +26,10 @@ def make_api_call(api_key: str, messages: List[Dict[str, str]], max_tokens: int,
         else:
             return {"title": "Error", "content": f"Failed to generate step. Error: {str(e)}", "next_action": "final_answer"}
 
-def generate_response(api_key: str, prompt: str) -> Tuple[List[Tuple[str, str, float]], float]:
+def generate_response(api_key: str, prompt: str) -> Generator[Tuple[List[Tuple[str, str]], float], None, None]:
     """
     Generate a response to the given prompt using a step-by-step reasoning approach.
+    This function is now a generator that yields each step as it's generated.
     """
     system_message = """You are an expert AI assistant that explains your reasoning step by step. For each step, provide a title that describes what you're doing in that step, along with the content. Decide if you need another step or if you're ready to give the final answer. Respond in JSON format with 'title', 'content', and 'next_action' (either 'continue' or 'final_answer') keys. USE AS MANY REASONING STEPS AS POSSIBLE. AT LEAST 3. BE AWARE OF YOUR LIMITATIONS AS AN LLM AND WHAT YOU CAN AND CANNOT DO. IN YOUR REASONING, INCLUDE EXPLORATION OF ALTERNATIVE ANSWERS. CONSIDER YOU MAY BE WRONG, AND IF YOU ARE WRONG IN YOUR REASONING, WHERE IT WOULD BE. FULLY TEST ALL OTHER POSSIBILITIES. YOU CAN BE WRONG. WHEN YOU SAY YOU ARE RE-EXAMINING, ACTUALLY RE-EXAMINE, AND USE ANOTHER APPROACH TO DO SO. DO NOT JUST SAY YOU ARE RE-EXAMINING. USE AT LEAST 3 METHODS TO DERIVE THE ANSWER. USE BEST PRACTICES."""
     
@@ -48,14 +49,20 @@ def generate_response(api_key: str, prompt: str) -> Tuple[List[Tuple[str, str, f
         thinking_time = time.time() - start_time
         total_thinking_time += thinking_time
         
-        steps.append((f"Step {step_count}: {step_data['title']}", step_data['content'], thinking_time))
+        step_title = f"Step {step_count}: {step_data['title']}"
+        step_content = step_data['content']
+        steps.append((step_title, step_content))
         messages.append({"role": "assistant", "content": json.dumps(step_data)})
         
+        # Yield the current conversation and total thinking time
+        yield steps, total_thinking_time
+        
         if step_data.get('next_action') == 'final_answer':
             break
         
         step_count += 1
 
+    # Request the final answer
     messages.append({"role": "user", "content": "Please provide the final answer based on your reasoning above."})
     
     start_time = time.time()
@@ -63,26 +70,33 @@ def generate_response(api_key: str, prompt: str) -> Tuple[List[Tuple[str, str, f
     thinking_time = time.time() - start_time
     total_thinking_time += thinking_time
     
-    steps.append(("Final Answer", final_data.get('content', 'No final answer provided.'), thinking_time))
-
-    return steps, total_thinking_time
+    steps.append(("Final Answer", final_data.get('content', 'No final answer provided.')))
+    
+    # Yield the final conversation and total thinking time
+    yield steps, total_thinking_time
 
-def generate_ui(api_key: str, prompt: str) -> Tuple[List[Tuple[str, str]], float]:
+def respond(api_key: str, message: str, history: List[Tuple[str, str]]) -> Generator[Tuple[List[Tuple[str, str]], str], None, None]:
     """
-    Generate the UI output based on the response to the given prompt.
+    Generator function to handle responses and yield updates for streaming.
     """
-    steps, total_time = generate_response(api_key, prompt)
-    conversation = []
-    for title, content, _ in steps:
-        if title.startswith("Step"):
-            conversation.append(("Assistant", f"**{title}**\n\n{content}"))
-        elif title == "Final Answer":
-            conversation.append(("Assistant", f"**{title}**\n\n{content}"))
-        else:
-            conversation.append(("Assistant", content))
-    return conversation, total_time
+    if not api_key:
+        yield history, "Please provide a valid Cerebras API key."
+        return
+
+    # Initialize the generator
+    response_generator = generate_response(api_key, message)
+    
+    for steps, total_time in response_generator:
+        conversation = history.copy()
+        for title, content in steps[len(conversation)//2:]:
+            if title.startswith("Step"):
+                conversation.append(("Assistant", f"**{title}**\n\n{content}"))
+            elif title == "Final Answer":
+                conversation.append(("Assistant", f"**{title}**\n\n{content}"))
+            else:
+                conversation.append(("Assistant", content))
+        yield conversation, f"**Total thinking time:** {total_time:.2f} seconds"
 
-# Gradio Blocks Interface with a Chatbot component and API key input
 def main():
     with gr.Blocks() as demo:
         gr.Markdown("# o1-like Chain of Thought - LLaMA-3.1 70B on Cerebras")
@@ -114,20 +128,6 @@ def main():
             interactive=False
         )
         
-        def respond(api_key, message, history):
-            if not api_key:
-                return history, "Please provide a valid Cerebras API key."
-
-            steps, total_time = generate_response(api_key, message)
-            for title, content, _ in steps:
-                if title.startswith("Step"):
-                    history.append(("Assistant", f"**{title}**\n\n{content}"))
-                elif title == "Final Answer":
-                    history.append(("Assistant", f"**{title}**\n\n{content}"))
-                else:
-                    history.append(("Assistant", content))
-            return history, f"**Total thinking time:** {total_time:.2f} seconds"
-
         submit_btn.click(
             fn=respond, 
             inputs=[api_key_input, user_input, chatbot], 
@@ -135,7 +135,7 @@ def main():
             queue=True
         )
         
-        # Optional: Allow pressing Enter to submit
+        # Allow pressing Enter to submit
         user_input.submit(
             fn=respond, 
             inputs=[api_key_input, user_input, chatbot], 
@@ -146,4 +146,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()