108 lines
4.5 KiB
Python
108 lines
4.5 KiB
Python
import requests
|
|
import logging
|
|
import os
|
|
import socket
|
|
import getpass
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class TrackingManager:
|
|
"""Gestiona la comunicación con el servidor de seguimiento web."""
|
|
|
|
def __init__(self, base_url: str = "http://localhost:8000"):
|
|
self.base_url = base_url
|
|
self.current_batch_id: Optional[str] = None
|
|
self.is_enabled = True
|
|
self._session = requests.Session() # Usar sesión para reutilizar conexiones
|
|
|
|
def start_batch(self, root_path: str, total_radicados: int, config: Optional[Dict] = None,
|
|
csv_output_path: Optional[str] = None):
|
|
"""Notifica el inicio de un nuevo lote al servidor."""
|
|
if not self.is_enabled:
|
|
return
|
|
|
|
try:
|
|
payload = {
|
|
"root_path": root_path,
|
|
"total_radicados": total_radicados,
|
|
"csv_output_path": csv_output_path,
|
|
"hostname": socket.gethostname(),
|
|
"username": getpass.getuser(),
|
|
"config": config or {}
|
|
}
|
|
response = self._session.post(f"{self.base_url}/batches", json=payload, timeout=5)
|
|
if response.status_code == 200:
|
|
self.current_batch_id = response.json().get("id")
|
|
logger.info(f"Tracking iniciado: Batch ID {self.current_batch_id}")
|
|
else:
|
|
logger.warning(f"Error al iniciar tracking: {response.status_code}")
|
|
except Exception as e:
|
|
logger.error(f"Error de conexión con el servidor de tracking: {e}")
|
|
|
|
def update_progress(self, current_index: int, success_count: int, error_count: int,
|
|
total_tokens: int, total_cost_usd: float, status: str = "processing"):
|
|
"""Actualiza el progreso del lote actual."""
|
|
if not self.is_enabled or not self.current_batch_id:
|
|
return
|
|
|
|
try:
|
|
payload = {
|
|
"current_radicado_index": current_index,
|
|
"success_count": success_count,
|
|
"error_count": error_count,
|
|
"total_tokens": total_tokens,
|
|
"total_cost_usd": total_cost_usd,
|
|
"status": status
|
|
}
|
|
self._session.patch(f"{self.base_url}/batches/{self.current_batch_id}", json=payload, timeout=1)
|
|
except Exception as e:
|
|
logger.debug(f"Error al actualizar tracking: {e}")
|
|
|
|
def log(self, message: str, level: str = "info", radicado: str = None):
|
|
"""Envía una entrada de log al servidor de tracking."""
|
|
if not self.is_enabled or not self.current_batch_id:
|
|
return
|
|
try:
|
|
payload = {"message": message, "level": level, "radicado": radicado}
|
|
self._session.post(
|
|
f"{self.base_url}/batches/{self.current_batch_id}/logs",
|
|
json=payload, timeout=1
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"Error al enviar log a tracking: {e}")
|
|
|
|
def report_phase(self, phase: str, radicado: str = None):
|
|
"""Reporta la fase actual de procesamiento (extract, normalize, csv, completed)."""
|
|
if not self.is_enabled or not self.current_batch_id:
|
|
return
|
|
try:
|
|
payload = {"current_phase": phase, "current_radicado": radicado}
|
|
self._session.patch(f"{self.base_url}/batches/{self.current_batch_id}", json=payload, timeout=1)
|
|
except Exception as e:
|
|
logger.debug(f"Error al reportar fase a tracking: {e}")
|
|
|
|
def report_error(self, error_message: str):
|
|
"""Reporta un error crítico al servidor."""
|
|
if not self.is_enabled or not self.current_batch_id:
|
|
return
|
|
|
|
try:
|
|
payload = {"last_error": error_message}
|
|
self._session.patch(f"{self.base_url}/batches/{self.current_batch_id}", json=payload, timeout=1)
|
|
except Exception as e:
|
|
logger.debug(f"Error al reportar error de tracking: {e}")
|
|
|
|
def complete_batch(self, success_count: int, error_count: int, total_tokens: int, total_cost_usd: float):
|
|
"""Notifica la finalización exitosa del lote."""
|
|
self.update_progress(
|
|
current_index=-1, # Indica finalización si se desea
|
|
success_count=success_count,
|
|
error_count=error_count,
|
|
total_tokens=total_tokens,
|
|
total_cost_usd=total_cost_usd,
|
|
status="completed"
|
|
)
|
|
self.current_batch_id = None
|