"""Estimador de costos y tokens. Encapsula la lógica de estimación previa al procesamiento. """ import os from typing import List, Dict, Any, Optional, TYPE_CHECKING if TYPE_CHECKING: from domain.service import InvoiceProcessingService from domain.profiles import get_profile from domain.prompts import PromptBuilder from domain.normalizer import normalize_to_schema from domain.costs import compute_cost_from_usage, sum_costs, accumulate_usage from domain.config import settings from domain.file_discovery import get_sorted_radicados, find_file_for_type_in_radicados, get_files_as_radicados from core.utils import collect_radicado_docs, collect_documents_flat class Estimator: def __init__(self, service: 'InvoiceProcessingService'): self.service = service def estimate( self, root_path: str, keywords: List[str] | None = None, schema: Dict[str, str] | None = None, selected_fields: List[str] | None = None, *, doc_types: Optional[List[Dict[str, str]]] = None, schema_by_type: Optional[Dict[str, Dict[str, str]]] = None, selected_paths: Optional[List[str]] = None, structure_mode: str = "radicados", ) -> Dict[str, Any]: """Calcula estimado de tokens y costos para toda la carpeta basado en el primer documento.""" if doc_types: return self._estimate_multidoc(root_path, doc_types, schema_by_type, selected_paths, structure_mode) return self._estimate_singledoc(root_path, keywords, schema, selected_fields, structure_mode) def _estimate_multidoc(self, root_path, doc_types, schema_by_type, selected_paths, structure_mode="radicados"): if structure_mode == "general": radicados = get_files_as_radicados(root_path) else: radicados = get_sorted_radicados(root_path) if not radicados: return {"total_docs": 0, "per_doc": {}, "estimated_total": {}} prof = get_profile('invoice') first_radicado_found = None 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, rad_used = find_file_for_type_in_radicados(root_path, radicados, kw) if not match: continue if first_radicado_found is None: first_radicado_found = rad_used sel_all = selected_paths or [] selected_paths_for_type = [p.split('/', 1)[1] for p in sel_all if p.startswith(f"{tname}/")] sch_t = (schema_by_type or {}).get(tname) or {} try: prompt_t = PromptBuilder.build_for_type(tname, selected_paths_for_type, sch_t, profile_name=prof.get('name', 'invoice')) except Exception: prompt_t = PromptBuilder.build(selected_paths_for_type, sch_t, profile_name=prof.get('name', 'invoice')) data, usage_ext = self.service.extract_full_from_pdf(match, instruction=prompt_t) self._accumulate_usage(agg_extract_usage, usage_ext) if isinstance(data, dict): data['radicado'] = rad_used else: data = {"radicado": rad_used, "data": data} norm = normalize_to_schema(data, paths=selected_paths_for_type, profile_name=prof.get('name', 'invoice')) if isinstance(norm, dict): self._accumulate_usage(agg_normalize_usage, norm.get("_usage")) return self._build_result(agg_extract_usage, agg_normalize_usage, len(radicados), first_radicado_found or radicados[0]) def _estimate_singledoc(self, root_path, keywords, schema, selected_fields, 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: return {"total_docs": 0, "per_doc": {}, "estimated_total": {}} first = entries[0] prof = get_profile('invoice') prompt = PromptBuilder.build(selected_fields or [], schema or {}) data, usage_ext = self.service.extract_full_from_pdf(first["path"], instruction=prompt) if isinstance(data, dict): data['radicado'] = first['radicado'] else: data = {"radicado": first['radicado'], "data": data} norm = normalize_to_schema(data, paths=selected_fields, profile_name=prof.get('name', 'invoice')) usage_norm = None if isinstance(norm, dict): usage_norm = norm.get("_usage") return self._build_result(usage_ext, usage_norm, len(entries), first["radicado"]) def _accumulate_usage(self, agg, new_usage): accumulate_usage(agg, new_usage) def estimate_from_file( self, file_path: str, doc_count: int = 1, selected_paths: Optional[List[str]] = None, tiff_page_range: Optional[tuple] = None, model: Optional[str] = None, ) -> Dict[str, Any]: """Estima tokens y costos procesando un único archivo subido desde el cliente. Si `model` se especifica y es distinto al modelo actual del extractor, crea un extractor temporal con ese modelo para obtener el uso real de tokens. """ from api.desktop_ingest import _service_for_model service = _service_for_model(self.service, model) prof = get_profile('invoice') paths = selected_paths or [] prompt = PromptBuilder.build(paths, {}, profile_name=prof.get('name', 'invoice')) data, usage_ext = service.extract_full_from_pdf( file_path, instruction=prompt, tiff_page_range=tiff_page_range, ) if isinstance(data, dict): data['radicado'] = 'sample' else: data = {"radicado": 'sample', "data": data} norm = normalize_to_schema(data, paths=paths, profile_name=prof.get('name', 'invoice')) usage_norm = None if isinstance(norm, dict): usage_norm = norm.get("_usage") result = self._build_result(usage_ext, usage_norm, doc_count, 'sample') # Enriquecer con modelo efectivo y conteo de páginas para el frontend effective_model = model or getattr(service.extractor, 'model', None) or getattr(self.service.extractor, 'model', None) result['extractor_model'] = effective_model result['normalizer_model'] = settings.api.normalizer_model or 'gpt-5.4-mini' # Estimar numero de paginas a partir de prompt_tokens de imagen: # GPT-4o cobra ~85 tok base + ~170 tok por tile 512x512. # Una A4 a 150 DPI (~1240x1754) genera ~6 tiles → ~1105 tokens. # Aproximamos: pages = max(1, round(image_tokens / 1000)) prompt_tokens = int((usage_ext or {}).get('prompt_tokens') or 0) text_prompt_est = len(prompt) // 4 # 1 token ≈ 4 chars image_tokens_est = max(0, prompt_tokens - text_prompt_est) pages_est = max(1, round(image_tokens_est / 1000)) result['pages_per_doc'] = pages_est return result def _build_result(self, usage_ext, usage_norm, total_docs, first_radicado): model_ext = getattr(self.service.extractor, "model", None) cost_ext = compute_cost_from_usage(usage_ext, "extract", model=model_ext) cost_norm = compute_cost_from_usage(usage_norm, "normalize", model=settings.api.normalizer_model) per_doc_cost = sum_costs(cost_ext, cost_norm) estimated_total = {k: round(v * total_docs, 6) for k, v in per_doc_cost.items()} return { "total_docs": total_docs, "first_radicado": first_radicado, "extractor_usage": usage_ext, "normalizer_usage": usage_norm, "per_doc_cost": per_doc_cost, "estimated_total_cost": estimated_total, }