Autoparts-DB/scripts/categorize_inventory.py

#!/usr/bin/env python3
"""Backfill inventory categories (A1).

Assigns ``inventory.category_id`` per tenant using the rule engine over the
Nexpart taxonomy (+ optional LLM fallback via ``ai_chat.classify_part``).

Usage (as root/postgres, from /home/Autopartes):
    python3 scripts/categorize_inventory.py --tenant 31 --limit 200
    python3 scripts/categorize_inventory.py --tenant 45 --limit 50 --only-uncategorized
    python3 scripts/categorize_inventory.py --all --batch 500
    python3 scripts/categorize_inventory.py --tenant 31 --no-llm --dry-run

Checkpoints progress (last processed inventory id per tenant) in a JSON file so
long backfills can be resumed.
"""

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
from services.product_categorizer import categorize_inventory  # noqa: E402

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


def list_tenants(only_id=None):
    conn = get_master_conn()
    cur = conn.cursor()
    if only_id:
        cur.execute(
            "SELECT id, name, db_name FROM tenants WHERE is_active = true AND id = %s",
            (only_id,),
        )
    else:
        cur.execute(
            "SELECT id, name, db_name FROM tenants WHERE is_active = true ORDER BY id"
        )
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows


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"  ! no se pudo guardar checkpoint: {e}")


def run_tenant(tenant_id, name, db_name, args, checkpoint):
    print(f"\n=== tenant {tenant_id} — {name} ({db_name}) ===")
    conn = get_tenant_conn(tenant_id)
    key = str(tenant_id)
    state = checkpoint.get(key, {})
    after_id = None if args.reset else state.get('after_id')
    total_scanned = 0
    total_persisted = 0
    total_by_source = {}

    try:
        while True:
            res = categorize_inventory(
                conn,
                limit=args.batch,
                only_uncategorized=args.only_uncategorized,
                use_llm=not args.no_llm,
                persist=not args.dry_run,
                create_categories=not args.no_create,
                after_id=after_id,
                sample_limit=args.show,
            )
            total_scanned += res['scanned']
            total_persisted += res['persisted']
            for k, v in res['by_source'].items():
                total_by_source[k] = total_by_source.get(k, 0) + v

            if args.show and res['sample']:
                for s in res['sample']:
                    print(f"  {s['part_number']:<22} {str(s['name'])[:34]:<34} "
                          f"-> {str(s['category_name']):<22} "
                          f"({s['source']}, {s['confidence']})")

            after_id = res['last_id']
            checkpoint[key] = {
                'after_id': after_id,
                'persisted': state.get('persisted', 0) + total_persisted,
                'updated_at': time.strftime('%Y-%m-%dT%H:%M:%S'),
            }
            if not args.dry_run:
                save_checkpoint(args.checkpoint, checkpoint)

            print(f"  batch: scanned={res['scanned']} persisted={res['persisted']} "
                  f"by_source={res['by_source']} last_id={after_id}")

            if res['scanned'] < args.batch:
                break
            if args.limit and total_scanned >= args.limit:
                break
        print(f"  total: scanned={total_scanned} persisted={total_persisted} "
              f"by_source={total_by_source}")
    finally:
        conn.close()

    return total_scanned, total_persisted


def main():
    p = argparse.ArgumentParser(description='Backfill inventory categories (A1).')
    p.add_argument('--tenant', type=int, help='tenant id (omit with --all)')
    p.add_argument('--all', action='store_true', help='process every active tenant')
    p.add_argument('--limit', type=int, default=0, help='max items per tenant (0 = all)')
    p.add_argument('--batch', type=int, default=500, help='rows per batch')
    p.add_argument('--only-uncategorized', action='store_true', default=True,
                   help='only fill items with category_id IS NULL (default)')
    p.add_argument('--all-items', dest='only_uncategorized', action='store_false',
                   help='consider every active item')
    p.add_argument('--no-llm', action='store_true', help='disable LLM fallback')
    p.add_argument('--no-create', action='store_true',
                   help='do not create missing categories')
    p.add_argument('--dry-run', action='store_true', help='do not persist')
    p.add_argument('--reset', action='store_true', help='ignore saved checkpoint')
    p.add_argument('--show', type=int, default=5, help='sample rows to print (0 = none)')
    p.add_argument('--checkpoint', default=DEFAULT_CHECKPOINT)
    args = p.parse_args()

    if not args.all and not args.tenant:
        p.error('--tenant ID o --all requerido')

    checkpoint = load_checkpoint(args.checkpoint)
    tenants = list_tenants(None if args.all else args.tenant)
    if not tenants:
        print('sin tenants para procesar')
        return 1

    grand_scanned = grand_persisted = 0
    for tid, name, db_name in tenants:
        scanned, persisted = run_tenant(tid, name, db_name, args, checkpoint)
        grand_scanned += scanned
        grand_persisted += persisted

    print(f"\nTOTAL: tenants={len(tenants)} scanned={grand_scanned} "
          f"persisted={grand_persisted}")
    return 0


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