Update app.py

app.py

@@ -1,12 +1,17 @@
 import gradio as gr
 import torch
 import time
-import spaces # Import the spaces library
+import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from db import save_test_result
+from db import init_db, save_test_result, get_test_history, get_test_details
+
+# --- Initialize Database ---
+db_initialized = init_db()
+if not db_initialized:
+    print("WARNING: Database initialization failed. Test history will not be saved.")
 
 # --- Configuration ---
-MODEL_ID = "Qwen/Qwen2.5-Math-1.5B" # Replace with actual ID if found
+MODEL_ID = "Qwen/Qwen2.5-Math-1.5B"  # Replace with actual ID if found
 # --- Load Model and Tokenizer ---
 print(f"Loading model: {MODEL_ID}")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -17,7 +22,6 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 print("Model loaded successfully.")
 
-
 # --- Generation Function (Returns response and token count) ---
 def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
     """Generate a response and return it along with the number of generated tokens."""
@@ -48,17 +52,6 @@ def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
         num_generated_tokens = len(output_ids)
         response = tokenizer.decode(output_ids, skip_special_tokens=True)
         print("Generation complete.")
-
-        save_test_result(
-            analysis_mode='',
-            system_prompt='',
-            input_content='', 
-            model_response= response, 
-            generation_time='',
-            tokens_generated='', 
-            temperature='',
-            top_p='',
-            max_length='')
         return response.strip(), num_generated_tokens
 
     except Exception as e:
@@ -73,7 +66,8 @@ def process_input(
     system_prompt, # Single system prompt from UI
     max_length,
     temperature,
-    top_p
+    top_p,
+    save_to_db=True  # New parameter to toggle database saving
 ):
     """Process inputs based on selected analysis mode using the provided system prompt."""
     print(f"GPU requested via decorator, starting processing in mode: {analysis_mode}")
@@ -117,11 +111,29 @@ def process_input(
         display_prompt = f"Selected Mode: {analysis_mode}\nSystem Prompt:\n{system_prompt}"
 
     print(f"Processing finished in {duration} seconds.")
+    
+    # Save to database if requested and if database is available
+    if save_to_db and db_initialized:
+        test_id = save_test_result(
+            analysis_mode=analysis_mode,
+            system_prompt=system_prompt,
+            input_content=user_content if user_content else "",
+            model_response=response,
+            generation_time=duration,
+            tokens_generated=generated_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            max_length=max_length
+        )
+        if test_id:
+            print(f"Test saved to database with ID: {test_id}")
+        else:
+            print("Failed to save test to database")
+    
     # Return all results including time and tokens
     return display_prompt, response, f"{duration} seconds", generated_tokens
 
 # --- System Prompts (Defaults only, UI will hold the editable version) ---
-
 DEFAULT_SYSTEM_PROMPT_FREQ = """You are an assistant that analyzes Rock-Paper-Scissors (RPS) player statistics. Your ONLY goal is to find the best single AI move to counter the player's MOST frequent move based on the provided frequency stats.
 
 Follow these steps EXACTLY. Do NOT deviate.
@@ -144,7 +156,6 @@ Step 4: State Final Recommendation.
 Base your analysis strictly on the provided frequencies and the stated RPS rules.
 """
 
-# Updated Markov system prompt - no changes to content
 DEFAULT_SYSTEM_PROMPT_MARKOV = """You are analyzing a Rock-Paper-Scissors (RPS) game using a Markov transition matrix.
 
 ### TRANSITION MATRIX:
@@ -186,7 +197,6 @@ Predicted Next Move: [Move with highest probability]
 Optimal Counter: [Move that beats the predicted move]
 """
 
-# New Behavior Analysis prompt
 DEFAULT_SYSTEM_PROMPT_BEHAVIOR = """You are an RPS assistant analyzing player behavior after wins, losses, and ties. Predict the player's next move and give counter strategy based on the Behavioral probabilities.
 
 **Behavioral Probabilities P(Change/not change | Win/Loss/Tie):**
@@ -224,74 +234,116 @@ DEFAULT_PLAYER_BEHAVIOR = "Player's Last Outcome: Win\nPlayer's Last Move: Rock"
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# {MODEL_ID} - RPS Strategy Tester")
-    gr.Markdown("Test model advice using Frequency Stats, Markov Predictions, or Win/Loss/Tie Behavior Analysis.")
-
-    # Mode Selector - now with three options
-    analysis_mode_selector = gr.Radio(
-        label="Select Analysis Mode",
-        choices=["Frequency Only", "Markov Prediction Only", "Behavior Analysis"],
-        value="Frequency Only" # Default mode
-    )
-
-    # --- Visible System Prompt Textbox ---
-    system_prompt_input = gr.Textbox(
-            label="System Prompt (Edit based on selected mode)",
-            value=DEFAULT_SYSTEM_PROMPT_FREQ, # Start with frequency prompt
-            lines=15
+    with gr.Tab("Model Testing"):
+        gr.Markdown(f"# {MODEL_ID} - RPS Strategy Tester with Test History")
+        gr.Markdown("Test model advice using Frequency Stats, Markov Predictions, or Win/Loss/Tie Behavior Analysis.")
+
+        # Mode Selector - now with three options
+        analysis_mode_selector = gr.Radio(
+            label="Select Analysis Mode",
+            choices=["Frequency Only", "Markov Prediction Only", "Behavior Analysis"],
+            value="Frequency Only" # Default mode
         )
 
-    # Input Sections (conditionally visible)
-    with gr.Group(visible=True) as frequency_inputs: # Visible by default
-        gr.Markdown("### Frequency Analysis Inputs")
-        player_stats_input = gr.Textbox(
-            label="Player Move Frequency Stats (Long-Term)", value=DEFAULT_PLAYER_STATS, lines=4,
-            info="Overall player move distribution."
+        # --- Visible System Prompt Textbox ---
+        system_prompt_input = gr.Textbox(
+                label="System Prompt (Edit based on selected mode)",
+                value=DEFAULT_SYSTEM_PROMPT_FREQ, # Start with frequency prompt
+                lines=15
+            )
+
+        # Input Sections (conditionally visible)
+        with gr.Group(visible=True) as frequency_inputs: # Visible by default
+            gr.Markdown("### Frequency Analysis Inputs")
+            player_stats_input = gr.Textbox(
+                label="Player Move Frequency Stats (Long-Term)", value=DEFAULT_PLAYER_STATS, lines=4,
+                info="Overall player move distribution."
+            )
+
+        with gr.Group(visible=False) as markov_inputs: # Hidden by default
+            gr.Markdown("### Markov Prediction Analysis Inputs")
+            gr.Markdown("*Use the System Prompt field to directly input your Markov analysis instructions.*")
+
+        # New behavior analysis inputs
+        with gr.Group(visible=False) as behavior_inputs:
+            gr.Markdown("### Win/Loss/Tie Behavior Analysis Inputs")
+            player_behavior_input = gr.Textbox(
+                label="Player's Last Outcome and Move", value=DEFAULT_PLAYER_BEHAVIOR, lines=4,
+                info="Enter the last outcome (Win/Loss/Tie) and move (Rock/Paper/Scissors)."
+            )
+
+        # General Inputs / Parameters / Outputs
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Generation Parameters")
+                max_length_slider = gr.Slider(minimum=50, maximum=1024, value=300, step=16, label="Max New Tokens")
+                temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Temperature")
+                top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P")
+                
+                # Add a checkbox to control saving to database
+                save_to_db_checkbox = gr.Checkbox(
+                    label="Save this test to database", 
+                    value=True,
+                    info="Store input and output in SQLite database for later reference"
+                )
+
+        submit_btn = gr.Button("Generate Response", variant="primary")
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Performance Metrics")
+                time_output = gr.Textbox(label="Generation Time", interactive=False)
+                tokens_output = gr.Number(label="Generated Tokens", interactive=False)
+            with gr.Column():
+                 gr.Markdown("""
+                #### Testing Tips
+                - Select the desired **Analysis Mode**.
+                - Fill in the inputs for the **selected mode only**.
+                - **Edit the System Prompt** above as needed for testing.
+                - Use low **Temperature** for factual analysis.
+                """)
+
+        with gr.Row():
+            final_prompt_display = gr.Textbox(
+                label="Formatted Input Sent to Model (via Chat Template)", lines=20
+            )
+            response_display = gr.Textbox(
+                label="Model Response", lines=20, show_copy_button=True
+            )
+
+    # Add a new tab for test history
+    with gr.Tab("Test History"):
+        gr.Markdown("### Saved Test Results")
+        
+        refresh_btn = gr.Button("Refresh History")
+        
+        # Display test history as a dataframe
+        test_history_df = gr.Dataframe(
+            headers=["Test ID", "Analysis Mode", "Timestamp", "Generation Time", "Tokens"],
+            label="Recent Tests",
+            interactive=False
         )
-
-    with gr.Group(visible=False) as markov_inputs: # Hidden by default
-        gr.Markdown("### Markov Prediction Analysis Inputs")
-        gr.Markdown("*Use the System Prompt field to directly input your Markov analysis instructions.*")
-
-    # New behavior analysis inputs
-    with gr.Group(visible=False) as behavior_inputs:
-        gr.Markdown("### Win/Loss/Tie Behavior Analysis Inputs")
-        player_behavior_input = gr.Textbox(
-            label="Player's Last Outcome and Move", value=DEFAULT_PLAYER_BEHAVIOR, lines=4,
-            info="Enter the last outcome (Win/Loss/Tie) and move (Rock/Paper/Scissors)."
-        )
-
-    # General Inputs / Parameters / Outputs
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("#### Generation Parameters")
-            max_length_slider = gr.Slider(minimum=50, maximum=1024, value=300, step=16, label="Max New Tokens")
-            temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Temperature")
-            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P")
-
-    submit_btn = gr.Button("Generate Response", variant="primary")
-
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("#### Performance Metrics")
-            time_output = gr.Textbox(label="Generation Time", interactive=False)
-            tokens_output = gr.Number(label="Generated Tokens", interactive=False)
-        with gr.Column():
-             gr.Markdown("""
-            #### Testing Tips
-            - Select the desired **Analysis Mode**.
-            - Fill in the inputs for the **selected mode only**.
-            - **Edit the System Prompt** above as needed for testing.
-            - Use low **Temperature** for factual analysis.
-            """)
-
-    with gr.Row():
-        final_prompt_display = gr.Textbox(
-            label="Formatted Input Sent to Model (via Chat Template)", lines=20
-        )
-        response_display = gr.Textbox(
-            label="Model Response", lines=20, show_copy_button=True
+        
+        # Add a number input to load a specific test
+        test_id_input = gr.Number(
+            label="Test ID", 
+            precision=0,
+            info="Enter a Test ID to load details"
         )
+        load_test_btn = gr.Button("Load Test")
+        
+        # Display test details
+        with gr.Group():
+            test_mode_display = gr.Textbox(label="Analysis Mode", interactive=False)
+            test_prompt_display = gr.Textbox(label="System Prompt", interactive=False, lines=8)
+            test_input_display = gr.Textbox(label="Input Content", interactive=False, lines=4)
+            test_response_display = gr.Textbox(label="Model Response", interactive=False, lines=8)
+            
+            with gr.Row():
+                test_time_display = gr.Number(label="Generation Time (s)", interactive=False)
+                test_tokens_display = gr.Number(label="Tokens Generated", interactive=False)
+                test_temp_display = gr.Number(label="Temperature", interactive=False)
+                test_topp_display = gr.Number(label="Top P", interactive=False)
 
     # --- Event Handlers ---
 
@@ -326,6 +378,34 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 system_prompt_input: gr.update(value=DEFAULT_SYSTEM_PROMPT_FREQ)
             }
 
+    # Function to update test history display
+    def update_test_history():
+        if db_initialized:
+            history = get_test_history(limit=20)
+            return [[h[0], h[1], h[2], h[3], h[4]] for h in history]
+        else:
+            return [["N/A", "Database Not Available", "N/A", 0, 0]]
+    
+    # Function to load test details
+    def load_test_details(test_id):
+        if not db_initialized:
+            return ["Database Not Available", "", "", "", 0, 0, 0, 0]
+        
+        test = get_test_details(test_id)
+        if test:
+            return [
+                test["analysis_mode"],
+                test["system_prompt"],
+                test["input_content"] or "",
+                test["model_response"],
+                test["generation_time"],
+                test["tokens_generated"],
+                test["temperature"],
+                test["top_p"]
+            ]
+        else:
+            return ["Test not found", "", "", "", 0, 0, 0, 0]
+
     # Link the radio button change to the UI update function
     analysis_mode_selector.change(
         fn=update_ui_visibility_and_prompt, # Use the combined update function
@@ -343,14 +423,33 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             system_prompt_input, # Pass the visible system prompt textbox
             max_length_slider,
             temperature_slider,
-            top_p_slider
+            top_p_slider,
+            save_to_db_checkbox  # Pass the checkbox value
         ],
         outputs=[
             final_prompt_display, response_display,
             time_output, tokens_output
-        ],
-        api_name="generate_rps_selectable_analysis_v4" # Updated api_name
+        ]
     )
+    
+    # Connect buttons for test history tab
+    refresh_btn.click(
+        update_test_history,
+        outputs=[test_history_df]
+    )
+    
+    load_test_btn.click(
+        load_test_details,
+        inputs=[test_id_input],
+        outputs=[
+            test_mode_display, test_prompt_display, test_input_display, 
+            test_response_display, test_time_display, test_tokens_display,
+            test_temp_display, test_topp_display
+        ]
+    )
+    
+    # Initialize history on page load
+    demo.load(update_test_history, outputs=[test_history_df])
 
 # --- Launch the demo ---
 if __name__ == "__main__":