Update app.py

app.py

@@ -7,8 +7,8 @@ from typing import List
 
 model = Llama(
     model_path=hf_hub_download(
-        repo_id=os.environ.get("REPO_ID", "Lyte/QuadConnect2.5-1.5B-v0.1.0b"), #"Lyte/QuadConnect2.5-0.5B-v0.0.9b"),#"Lyte/QuadConnect2.5-0.5B-v0.0.8b"), #"Lyte/QuadConnect2.5-0.5B-v0.0.6b"), #"Lyte/QuadConnect-Llama-1B-v0.0.7b"),#"
-        filename=os.environ.get("MODEL_FILE", "unsloth.Q8_0.gguf"), #"quadconnect.Q8_0.gguf"),
+        repo_id=os.environ.get("REPO_ID", "Lyte/QuadConnect2.5-1.5B-v0.1.0b"),
+        filename=os.environ.get("MODEL_FILE", "unsloth.Q8_0.gguf"),
     ),
     n_ctx=16384
 )
@@ -386,71 +386,114 @@ def create_interface():
             game_state = game.format_game_state()
             print(game_state)
             
-            # Get AI response with user-defined temperature
-            response = model.create_chat_completion(
-                messages=[
-                    {"role": "system", "content": SYSTEM_PROMPT},
-                    {"role": "user", "content": game_state}
-                ],
-                temperature=temperature,
-                top_p=0.95,
-                max_tokens=1024
-            )
-            
-            ai_response = response['choices'][0]['message']['content']
-            print(ai_response)
-            
-            # Extract reasoning and move
-            try:
-                reasoning = ai_response.split("<reasoning>")[1].split("</reasoning>")[0].strip()
-                move_str = extract_xml_move(ai_response)
+            # 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>'
                 
-                if not move_str:
-                    raise ValueError("Invalid move format from AI")
+                # 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
+                )
                 
-                ai_col = game.parse_ai_move(move_str)
+                # Process stream chunks
+                for chunk in response_stream:
+                    if 'choices' in chunk and len(chunk['choices']) > 0:
+                        content = chunk['choices'][0]['delta'].get('content', '')
+                        if 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>")
+                                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>'
+                            
+                            current_reasoning_html += '</div></div>'
+                            
+                            # Yield the updated HTML
+                            yield [
+                                render_board(game.board),
+                                "AI is thinking...",
+                                current_reasoning_html
+                            ]
                 
-                if ai_col == -1:
-                    raise ValueError("Invalid move format from AI")
+                # 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."
                 
-                # Format reasoning for display
-                reasoning_html = f'''
+                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}</p>
-                        <p><strong>📍 Move chosen:</strong> Column {move_str.upper()}</p>
+                        <p>{reasoning_content}</p>
+                        <p><strong>📍 Move chosen:</strong> Column {move_str.upper() if move_str else "?"}</p>
                     </div>
                 </div>
                 '''
                 
-                success, _ = game.make_move(ai_col)
-                if success:
-                    # Check for AI winner
-                    winner = game.check_winner()
-                    if winner == 2:
-                        game.game_over = True
-                        return [
-                            render_board(game.board),
-                            "🤖 AI wins! Better luck next time!",
-                            reasoning_html
-                        ]
-                else:
-                    return [
-                        render_board(game.board),
-                        "AI made invalid move! You win by default!",
-                        '<div id="ai-reasoning">AI made an invalid move!</div>'
-                    ]
-            except Exception as e:
+                # 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
+                                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
+                
+                # If we get here, something went wrong with the AI move
                 game.game_over = True
-                return [
+                yield [
                     render_board(game.board),
-                    "AI error occurred! You win by default!",
-                    f'<div id="ai-reasoning">Error: {str(e)}</div>'
+                    "AI made an invalid move! You win by default!",
+                    '<div id="ai-reasoning">AI made an invalid move or provided an invalid response!</div>'
                 ]
             
-            game.current_player = 1
-            return [render_board(game.board), "Your turn!", reasoning_html]
+            # Return the streaming generator
+            return stream_response()
 
         def reset_game():
             game.board = np.zeros((6, 7))