Files
IPSExtractor/application/batch_processor.py
T

592 lines
27 KiB
Python

"""Procesador por lotes para la extracción de facturas.
Encapsula la lógica de iteración sobre radicados, extracción, normalización y exportación a CSV.
"""
import os
import logging
import asyncio
import openai
from typing import List, Dict, Any, Optional, Tuple, Callable, Awaitable, TYPE_CHECKING
if TYPE_CHECKING:
from application.service import InvoiceProcessingService
from application.profiles import get_profile
from application.prompts import PromptBuilder
from application.normalizer import normalize_to_schema
from application.costs import compute_cost_from_usage
from application.file_discovery import get_sorted_radicados, find_file_for_type_in_radicados, get_files_as_radicados
from application.csv_manager import CSVManager
from application.config import settings
from core.utils import collect_radicado_docs, find_first_matching_file, collect_documents_flat
from core.csv_utils import rows_from_data
logger = logging.getLogger(__name__)
class BatchProcessor:
def __init__(self, service: 'InvoiceProcessingService'):
"""
Args:
service: Instancia de InvoiceProcessingService para delegar la extracción individual.
"""
self.service = service
self.csv_manager = CSVManager()
self._pause_event = asyncio.Event()
self._pause_event.set() # No pausado inicialmente
self._stop_requested = False
def process_root(
self,
root_path: str,
keywords: List[str] | None = None,
*,
doc_types: Optional[List[Dict[str, str]]] = None,
fields: Optional[List[str]] = None,
on_progress: Optional[callable] = None,
csv_output_path: Optional[str] = None,
csv_header_override: Optional[List[str]] = None,
instruction_override: Optional[str] = None,
on_usage: Optional[Callable[[str, str, Dict[str, Any] | None, Dict[str, Any] | None], None]] = None,
on_error_retry: Optional[Callable[[Exception], Awaitable[bool]]] = None,
alias_map: Optional[Dict[str, str]] = None,
tiff_page_range: Optional[Tuple[int, int]] = None,
structure_mode: str = "radicados",
) -> List[Dict[str, Any]]:
"""Procesa radicados con selección de campos y CSV incremental.
Delega en métodos específicos para multi-documento o single-document.
"""
# Resetear estado de pausa/parada
self._pause_event.set()
self._stop_requested = False
# 1. Preparación de Schemas (Fase 1 Multi-doc)
if doc_types and not fields:
return self._prepare_multidoc_schemas(root_path, doc_types, instruction_override, tiff_page_range, structure_mode)
# 2. Procesamiento Multi-doc (Fase 2) - ASYNC
if doc_types and fields:
return asyncio.run(self._process_multidoc_async(
root_path=root_path,
doc_types=doc_types,
fields=fields,
csv_output_path=csv_output_path,
csv_header_override=csv_header_override,
alias_map=alias_map,
tiff_page_range=tiff_page_range,
on_progress=on_progress,
on_usage=on_usage,
on_error_retry=on_error_retry,
structure_mode=structure_mode
))
# 3. Procesamiento Single-doc (Retrocompatibilidad)
# Si no hay fields, es la fase de detección de esquema (secuencial, solo 1 doc)
if not fields:
return self._process_singledoc(
root_path=root_path,
keywords=keywords,
fields=fields,
csv_output_path=csv_output_path,
csv_header_override=csv_header_override,
instruction_override=instruction_override,
alias_map=alias_map,
tiff_page_range=tiff_page_range,
on_progress=on_progress,
on_usage=on_usage,
structure_mode=structure_mode
)
else:
# Si hay fields, es procesamiento masivo -> ASYNC
return asyncio.run(self._process_singledoc_async(
root_path=root_path,
keywords=keywords,
fields=fields,
csv_output_path=csv_output_path,
csv_header_override=csv_header_override,
instruction_override=instruction_override,
alias_map=alias_map,
tiff_page_range=tiff_page_range,
on_progress=on_progress,
on_usage=on_usage,
on_error_retry=on_error_retry,
structure_mode=structure_mode
))
def _prepare_multidoc_schemas(self, root_path, doc_types, instruction, tiff_page_range, structure_mode):
schema_by_type, first_full_by_type, first_radicado = self.service.build_schema_from_first_for_types(
root_path,
doc_types,
instruction=instruction,
tiff_page_range=tiff_page_range,
structure_mode=structure_mode
)
return [{
"_needs_fields": True,
"schema_by_type": schema_by_type,
"first_full_by_type": first_full_by_type,
"first_radicado": first_radicado,
}]
async def _process_multidoc_async(
self, root_path, doc_types, fields, csv_output_path,
csv_header_override, alias_map, tiff_page_range, on_progress, on_usage, on_error_retry=None,
structure_mode="radicados"
):
csv_header = self._prepare_csv_header(fields, alias_map, csv_header_override)
if structure_mode == "general":
radicados = get_files_as_radicados(root_path)
else:
radicados = get_sorted_radicados(root_path)
if not radicados and not os.path.isdir(root_path):
raise ValueError(f"La ruta raíz no es un directorio: {root_path}")
total = len(radicados)
prof = get_profile('invoice')
# Limitar concurrencia (Configurable)
limit = settings.processing.concurrency_limit
if limit < 1: limit = 1
sem = asyncio.Semaphore(limit)
# Para mantener el orden de escritura en CSV, usamos un Lock
csv_lock = asyncio.Lock()
# Contador thread-safe para progreso
completed_count = 0
progress_lock = asyncio.Lock()
async def process_radicado(idx, rad):
nonlocal completed_count
async with sem:
while True: # Bucle para permitir reintentos tras pausa
await self._pause_event.wait()
if self._stop_requested:
return {"radicado": rad, "status": "stopped"}
try:
# Determinar ID visual y ruta
if structure_mode == "general":
rad_id = os.path.splitext(rad)[0]
rad_path = os.path.join(root_path, rad) # Es el archivo
else:
rad_id = rad
rad_path = os.path.join(root_path, rad) # Es la carpeta
logger.info(f"Procesando radicado {rad_id} (multi) [{idx}/{total}]", extra={"radicado": rad_id})
fused_data = {"radicado": rad_id}
agg_extract_usage = {"prompt_tokens": 0, "completion_tokens": 0, "prompt_tokens_details": {"cached_tokens": 0}}
agg_normalize_usage = {"prompt_tokens": 0, "completion_tokens": 0, "prompt_tokens_details": {"cached_tokens": 0}}
for dt in doc_types:
tname = str(dt.get("name") or "").strip() or "Tipo"
kw = str(dt.get("keyword") or "").strip()
match = None
if structure_mode == "general":
# En modo general, el archivo ES el documento.
# Opcional: filtrar si el nombre coincide con keyword, pero usualmente
# si el usuario elige "General" y "Tipos", quiere aplicar los tipos a ese archivo.
# Pero si tiene varios tipos (Factura, Anexo), no tiene sentido aplicarlos todos al mismo archivo
# a menos que quiera intentar extraer cosas distintas.
# Asumiremos que el archivo aplica para el tipo.
match = rad_path
else:
try:
files = sorted(await asyncio.to_thread(os.listdir, rad_path))
except Exception:
files = []
match_name = find_first_matching_file(files, kw) if kw else None
match = os.path.join(rad_path, match_name) if match_name else None
if not match:
logger.info(f"[{rad_id}] Tipo '{tname}' sin archivo coincidente (kw='{kw}')")
fused_data[tname] = {}
continue
# Extraer
selected_paths_for_type = [p.split('/', 1)[1] for p in (fields or []) if p.startswith(f"{tname}/")]
try:
prompt_t = PromptBuilder.build_for_type(tname, selected_paths_for_type, {}, profile_name=prof.get('name', 'invoice'))
except Exception:
prompt_t = PromptBuilder.build(selected_paths_for_type, {}, profile_name=prof.get('name', 'invoice'))
# Llamada ASYNC
data, usage_ext = await self.service.extract_full_from_pdf_async(match, instruction=prompt_t, tiff_page_range=tiff_page_range)
self._accumulate_usage(agg_extract_usage, usage_ext)
# Normalizar (CPU bound, mover a thread)
if isinstance(data, dict):
data['radicado'] = rad_id
else:
data = {"radicado": rad_id, "data": data}
# Normalización en thread aparte
norm = await asyncio.to_thread(
normalize_to_schema,
data,
paths=selected_paths_for_type,
profile_name=prof.get('name', 'invoice')
)
normalized_obj = data
n_usage = None
if isinstance(norm, dict):
n_usage = norm.get("_usage")
normalized_obj = norm.get("normalized") or data
self._accumulate_usage(agg_normalize_usage, n_usage)
fused_data[tname] = normalized_obj or {}
fused_data["_source_pdf"] = match
# Report usage
self._report_usage(on_usage, rad_id, "extract", agg_extract_usage)
self._report_usage(on_usage, rad_id, "normalize", agg_normalize_usage)
# CSV Export protegido por lock
if csv_output_path:
async with csv_lock:
await asyncio.to_thread(
self._write_csv_multidoc,
csv_output_path, fused_data, fields, doc_types,
prof, csv_header, alias_map, rad_id
)
return fused_data
except openai.RateLimitError as e:
if on_error_retry:
# Pausar todos los demás
self._pause_event.clear()
try:
should_retry = await on_error_retry(e)
if should_retry:
self._pause_event.set()
continue # Reintentar este radicado
else:
self._stop_requested = True
self._pause_event.set()
return self._handle_error_async(rad, e, csv_output_path, csv_header)
except Exception as e_retry:
self._pause_event.set()
return self._handle_error_async(rad, e_retry, csv_output_path, csv_header)
else:
return self._handle_error_async(rad, e, csv_output_path, csv_header)
except Exception as e:
return self._handle_error_async(rad, e, csv_output_path, csv_header)
finally:
async with progress_lock:
completed_count += 1
self._notify_progress(on_progress, rad, completed_count, total)
break # Salir del bucle True si no hubo RateLimitError o se manejó
tasks = [process_radicado(i, r) for i, r in enumerate(radicados, start=1)]
resultados = await asyncio.gather(*tasks)
return list(resultados)
def _process_singledoc(
self, root_path, keywords, fields, csv_output_path,
csv_header_override, instruction_override, alias_map, tiff_page_range, on_progress, on_usage, on_error_retry=None,
structure_mode="radicados"
):
if structure_mode == "general":
entries = collect_documents_flat(root_path, keywords)
else:
entries = collect_radicado_docs(root_path, keywords or [])
if not entries:
logger.info("No se encontraron radicados con PDFs que coincidan.")
return []
# Primer documento para schema
first = entries[0]
prof = get_profile('invoice')
inst0 = instruction_override if instruction_override is not None else prof["default_instruction"]
schema, first_full = self.service.build_schema_from_first(first["path"], instruction=inst0, tiff_page_range=tiff_page_range)
selected_fields = fields or []
if not selected_fields:
return [{
"_needs_fields": True,
"schema": schema,
"first_full": first_full,
"radicado": first["radicado"],
"_source_pdf": first["path"],
}]
# Fallback a procesamiento secuencial/async si se llama directamente
return asyncio.run(self._process_singledoc_async(
root_path, keywords, fields, csv_output_path, csv_header_override,
instruction_override, alias_map, tiff_page_range, on_progress, on_usage, on_error_retry, structure_mode
))
async def _process_singledoc_async(
self, root_path, keywords, fields, csv_output_path,
csv_header_override, instruction_override, alias_map, tiff_page_range, on_progress, on_usage, on_error_retry=None,
structure_mode="radicados"
):
# Recolección de archivos (IO bloqueante)
if structure_mode == "general":
entries = await asyncio.to_thread(collect_documents_flat, root_path, keywords)
else:
entries = await asyncio.to_thread(collect_radicado_docs, root_path, keywords or [])
if not entries:
logger.info("No se encontraron radicados con PDFs que coincidan.")
return []
selected_fields = fields or []
prompt = PromptBuilder.build(selected_fields, {})
csv_header = self._prepare_csv_header(selected_fields, alias_map, csv_header_override)
prof = get_profile('invoice')
total = len(entries)
# Limitar concurrencia (Configurable)
limit = settings.processing.concurrency_limit
if limit < 1: limit = 1
sem = asyncio.Semaphore(limit)
csv_lock = asyncio.Lock()
progress_lock = asyncio.Lock()
completed_count = 0
async def process_entry(idx, entry):
nonlocal completed_count
async with sem:
while True: # Bucle para permitir reintentos tras pausa
await self._pause_event.wait()
if self._stop_requested:
return {"radicado": entry['radicado'], "status": "stopped"}
try:
logger.info(f"Procesando radicado {entry['radicado']} [{idx}/{total}]", extra={"radicado": entry['radicado'], "pdf": entry['path']})
# Async Extract
data, usage_ext = await self.service.extract_full_from_pdf_async(entry["path"], instruction=prompt, tiff_page_range=tiff_page_range)
self._report_usage(on_usage, entry['radicado'], "extract", usage_ext)
if isinstance(data, dict):
data['radicado'] = entry['radicado']
else:
data = {"radicado": entry['radicado'], "data": data}
data["_source_pdf"] = entry["path"]
# Normalizar (Thread)
norm = await asyncio.to_thread(
normalize_to_schema,
data, paths=selected_fields, profile_name=prof.get('name', 'invoice')
)
if isinstance(norm, dict):
n_usage = norm.get("_usage")
self._report_usage(on_usage, entry['radicado'], "normalize", n_usage)
normalized_obj = norm.get("normalized")
if isinstance(normalized_obj, dict):
data = normalized_obj
data['radicado'] = entry['radicado']
data["_source_pdf"] = entry["path"]
if csv_output_path:
async with csv_lock:
await asyncio.to_thread(
self._write_csv_singledoc,
csv_output_path, data, selected_fields, prof, csv_header, alias_map, entry['radicado']
)
return data
except openai.RateLimitError as e:
if on_error_retry:
# Pausar todos los demás
self._pause_event.clear()
try:
should_retry = await on_error_retry(e)
if should_retry:
self._pause_event.set()
continue # Reintentar este radicado
else:
self._stop_requested = True
self._pause_event.set()
return self._handle_error_async(entry['radicado'], e, csv_output_path, csv_header, entry['path'])
except Exception as e_retry:
self._pause_event.set()
return self._handle_error_async(entry['radicado'], e_retry, csv_output_path, csv_header, entry['path'])
else:
return self._handle_error_async(entry['radicado'], e, csv_output_path, csv_header, entry['path'])
except Exception as e:
return self._handle_error_async(entry['radicado'], e, csv_output_path, csv_header, entry['path'])
finally:
async with progress_lock:
completed_count += 1
self._notify_progress(on_progress, entry['radicado'], completed_count, total)
break # Salir del bucle si no hubo RateLimitError o se manejó
tasks = [process_entry(i, e) for i, e in enumerate(entries, start=1)]
resultados = await asyncio.gather(*tasks)
return list(resultados)
# --- Helpers ---
def _prepare_csv_header(self, fields, alias_map, override):
if override:
return override
header_base = ['radicado'] + list(fields or [])
if not alias_map:
return header_base
alias_map2 = dict(alias_map)
# Expandir alias items
for orig_key, alias_value in list(alias_map.items()):
if "/*/" in orig_key:
sub = orig_key.split("/*/")[-1]
if sub and sub not in alias_map2:
alias_map2[sub] = alias_value
raw_header = [alias_map2.get(h, h) for h in header_base]
# Deduplicar conservando orden
seen = set()
final_header = []
for h in raw_header:
if h not in seen:
final_header.append(h)
seen.add(h)
return final_header
def _notify_progress(self, callback, rad, idx, total):
if callback:
try:
callback(rad, idx, total)
except Exception:
pass
def _accumulate_usage(self, agg, new_usage):
if not isinstance(new_usage, dict): return
agg["prompt_tokens"] += int(new_usage.get("prompt_tokens") or 0)
agg["completion_tokens"] += int(new_usage.get("completion_tokens") or 0)
details = new_usage.get("prompt_tokens_details") or {}
agg["prompt_tokens_details"]["cached_tokens"] += int(details.get("cached_tokens") or 0)
def _report_usage(self, callback, rad, phase, usage):
if callback and usage:
try:
model = None
if phase == "extract":
model = getattr(self.service.extractor, "model", None)
elif phase == "normalize":
# El normalizador está hardcodeado a gpt-4.1-mini en normalizer.py
model = "gpt-4.1-mini"
callback(phase, rad, usage, compute_cost_from_usage(usage, phase, model=model))
except Exception:
pass
def _write_csv_multidoc(self, path, data, fields, doc_types, prof, header, alias_map, rad):
try:
agg_base = prof.get('aggregate_list_paths') or None
agg_paths_expanded = []
if agg_base:
for base in agg_base:
for dt in doc_types:
tname = str(dt.get("name") or "").strip() or "Tipo"
agg_paths_expanded.append(f"{tname}/{base}")
rows, computed_header = rows_from_data(
data,
aggregate_list_paths=agg_paths_expanded or None,
allowed_paths=fields,
)
final_header = header if header else computed_header
final_rows = self._apply_alias_to_rows(rows, alias_map)
self.csv_manager.append_rows(path, final_rows, final_header)
except Exception:
logger.exception(f"Error escribiendo CSV para radicado {rad}", extra={"radicado": rad})
def _write_csv_singledoc(self, path, data, fields, prof, header, alias_map, rad):
try:
rows, computed_header = rows_from_data(
data,
aggregate_list_paths=(prof.get('aggregate_list_paths') or None),
allowed_paths=fields,
)
final_header = header if header else computed_header
final_rows = self._apply_alias_to_rows(rows, alias_map)
self.csv_manager.append_rows(path, final_rows, final_header)
except Exception:
logger.exception(f"Error escribiendo CSV para radicado {rad}", extra={"radicado": rad})
def _apply_alias_to_rows(self, rows, alias_map):
if not alias_map: return rows
alias_map2 = dict(alias_map)
for orig_key, alias_value in list(alias_map.items()):
if "/*/" in orig_key:
sub = orig_key.split("/*/")[-1]
if sub and sub not in alias_map2:
alias_map2[sub] = alias_value
aliased_rows = []
for r in rows:
if isinstance(r, dict):
new_r = {}
for k, v in r.items():
new_k = alias_map2.get(k, k)
if new_k in new_r:
# Lógica de fusión solicitada:
# 1. Si son iguales, dejar uno.
# 2. Si son diferentes, separar por coma.
prev_val = new_r[new_k]
# Normalizar a string para comparación y fusión
s_prev = str(prev_val).strip() if prev_val is not None else ""
s_curr = str(v).strip() if v is not None else ""
if s_prev == s_curr:
continue
if not s_prev:
new_r[new_k] = v
elif not s_curr:
pass
else:
new_r[new_k] = f"{s_prev}, {s_curr}"
else:
new_r[new_k] = v
aliased_rows.append(new_r)
else:
aliased_rows.append(r)
return aliased_rows
def _handle_error(self, resultados, rad, e, csv_path, header, pdf_path=None):
error_obj = {
"radicado": rad,
"error": str(e)
}
if pdf_path: error_obj["_source_pdf"] = pdf_path
resultados.append(error_obj)
logger.exception(f"Error procesando radicado {rad}", extra={"radicado": rad})
if csv_path:
header_err = header if header else ["radicado", "error"]
self.csv_manager.append_rows(csv_path, [error_obj], header_err)
def _handle_error_async(self, rad, e, csv_path, header, pdf_path=None):
error_obj = {
"radicado": rad,
"error": str(e)
}
if pdf_path: error_obj["_source_pdf"] = pdf_path
logger.exception(f"Error procesando radicado {rad}", extra={"radicado": rad})
if csv_path:
header_err = header if header else ["radicado", "error"]
try:
self.csv_manager.append_rows(csv_path, [error_obj], header_err)
except Exception:
pass
return error_obj