aduc_framework/managers/llama_multimodal_manager.py

# aduc_framework/managers/llama_multimodal_manager.py
#
# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
#
# Versão 1.3.0 (Comprehensive Logging & Stable Attention)
#
# Manager especialista em conversas multimodais. Esta versão adiciona
# logging detalhado em todos os pontos críticos para garantir que não haja
# falhas silenciosas e para fornecer visibilidade total do fluxo de dados.

import yaml
import os
import logging
import torch
from PIL import Image
from typing import List, Dict, Any, Optional

from transformers import MllamaForConditionalGeneration, MllamaProcessor
from huggingface_hub import HfFolder

from ..tools.hardware_manager import hardware_manager

# Logger específico para este manager
logger = logging.getLogger(__name__)

class LlamaMultiModalManager:
    """
    Gerencia uma única instância do Llama 3.2 Vision, com logging robusto
    para rastrear todo o ciclo de vida de uma requisição.
    """
    def __init__(self, config: dict):
        logger.debug("LLAMA_MANAGER: Iniciando __init__.")
        self.hf_token = os.getenv("HF_TOKEN") or HfFolder.get_token()
        if not self.hf_token:
            raise ValueError("Token da Hugging Face não encontrado. Faça login via `huggingface-cli login` ou defina a variável de ambiente HF_TOKEN.")

        self.model_id = config['model_id']
        self.max_new_tokens = config.get('max_new_tokens', 2048)
        self.temperature = config.get('temperature', 0.6)
        self.top_p = config.get('top_p', 0.9)
        self.max_image_size = (1120, 1120)

        logger.info(f"LLAMA_MANAGER: Carregando processador do modelo: {self.model_id}...")
        self.processor = MllamaProcessor.from_pretrained(self.model_id, token=self.hf_token)
        logger.info("LLAMA_MANAGER: Processador carregado.")

        logger.info(f"LLAMA_MANAGER: Carregando modelo base: {self.model_id}...")
        self.model = MllamaForConditionalGeneration.from_pretrained(
            self.model_id,
            attn_implementation="flex_attention
            torch_dtype=torch.bfloat16,
            device_map="auto"
            #token=self.hf_token
        )
        logger.info("LLAMA_MANAGER: Modelo carregado e mapeado para o dispositivo.")
        logger.debug("LLAMA_MANAGER: __init__ concluído.")

    def _preprocess_image(self, image: Image.Image) -> Image.Image:
        """Garante que a imagem esteja no formato RGB e dentro do tamanho máximo."""
        logger.debug(f"Pré-processando imagem. Tamanho original: {image.size}, Modo: {image.mode}")
        img = image.convert("RGB")
        if img.width > self.max_image_size[0] or img.height > self.max_image_size[1]:
            original_size = img.size
            img.thumbnail(self.max_image_size, Image.Resampling.LANCZOS)
            logger.debug(f"Imagem redimensionada de {original_size} para {img.size}.")
        return img

    @torch.inference_mode()
    def process_turn(self, prompt_text: str, image_list: Optional[List[Image.Image]] = None) -> str:
        """
        Ponto de entrada para processar uma requisição. Lida com a orquestração
        interna e o tratamento de exceções.
        """
        logger.info(f"LLAMA_MANAGER: Recebido novo turno. Comprimento do prompt: {len(prompt_text)}, Imagens: {len(image_list) if image_list else 0}.")
        image_list = image_list or []
        
        try:
            # Prepara a imagem (se houver)
            last_image_processed = self._preprocess_image(image_list[0]) if image_list else None
            
            # Gera a resposta
            assistant_response_text = self._generate_response(prompt_text, last_image_processed)
            
            logger.info("LLAMA_MANAGER: Turno processado com sucesso.")
            return assistant_response_text

        except Exception as e:
            logger.error(f"LLAMA_MANAGER: ERRO BRUTO DURANTE O PROCESSAMENTO DO TURNO: {e}", exc_info=True)
            # Propaga a exceção para que as camadas superiores (Composer, Orchestrator) possam tratá-la.
            raise e

    def _generate_response(self, prompt_str: str, image: Optional[Image.Image] = None) -> str:
        """
        Função de inferência interna: processa, gera e decodifica.
        """
        # 1. Processamento da Entrada
        logger.debug("---> LLAMA_MANAGER: Etapa de Processamento da Entrada Iniciada.")
        logger.debug(f"Texto do prompt recebido (primeiros 500 chars):\n---\n{prompt_str[:500]}\n---")
        
        inputs = self.processor(
            text=prompt_str,
            images=[image] if image else None,
            return_tensors="pt"
        ).to(self.model.device)
        logger.debug(f"Entrada processada e movida para o dispositivo: {self.model.device}. Shape dos input_ids: {inputs['input_ids'].shape}")

        # 2. Geração do Modelo
        logger.debug("---> LLAMA_MANAGER: Etapa de Geração do Modelo Iniciada.")
        generate_ids = self.model.generate(
            **inputs,
            max_new_tokens=self.max_new_tokens,
            do_sample=True,
            temperature=self.temperature,
            top_p=self.top_p,
        )
        logger.debug(f"Geração concluída. Shape dos IDs de saída: {generate_ids.shape}")

        # 3. Decodificação da Saída
        logger.debug("---> LLAMA_MANAGER: Etapa de Decodificação da Saída Iniciada.")
        input_ids_len = inputs['input_ids'].shape[1]
        output_ids = generate_ids[:, input_ids_len:]
        
        response_text_raw = self.processor.batch_decode(
            output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]
        
        logger.debug(f"<--- LLAMA_MANAGER: Resposta bruta decodificada (antes da limpeza):\n---\n{response_text_raw}\n---")
        
        return response_text_raw.strip()


# --- Placeholder e Instanciação Singleton ---
class LlamaMultiModalPlaceholder:
    def __init__(self, reason: str = "Motivo desconhecido"):
        logger.error(f"LlamaMultiModalManager não inicializado. Razão: {reason}. Placeholder em uso.")
        self.reason = reason
    def process_turn(self, *args, **kwargs):
        return json.dumps({"error": f"Especialista Llama MultiModal indisponível. Razão: {self.reason}"})

try:
    with open("config.yaml", 'r') as f:
        config = yaml.safe_load(f)
    llama_config = config['specialists'].get('llama_multimodal')
    if llama_config and llama_config.get('gpus_required', 0) > 0:
        hardware_manager.allocate_gpus('LlamaMultiModal', llama_config['gpus_required'])
        llama_multimodal_manager_singleton = LlamaMultiModalManager(config=llama_config)
        logger.info("Especialista Llama MultiModal (Stateful) pronto.")
    else:
        llama_multimodal_manager_singleton = LlamaMultiModalPlaceholder("Não habilitado ou sem gpus_required na config.yaml")
except Exception as e:
    logger.critical(f"Falha CRÍTICA ao inicializar o LlamaMultiModalManager: {e}", exc_info=True)
    llama_multimodal_manager_singleton = LlamaMultiModalPlaceholder(reason=str(e))