hendrik/lib/moop.py

#  moop.py — Model Option (moop).
#
#  Penyimpanan model/API key/provider dipindah dari config.yaml ke database
#  sqlite. Modul ini bertanggung jawab atas:
#    * inisialisasi schema (referensi: plan/moop.sql)
#    * CRUD model set, provider (moop_api), key, dan model
#    * default model set per tipe (llm/embedding/imagegen/imagevision)
#    * resolve "chain" kandidat (base_url, model, api_key) untuk auto-switch
#
#  Prioritas auto-switch di dalam satu model set:
#    1. Key  (urut berdasarkan priority, NULL = disabled)
#    2. Setelah semua key gagal baru pindah Model/Provider berikutnya.
#  Setiap kandidat dalam chain membawa (api_id, model_id) agar pemanggil
#  bisa membedakan rotasi key vs pindah provider.

import os
import sqlite3
import uuid

import config

MODEL_TYPES = ("llm", "embedding", "imagegen", "imagevision")

_SCHEMA = """
PRAGMA foreign_keys = ON;

CREATE TABLE IF NOT EXISTS moop_model_set (
    id      BLOB PRIMARY KEY NOT NULL,
    name    TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS moop_model_type (
    id      BLOB PRIMARY KEY NOT NULL,
    type    TEXT NOT NULL UNIQUE
);

CREATE TABLE IF NOT EXISTS moop_model_set_type (
    id      BLOB PRIMARY KEY NOT NULL,
    "set"   BLOB NOT NULL,
    type    BLOB NOT NULL,
    FOREIGN KEY ("set")  REFERENCES moop_model_set(id)   ON UPDATE CASCADE ON DELETE CASCADE,
    FOREIGN KEY (type)   REFERENCES moop_model_type(id)  ON UPDATE CASCADE ON DELETE CASCADE,
    UNIQUE ("set", type)
);

CREATE TABLE IF NOT EXISTS moop_model_type_default (
    id          BLOB PRIMARY KEY NOT NULL,
    type        BLOB NOT NULL UNIQUE,
    set_type    BLOB NOT NULL UNIQUE,
    FOREIGN KEY (type)      REFERENCES moop_model_type(id)      ON UPDATE CASCADE ON DELETE CASCADE,
    FOREIGN KEY (set_type)  REFERENCES moop_model_set_type(id)  ON UPDATE CASCADE ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS moop_api (
    id          BLOB PRIMARY KEY NOT NULL,
    name        TEXT NOT NULL,
    baseurl     TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS moop_key (
    id          BLOB PRIMARY KEY NOT NULL,
    key         TEXT NOT NULL,
    priority    INTEGER NULL,
    api         BLOB NOT NULL,
    FOREIGN KEY (api) REFERENCES moop_api(id) ON UPDATE CASCADE ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS moop_model (
    id          BLOB PRIMARY KEY NOT NULL,
    api         BLOB NOT NULL,
    "set"       BLOB NOT NULL,
    model       TEXT NOT NULL,
    priority    INTEGER NULL,
    FOREIGN KEY (api)    REFERENCES moop_api(id)        ON UPDATE CASCADE ON DELETE CASCADE,
    FOREIGN KEY ("set")  REFERENCES moop_model_set(id)  ON UPDATE CASCADE ON DELETE CASCADE
);

CREATE INDEX IF NOT EXISTS idx_moop_key_api              ON moop_key(api);
CREATE INDEX IF NOT EXISTS idx_moop_key_api_priority     ON moop_key(api, priority);

CREATE INDEX IF NOT EXISTS idx_moop_model_api            ON moop_model(api);
CREATE INDEX IF NOT EXISTS idx_moop_model_set            ON moop_model("set");
CREATE INDEX IF NOT EXISTS idx_moop_model_api_priority   ON moop_model(api, priority);

CREATE INDEX IF NOT EXISTS idx_moop_model_set_type_set   ON moop_model_set_type("set");
CREATE INDEX IF NOT EXISTS idx_moop_model_set_type_type  ON moop_model_set_type(type);
"""


def moop_db_path():
    return os.path.expanduser(config.moop_db_path)


def _id():
    return uuid.uuid4().bytes


def _connect(path=None):
    p = os.path.expanduser(path or moop_db_path())
    parent = os.path.dirname(p)
    if parent:
        os.makedirs(parent, exist_ok=True)
    conn = sqlite3.connect(p, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    try:
        # Self-healing: buat schema + seed tipe kalau DB belum pernah dibuat.
        if not conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='moop_model_type'"
        ).fetchone():
            conn.executescript(_SCHEMA)
            for t in MODEL_TYPES:
                conn.execute(
                    "INSERT INTO moop_model_type (id, type) VALUES (?, ?)",
                    (_id(), t),
                )
            conn.commit()
    except sqlite3.OperationalError:
        pass
    return conn


# ---------------------------------------------------------------- init ---

def init_db(path=None):
    """Buat schema kalau belum ada, lalu seed tipe model yang didukung."""
    conn = _connect(path)
    try:
        conn.executescript(_SCHEMA)
        for t in MODEL_TYPES:
            conn.execute(
                "INSERT OR IGNORE INTO moop_model_type (id, type) VALUES (?, ?)",
                (_id(), t),
            )
        conn.commit()
    finally:
        conn.close()


def ensure_defaults(path=None):
    """Untuk tipe yang belum punya default set, ambil set+type pertama."""
    conn = _connect(path)
    try:
        _ensure_defaults_conn(conn)
        conn.commit()
    finally:
        conn.close()


def _ensure_defaults_conn(conn):
    types = conn.execute("SELECT id FROM moop_model_type").fetchall()
    for t in types:
        tid = t["id"]
        has = conn.execute(
            "SELECT id FROM moop_model_type_default WHERE type = ?", (tid,)
        ).fetchone()
        if has:
            continue
        st = conn.execute(
            'SELECT st.id FROM moop_model_set_type st '
            'WHERE st.type = ? ORDER BY st.rowid LIMIT 1',
            (tid,),
        ).fetchone()
        if st:
            conn.execute(
                "INSERT OR IGNORE INTO moop_model_type_default (id, type, set_type) "
                "VALUES (?, ?, ?)",
                (_id(), tid, st["id"]),
            )


# ---------------------------------------------------------------- sets ---

def _set_types_conn(conn, set_id):
    return [
        r["type"]
        for r in conn.execute(
            'SELECT mt.type AS type FROM moop_model_set_type st '
            'JOIN moop_model_type mt ON mt.id = st.type '
            'WHERE st."set" = ? ORDER BY mt.type',
            (set_id,),
        ).fetchall()
    ]


def _set_defaults_conn(conn, set_id):
    return [
        r["type_name"]
        for r in conn.execute(
            'SELECT mt.type AS type_name FROM moop_model_type_default d '
            'JOIN moop_model_type mt ON mt.id = d.type '
            'JOIN moop_model_set_type st ON st.id = d.set_type '
            'WHERE st."set" = ?',
            (set_id,),
        ).fetchall()
    ]


def list_sets(path=None):
    conn = _connect(path)
    try:
        result = []
        for r in conn.execute(
            "SELECT id, name FROM moop_model_set ORDER BY name COLLATE NOCASE"
        ).fetchall():
            set_id = r["id"]
            result.append({
                "id": set_id,
                "name": r["name"],
                "types": _set_types_conn(conn, set_id),
                "default_for": _set_defaults_conn(conn, set_id),
            })
        return result
    finally:
        conn.close()


def get_set(path=None, set_id=None):
    if set_id is None:
        return None
    conn = _connect(path)
    try:
        r = conn.execute(
            "SELECT id, name FROM moop_model_set WHERE id = ?", (set_id,)
        ).fetchone()
        if not r:
            return None
        return {
            "id": r["id"],
            "name": r["name"],
            "types": _set_types_conn(conn, set_id),
            "default_for": _set_defaults_conn(conn, set_id),
        }
    finally:
        conn.close()


def create_set(path=None, name=""):
    conn = _connect(path)
    try:
        set_id = _id()
        conn.execute(
            "INSERT INTO moop_model_set (id, name) VALUES (?, ?)",
            (set_id, (name or "").strip()),
        )
        conn.commit()
        return set_id
    finally:
        conn.close()


def rename_set(path=None, set_id=None, name=""):
    if set_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            "UPDATE moop_model_set SET name = ? WHERE id = ?",
            ((name or "").strip(), set_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_set(path=None, set_id=None):
    if set_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute("DELETE FROM moop_model_set WHERE id = ?", (set_id,))
        conn.commit()
        _ensure_defaults_conn(conn)
        conn.commit()
    finally:
        conn.close()


# --------------------------------------------------------------- types ---

def _type_id(conn, type_name):
    r = conn.execute(
        "SELECT id FROM moop_model_type WHERE type = ?", (type_name,)
    ).fetchone()
    if r:
        return r["id"]
    tid = _id()
    conn.execute(
        "INSERT INTO moop_model_type (id, type) VALUES (?, ?)", (tid, type_name)
    )
    return tid


def add_type(path=None, set_id=None, type_name=""):
    if set_id is None:
        return
    conn = _connect(path)
    try:
        tid = _type_id(conn, type_name)
        exists = conn.execute(
            'SELECT id FROM moop_model_set_type WHERE "set" = ? AND type = ?',
            (set_id, tid),
        ).fetchone()
        if not exists:
            conn.execute(
                'INSERT INTO moop_model_set_type (id, "set", type) VALUES (?, ?, ?)',
                (_id(), set_id, tid),
            )
        conn.commit()
        _ensure_defaults_conn(conn)
        conn.commit()
    finally:
        conn.close()


def remove_type(path=None, set_id=None, type_name=""):
    if set_id is None:
        return
    conn = _connect(path)
    try:
        tid = conn.execute(
            "SELECT id FROM moop_model_type WHERE type = ?", (type_name,)
        ).fetchone()
        if not tid:
            return
        conn.execute(
            'DELETE FROM moop_model_set_type WHERE "set" = ? AND type = ?',
            (set_id, tid["id"]),
        )
        conn.commit()
        _ensure_defaults_conn(conn)
        conn.commit()
    finally:
        conn.close()


# -------------------------------------------------------------- default ---

def set_default(path=None, type_name="", set_id=None):
    """Set `set_id` sebagai default untuk `type_name`. Jika set belum punya
    type tersebut, type otomatis ditambahkan."""
    if set_id is None or not type_name:
        return
    conn = _connect(path)
    try:
        tid = _type_id(conn, type_name)
        st = conn.execute(
            'SELECT id FROM moop_model_set_type WHERE "set" = ? AND type = ?',
            (set_id, tid),
        ).fetchone()
        if not st:
            st_id = _id()
            conn.execute(
                'INSERT INTO moop_model_set_type (id, "set", type) VALUES (?, ?, ?)',
                (st_id, set_id, tid),
            )
        else:
            st_id = st["id"]
        conn.execute("DELETE FROM moop_model_type_default WHERE type = ?", (tid,))
        conn.execute(
            "INSERT INTO moop_model_type_default (id, type, set_type) VALUES (?, ?, ?)",
            (_id(), tid, st_id),
        )
        conn.commit()
    finally:
        conn.close()


def get_default_set_id(path=None, type_name=""):
    conn = _connect(path)
    try:
        r = conn.execute(
            'SELECT s.id AS id FROM moop_model_type_default d '
            'JOIN moop_model_type mt ON mt.id = d.type '
            'JOIN moop_model_set_type st ON st.id = d.set_type '
            'JOIN moop_model_set s ON s.id = st."set" '
            'WHERE mt.type = ?',
            (type_name,),
        ).fetchone()
        return r["id"] if r else None
    finally:
        conn.close()


# -------------------------------------------------------------- provider ---

def list_apis(path=None):
    conn = _connect(path)
    try:
        rows = conn.execute(
            "SELECT a.id, a.name, a.baseurl, "
            "(SELECT COUNT(*) FROM moop_key k WHERE k.api = a.id) AS key_count, "
            "(SELECT COUNT(*) FROM moop_model m WHERE m.api = a.id) AS model_count "
            "FROM moop_api a ORDER BY a.name COLLATE NOCASE"
        ).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def get_api(path=None, api_id=None):
    if api_id is None:
        return None
    conn = _connect(path)
    try:
        r = conn.execute(
            "SELECT id, name, baseurl FROM moop_api WHERE id = ?", (api_id,)
        ).fetchone()
        return dict(r) if r else None
    finally:
        conn.close()


def create_api(path=None, name="", baseurl=""):
    conn = _connect(path)
    try:
        api_id = _id()
        conn.execute(
            "INSERT INTO moop_api (id, name, baseurl) VALUES (?, ?, ?)",
            (api_id, (name or "").strip(), (baseurl or "").strip()),
        )
        conn.commit()
        return api_id
    finally:
        conn.close()


def update_api(path=None, api_id=None, name="", baseurl=""):
    if api_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            "UPDATE moop_api SET name = ?, baseurl = ? WHERE id = ?",
            ((name or "").strip(), (baseurl or "").strip(), api_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_api(path=None, api_id=None):
    if api_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute("DELETE FROM moop_api WHERE id = ?", (api_id,))
        conn.commit()
    finally:
        conn.close()


# ------------------------------------------------------------------ keys ---

def list_keys(path=None, api_id=None):
    conn = _connect(path)
    try:
        rows = conn.execute(
            "SELECT id, key, priority FROM moop_key WHERE api = ? "
            "ORDER BY priority IS NULL ASC, priority ASC, rowid ASC",
            (api_id,),
        ).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def add_key(path=None, api_id=None, key="", priority=None):
    if api_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            "INSERT INTO moop_key (id, key, priority, api) VALUES (?, ?, ?, ?)",
            (_id(), key or "", priority, api_id),
        )
        conn.commit()
    finally:
        conn.close()


def update_key(path=None, key_id=None, key="", priority=None):
    if key_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            "UPDATE moop_key SET key = ?, priority = ? WHERE id = ?",
            (key or "", priority, key_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_key(path=None, key_id=None):
    if key_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute("DELETE FROM moop_key WHERE id = ?", (key_id,))
        conn.commit()
    finally:
        conn.close()


# ----------------------------------------------------------------- models ---

def list_models(path=None, set_id=None):
    conn = _connect(path)
    try:
        rows = conn.execute(
            'SELECT m.id AS model_id, m.model, m.priority, '
            'a.id AS api_id, a.name AS provider, a.baseurl AS baseurl '
            'FROM moop_model m JOIN moop_api a ON a.id = m.api '
            'WHERE m."set" = ? '
            'ORDER BY m.priority IS NULL ASC, m.priority ASC, m.rowid ASC',
            (set_id,),
        ).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def add_model(path=None, set_id=None, api_id=None, model="", priority=None):
    if set_id is None or api_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            'INSERT INTO moop_model (id, api, "set", model, priority) '
            "VALUES (?, ?, ?, ?, ?)",
            (_id(), api_id, set_id, model or "", priority),
        )
        conn.commit()
    finally:
        conn.close()


def update_model(path=None, model_id=None, model="", priority=None):
    if model_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute(
            "UPDATE moop_model SET model = ?, priority = ? WHERE id = ?",
            (model or "", priority, model_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_model(path=None, model_id=None):
    if model_id is None:
        return
    conn = _connect(path)
    try:
        conn.execute("DELETE FROM moop_model WHERE id = ?", (model_id,))
        conn.commit()
    finally:
        conn.close()


# ----------------------------------------------------------------- chain ---

def resolve_chain(path=None, set_id=None):
    """Bangun chain kandidat untuk satu model set.

    Urutan: model diurutkan berdasarkan priority (NULL = disabled), lalu untuk
    tiap model, semua key provider-nya (priority, NULL = disabled). Provider
    tanpa key tetap jadi 1 kandidat dengan api_key kosong (= no bearer).
    Setiap kandidat membawa (api_id, model_id) untuk membedakan rotasi key."""
    if set_id is None:
        return []
    conn = _connect(path)
    try:
        models = conn.execute(
            'SELECT m.id AS model_id, m.model, '
            'a.id AS api_id, a.name AS provider, a.baseurl AS base_url '
            'FROM moop_model m JOIN moop_api a ON a.id = m.api '
            'WHERE m."set" = ? AND m.priority IS NOT NULL '
            'ORDER BY m.priority ASC, m.rowid ASC',
            (set_id,),
        ).fetchall()
        candidates = []
        for m in models:
            keys = conn.execute(
                "SELECT key FROM moop_key WHERE api = ? AND priority IS NOT NULL "
                "ORDER BY priority ASC, rowid ASC",
                (m["api_id"],),
            ).fetchall()
            if keys:
                for k in keys:
                    candidates.append({
                        "base_url": m["base_url"],
                        "model": m["model"],
                        "api_key": k["key"],
                        "api_id": m["api_id"],
                        "model_id": m["model_id"],
                        "provider": m["provider"],
                    })
            else:
                candidates.append({
                    "base_url": m["base_url"],
                    "model": m["model"],
                    "api_key": "",
                    "api_id": m["api_id"],
                    "model_id": m["model_id"],
                    "provider": m["provider"],
                })
        return candidates
    finally:
        conn.close()


def default_chain(path=None, type_name=""):
    """Chain untuk default model set dari tipe tertentu."""
    set_id = get_default_set_id(path, type_name)
    if not set_id:
        return []
    return resolve_chain(path, set_id)


def resolve_provider(path=None, base_url="", model=""):
    """Cari nama provider untuk (base_url, model) di chain default semua tipe."""
    base_url = (base_url or "").rstrip("/")
    for t in MODEL_TYPES:
        for cand in default_chain(path, t):
            if (cand["base_url"] or "").rstrip("/") == base_url and cand["model"] == model:
                return cand["provider"]
    return None


def configure_client(llm_client, path=None):
    """Isi chain llm + imagevision ke LLMClient berdasarkan default set.
    Mengembalikan True jika chain default 'llm' tersedia."""
    p = os.path.expanduser(path or moop_db_path())
    init_db(p)
    chains = {
        "llm": default_chain(p, "llm"),
        "imagevision": default_chain(p, "imagevision"),
    }
    llm_client.set_chains(chains)
    return bool(chains["llm"])


def first_endpoint(path=None, type_name=""):
    """Kandidat pertama (base_url, model, api_key) dari chain default tipe.
    Mengembalikan None jika tidak ada."""
    chain = default_chain(path, type_name)
    if not chain:
        return None
    c = chain[0]
    return c["base_url"], c["model"], c.get("api_key")


def embedding_endpoint(path=None):
    """Endpoint embedding (url, model, api_key) dari default chain 'embedding'."""
    return first_endpoint(path, "embedding")


def imagegen_endpoint(path=None):
    """Endpoint image generation (url, model, api_key) dari default chain 'imagegen'."""
    return first_endpoint(path, "imagegen")