app.py

"""PromptClimb — interactive BYOK demo of the prompt hill-climber.

Visitors paste their OWN provider key (BYOK): it's used only for the duration of
one run and never stored. The hill-climb is the real promptclimb package
(installed from GitHub), driven over a fixed sentiment-classification task so we
never execute visitor-supplied code on the Space.

Providers: OpenAI / Anthropic native backends; Gemini + NVIDIA via their
OpenAI-compatible endpoints (no promptclimb change — just OPENAI_BASE_URL + key);
Custom for any other OpenAI-compatible server (incl. a *publicly reachable*
Ollama/vLLM — a hosted Space can't reach your localhost).
"""
import os
import io
import json
import queue
import tempfile
import threading
import contextlib

import gradio as gr

# ponytail: hard caps keep a BYOK run to cents and bound what a careless visitor
# can spend on their own key. Raise if you self-fund a bigger demo.
MAX_ITERS = 8
MAX_CASES = 12

# provider -> how to drive promptclimb. prefix picks the backend (proposer.py
# routes anthropic:/ollama:/else->openai); base_url (for openai-compatible
# providers) goes in OPENAI_BASE_URL, which the openai SDK reads when the client
# is built with base_url=None (promptclimb's openai: path).
PROVIDERS = {
    "OpenAI": dict(prefix="openai", base_url=None, key_env="OPENAI_API_KEY",
                   model="gpt-4o-mini", hint="sk-..."),
    "Anthropic": dict(prefix="anthropic", base_url=None, key_env="ANTHROPIC_API_KEY",
                      model="claude-haiku-4-5", hint="sk-ant-..."),
    "Google Gemini": dict(prefix="openai",
                          base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
                          key_env="OPENAI_API_KEY", model="gemini-2.0-flash", hint="AIza..."),
    "NVIDIA (Nemotron)": dict(prefix="openai", base_url="https://integrate.api.nvidia.com/v1",
                              key_env="OPENAI_API_KEY",
                              model="nvidia/nvidia-nemotron-nano-9b-v2", hint="nvapi-..."),
    "Custom (OpenAI-compatible)": dict(prefix="openai", base_url="", key_env="OPENAI_API_KEY",
                                       model="", hint="key (or anything for keyless servers)"),
}

DEFAULT_PROMPT = "Classify the sentiment of the following text as positive, negative, or neutral."
DEFAULT_CASES = """I love this product! => positive
This is the worst thing I ever bought. => negative
It arrived on a Tuesday. => neutral
Absolutely fantastic, exceeded expectations. => positive
I'm so disappointed and frustrated. => negative
The meeting is scheduled for noon. => neutral"""


def parse_cases(text: str) -> list[dict]:
    cases = []
    for line in text.strip().splitlines():
        if "=>" in line:
            inp, exp = line.rsplit("=>", 1)
            if inp.strip() and exp.strip():
                cases.append({"input": inp.strip(), "expected": exp.strip().lower()})
    return cases


def make_scorer(model: str):
    from promptclimb import call_model

    def score(prompt: str, cases: list[dict]) -> float:
        correct = 0
        for c in cases:
            out = call_model(prompt + "\n\nText: " + c["input"], model).strip().lower()
            if c["expected"] in out:  # tolerant: label appears anywhere in output
                correct += 1
        return correct / len(cases) if cases else 0.0

    return score


def run_climb(provider, api_key, model_name, base_url, prompt_text, cases_text, iterations):
    cfg = PROVIDERS[provider]
    is_custom = provider.startswith("Custom")
    if not model_name.strip():
        yield "⚠️ Enter a model name.", ""
        return
    if not api_key.strip() and not is_custom:
        yield "⚠️ Paste your API key first (BYOK — used only for this run, never stored).", ""
        return
    cases = parse_cases(cases_text)[:MAX_CASES]
    if len(cases) < 2:
        yield "⚠️ Need at least 2 test cases. Format: `text => label` (one per line).", ""
        return
    iterations = max(1, min(int(iterations), MAX_ITERS))

    # BYOK: set this run's provider env from scratch so nothing leaks between runs.
    for k in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "OPENAI_BASE_URL"):
        os.environ.pop(k, None)
    os.environ[cfg["key_env"]] = api_key.strip() or "no-key-required"
    eff_base = base_url.strip() if is_custom else cfg["base_url"]
    if eff_base:
        os.environ["OPENAI_BASE_URL"] = eff_base
    model = f"{cfg['prefix']}:{model_name.strip()}"

    workdir = tempfile.mkdtemp()
    prompt_path = os.path.join(workdir, "prompt.txt")
    gold_dir = os.path.join(workdir, "gold")
    os.makedirs(gold_dir)
    with open(prompt_path, "w") as f:
        f.write(prompt_text)
    for i, c in enumerate(cases):
        with open(os.path.join(gold_dir, f"case_{i:02d}.json"), "w") as f:
            json.dump(c, f)

    from promptclimb.climber import HillClimber

    q: queue.Queue = queue.Queue()

    class QWriter(io.TextIOBase):
        def write(self, s):
            if s:
                q.put(s)
            return len(s)

    result: dict = {}

    def worker():
        try:
            with contextlib.redirect_stdout(QWriter()):
                hc = HillClimber(
                    prompt_path, make_scorer(model), gold_dir,
                    model=model, output_dir=os.path.join(workdir, "results"),
                )
                result["r"] = hc.run(max_iterations=iterations)
        except Exception as e:  # surface auth/quota/model errors to the visitor
            q.put(f"\n❌ {type(e).__name__}: {e}\n")
        finally:
            q.put(None)

    threading.Thread(target=worker, daemon=True).start()

    log = f"provider: {provider}  |  model: {model}\n"
    yield log, ""
    while True:
        item = q.get()
        if item is None:
            break
        log += item
        yield log, ""
    r = result.get("r")
    best = r.best_prompt if r else ""
    if r:
        log += f"\n✅ {r.start_score:.0%} → {r.best_score:.0%}  ({r.n_keeps} keeps / {r.n_iterations} iters)"
    yield log, best


with gr.Blocks(title="PromptClimb") as demo:
    gr.Markdown(
        "# 🧗 PromptClimb\n"
        "Automatically improve an LLM prompt by hill-climbing against a scored test set.\n\n"
        "**BYOK** — paste your own provider key. It's used only for your run and never stored. "
        f"Capped at {MAX_ITERS} iterations / {MAX_CASES} cases so a run costs cents. "
        "[Source on GitHub](https://github.com/treebird7/promptclimb).\n\n"
        "_Ollama / LM Studio are localhost-only — a hosted Space can't reach your machine. "
        "Run this Space locally, or expose a public OpenAI-compatible URL and use **Custom**._"
    )
    with gr.Row():
        provider = gr.Dropdown(list(PROVIDERS), value="OpenAI", label="Provider")
        api_key = gr.Textbox(label="API key (BYOK)", type="password", placeholder="sk-...")
        iterations = gr.Slider(1, MAX_ITERS, value=5, step=1, label="Max iterations")
    with gr.Row():
        model_name = gr.Textbox(label="Model", value="gpt-4o-mini")
        base_url = gr.Textbox(label="Base URL (Custom only)", visible=False,
                              placeholder="https://your-server/v1")
    prompt_text = gr.Textbox(label="Starting prompt", value=DEFAULT_PROMPT, lines=3)
    cases_text = gr.Textbox(
        label="Test cases  (text => label, one per line)", value=DEFAULT_CASES, lines=7
    )
    run_btn = gr.Button("Climb 🧗", variant="primary")
    log_out = gr.Textbox(label="Hill-climb log", lines=16)
    best_out = gr.Textbox(label="Best prompt found", lines=4)

    def on_provider(p):
        c = PROVIDERS[p]
        return (
            gr.update(value=c["model"]),
            gr.update(visible=p.startswith("Custom"), value=c["base_url"] or ""),
            gr.update(placeholder=c["hint"]),
        )

    provider.change(on_provider, provider, [model_name, base_url, api_key])

    # concurrency_limit=1: one run at a time so BYOK keys/base-urls never cross
    # between concurrent visitors via the shared process env.
    run_btn.click(
        run_climb,
        [provider, api_key, model_name, base_url, prompt_text, cases_text, iterations],
        [log_out, best_out],
        concurrency_limit=1,
    )

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