ragflow_preprocess/agents/vision.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Agent LLM pour l'analyse visuelle d'images
"""

import base64
import json
import requests
from typing import List, Optional, Union

from .base import LLMBaseAgent

class VisionAgent(LLMBaseAgent):
    """
    Agent LLM spécialisé dans l'analyse d'images
    """
    
    def __init__(self, model_name: str = "llama3.2-vision:90b", **config):
        """
        Initialise l'agent de vision
        
        Args:
            model_name (str): Nom du modèle de vision (défaut: llama3.2-vision:90b)
            **config: Paramètres de configuration supplémentaires
        """
        super().__init__(model_name, **config)
        
        # Définir les prompts par défaut selon la langue
        self.prompts = {
            "fr": {
                "schéma": "Décris en détail ce schéma. Quelle information principale est-elle représentée ? "
                          "Comment les éléments sont-ils organisés ? Quelle est la signification des différentes parties ?",
                
                "tableau": "Analyse ce tableau. Quel type de données contient-il ? "
                           "Quelles sont les informations importantes ? Résume son contenu et sa structure.",
                
                "formule": "Analyse cette formule ou équation mathématique. "
                           "Que représente-t-elle ? Quelle est sa signification et son application ?",
                
                "texte": "Lis et résume ce texte. Quel est le contenu principal ? "
                         "Quels sont les points importants à retenir ?",
                
                "autre": "Décris en détail ce que tu vois sur cette image. "
                         "Quel est le contenu principal ? Quelle information est présentée ?"
            },
            "en": {
                "schéma": "Describe this diagram in detail. What main information is represented? "
                          "How are the elements organized? What is the meaning of the different parts?",
                
                "tableau": "Analyze this table. What kind of data does it contain? "
                           "What is the important information? Summarize its content and structure.",
                
                "formule": "Analyze this mathematical formula or equation. "
                           "What does it represent? What is its meaning and application?",
                
                "texte": "Read and summarize this text. What is the main content? "
                         "What are the important points to remember?",
                
                "autre": "Describe in detail what you see in this image. "
                         "What is the main content? What information is presented?"
            }
        }
    
    def generate(self, prompt: Optional[str] = "", images: Optional[List[bytes]] = None, 
                 selection_type: str = "autre", context: Optional[str] = "") -> str:
        """
        Génère une description ou une analyse d'une image
        
        Args:
            prompt (str, optional): Prompt personnalisé (si vide, utilise le prompt par défaut)
            images (List[bytes], optional): Liste d'images en bytes
            selection_type (str): Type de contenu ("schéma", "tableau", "formule", "texte", "autre")
            context (str, optional): Contexte textuel associé à la sélection
            
        Returns:
            str: La réponse générée par le modèle
        """
        if not images:
            return "Aucune image fournie pour l'analyse."
        
        # Utiliser le prompt par défaut si aucun n'est fourni
        if not prompt:
            lang = self.config.get("language", "fr")
            prompts_dict = self.prompts.get(lang, self.prompts["fr"])
            prompt = prompts_dict.get(selection_type, prompts_dict["autre"])
            
            # Ajouter le contexte si disponible
            if context:
                prompt = f"Contexte: {context}\n\n{prompt}"
        
        try:
            # Pour chaque image, encoder en base64
            base64_images = []
            for image in images:
                if isinstance(image, bytes):
                    base64_image = base64.b64encode(image).decode("utf-8")
                    base64_images.append(base64_image)
            
            # Construire la payload pour l'API Ollama
            payload = {
                "model": self.model_name,
                "prompt": prompt,
                "images": base64_images,
                "options": {
                    "temperature": self.config.get("temperature", 0.2),
                    "top_p": self.config.get("top_p", 0.95),
                    "top_k": self.config.get("top_k", 40),
                    "num_predict": self.config.get("max_tokens", 1024)
                }
            }
            
            # Dans une implémentation réelle, envoyer la requête à l'API Ollama
            # response = requests.post(f"{self.endpoint}/api/generate", json=payload)
            # json_response = response.json()
            # return json_response.get("response", "")
            
            # Pour cette démonstration, retourner une réponse simulée
            if selection_type == "schéma":
                return "Le schéma présenté illustre un processus structuré avec plusieurs étapes interconnectées. " \
                       "On peut observer une organisation hiérarchique des éléments, avec des flèches indiquant " \
                       "le flux d'information ou la séquence d'opérations. Les différentes composantes sont " \
                       "clairement délimitées et semblent représenter un workflow ou un système de classification."
            
            elif selection_type == "tableau":
                return "Ce tableau contient plusieurs colonnes et rangées de données structurées. " \
                       "Il présente une organisation systématique d'informations, probablement des " \
                       "valeurs numériques ou des catégories. Les en-têtes indiquent le type de données " \
                       "dans chaque colonne, et l'ensemble forme une matrice cohérente d'informations liées."
            
            elif selection_type == "formule":
                return "Cette formule mathématique représente une relation complexe entre plusieurs variables. " \
                       "Elle utilise divers opérateurs et symboles mathématiques pour exprimer un concept ou " \
                       "une règle. La structure suggère qu'il s'agit d'une équation importante dans son domaine, " \
                       "possiblement liée à un phénomène physique ou à un modèle théorique."
            
            else:
                return "L'image montre un contenu visuel structuré qui présente des informations importantes " \
                       "dans le contexte du document. Les éléments visuels sont organisés de manière à faciliter " \
                       "la compréhension d'un concept ou d'un processus spécifique. La qualité et la disposition " \
                       "des éléments suggèrent qu'il s'agit d'une représentation professionnelle destinée à " \
                       "communiquer efficacement l'information."
        
        except Exception as e:
            return f"Erreur lors de l'analyse de l'image: {str(e)}"