"""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