scripts/02_coletar_corpus.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Coleta do corpus v2 — UM DATASET POR VEZ, cada fonte em SUBPROCESSO isolado
(estado limpo da biblioteca datasets; RAM devolvida ao SO; erro de uma fonte
nunca contamina a seguinte). Streaming=True (+trust_remote_code onde aceito),
18 fontes, dedup global, multimodal em base64. Cache temporário limpo após
cada fonte."""
import base64
import json
import os
import subprocess
import sys

sys.path.insert(0, "/home/z/my-project/aurora/src")

from aurora.dados.streaming import StreamingCorpus, limpar_cache_temp

SAIDA = "/home/z/my-project/aurora/cache_dados/corpus_v2.jsonl"
DIR_PARCIAL = "/home/z/my-project/aurora/cache_dados/parciais"
RUNNER = "/home/z/my-project/aurora/scripts/02a_coletar_fonte.py"


def _hash_texto(t: str) -> str:
    import hashlib
    return hashlib.sha1(t.encode("utf-8", "ignore")).hexdigest()[:16]


def main():
    os.makedirs(os.path.dirname(SAIDA), exist_ok=True)
    os.makedirs("/home/z/my-project/aurora/telemetria_out", exist_ok=True)
    os.makedirs(DIR_PARCIAL, exist_ok=True)
    limpar_cache_temp()
    plano_total = StreamingCorpus().PLANO
    eventos = []
    vistos: set[str] = set()
    # v5 — RETOMÁVEL: corpus escrito INCREMENTALMENTE (append por fonte);
    # na retomada, recarrega o acumulado e pula fontes já OK (estado persistido)
    registros = []
    if os.path.exists(SAIDA):
        with open(SAIDA, encoding="utf-8") as f:
            for linha in f:
                try:
                    reg = json.loads(linha)
                except Exception:
                    continue
                chave = _hash_texto(reg.get("texto") or
                                    ((reg.get("entrada") or "") + "␟" +
                                     (reg.get("saida") or "")))
                if chave not in vistos:
                    vistos.add(chave)
                    registros.append(reg)
        print(f"retomada: {len(registros)} registros acumulados")
    for i, plano in enumerate(plano_total):
        fonte = plano["id"]
        cam_estado_prev = os.path.join(DIR_PARCIAL, f"fonte_{i:02d}.estado.json")
        if os.path.exists(cam_estado_prev):
            ev_prev = json.load(open(cam_estado_prev, encoding="utf-8"))
            if ev_prev.get("status") == "OK" and ev_prev.get("n", 0) > 0:
                eventos.append(ev_prev)
                print(f"  [SKIP] {fonte} (já coletada: n={ev_prev['n']})")
                continue
        if plano["via"] == "degradado":
            # mensagem determinística gerada pelo próprio módulo de streaming
            sub = StreamingCorpus()
            list(sub._iterar_fonte(plano))          # registra evento degradado
            eventos.extend([vars(ev) for ev in sub.eventos])
            continue
        cam_parcial = os.path.join(DIR_PARCIAL, f"fonte_{i:02d}.jsonl")
        cam_estado = os.path.join(DIR_PARCIAL, f"fonte_{i:02d}.estado.json")
        try:
            r = subprocess.run(
                [sys.executable, RUNNER, "--indice", str(i),
                 "--saida", cam_parcial, "--estado", cam_estado],
                capture_output=True, text=True, timeout=900,
                env={**os.environ, "PYTHONPATH": "/home/z/my-project/aurora/src"})
            r_erro = (r.stderr or "sem resposta")[-180:]
        except subprocess.TimeoutExpired:
            r_erro = "TIMEOUT 900s (fonte lenta — pulada)"
        except Exception as e:                     # rede/Hub: nunca para o lote
            r_erro = f"{type(e).__name__}: {e}"[-180:]
        if os.path.exists(cam_estado):
            ev = json.load(open(cam_estado, encoding="utf-8"))
        else:
            ev = {"fonte": fonte, "status": "ERRO", "detalhe": r_erro, "n": 0}
        # v5 — dedup/append durável ÚNICO (abaixo); contagem de dedup por leitura
        n_dedup = 0
        if os.path.exists(cam_parcial):
            with open(cam_parcial, encoding="utf-8") as f:
                for linha in f:
                    try:
                        reg = json.loads(linha)
                    except Exception:
                        continue
                    if reg.get("texto"):
                        chave = _hash_texto(reg["texto"])
                    else:
                        chave = _hash_texto((reg.get("entrada") or "") + "␟" +
                                             (reg.get("saida") or ""))
                    if chave in vistos:
                        n_dedup += 1
        ev["dedup_descartados"] = n_dedup
        eventos.append(ev)
        print(f"  [{ev['status']}] {fonte} (n={ev['n']}, dedup={n_dedup})"
              + (f" {ev['detalhe'][:80]}" if ev.get("detalhe") else ""))
        # v5 — append DURÁVEL ao corpus (morte do processo não perde progresso)
        novos = 0
        if os.path.exists(cam_parcial) and ev.get("status") in ("OK", "PARCIAL"):
            with open(cam_parcial, encoding="utf-8") as f, \
                    open(SAIDA, "a", encoding="utf-8") as fa:
                for linha in f:
                    try:
                        reg = json.loads(linha)
                    except Exception:
                        continue
                    if plano["via"] in ("tts_marcos",):
                        reg.setdefault("imagem", None)
                        reg.setdefault("audio", None)
                    if reg.get("texto"):
                        chave = _hash_texto(reg["texto"])
                    else:
                        chave = _hash_texto((reg.get("entrada") or "") + "␟" +
                                             (reg.get("saida") or ""))
                    if chave in vistos:
                        continue
                    vistos.add(chave)
                    if reg.get("imagem") is None:
                        reg.pop("imagem", None)
                    if reg.get("audio") is None:
                        reg.pop("audio", None)
                    registros.append(reg)
                    fa.write(json.dumps(reg, ensure_ascii=False) + "\n")
                    novos += 1
        os.remove(cam_parcial) if os.path.exists(cam_parcial) else None
        limpar_cache_temp()            # um dataset por vez → disco sob controle
        import gc
        gc.collect()
    # escreve corpus final
    n_multimodal = 0
    with open(SAIDA, "w", encoding="utf-8") as f:
        for r in registros:
            if r.get("imagem") or r.get("audio"):
                n_multimodal += 1
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    cont = {}
    for r in registros:
        cont[r["tarefa"]] = cont.get(r["tarefa"], 0) + 1
    print("== contagem por tarefa ==", cont)
    print(f"total: {len(registros)} registros ({n_multimodal} multimodais) → {SAIDA}")
    with open("/home/z/my-project/aurora/telemetria_out/eventos_coleta_v2.json",
              "w", encoding="utf-8") as f:
        json.dump(eventos, f, ensure_ascii=False, indent=2)
    limpar_cache_temp()


if __name__ == "__main__":
    main()