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/khtst/src")

from khtst.dados.streaming import (StreamingCorpus, limpar_cache_temp,
                                   _chave_dedup)  # v23: fonte ÚNICA da chave

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

# v23 (DRY/T-02): as TRÊS rotinas de dedup agora usam a MESMA chave
# (_chave_dedup — texto+entrada+saida+mídia, NFC). Antes: streaming usava
# `texto or saida or entrada` e este script `texto or entrada+␟+saida` —
# duas semânticas, 140+ registros divergentes (medido na auditoria).


def main():
    os.makedirs(os.path.dirname(SAIDA), exist_ok=True)
    os.makedirs("/home/z/my-project/khtst/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 = _chave_dedup(reg)
                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/khtst/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
                    chave = _chave_dedup(reg)
                    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()
    # v12 — FONTE LOCAL (órfão conectado: dados/documentos.py): arquivos
    # .txt/.md/.csv/.pdf/.docx/.xlsx soltos em cache_dados/documentos_locais
    # entram como registros lm/texto (se existirem; sem rede, sem erro)
    try:
        from khtst.dados.documentos import texto_de_arquivo
        DIR_LOC = os.path.join(os.path.dirname(SAIDA), "documentos_locais")
        if os.path.isdir(DIR_LOC):
            for nome in sorted(os.listdir(DIR_LOC)):
                cam = os.path.join(DIR_LOC, nome)
                if not os.path.isfile(cam):
                    continue
                try:
                    texto = texto_de_arquivo(cam)
                except Exception:
                    continue            # formato não suportado: pula
                texto = " ".join((texto or "").split())
                if len(texto) < 200:
                    continue
                reg = {"tarefa": "lm", "texto": texto[:8000]}
                chave = _hash_texto(texto[:8000])
                if chave in vistos:
                    continue
                vistos.add(chave)
                registros.append(reg)
                print(f"  [LOCAL] {nome}: {len(texto)} caracteres → lm")
    except Exception as e:
        print(f"  [LOCAL] indisponível: {type(e).__name__}: {e}")
    # 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/khtst/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()