app.py

import gradio as gr
import pandas as pd
import traceback
import gc
import torch
import json

from cognitive_mapping_probe.orchestrator_seismograph import run_seismic_analysis
from cognitive_mapping_probe.auto_experiment import run_auto_suite, get_curated_experiments
from cognitive_mapping_probe.prompts import RESONANCE_PROMPTS
from cognitive_mapping_probe.utils import dbg

theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue").set(body_background_fill="#f0f4f9", block_background_fill="white")

def cleanup_memory():
    """Eine zentrale Funktion zum Aufräumen des Speichers nach einem Lauf."""
    dbg("Cleaning up memory...")
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    dbg("Memory cleanup complete.")

def run_single_analysis_display(*args, progress=gr.Progress(track_tqdm=True)):
    """Wrapper für ein einzelnes manuelles Experiment mit robuster Fehlerbehandlung."""
    try:
        results = run_seismic_analysis(*args, progress_callback=progress)
        stats, deltas = results.get("stats", {}), results.get("state_deltas", [])
        df = pd.DataFrame({"Internal Step": range(len(deltas)), "State Change (Delta)": deltas})
        stats_md = f"### Statistical Signature\n- **Mean Delta:** {stats.get('mean_delta', 0):.4f}\n- **Std Dev Delta:** {stats.get('std_delta', 0):.4f}\n- **Max Delta:** {stats.get('max_delta', 0):.4f}\n"
        serializable_results = json.dumps(results, indent=2, default=str)
        return f"{results.get('verdict', 'Error')}\n\n{stats_md}", df, serializable_results
    except Exception:
        # Im Fehlerfall, gib für jede Komponente einen leeren, aber typ-korrekten Wert zurück.
        error_message = f"### ❌ Analysis Failed\n```\n{traceback.format_exc()}\n```"
        return error_message, pd.DataFrame(), "{}"
    finally:
        cleanup_memory()

PLOT_PARAMS = {
    "x": "Step", "y": "Delta", "color": "Experiment",
    "title": "Comparative Cognitive Dynamics", "color_legend_title": "Experiment Runs",
    "color_legend_position": "bottom", "show_label": True, "height": 400, "interactive": True
}

def run_auto_suite_display(model_id, num_steps, seed, experiment_name, progress=gr.Progress(track_tqdm=True)):
    """
    Wrapper für die automatisierte Experiment-Suite mit robuster Fehlerbehandlung.
    """
    try:
        summary_df, plot_df, all_results = run_auto_suite(model_id, int(num_steps), int(seed), experiment_name, progress)
        new_plot = gr.LinePlot(value=plot_df, **PLOT_PARAMS)
        serializable_results = json.dumps(all_results, indent=2, default=str)
        return summary_df, new_plot, serializable_results
    except Exception:
        # KORREKTUR: Gib für jede Komponente einen typ-korrekten, leeren Wert zurück.
        # Insbesondere für die JSON-Komponente einen leeren JSON-String.
        empty_plot = gr.LinePlot(value=pd.DataFrame(), **PLOT_PARAMS)
        error_string_for_json = json.dumps({"error": traceback.format_exc()}, indent=2)
        return pd.DataFrame([{"Error": "Experiment failed. See Raw JSON."}]), empty_plot, error_string_for_json
    finally:
        cleanup_memory()

with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.3") as demo:
    gr.Markdown("# 🧠 Cognitive Seismograph 2.3: Advanced Experiment Suite")

    with gr.Tabs():
        with gr.TabItem("🔬 Manual Single Run"):
            # ... (UI unverändert)
            gr.Markdown("Run a single experiment with manual parameters to explore hypotheses.")
            with gr.Row(variant='panel'):
                with gr.Column(scale=1):
                    gr.Markdown("### 1. General Parameters")
                    manual_model_id = gr.Textbox(value="google/gemma-3-1b-it", label="Model ID")
                    manual_prompt_type = gr.Radio(choices=list(RESONANCE_PROMPTS.keys()), value="resonance_prompt", label="Prompt Type")
                    manual_seed = gr.Slider(1, 1000, 42, step=1, label="Seed")
                    manual_num_steps = gr.Slider(50, 1000, 300, step=10, label="Number of Internal Steps")
                    gr.Markdown("### 2. Modulation Parameters")
                    manual_concept = gr.Textbox(label="Concept to Inject", placeholder="e.g., 'calmness' (leave blank for baseline)")
                    manual_strength = gr.Slider(0.0, 5.0, 1.5, step=0.1, label="Injection Strength")
                    manual_run_btn = gr.Button("Run Single Analysis", variant="primary")
                with gr.Column(scale=2):
                    gr.Markdown("### Single Run Results")
                    manual_verdict = gr.Markdown("Analysis results will appear here.")
                    manual_plot = gr.LinePlot(x="Internal Step", y="State Change (Delta)", title="Internal State Dynamics", show_label=True, height=400, interactive=True)
                    with gr.Accordion("Raw JSON Output", open=False):
                        manual_raw_json = gr.JSON()
            manual_run_btn.click(
                fn=run_single_analysis_display,
                inputs=[manual_model_id, manual_prompt_type, manual_seed, manual_num_steps, manual_concept, manual_strength],
                outputs=[manual_verdict, manual_plot, manual_raw_json]
            )

        with gr.TabItem("🚀 Automated Suite"):
            # ... (UI unverändert)
            gr.Markdown("Run a predefined, curated suite of experiments and visualize the results comparatively.")
            with gr.Row(variant='panel'):
                with gr.Column(scale=1):
                    gr.Markdown("### Auto-Experiment Parameters")
                    auto_model_id = gr.Textbox(value="google/gemma-3-4b-it", label="Model ID")
                    auto_num_steps = gr.Slider(50, 1000, 300, step=10, label="Steps per Run")
                    auto_seed = gr.Slider(1, 1000, 42, step=1, label="Seed")
                    auto_experiment_name = gr.Dropdown(choices=list(get_curated_experiments().keys()), value="Therapeutic Intervention (4B-Model)", label="Curated Experiment Protocol")
                    auto_run_btn = gr.Button("Run Curated Auto-Experiment", variant="primary")
                with gr.Column(scale=2):
                    gr.Markdown("### Suite Results Summary")
                    auto_plot_output = gr.LinePlot(**PLOT_PARAMS)
                    auto_summary_df = gr.DataFrame(label="Comparative Statistical Signature", wrap=True)
                    with gr.Accordion("Raw JSON for all runs", open=False):
                        auto_raw_json = gr.JSON()
            auto_run_btn.click(
                fn=run_auto_suite_display,
                inputs=[auto_model_id, auto_num_steps, auto_seed, auto_experiment_name],
                outputs=[auto_summary_df, auto_plot_output, auto_raw_json]
            )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)