Update app.py

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]:
@@ -12,6 +12,7 @@ def make_api_call(api_key: str, messages: List[Dict[str, str]], max_tokens: int,
     client = Cerebras(api_key=api_key)
     
     try:
+        start_time = time.time()
         response = client.chat.completions.create(
             model="llama3.1-70b",
             messages=messages,
@@ -19,16 +20,31 @@ def make_api_call(api_key: str, messages: List[Dict[str, str]], max_tokens: int,
             temperature=0.2,
             response_format={"type": "json_object"}
         )
-        return json.loads(response.choices[0].message.content)
+        end_time = time.time()
+        
+        content = json.loads(response.choices[0].message.content)
+        
+        # Calculate tokens per second
+        total_tokens = response.usage.total_tokens
+        elapsed_time = end_time - start_time
+        tokens_per_second = total_tokens / elapsed_time if elapsed_time > 0 else 0
+        
+        content['token_info'] = {
+            'total_tokens': total_tokens,
+            'tokens_per_second': tokens_per_second
+        }
+        
+        return content
     except Exception as e:
         if is_final_answer:
             return {"title": "Error", "content": f"Failed to generate final answer. Error: {str(e)}"}
         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, int, 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."""
     
@@ -41,6 +57,8 @@ def generate_response(api_key: str, prompt: str) -> Tuple[List[Tuple[str, str, f
     steps = []
     step_count = 1
     total_thinking_time = 0
+    total_tokens = 0
+    total_tokens_per_second = 0
     
     while True:
         start_time = time.time()
@@ -48,14 +66,24 @@ 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))
+        token_info = step_data.pop('token_info', {'total_tokens': 0, 'tokens_per_second': 0})
+        total_tokens += token_info['total_tokens']
+        total_tokens_per_second += token_info['tokens_per_second']
+        
+        step_title = f"Step {step_count}: {step_data['title']}"
+        step_content = f"{step_data['content']}\n\n**Cerebras LLM Call Duration: {thinking_time:.2f} seconds**\n**Tokens: {token_info['total_tokens']}, Tokens/s: {token_info['tokens_per_second']:.2f}**"
+        steps.append((step_title, step_content))
         messages.append({"role": "assistant", "content": json.dumps(step_data)})
         
+        # Yield the current conversation, total thinking time, total tokens, and average tokens per second
+        yield steps, total_thinking_time, total_tokens, total_tokens_per_second / step_count if step_count > 0 else 0
+        
         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 +91,36 @@ 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
+    token_info = final_data.pop('token_info', {'total_tokens': 0, 'tokens_per_second': 0})
+    total_tokens += token_info['total_tokens']
+    total_tokens_per_second += token_info['tokens_per_second']
+    
+    final_content = f"{final_data.get('content', 'No final answer provided.')}\n\n**Final answer thinking time: {thinking_time:.2f} seconds**\n**Tokens: {token_info['total_tokens']}, Tokens/s: {token_info['tokens_per_second']:.2f}**"
+    steps.append(("Final Answer", final_content))
+    
+    # Yield the final conversation, total thinking time, total tokens, and average tokens per second
+    yield steps, total_thinking_time, total_tokens, total_tokens_per_second / (step_count + 1)
 
-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, total_tokens, avg_tokens_per_second in response_generator:
+        conversation = history.copy()
+        for title, content in steps[len(conversation):]:
+            if title.startswith("Step") or title == "Final Answer":
+                conversation.append((title, content))
+            else:
+                conversation.append((title, content))
+        yield conversation, f"**Total thinking time:** {total_time:.2f} seconds\n**Total tokens:** {total_tokens}\n**Average tokens/s:** {avg_tokens_per_second:.2f}"
 
-# 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")
@@ -109,25 +147,11 @@ def main():
             submit_btn = gr.Button("Submit")
         
         thinking_time_display = gr.Textbox(
-            label="Total Thinking Time", 
+            label="Performance Metrics", 
             value="", 
             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 +159,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],