183 lines
8.0 KiB
Python
183 lines
8.0 KiB
Python
"""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,
|
|
}
|
|
|