photoanalyzer/compare_models.py

#!/usr/bin/env python3
"""
Side-by-side model comparison for photo_analyzer.

Reads already-analyzed photos from the DB (the stored result — no re-run, no
extra cost on the model that produced them) and runs a second model (default:
Gemini 2.5 Flash via its OpenAI-compatible endpoint) live on the same images,
then prints the two results next to each other plus measured token cost.

Reuses photo_analyzer's prepare_image, ANALYSIS_PROMPT, get_db and env loader —
no logic is duplicated.

Usage:
    GEMINI_API_KEY=... python compare_models.py -n 10
    python compare_models.py --selftest          # offline check, no API call

Put GEMINI_API_KEY in photo_analyzer.env (the same file the main pipeline reads).
"""
import argparse
import json
import os
import random
import sys
from collections import defaultdict
from pathlib import Path

from openai import OpenAI

import photo_analyzer as pa

# Gemini's OpenAI-compatible endpoint — the existing OpenAI client just works.
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
DEFAULT_MODEL   = "gemini-2.5-flash"

# $/million tokens, for the live model's measured cost. Update if you change model.
PRICES = {  # (input, output)
    "gemini-2.5-flash":      (0.30, 2.50),
    "gemini-2.5-flash-lite": (0.10, 0.40),
    "gemini-3.5-flash":      (1.50, 9.00),
}


def _strip_json_fences(raw: str) -> str:
    """Same fence-stripping photo_analyzer uses — models often wrap JSON in ```."""
    raw = raw.strip()
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
        raw = raw.strip()
    return raw


def analyze_with(client: OpenAI, model: str, path: Path) -> tuple[dict, dict]:
    """Run `model` on one image. Returns (parsed_result, usage_dict)."""
    b64, mime = pa.prepare_image(path)
    resp = client.chat.completions.create(
        model=model,
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}},
                {"type": "text", "text": pa.ANALYSIS_PROMPT},
            ],
        }],
    )
    result = json.loads(_strip_json_fences(resp.choices[0].message.content))
    u = resp.usage
    usage = {"prompt": u.prompt_tokens if u else 0,
             "completion": u.completion_tokens if u else 0}
    return result, usage


def fetch_analyzed(conn, limit: int) -> list:
    """
    Randomly sample already-analyzed photos, spread across folders. Groups rows by
    parent directory, shuffles, then round-robins one per folder so the sample
    covers as many different folders as possible. Only picked files are checked on
    disk (cheap), not the whole library.
    """
    # "analyzed already" = has a description, regardless of status — once EXIF is
    # written the status flips from 'analyzed' to 'exif_written', but the result stays.
    rows = conn.execute(
        """SELECT path, description, tags, people_count, mood, setting, approx_year
           FROM photos WHERE description IS NOT NULL AND description != ''""").fetchall()
    buckets = defaultdict(list)
    for r in rows:
        buckets[str(Path(r["path"]).parent)].append(r)
    folders = list(buckets.keys())
    random.shuffle(folders)
    for f in folders:
        random.shuffle(buckets[f])

    out, exhausted = [], False
    while len(out) < limit and not exhausted:
        exhausted = True
        for f in folders:                       # one per folder per pass → spread
            if not buckets[f]:
                continue
            exhausted = False
            r = buckets[f].pop()
            if Path(r["path"]).exists():
                out.append(r)
                if len(out) >= limit:
                    break
    return out


def _fmt(label: str, desc, tags, people, mood) -> str:
    tags_s = ", ".join(tags) if isinstance(tags, list) else (tags or "")
    return (f"  {label}\n"
            f"    desc:   {desc}\n"
            f"    tags:   {tags_s}\n"
            f"    people: {people}   mood: {mood}")


def selftest():
    assert _strip_json_fences('```json\n{"a":1}\n```') == '{"a":1}'
    assert _strip_json_fences('{"a":1}') == '{"a":1}'
    assert _strip_json_fences('```\n{"b":2}\n```') == '{"b":2}'
    print("selftest OK")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("-n", "--num", type=int, default=10, help="photos to compare")
    ap.add_argument("--model", default=DEFAULT_MODEL, help="live model to compare against the stored results")
    ap.add_argument("--db", default="photo_analysis.db")
    ap.add_argument("--selftest", action="store_true", help="run offline check and exit")
    args = ap.parse_args()

    if args.selftest:
        selftest()
        return

    pa.load_env_file()  # picks up GEMINI_API_KEY from photo_analyzer.env
    key = os.environ.get("GEMINI_API_KEY")
    if not key:
        sys.exit("GEMINI_API_KEY not set (add it to photo_analyzer.env or export it).")

    conn = pa.get_db(args.db)
    rows = fetch_analyzed(conn, args.num)
    if not rows:
        sys.exit("No analyzed photos found in the DB to compare against.")

    client = OpenAI(api_key=key, base_url=GEMINI_BASE_URL)
    p_in, p_out = PRICES.get(args.model, (0.0, 0.0))
    tot_in = tot_out = 0
    ok = 0

    for r in rows:
        path = Path(r["path"])
        print(f"\n=== {path.parent.name}/{path.name} ===")
        # stored result (already in the DB — no re-run)
        try:
            stored_tags = json.loads(r["tags"]) if r["tags"] else []
        except Exception:
            stored_tags = r["tags"]
        print(_fmt("stored (DB)", r["description"], stored_tags,
                   r["people_count"], r["mood"]))
        # live call to the comparison model
        try:
            res, usage = analyze_with(client, args.model, path)
            tot_in += usage["prompt"]; tot_out += usage["completion"]; ok += 1
            print(_fmt(f"{args.model} (live)", res.get("description"),
                       res.get("tags", []), res.get("people_count"), res.get("mood")))
        except Exception as e:
            print(f"  {args.model} (live)\n    ERROR: {e}")

    if ok:
        cost = tot_in / 1e6 * p_in + tot_out / 1e6 * p_out
        print(f"\n— {args.model}: {ok} photos, {tot_in:,} in + {tot_out:,} out tokens, "
              f"measured cost ${cost:.4f} (${cost / ok * 1000:.2f}/1k photos)")


if __name__ == "__main__":
    main()