app_v2.py

"""
Molecule Maestro – Gradio (≥4) + OpenAI ≥ 1.0  
**Thermo‑Reaction Mode** with *50 % carry‑over shuffle*: when you hit **Reshuffle palette**, half of the current reactants stay, half are refreshed.

### Quick start
```bash
pip install "gradio>=4.0" "openai>=1.0.0"
export OPENAI_API_KEY="sk‑…"
python molecule_maestro_gradio.py
```
Open <http://localhost:7860>.
"""

from __future__ import annotations
import json, random
from typing import List, Dict
from openai import OpenAI
import gradio as gr

client = OpenAI()

CATEGORY_EMOJI = {
    "acid": "🔴", "base": "🔵", "solvent": "🟢", "gas": "🟡", "oxidizer": "🟣", "other": "⚪",
}

SYSTEM_PROMPT_PALETTE = (
    "You are a helpful chemistry assistant. Provide a JSON array named 'palette' with exactly 12 reactants used in school labs. "
    "Each item: {name, category} (acid|base|solvent|gas|oxidizer|other)."
)

SYSTEM_PROMPT_REACTION = (
    "You are a chemistry mentor. The user provides 1‑5 reactants, optional hints, and conditions. "
    "Return up to five chemical reactions in **strict JSON**:\n"
    "{\n  'reactions': [\n    { 'equation':'string', 'conditions to occur':'string', 'thermo':'exothermic|endothermic', 'deltaH_kJ':number, 'explanation':'string (≤25 words)'}\n    … up to 5 …\n  ]\n}\n"
    "Reactions must be single‑step, common, school‑level. If none plausible, 'reactions' is empty."
)


def chat_completion(prompt: str, system: str, temperature: float = 0.6) -> str:
    resp = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "system", "content": system}, {"role": "user", "content": prompt}],
        temperature=temperature,
    )
    return resp.choices[0].message.content.strip()


# — Palette helpers —

def generate_palette() -> List[Dict[str, str]]:
    """Ask GPT for 12 reactants; fall back to static list on parse issues."""
    raw = chat_completion("Give me the palette", SYSTEM_PROMPT_PALETTE, 0.4)
    try:
        return json.loads(raw)["palette"]
    except Exception:
        return [
            {"name": "HCl", "category": "acid"}, {"name": "NaOH", "category": "base"}, {"name": "Ethanol", "category": "solvent"},
            {"name": "Acetic Acid", "category": "acid"}, {"name": "NH3", "category": "base"}, {"name": "H2O", "category": "solvent"},
            {"name": "H2", "category": "gas"}, {"name": "O2", "category": "gas"}, {"name": "KMnO4", "category": "oxidizer"},
            {"name": "H2SO4", "category": "acid"}, {"name": "NaCl", "category": "other"}, {"name": "CuSO4", "category": "other"},
        ]

palette: List[Dict[str, str]] = generate_palette()


def palette_labels() -> List[str]:
    return [f"{CATEGORY_EMOJI.get(it['category'], '⚪')} {it['name']}" for it in palette]


# — Reaction evaluation —

def build_prompt(reactants: List[str], t: float, p: float, cat: str, solv: str, user_hint: str) -> str:
    return (
        f"Reactants: {', '.join(reactants)}. Temperature: {t} °C. Pressure: {p} atm. "
        f"Catalyst: {cat or 'none'}. Solvent: {solv or 'none'}. User hint: {user_hint or 'none'}."
    )


def evaluate_reaction(reactants: List[str], t: float, p: float, cat: str, solv: str, hint: str):
    if not reactants:
        return "⚠️ Select at least one reactant."

    raw = chat_completion(build_prompt(reactants, t, p, cat, solv, hint), SYSTEM_PROMPT_REACTION, 0.3)
    try:
        data = json.loads(raw)
        reactions = data.get("reactions", [])[:5]
        if not reactions:
            return "No simple reactions predicted. Try other reactants or conditions."
        bullets = [
            f"- **{r['equation']}** – {r['thermo']} (ΔH ≈ {r['deltaH_kJ']} kJ) – {r['explanation']}"
            for r in reactions
        ]
        return "\n".join(bullets)
    except Exception:
        return f"Parse err 🤖\n```\n{raw}\n```"


# — Gradio UI —
with gr.Blocks() as demo:
    gr.Markdown("# 🧪 Molecule Maestro – Thermo‑Reaction Mode")

    with gr.Row():
        with gr.Column():
            reactant_group = gr.CheckboxGroup(choices=palette_labels(), label="Reactants (emoji‑coded)")
            shuffle_btn = gr.Button("🔀 Reshuffle palette", variant="secondary")
        with gr.Column():
            temp = gr.Slider(-20, 200, value=25, step=1, label="Temperature (°C)")
            press = gr.Slider(0.5, 10, value=1, step=0.1, label="Pressure (atm)")
            catalyst = gr.Textbox(label="Catalyst", placeholder="optional")
            solvent = gr.Textbox(label="Solvent", placeholder="optional")
            user_hint = gr.Textbox(label="Notes / hypotheses (optional)")
            run_btn = gr.Button("▶︎ Suggest Reactions")

    result_md = gr.Markdown("—")

    # — Callbacks —
    def _shuffle():
        """Keep 50 % of current reactants, replace the rest with fresh ones."""
        global palette
        keep_n = len(palette) // 2  # 6 if palette has 12
        kept = random.sample(palette, keep_n)
        kept_names = {item['name'] for item in kept}

        # Generate until we get enough new unique reactants
        fresh: List[Dict[str, str]] = []
        attempts = 0
        while len(fresh) < (12 - keep_n) and attempts < 5:
            candidate_batch = generate_palette()
            for item in candidate_batch:
                if item['name'] not in kept_names and item['name'] not in {f['name'] for f in fresh}:
                    fresh.append(item)
                    if len(fresh) == 12 - keep_n:
                        break
            attempts += 1

        palette = kept + fresh
        random.shuffle(palette)
        return gr.update(choices=palette_labels(), value=[]), "—"

    def _run(labels, t, p, cat, solv, hint):
        names = [lbl.split(' ', 1)[1] for lbl in labels]
        return evaluate_reaction(names, t, p, cat, solv, hint)

    shuffle_btn.click(_shuffle, None, [reactant_group, result_md])
    run_btn.click(_run, [reactant_group, temp, press, catalyst, solvent, user_hint], result_md)

if __name__ == "__main__":
    demo.launch()