133 lines
6.1 KiB
Python
133 lines
6.1 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 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, sum_costs
|
|
from application.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):
|
|
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 _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)
|
|
# El normalizador está hardcodeado a gpt-4.1-mini en normalizer.py
|
|
cost_norm = compute_cost_from_usage(usage_norm, "normalize", model="gpt-4.1-mini")
|
|
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,
|
|
}
|
|
|