app.py

import streamlit as st
from google import genai
from google.genai import types
from PIL import Image
import json
import logging
import re
from typing import Optional, Generator, Any, Tuple
import sys
from pathlib import Path

# Configuration du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(Path('app.log'))
    ]
)
logger = logging.getLogger(__name__)


class LatexFormatter:
    """Classe améliorée pour le formatage LaTeX"""
    
    @staticmethod
    def cleanup_latex_fractions(text: str) -> str:
        """Nettoie et formate correctement les fractions LaTeX"""
        # Remplace \frac{x}{y} par \frac{x}{y}
        text = re.sub(r'\\frac\{([^}]*)\}\{([^}]*)\}', r'$$\\frac{\1}{\2}$$', text)
        # Gère les cas spéciaux comme \left et \right
        text = re.sub(r'\\left\(', r'\\left(', text)
        text = re.sub(r'\\right\)', r'\\right)', text)
        return text

    @staticmethod
    def format_inline_math(text: str) -> str:
        """Formate les expressions mathématiques en ligne"""
        # Ajoute des espaces autour des expressions inline
        text = re.sub(r'(?<!\\)\$([^$]+)(?<!\\)\$', r' $\1$ ', text)
        return text

    @staticmethod
    def format_display_math(text: str) -> str:
        """Formate les expressions mathématiques en mode display"""
        # Ajoute des sauts de ligne autour des expressions display
        text = re.sub(r'\$\$(.*?)\$\$', r'\n\n$$\1$$\n\n', text, flags=re.DOTALL)
        return text

    @staticmethod
    def enhance_latex_display(text: str) -> str:
        """Améliore globalement l'affichage LaTeX"""
        # Prétraitement
        text = text.replace('\\[', '$$').replace('\\]', '$$')
        text = text.replace('\\(', '$').replace('\\)', '$')
        
        # Nettoyage des fractions
        text = LatexFormatter.cleanup_latex_fractions(text)
        
        # Formatage du mode inline
        text = LatexFormatter.format_inline_math(text)
        
        # Formatage du mode display
        text = LatexFormatter.format_display_math(text)
        
        # Assure que chaque ligne d'équation est bien espacée
        text = re.sub(r'(\n\s*\$\$[^$]+\$\$)', r'\n\n\1\n\n', text)
        
        return text


    



class GeminiClient:
    """Classe pour gérer les interactions avec l'API Gemini"""
    def __init__(self, api_key: str):
        self.client = None
        self.init_client(api_key)
        
    def init_client(self, api_key: str) -> None:
        """Initialise le client Gemini"""
        try:
            self.client = genai.Client(
                api_key=api_key,
                http_options={'api_version': 'v1alpha'}
            )
        except Exception as e:
            logger.error(f"Erreur d'initialisation du client Gemini: {e}")
            raise RuntimeError(f"Impossible d'initialiser le client Gemini: {e}")

    def analyze_image(self, image: Image.Image, prompt: str, model_name: str) -> Generator:
        """Analyse une image avec Gemini"""
        if not self.client:
            raise RuntimeError("Client Gemini non initialisé")
            
        try:
            response = self.client.models.generate_content_stream(
                model=model_name,
                config={'thinking_config': {'include_thoughts': True}},
                contents=[
                    image,
                    prompt + " Si ta réponse contient des expressions mathématiques, utilise la notation LaTeX avec les délimiteurs appropriés ($ pour inline, $$ pour les blocs)."
                ]
            )
            return response
        except Exception as e:
            logger.error(f"Erreur lors de l'analyse de l'image: {e}")
            raise

def format_text_with_latex(text: str) -> str:
    """Formate le texte en appliquant les améliorations LaTeX"""
    formatter = LatexFormatter()
    text = formatter.format_inline_latex(text)
    text = formatter.format_block_latex(text)
    text = formatter.enhance_latex_display(text)
    return text

def stream_response(container, response: Generator) -> None:
    """Fonction de streaming améliorée avec meilleur support LaTeX"""
    
    # Configuration CSS améliorée pour le rendu LaTeX
    st.markdown("""
        <style>
        /* Amélioration du rendu LaTeX */
        .katex { 
            font-size: 1.2em !important; 
            padding: 0.2em 0;
        }
        .katex-display { 
            margin: 1.5em 0 !important;
            overflow: auto hidden;
            background: rgba(248, 249, 250, 0.05);
            padding: 0.5em;
            border-radius: 4px;
        }
        /* Amélioration de l'espacement des fractions */
        .katex .frac-line {
            border-bottom-width: 0.08em;
        }
        .katex .mfrac .frac-line {
            margin: 0.1em 0;
        }
        /* Meilleur rendu sur mobile */
        @media (max-width: 768px) {
            .katex { font-size: 1.1em !important; }
            .katex-display { 
                padding: 0.3em;
                margin: 1em 0 !important;
            }
        }
        </style>
    """, unsafe_allow_html=True)
        
    try:
        for chunk in response:
            logger.debug(f"Chunk reçu: {chunk}")
            
            if not isinstance(chunk, (dict, types.GenerateContentResponse)):
                logger.warning(f"Format de chunk invalide reçu: {type(chunk)}")
                continue
                
            try:
                candidates = getattr(chunk, 'candidates', None)
                if not candidates or not len(candidates):
                    logger.warning("Pas de candidats dans le chunk")
                    continue
                    
                content = getattr(candidates[0], 'content', None)
                if not content:
                    logger.warning("Pas de contenu dans le premier candidat")
                    continue
                    
                parts = getattr(content, 'parts', [])
                for part in parts:
                    has_thought = False
                    try:
                        has_thought = hasattr(part, 'thought') and part.thought
                    except Exception as e:
                        logger.warning(f"Erreur lors de la vérification de thought: {e}")
                    
                    text = getattr(part, 'text', '')
                    if not text:
                        continue
                    
                    # Formatage LaTeX du texte
                    formatted_text = format_text_with_latex(text)
                        
                    if has_thought:
                        if mode != "thinking":
                            if thinking_placeholder is None:
                                with container.expander("Voir le raisonnement", expanded=False):
                                    thinking_placeholder = st.empty()
                            mode = "thinking"
                        thinking_text += formatted_text
                        thinking_placeholder.markdown(thinking_text)
                    else:
                        if mode != "answering":
                            if answer_placeholder is None:
                                answer_placeholder = container.empty()
                                container.subheader("Réponse")
                            mode = "answering"
                        answer_text += formatted_text
                        answer_placeholder.markdown(answer_text)
                        
            except json.JSONDecodeError as e:
                logger.error(f"Erreur de décodage JSON: {e}")
                continue
            except Exception as e:
                logger.error(f"Erreur lors du traitement d'un chunk: {e}")
                logger.debug(f"Contenu du chunk problématique: {chunk}")
                continue
                
    except Exception as e:
        logger.error(f"Erreur fatale dans le streaming de la réponse: {e}")
        if not answer_text and not thinking_text:
            container.error("Une erreur est survenue lors de l'analyse de l'image. Veuillez réessayer.")
        raise
    finally:
        if not answer_text and not thinking_text:
            container.warning("Aucune réponse n'a pu être générée. Veuillez réessayer.")

def validate_image(uploaded_file) -> Optional[Image.Image]:
    """Valide et ouvre une image téléchargée"""
    try:
        image = Image.open(uploaded_file)
        return image
    except Exception as e:
        logger.error(f"Erreur lors de l'ouverture de l'image: {e}")
        st.error("L'image n'a pas pu être ouverte. Veuillez vérifier le format.")
        return None

def main():
    st.set_page_config(
        page_title="Mariam M-0",
        page_icon="🔍",
        layout="wide"
    )
    
    st.title("Mariam M-0")
    
    # Configuration supplémentaire pour LaTeX
    st.markdown("""
        <style>
        /* Styles globaux pour LaTeX */
        .latex-container { 
            margin: 1em 0;
            padding: 0.5em;
            background: #f9f9f9;
            border-radius: 4px;
        }
        /* Amélioration de l'espacement des équations */
        .katex-display { 
            margin: 1.5em 0 !important; 
            padding: 0.5em 0;
        }
        /* Amélioration du rendu sur mobile */
        @media (max-width: 768px) {
            .katex { font-size: 1em !important; }
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Récupération de la clé API
    try:
        api_key = st.secrets["GEMINI_API_KEY"]
    except Exception as e:
        logger.error(f"Erreur dans la récupération des secrets: {e}")
        st.error("Erreur: Impossible d'accéder aux secrets de l'application.")
        return
    
    # Initialisation du client
    try:
        gemini_client = GeminiClient(api_key)
    except Exception as e:
        st.error(f"Erreur lors de l'initialisation du client Gemini: {e}")
        return
    
    # Interface utilisateur
    uploaded_file = st.file_uploader(
        "Choisissez une image géométrique",
        type=['png', 'jpg', 'jpeg'],
        help="Formats supportés: PNG, JPG, JPEG"
    )
    
    if uploaded_file:
        image = validate_image(uploaded_file)
        if image:
            st.image(image, caption="Image téléchargée", use_container_width=True)
            
            model_name = "gemini-2.0-flash-thinking-exp-01-21"
            prompt = "Résous cet exercice mathématique. La réponse doit être bien présentée et espacée pour faciliter la lecture. Réponds en français et utilise la notation LaTeX pour toutes les expressions mathématiques."
            
            if st.button("Analyser l'image", type="primary"):
                response_container = st.container()
                
                with st.spinner("Analyse en cours..."):
                    try:
                        response = gemini_client.analyze_image(image, prompt, model_name)
                        stream_response(response_container, response)
                    except Exception as e:
                        logger.error(f"Erreur lors de l'analyse: {e}", exc_info=True)
                        st.error("Une erreur est survenue lors de l'analyse. Veuillez réessayer.")

if __name__ == "__main__":
    main()