Add venice model to Brave Browser profile

add_venice_models.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "httpx>=0.27",
#   "cryptography>=42",
# ]
# ///
"""
Add all Venice.ai text models as custom models to Brave Leo across all profiles.

Usage:
  uv run add_venice_models.py [--api-key KEY] [--dry-run]

Brave must be closed before running, or the Preferences file will be overwritten
by the browser on exit.
"""

import argparse
import hashlib
import json
import shutil
import sys
from datetime import datetime
from pathlib import Path

VENICE_MODELS_URL = "https://api.venice.ai/api/v1/models?type=text"
VENICE_BASE_URL = "https://api.venice.ai/api/v1/chat/completions"
BRAVE_BASE = (
    Path.home() / "Library/Application Support/BraveSoftware/Brave-Browser-Nightly"
)

PROFILE_DIRS = [
    "Default",
    "Profile 1",
    "Profile 2",
    "Profile 3",
    "Profile 4",
    "Profile 5",
]


def make_key(model_id: str) -> str:
    """Generate a stable 8-char hex key matching Brave's custom:XXXXXXXX format."""
    return "custom:" + hashlib.md5(f"venice:{model_id}".encode()).hexdigest()[:8]


def _brave_encryption_key() -> bytes:
    """
    Derive the AES-128 key Brave uses for OSCrypt on macOS.
    Keychain stores the password as a base64 string; it is used AS-IS (raw string bytes).
    AES key = PBKDF2-HMAC-SHA1(password_bytes, salt=b'saltysalt', iterations=1003, dklen=16).
    """
    import subprocess
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    from cryptography.hazmat.primitives import hashes

    result = subprocess.run(
        [
            "security",
            "find-generic-password",
            "-s",
            "Brave Safe Storage",
            "-a",
            "Brave",
            "-w",
        ],
        capture_output=True,
        text=True,
        check=True,
    )
    # Use the raw string value from Keychain, not base64-decoded
    password = result.stdout.strip().encode()
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA1(), length=16, salt=b"saltysalt", iterations=1003
    )
    return kdf.derive(password)


def encrypt_api_key(plaintext: str) -> str:
    """
    Encrypt an API key the same way Brave's OSCrypt does on macOS:
      ciphertext = b'v10' + AES-128-CBC(key, iv=b' '*16, plaintext_pkcs7_padded)
    Returns base64-encoded result (as stored in Preferences).
    """
    import base64
    import os
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
    from cryptography.hazmat.primitives import padding

    key = _brave_encryption_key()
    iv = b" " * 16  # Chromium OSCrypt macOS uses a fixed IV of 16 spaces

    padder = padding.PKCS7(128).padder()
    padded = padder.update(plaintext.encode()) + padder.finalize()

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    enc = cipher.encryptor()
    ciphertext = enc.update(padded) + enc.finalize()

    return base64.b64encode(b"v10" + ciphertext).decode()


def fetch_venice_models(api_key: str | None) -> list[dict]:
    import httpx

    headers = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    resp = httpx.get(VENICE_MODELS_URL, headers=headers, timeout=15)
    resp.raise_for_status()
    data = resp.json()
    return [m for m in data.get("data", []) if m.get("type") == "text"]


import re as _re

# Tokens stripped from the right when computing a model family name
_STRIP_TOKEN = _re.compile(
    r"^("
    r"\d{4,8}"  # date stamps / long numeric versions
    r"|\d+\.\d+(\.\d+)*"  # dotted versions: 3.3 / 4.5 / 4.6
    r"|\d+"  # bare integers: 4, 5, 54, 70, 235
    r"|a\d+b?"  # expert counts: a22b, a3b, a10b
    r"|\d+b"  # param counts: 70b, 405b, 9b
    r"|fp\d+|int\d+|bf\d+"  # quantisation
    r"|fast|turbo|mini|pro|plus|flash|preview|nano|large|small|medium|big"
    r"|thinking|reasoning|instruct|chat|it|next"
    r"|multi.agent|codex|coder|role.play|uncensored|heretic"
    r"|cascade"
    r"|p"  # single-letter e2ee suffix
    r")$",
    _re.IGNORECASE,
)


def model_family(model_id: str) -> str:
    """
    Canonical family name: strip the e2ee- prefix first, then strip all
    trailing version/size/qualifier tokens.

    Examples:
      e2ee-glm-4-7-p        → glm
      e2ee-qwen3-30b-a3b-p  → qwen3
      zai-org-glm-5-1       → zai-org-glm  (different namespace, kept separate)
      claude-opus-4-6-fast  → claude-opus
      minimax-m27           → minimax-m27  (opaque name, no strippable suffix)
      qwen3-5-9b            → qwen3
    """
    mid = model_id
    # Normalise: treat e2ee-* and plain counterpart as the same family
    if mid.startswith("e2ee-"):
        mid = mid[5:]  # strip "e2ee-"

    parts = mid.split("-")
    while len(parts) > 1 and _STRIP_TOKEN.match(parts[-1]):
        parts.pop()

    return "-".join(parts) if parts else model_id


def _version_tuple(model_id: str) -> tuple:
    """
    Extract a comparable version tuple from a model ID for tiebreaking.
    Extracts all numeric sequences found in the ID (after stripping e2ee- prefix).
    Higher numbers = newer version.
    e2ee-glm-5 → (5,)  >  e2ee-glm-4-7-flash → (4, 7)
    openai-gpt-4o → (4,)  >  openai-gpt-4o-mini → (4,) [same, falls to shorter = base wins]
    """
    mid = model_id
    if mid.startswith("e2ee-"):
        mid = mid[5:]
    nums = [int(x) for x in _re.findall(r"\d+", mid)]
    return tuple(nums)


def model_sort_key(model: dict) -> tuple:
    """
    Comparison key for picking the best model within a family.
    Priority (descending):
      1. e2ee variant preferred over non-e2ee (privacy-first)
      2. Higher `created` timestamp
      3. Version tuple (numerically higher = newer)
      4. Shorter ID as tiebreaker (base model > qualified variant, e.g. gpt-4o > gpt-4o-mini)
    """
    mid = model["id"]
    is_e2ee = 1 if mid.startswith("e2ee-") else 0
    return (is_e2ee, model.get("created", 0), _version_tuple(mid), -len(mid))


def deduplicate_models(models: list[dict]) -> list[dict]:
    """
    Per family: keep the single best model (e2ee > newest > longest id).
    Additionally collapse minimax-mXX variants to keep only the newest.
    """

    # --- minimax special case: treat all minimax-m* as one family ---
    def canonical_family(m: dict) -> str:
        fam = model_family(m["id"])
        if fam.startswith("minimax-"):
            return "minimax"
        return fam

    best: dict[str, dict] = {}
    for m in models:
        fam = canonical_family(m)
        if fam not in best or model_sort_key(m) > model_sort_key(best[fam]):
            best[fam] = m

    # Return in original API order (stable); emit each winner exactly once
    seen: set[str] = set()
    result = []
    for m in models:
        fam = canonical_family(m)
        if fam not in seen and best[fam]["id"] == m["id"]:
            seen.add(fam)
            result.append(m)
    return result


def build_custom_model(model: dict, api_key: str, encrypted_key: str) -> dict:
    spec = model.get("model_spec", {})
    ctx = spec.get("availableContextTokens", 4096)
    vision = spec.get("capabilities", {}).get("supportsVision", False)
    tools = spec.get("capabilities", {}).get("supportsFunctionCalling", False)
    model_id = model["id"]
    name = spec.get("name", model_id)

    # Strip e2ee- prefix from label (kept in model_request_name for the API call)
    display_name = _re.sub(r"^e2ee[\s\-]", "", name, flags=_re.IGNORECASE)

    return {
        "api_key": encrypted_key,
        "context_size": ctx,
        "endpoint_url": VENICE_BASE_URL,
        "key": make_key(model_id),
        "label": f"Venice: {display_name}",
        "model_request_name": model_id,
        "supports_tools": tools,
        "vision_support": vision,
    }


def is_venice_model(m: dict) -> bool:
    """Detect any previously injected Venice model regardless of label format."""
    return (
        m.get("label", "").startswith("Venice:")
        or m.get("endpoint_url", "") == VENICE_BASE_URL
    )


def update_profile(prefs_path: Path, new_models: list[dict], dry_run: bool) -> bool:
    if not prefs_path.exists():
        return False

    with open(prefs_path) as f:
        prefs = json.load(f)

    ai_chat = prefs.setdefault("brave", {}).setdefault("ai_chat", {})
    existing: list[dict] = ai_chat.get("custom_models", [])

    # Drop ALL previously injected Venice models (stale, old dedup set, duplicates)
    non_venice = [m for m in existing if not is_venice_model(m)]
    removed = len(existing) - len(non_venice)

    final = non_venice + new_models
    ai_chat["custom_models"] = final

    if dry_run:
        print(
            f"  [dry-run] {removed} venice models removed, "
            f"{len(new_models)} added, "
            f"{len(non_venice)} non-venice kept → {len(final)} total"
        )
        return True

    # Backup
    backup = prefs_path.with_suffix(f".bak.{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    shutil.copy2(prefs_path, backup)

    with open(prefs_path, "w") as f:
        json.dump(prefs, f, separators=(",", ":"))

    print(
        f"  {removed} removed, {len(new_models)} added, "
        f"{len(non_venice)} non-venice kept → {len(final)} total. Backup: {backup.name}"
    )
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Inject Venice.ai models into Brave Leo profiles"
    )
    parser.add_argument(
        "--api-key", default="", help="Venice API key (stored in Preferences)"
    )
    parser.add_argument(
        "--dry-run", action="store_true", help="Preview changes without writing"
    )
    args = parser.parse_args()

    print("Fetching Venice.ai text models...")
    try:
        models = fetch_venice_models(args.api_key or None)
    except Exception as e:
        print(f"ERROR fetching models: {e}", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(models)} text models from Venice.ai")

    deduped = deduplicate_models(models)
    print(f"After dedup (latest per family): {len(deduped)} models")
    if args.dry_run:
        for m in deduped:
            fam = model_family(m["id"])
            print(f"  {m['id']:50s}  (family: {fam})")

    if args.api_key and not args.dry_run:
        print("Encrypting API key via Brave Safe Storage...")
        try:
            encrypted_key = encrypt_api_key(args.api_key)
        except Exception as e:
            print(
                f"WARNING: could not encrypt API key ({e}), storing plaintext",
                file=sys.stderr,
            )
            encrypted_key = args.api_key
    else:
        encrypted_key = args.api_key  # dry-run: plaintext is fine for preview

    custom_models = [
        build_custom_model(m, args.api_key, encrypted_key) for m in deduped
    ]

    updated = 0
    for profile in PROFILE_DIRS:
        prefs_path = BRAVE_BASE / profile / "Preferences"
        if not prefs_path.exists():
            continue
        print(f"\nProfile: {profile}")
        if update_profile(prefs_path, custom_models, args.dry_run):
            updated += 1

    print(f"\nDone. Updated {updated} profile(s).")
    if not args.dry_run:
        print("NOTE: Make sure Brave is closed before running this script,")
        print("      otherwise Brave will overwrite the Preferences on exit.")


if __name__ == "__main__":
    main()