app.py

import os
import time
import requests
from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import openai

# ==============================
# CONFIGURACIÓN DE CLASIFICADORES
# ==============================

CLASIFICADORES = {
    "abaco": {
        "candidate_labels": [
            'Otros No Alimentos', 'Alimentos de hoteles, restaurantes o casinos', 'Juguetes', 'Muebles',
            'Pequeño Aparato', 'Textiles', 'Bebidas Alcohólicas', 'Dulces y Postres', 'Papelería',
            'Bebidas Azucaradas', 'Carnes', 'Panadería de sal', 'Materiales de construcción', 'Hogar',
            'Otras bebidas', 'Grasas tipo 3 (saturadas)', 'Otros Alimentos', 'Cereales', 'Frutas', 'Personal',
            'Gama Blanca', 'Gama Marron', 'Frutos secos y semillas', 'Paquetes/snacks', 'Agua',
            'Productos lácteos grasa entera', 'Dulces y postres', 'Otras fórmulas especiales',
            'Leches Enteras', 'Fórmulas infantiles', 'Leguminosas Secas', 'Tuberculos', 'Raíz',
            'Plátano', 'Verduras', 'Panadería dulce', 'Huevos', 'Azúcares Simples',
            'Medicamentos sin fórmula médica', 'Suplementos nutricionales', 'Alimentos para Mascotas',
            'Grasas tipo 2 (poliinsaturadas)', 'Grasas tipo 1 (monoinsaturadas)'
        ],
        "categoria_tipo": {
            'Sin Categoria': 'Sin Categoria',
            'Otros No Alimentos': 'other', 'Alimentos de hoteles, restaurantes o casinos': 'food',
            'Juguetes': 'other', 'Muebles': 'other', 'Pequeño Aparato': 'other', 'Textiles': 'other',
            'Bebidas Alcohólicas': 'food', 'Dulces y Postres': 'food', 'Papelería': 'other',
            'Bebidas Azucaradas': 'food', 'Carnes': 'food', 'Panadería de sal': 'food',
            'Materiales de construcción': 'other', 'Hogar': 'other', 'Otras bebidas': 'food',
            'Grasas tipo 3 (saturadas)': 'food', 'Otros Alimentos': 'food', 'Cereales': 'food',
            'Frutas': 'food', 'Personal': 'other', 'Gama Blanca': 'other', 'Gama Marron': 'other',
            'Frutos secos y semillas': 'food', 'Paquetes/snacks': 'food', 'Agua': 'food',
            'Productos lácteos grasa entera': 'food', 'Dulces y postres': 'food',
            'Otras fórmulas especiales': 'food', 'Leches Enteras': 'food',
            'Fórmulas infantiles': 'food', 'Leguminosas Secas': 'food', 'Tuberculos': 'food',
            'Raíz': 'food', 'Plátano': 'food', 'Verduras': 'food', 'Panadería dulce': 'food',
            'Huevos': 'food', 'Azúcares Simples': 'food',
            'Medicamentos sin fórmula médica': 'other', 'Suplementos nutricionales': 'food',
            'Alimentos para Mascotas': 'other', 'Grasas tipo 2 (poliinsaturadas)': 'food',
            'Grasas tipo 1 (monoinsaturadas)': 'food'
        }
    },
    "mexico": {
        "candidate_labels": [
            'No Comestible', 'Fruta y verdura', 'Pan', 'Cereales y Leguminosas', 'Bebidas Saborizadas',
            'Azucares', 'Abarrotes', 'Proteina animal', 'Alimento Preparado',
            'Refrigerados y Congelados', 'Lactéos', 'Agua', 'Jugo y Néctar',
            'Leguminosa', 'Grasas Animales', 'Confiteria', 'Suplementos',
            'Salsas', 'Café y Té', 'Sabores listos para comer', 'Aceites Vegetales',
            'Raices, Tuberculos', 'Semillas y Nueces', 'Productos veganos'
        ],
        "categoria_tipo": {
            'No Comestible': 'other', 'Fruta y verdura': 'food', 'Pan': 'food',
            'Cereales y Leguminosas': 'food', 'Bebidas Saborizadas': 'food', 'Azucares': 'food',
            'Abarrotes': 'food', 'Proteina animal': 'food', 'Alimento Preparado': 'food',
            'Refrigerados y Congelados': 'food', 'Lactéos': 'food', 'Agua': 'food',
            'Jugo y Néctar': 'food', 'Leguminosa': 'food', 'Grasas Animales': 'food',
            'Confiteria': 'food', 'Suplementos': 'food', 'Salsas': 'food', 'Café y Té': 'food',
            'Sabores listos para comer': 'food', 'Aceites Vegetales': 'food',
            'Raices, Tuberculos': 'food', 'Semillas y Nueces': 'food', 'Productos veganos': 'food'
        }
    }
}

# ==============================
# CONFIGURACIÓN APP Y API
# ==============================

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

HF_TOKEN = os.getenv("HG_TOKEN")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
openai.api_key = OPENAI_API_KEY

MODEL_1_URL = "https://api-inference.huggingface.co/models/facebook/bart-large-mnli?wait_for_model=true"
MODEL_2_URL = "https://api-inference.huggingface.co/models/joeddav/xlm-roberta-large-xnli?wait_for_model=true"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}

# ==============================
# FUNCIONES DE UTILIDAD
# ==============================

def query_huggingface_api(sequence, labels, api_url):
    payload = {
        "inputs": sequence,
        "parameters": {"candidate_labels": labels}
    }
    response = requests.post(api_url, headers=HEADERS, json=payload)
    if response.status_code != 200:
        raise Exception(f"Error al invocar Hugging Face API: {response.status_code}, {response.text}")
    return response.json()

def query_openai_classification(sequence, labels):
    prompt = (
        f"Dado el siguiente texto: '{sequence}', selecciona la categoría más adecuada entre: {', '.join(labels)}. "
        f"Devuelve únicamente el nombre exacto de la categoría seleccionada."
    )
    try:
        # Usar la nueva API de openai >= 1.0.0
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "Eres un clasificador experto."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.0,
            max_tokens=20
        )
        choice = response.choices[0].message.content.strip()
        # Validar que la respuesta esté entre las etiquetas
        if choice in labels:
            return {"labels": [choice], "scores": [1.0]}
        else:
            # A veces el modelo puede devolver texto adicional, intentar limpiar
            for label in labels:
                if label.lower() in choice.lower():
                    return {"labels": [label], "scores": [1.0]}
            raise Exception(f"Respuesta inválida del modelo OpenAI: '{choice}'")
    except Exception as e:
        raise Exception(f"Error en modelo OpenAI: {str(e)}")

def clasificar_con_modelos(sequence, labels):
    errores = []
    modelos = [
        (MODEL_1_URL, "Modelo 1"),
        (MODEL_2_URL, "Modelo 2"),
        ("openai", "Modelo 3 OpenAI")
    ]
    for url, nombre in modelos:
        for intento in range(2):
            try:
                if url == "openai":
                    resultado = query_openai_classification(sequence, labels)
                else:
                    resultado = query_huggingface_api(sequence, labels, url)
                return resultado, nombre
            except Exception as e:
                error_text = f"{nombre} intento {intento+1}: {str(e)}"
                if ("Model too busy" in error_text or "Resource temporarily unavailable" in error_text or "Model" in error_text and "is currently loading" in error_text) and intento == 0: # Añadido "is currently loading"
                    time.sleep(5)
                    continue
                errores.append(error_text)
                break # Salir del bucle de intentos para este modelo si falla y no es reintentable
    # Si todos los modelos fallan para este chunk de etiquetas
    raise Exception("Fallo en todos los modelos para este chunk de etiquetas:\n" + "\n".join(errores))

def split_list(lst, max_size):
    for i in range(0, len(lst), max_size):
        yield lst[i:i + max_size]

# ==============================
# ENDPOINT PRINCIPAL
# ==============================

@app.get("/categorize")
def categorizar_producto(
    cua_master: str = Query(..., description="Selección de clasificador (abaco/mexico)"),
    typology: str = Query(..., description="Valor de typology a clasificar")
):
    cua_master = cua_master.lower()
    if cua_master not in CLASIFICADORES:
        raise HTTPException(
            status_code=400,
            detail=f"Clasificador no válido. Opciones: {', '.join(CLASIFICADORES.keys())}"
        )

    config = CLASIFICADORES[cua_master]
    best_score = -1.0 # Inicializar a un flotante para consistencia
    best_label = "Sin Categoria"
    modelo_utilizado = "Ninguno" # Un default más informativo

    # Validar que los tokens estén presentes si se van a usar
    if not HF_TOKEN and (MODEL_1_URL or MODEL_2_URL): # Si se usan modelos HF
        # Podrías lanzar un error aquí o loguear una advertencia severa
        # Por ahora, permitiremos que falle en query_huggingface_api si es el caso
        print("Advertencia: HG_TOKEN no está configurado, los modelos de Hugging Face pueden fallar.")
        # raise HTTPException(status_code=503, detail="Configuración del servidor incompleta: HG_TOKEN faltante.")

    if not OPENAI_API_KEY: # Si se usa OpenAI
        # Similar al anterior
        print("Advertencia: OPENAI_API_KEY no está configurado, el modelo de OpenAI puede fallar.")
        # raise HTTPException(status_code=503, detail="Configuración del servidor incompleta: OPENAI_API_KEY faltante.")


    typology_stripped = typology.strip()
    if not typology_stripped:
        raise HTTPException(status_code=400, detail="El parámetro 'typology' no puede estar vacío.")

    # Lógica para intentar clasificar con los modelos disponibles
    # Si un chunk de etiquetas da un resultado, ese se usa.
    # Si un modelo falla en un chunk, se intenta el siguiente modelo para ESE MISMO CHUNK.
    # Si todos los modelos fallan para un chunk, se lanza una excepción y la clasificación general falla.
    # Si un chunk es exitoso, se actualiza best_score y best_label.
    # El objetivo es encontrar la mejor clasificación entre todos los chunks.

    # Almacenar los resultados de cada chunk si es necesario, o simplemente el mejor hasta ahora.
    # La lógica original actualiza best_score y best_label iterativamente.

    try:
        # Itera sobre los chunks de etiquetas candidatas
        for sublist_labels in split_list(config["candidate_labels"], 10):
            if not sublist_labels: # En caso de que split_list devuelva una lista vacía
                continue
            
            try:
                # Intenta clasificar este chunk con la cascada de modelos
                result, modelo = clasificar_con_modelos(typology_stripped, sublist_labels)
                
                # Asegurarse que el resultado tiene la estructura esperada
                if result and "scores" in result and result["scores"] and \
                   "labels" in result and result["labels"]:
                    
                    current_score = result["scores"][0]
                    current_label = result["labels"][0]

                    if current_score > best_score:
                        best_score = current_score
                        best_label = current_label
                        modelo_utilizado = modelo
                else:
                    # Esto no debería ocurrir si las funciones query_* devuelven el formato correcto
                    # o lanzan una excepción. Pero es una salvaguarda.
                    print(f"Advertencia: Resultado malformado o vacío del clasificador para el chunk: {sublist_labels} con tipología '{typology_stripped}'. Resultado: {result}")
                    # Continuar al siguiente chunk o manejar como un error si es crítico
                    # Por ahora, lo ignoramos y no actualizamos best_score

            except Exception as e_chunk:
                # Si clasificar_con_modelos lanza una excepción (todos los modelos fallaron para este chunk)
                # Se podría decidir si continuar con el siguiente chunk o fallar toda la clasificación.
                # La lógica original parece implicar que un fallo en un chunk no detiene los demás,
                # sino que se busca el "best_score" global.
                # Sin embargo, clasificar_con_modelos lanzará una excepción que detendría este bucle.
                # Para que continúe y pruebe otros chunks (y tome el mejor de los exitosos),
                # necesitaríamos capturar la excepción aquí y decidir.
                print(f"Error procesando chunk de etiquetas {sublist_labels} para '{typology_stripped}': {e_chunk}. Continuando con el siguiente chunk si hay.")
                # Si queremos que la clasificación general falle si un chunk falla, entonces no necesitamos este try-except interno.
                # La lógica original tenía el try-except principal más afuera.
                # Vamos a mantener la lógica de que si un chunk falla completamente (todos los modelos),
                # la clasificación general para ese typology podría fallar o continuar.
                # Si se quiere que falle toda la llamada a /categorize:
                # raise HTTPException(status_code=500, detail=f"Error clasificando chunk: {str(e_chunk)}")
                # Si se quiere continuar y ver si otros chunks dan resultado:
                continue # Pasa al siguiente sublist_labels

        # Después de probar todos los chunks:
        if best_score == -1.0 and best_label == "Sin Categoria":
            # Esto significa que o no hubo chunks, o todos los chunks fallaron de tal manera
            # que no se actualizó best_score (por ejemplo, si capturamos excepciones por chunk y continuamos)
            # o si la lista de candidate_labels estaba vacía.
            # En este punto, si no hay una mejor etiqueta, se queda "Sin Categoria".
            # Si se quiere lanzar un error si no se clasificó:
            # raise HTTPException(status_code=404, detail=f"No se pudo clasificar la tipología '{typology_stripped}'.")
            pass # Se devolverá "Sin Categoria"

        etiqueta_a = config["categoria_tipo"].get(best_label, "Sin Categoria")
        
        # Asegurar que probabilidad sea un número antes de redondear
        probabilidad_final = 0.0
        if isinstance(best_score, (int, float)) and best_score > -1.0 : # Solo si es un score válido
            probabilidad_final = round(best_score * 100, 2)
        
        return {
            "eatc-odd_typology_b": best_label,
            "eatc-odd_typology_a": etiqueta_a,
            "probabilidad": probabilidad_final,
            "modelo_usado": modelo_utilizado
        }

    except Exception as e:
        # Este es el try-except principal que estaba en tu código original.
        # Si `clasificar_con_modelos` lanza una excepción (todos los modelos fallaron para un chunk),
        # y no la capturamos dentro del bucle de chunks, llegará aquí.
        print(f"Error general en la clasificación para typology '{typology_stripped}': {str(e)}") # Loguear el error
        raise HTTPException(status_code=500, detail=f"Error en clasificación: {str(e)}")