Update app.py

app.py

@@ -4,6 +4,7 @@ from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 import numpy as np
 from typing import List
+import time
 
 model = Llama(
     model_path=hf_hub_download(
@@ -47,6 +48,16 @@ def extract_xml_move(text: str) -> str:
         return match.group(1)
     return ""
 
+def extract_reasoning(text: str) -> str:
+    """
+    Extract the reasoning part from the XML response.
+    """
+    import re
+    match = re.search(r'<reasoning>(.*?)</reasoning>', text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    return ""
+
 def convert_moves_to_coordinate_list(moves_list: List[str]) -> str:
     """
     Converts a list of moves to a coordinate list representation.
@@ -291,6 +302,9 @@ def create_interface():
         margin: 15px 0;
         font-family: monospace;
         min-height: 100px;
+        max-height: 400px;
+        overflow-y: auto;
+        white-space: pre-wrap;
     }
     .reasoning-box {
         border-left: 4px solid #2196F3;
@@ -313,6 +327,10 @@ def create_interface():
     div.svelte-1nguped {
       display: block;
     }
+    .thinking-indicator {
+      color: #ffc107;
+      font-style: italic;
+    }
     """
 
     with gr.Blocks(css=css) as interface:
@@ -352,7 +370,7 @@ def create_interface():
                 info="Lower values make AI more deterministic, higher values more creative"
             )
 
-        def handle_move(col, temperature=0.8):
+        def handle_player_move(col, temperature=0.8):
             if game.game_over:
                 return [
                     render_board(game.board),
@@ -386,114 +404,145 @@ def create_interface():
             game_state = game.format_game_state()
             print(game_state)
             
-            # Create streaming response generator
-            def stream_response():
-                # Initialize streaming response
-                reasoning_html = '<div id="ai-reasoning"><div class="reasoning-box"><p><strong>🤔 Reasoning:</strong></p><p>'
-                
-                # Initialize variables to store the full response
-                full_response = ""
-                reasoning_content = ""
-                move_str = ""
-                
-                # Set up streaming
-                response_stream = model.create_chat_completion(
-                    messages=[
-                        {"role": "system", "content": SYSTEM_PROMPT},
-                        {"role": "user", "content": game_state}
-                    ],
-                    temperature=temperature,
-                    top_p=0.95,
-                    max_tokens=1024,
-                    stream=True
-                )
-                
-                # Process stream chunks
-                for chunk in response_stream:
-                    if 'choices' in chunk and len(chunk['choices']) > 0:
-                        content = chunk['choices'][0]['delta'].get('content', '')
+            # Prepare for AI's turn
+            reasoning_html = '''
+            <div id="ai-reasoning">
+                <div class="reasoning-box">
+                    <p><strong>🤔 Reasoning:</strong></p>
+                    <p class="thinking-indicator">AI is thinking...</p>
+                </div>
+            </div>
+            '''
+            
+            return [
+                render_board(game.board),
+                "AI is thinking...",
+                reasoning_html
+            ]
+        
+        def process_ai_response(temperature=0.8):
+            if game.game_over or game.current_player != 2:
+                return [
+                    render_board(game.board),
+                    "Your turn!" if not game.game_over else "Game is over! Click New Game to play again.",
+                    reasoning_display.value
+                ]
+            
+            game_state = game.format_game_state()
+            
+            # Start streaming response
+            full_response = ""
+            reasoning_content = ""
+            move_str = ""
+            
+            # Begin HTML
+            reasoning_html = '''
+            <div id="ai-reasoning">
+                <div class="reasoning-box">
+                    <p><strong>🤔 Reasoning:</strong></p>
+                    <p>
+            '''
+            
+            # Create streaming session
+            response_stream = model.create_chat_completion(
+                messages=[
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": game_state}
+                ],
+                temperature=temperature,
+                top_p=0.95,
+                max_tokens=1024,
+                stream=True
+            )
+            
+            # Process stream
+            collected_text = ""
+            for chunk in response_stream:
+                if 'choices' in chunk and len(chunk['choices']) > 0:
+                    delta = chunk['choices'][0]['delta']
+                    if 'content' in delta:
+                        content = delta['content']
                         if content:
+                            collected_text += content
                             full_response += content
                             
-                            # Try to extract reasoning as it comes
-                            if "<reasoning>" in full_response and "</reasoning>" in full_response:
-                                parts = full_response.split("<reasoning>")
+                            # Extract and update reasoning
+                            if "<reasoning>" in collected_text and "</reasoning>" not in collected_text:
+                                # Still collecting reasoning
+                                parts = collected_text.split("<reasoning>")
                                 if len(parts) > 1:
-                                    reasoning_parts = parts[1].split("</reasoning>")
-                                    if len(reasoning_parts) > 0:
-                                        reasoning_content = reasoning_parts[0].strip()
-                            
-                            # Update the HTML with the current reasoning content
-                            current_reasoning_html = reasoning_html + reasoning_content + '</p>'
-                            
-                            # Try to extract move as it comes
-                            temp_move_str = extract_xml_move(full_response)
-                            if temp_move_str:
-                                move_str = temp_move_str
-                                current_reasoning_html += f'<p><strong>📍 Move chosen:</strong> Column {move_str.upper()}</p>'
+                                    reasoning_content = parts[1]
+                                    current_html = reasoning_html + reasoning_content + '</p>'
+                                    if not move_str:
+                                        current_html += '</div></div>'
+                                    yield [
+                                        render_board(game.board),
+                                        "AI is thinking...",
+                                        current_html
+                                    ]
                             
-                            current_reasoning_html += '</div></div>'
+                            # Check for complete reasoning
+                            if "<reasoning>" in collected_text and "</reasoning>" in collected_text:
+                                reasoning_content = extract_reasoning(collected_text)
                             
-                            # Yield the updated HTML
-                            yield [
-                                render_board(game.board),
-                                "AI is thinking...",
-                                current_reasoning_html
-                            ]
-                
-                # Final response with complete reasoning and move
-                if not reasoning_content:
-                    try:
-                        reasoning_content = full_response.split("<reasoning>")[1].split("</reasoning>")[0].strip()
-                    except:
-                        reasoning_content = "AI couldn't provide clear reasoning."
-                
-                if not move_str:
-                    move_str = extract_xml_move(full_response)
-                
-                final_reasoning_html = f'''
-                <div id="ai-reasoning">
-                    <div class="reasoning-box">
-                        <p><strong>🤔 Reasoning:</strong></p>
-                        <p>{reasoning_content}</p>
-                        <p><strong>📍 Move chosen:</strong> Column {move_str.upper() if move_str else "?"}</p>
-                    </div>
-                </div>
-                '''
-                
-                # Make AI move
-                if move_str:
-                    ai_col = game.parse_ai_move(move_str)
-                    
-                    if ai_col != -1:
-                        success, _ = game.make_move(ai_col)
-                        if success:
-                            # Check for AI winner
-                            winner = game.check_winner()
-                            if winner == 2:
-                                game.game_over = True
+                            # Check for move
+                            new_move = extract_xml_move(collected_text)
+                            if new_move and not move_str:
+                                move_str = new_move
+                                current_html = reasoning_html + reasoning_content + '</p>'
+                                current_html += f'<p><strong>📍 Move chosen:</strong> Column {move_str.upper()}</p>'
+                                current_html += '</div></div>'
                                 yield [
                                     render_board(game.board),
-                                    "🤖 AI wins! Better luck next time!",
-                                    final_reasoning_html
+                                    "AI is making a move...",
+                                    current_html
                                 ]
-                                return
-                            else:
-                                # AI made a valid move but didn't win
-                                game.current_player = 1
-                                yield [render_board(game.board), "Your turn!", final_reasoning_html]
-                                return
+            
+            # Final formatting
+            final_reasoning_html = f'''
+            <div id="ai-reasoning">
+                <div class="reasoning-box">
+                    <p><strong>🤔 Reasoning:</strong></p>
+                    <p>{reasoning_content if reasoning_content else "AI couldn't provide clear reasoning."}</p>
+                    <p><strong>📍 Move chosen:</strong> Column {move_str.upper() if move_str else "?"}</p>
+                </div>
+            </div>
+            '''
+            
+            # Make AI move
+            if move_str:
+                ai_col = game.parse_ai_move(move_str)
                 
-                # If we get here, something went wrong with the AI move
-                game.game_over = True
-                yield [
-                    render_board(game.board),
-                    "AI made an invalid move! You win by default!",
-                    '<div id="ai-reasoning">AI made an invalid move or provided an invalid response!</div>'
-                ]
+                if ai_col != -1:
+                    success, _ = game.make_move(ai_col)
+                    if success:
+                        # Check for AI winner
+                        winner = game.check_winner()
+                        if winner == 2:
+                            game.game_over = True
+                            yield [
+                                render_board(game.board),
+                                "🤖 AI wins! Better luck next time!",
+                                final_reasoning_html
+                            ]
+                            return
+                        else:
+                            # AI made a valid move but didn't win
+                            game.current_player = 1
+                            yield [
+                                render_board(game.board), 
+                                "Your turn!", 
+                                final_reasoning_html
+                            ]
+                            return
             
-            # Return the streaming generator
-            return stream_response()
+            # If we get here, something went wrong with the AI move
+            game.game_over = True
+            yield [
+                render_board(game.board),
+                "AI made an invalid move! You win by default!",
+                '<div id="ai-reasoning">AI made an invalid move or provided an invalid response!</div>'
+            ]
 
         def reset_game():
             game.board = np.zeros((6, 7))
@@ -509,13 +558,20 @@ def create_interface():
 
         # Event handlers
         for i, btn in enumerate(col_buttons):
+            # First handle player move
             btn.click(
-                fn=handle_move,
+                fn=handle_player_move,
                 inputs=[
                     gr.Number(value=i, visible=False),
                     temperature_slider
                 ],
-                outputs=[board_display, status, reasoning_display]
+                outputs=[board_display, status, reasoning_display],
+                queue=False  # Process player move immediately
+            ).then(  # Then process AI response with streaming
+                fn=process_ai_response,
+                inputs=[temperature_slider],
+                outputs=[board_display, status, reasoning_display],
+                streaming=True  # Enable streaming for AI response
             )
         
         reset_btn.click(