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.
# Opera de forma stateless, recebendo prompts completos a cada chamada.

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

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. Opera de forma stateless.
    """
    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', 8096)
        self.temperature = config.get('temperature', 0.9)
        self.top_p = config.get('top_p', 0.9)
        self.max_image_size = (480, 480)

        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}...")
        # A implementação de atenção padrão é usada para máxima compatibilidade.
        self.model = MllamaForConditionalGeneration.from_pretrained(
            self.model_id,
            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)
            processed_image = self._preprocess_image(image_list[0]) if image_list else None
            
            # Gera a resposta
            assistant_response_text = self._generate_response(prompt_text, processed_image)
            
            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)
            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 (Stateless) 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))