Autoparts-DB/scripts/enrich_descriptions.py

#!/usr/bin/env python3
"""Generate short Spanish descriptions for empty catalog descriptions.

Two targets:

  inventory (tenant DB)      — 95% of rows have ``description`` empty
                               (Rached: 6851/6851; Capi: 3766/14568).
  supplier_catalog (master)  — 17,246/17,246 have no description at all.

Description format (documented, enforced by the prompt):

  * one short sentence, 8–20 words, plain Mexican Spanish;
  * starts with the part type, then the brand/compatibility qualifier already
    present in the name (never invent a vehicle, price, stock or warranty);
  * no emojis, no markdown, no quotes, no trailing period duplication.

The LLM call batches ``--batch-llm`` items per request through the same
``ai_chat`` backends (QWEN first, OpenRouter free models as fallback).  When
the provider does not answer, the item is skipped — nothing is written and no
exception is raised.

Quality gate: nothing is written unless ``--apply``.  Pass ``--review-out
FILE`` to dump every generated candidate as JSONL so a human can review the
batch before applying.  A checkpoint (last processed id per target) lets long
runs resume.

Usage (as root, from /home/Autopartes):
    python3 scripts/enrich_descriptions.py --tenant 31 --target inventory --limit 50
    python3 scripts/enrich_descriptions.py --tenant 31 --target inventory \
        --limit 200 --review-out /tmp/desc_rached.jsonl
    python3 scripts/enrich_descriptions.py --target supplier_catalog --limit 100 --dry-run
    python3 scripts/enrich_descriptions.py --tenant 31 --target inventory --apply --limit 200
"""

import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'pos'))

from tenant_db import get_tenant_conn, get_master_conn  # noqa: E402

DEFAULT_CHECKPOINT = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), '..', 'backups',
    'enrich_descriptions_checkpoint.json'
)

SYSTEM_PROMPT = (
    "Eres redactor de catálogo de una refaccionaria mexicana. Escribes "
    "descripciones de autopartes claras y breves en español. NUNCA inventas "
    "compatibilidad de vehículos, precios, garantías ni stock: solo describes "
    "lo que se te da. Responde SOLO con JSON válido, sin texto adicional."
)

USER_PROMPT = (
    "Para cada artículo genera una descripción de 8 a 20 palabras en español "
    "mexicano que empiece con el tipo de pieza. Prohibido: emojis, markdown, "
    "precios, garantías, años de vehículo que no aparezcan en el nombre.\n"
    "Devuelve un arreglo JSON: "
    '[{"id": <id>, "descripcion": "<texto>"}, ...]\n'
    "Artículos:\n__ITEMS__"
)


def _load_checkpoint(path):
    try:
        with open(path) as f:
            return json.load(f)
    except Exception:
        return {}


def _save_checkpoint(path, data):
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, 'w') as f:
            json.dump(data, f, indent=2)
    except Exception as e:
        print(f"  ! checkpoint no guardado: {e}")


def _llm_descriptions(items, timeout=60):
    """Return {id: description} via ai_chat backends; {} on any failure."""
    try:
        from services import ai_chat
    except Exception:
        return {}
    payload = json.dumps(
        [{"id": it["id"], "nombre": it["label"]} for it in items],
        ensure_ascii=False,
    )
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT.replace("__ITEMS__", payload)},
    ]
    backends = []
    if getattr(ai_chat, "QWEN_ENABLED", False):
        backends.append((ai_chat.QWEN_CHAT_URL, ai_chat.QWEN_API_KEY, ai_chat.QWEN_MODEL))
    if getattr(ai_chat, "OPENROUTER_API_KEY", None):
        backends.append((ai_chat.OPENROUTER_URL, ai_chat.OPENROUTER_API_KEY, ai_chat.MODEL))
    for url, key, model_id in backends:
        try:
            ai_chat._validate_model(model_id)
        except Exception:
            continue
        try:
            res = ai_chat._post_chat_completion(
                url, key, model_id, messages,
                max_tokens=1400, temperature=0.2, timeout=timeout)
        except Exception:
            res = None
        if not res:
            continue
        content = (res.get("content") or "").strip()
        if content.startswith("```"):
            content = "\n".join(content.split("\n")[1:-1])
        try:
            parsed = json.loads(content)
        except Exception:
            continue
        if isinstance(parsed, dict):
            parsed = parsed.get("items") or parsed.get("data") or []
        out = {}
        for row in parsed or []:
            if not isinstance(row, dict):
                continue
            rid = row.get("id")
            text = (row.get("descripcion") or row.get("description") or "").strip()
            if rid is not None and text:
                out[int(rid)] = text[:300]
        if out:
            return out
    return {}


def _iter_inventory(conn, after_id, limit, batch):
    cur = conn.cursor()
    where = ["i.is_active = true",
             "(i.description IS NULL OR i.description = '')"]
    params = []
    if after_id is not None:
        where.append("i.id > %s")
        params.append(after_id)
    cur.execute(f"""
        SELECT i.id, i.name, i.brand, c.name
        FROM inventory i
        LEFT JOIN categories c ON c.id = i.category_id
        WHERE {' AND '.join(where)}
        ORDER BY i.id
        LIMIT %s
    """, params + [batch])
    rows = cur.fetchall()
    cur.close()
    return rows


def _iter_supplier(conn, after_id, batch):
    cur = conn.cursor()
    where = ["is_active = true", "(description IS NULL OR description = '')"]
    params = []
    if after_id is not None:
        where.append("id > %s")
        params.append(after_id)
    cur.execute(f"""
        SELECT id, supplier_name, sku, name, category
        FROM supplier_catalog
        WHERE {' AND '.join(where)}
        ORDER BY id
        LIMIT %s
    """, params + [batch])
    rows = cur.fetchall()
    cur.close()
    return rows


def run(args, state):
    master = get_master_conn() if args.target in ('supplier_catalog', 'both') else None
    tenant = get_tenant_conn(args.tenant) if args.target in ('inventory', 'both') else None
    review = open(args.review_out, 'a') if args.review_out else None
    try:
        total = {'scanned': 0, 'generated': 0, 'written': 0, 'llm_fail': 0}
        targets = []
        if args.target in ('inventory', 'both'):
            targets.append('inventory')
        if args.target in ('supplier_catalog', 'both'):
            targets.append('supplier_catalog')

        for target in targets:
            key = f"{target}:{args.tenant or 'master'}"
            after_id = None if args.reset else state.get(key, {}).get('after_id')
            conn = tenant if target == 'inventory' else master
            if conn is None:
                continue
            print(f"\n=== {target} (tenant={args.tenant}) after_id={after_id} ===")
            while True:
                if target == 'inventory':
                    rows = _iter_inventory(conn, after_id, args.limit, args.batch)
                else:
                    rows = _iter_supplier(conn, after_id, args.batch)
                if not rows:
                    break

                # Build LLM request payload
                items = []
                for r in rows:
                    if target == 'inventory':
                        rid, name, brand, cat = r
                        label = ", ".join(x for x in [name, brand, cat] if x)
                    else:
                        rid, sup, sku, name, cat = r
                        label = ", ".join(x for x in [name, cat, sup,
                                                      f"SKU {sku}"] if x)
                    items.append({'id': rid, 'label': label})

                got = _llm_descriptions(items)
                if not got:
                    total['llm_fail'] += 1
                    print("  LLM sin respuesta — se detiene el lote (degradado, "
                          "checkpoint no avanza)")
                    break

                for it in items:
                    rid = it['id']
                    text = got.get(rid)
                    total['scanned'] += 1
                    if not text:
                        continue
                    total['generated'] += 1
                    if review:
                        review.write(json.dumps(
                            {'target': target, 'id': rid, 'text': text,
                             'source_label': it['label']},
                            ensure_ascii=False) + "\n")
                        review.flush()
                    if args.apply:
                        c = conn.cursor()
                        try:
                            if target == 'inventory':
                                c.execute(
                                    "UPDATE inventory SET description = %s "
                                    "WHERE id = %s AND (description IS NULL OR description = '')",
                                    (text, rid))
                            else:
                                c.execute(
                                    "UPDATE supplier_catalog SET description = %s "
                                    "WHERE id = %s AND (description IS NULL OR description = '')",
                                    (text, rid))
                            total['written'] += c.rowcount
                            conn.commit()
                        finally:
                            c.close()
                    if total['generated'] <= args.show:
                        print(f"  {rid:<8} {it['label'][:44]:<44} -> {text[:70]}")

                after_id = rows[-1][0]
                state[key] = {'after_id': after_id, 'written': total['written'],
                              'updated_at': time.strftime('%Y-%m-%dT%H:%M:%S')}
                if args.apply:
                    _save_checkpoint(args.checkpoint, state)
                if len(rows) < args.batch:
                    break
                if args.limit and total['scanned'] >= args.limit:
                    break

        print(f"\nresumen: {total} apply={args.apply}")
        return 0
    finally:
        if review:
            review.close()
        if tenant is not None:
            tenant.close()
        if master is not None:
            master.close()


def main():
    p = argparse.ArgumentParser(description='Generate short Spanish descriptions.')
    p.add_argument('--tenant', type=int)
    p.add_argument('--target', choices=['inventory', 'supplier_catalog', 'both'],
                   default='inventory')
    p.add_argument('--limit', type=int, default=0)
    p.add_argument('--batch', type=int, default=10, help='items per LLM request')
    p.add_argument('--apply', action='store_true', help='write descriptions')
    p.add_argument('--dry-run', action='store_true', default=True)
    p.add_argument('--review-out', default=None, help='JSONL file with candidates')
    p.add_argument('--reset', action='store_true')
    p.add_argument('--show', type=int, default=5)
    p.add_argument('--checkpoint', default=DEFAULT_CHECKPOINT)
    args = p.parse_args()

    if args.apply:
        args.dry_run = False
    if args.target in ('inventory', 'both') and not args.tenant:
        p.error('--tenant requerido para target inventory')

    state = _load_checkpoint(args.checkpoint)
    return run(args, state)


if __name__ == '__main__':
    sys.exit(main())