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, Dict
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"""
        # Améliore le rendu des fractions
        text = re.sub(r'\\frac\{([^}]*)\}\{([^}]*)\}', r'$$\\frac{\1}{\2}$$', text)
        # Gère les délimiteurs left/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"""
        # Améliore l'espacement autour des expressions inline
        text = re.sub(r'(?<!\\)\$([^$]+)(?<!\\)\$', r' $\1$ ', text)
        # Gère les cas spéciaux où il y a déjà des $$ mais en 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 et de l'espacement
        text = re.sub(r'\$\$(.*?)\$\$', r'\n\n$$\1$$\n\n', text, flags=re.DOTALL)
        return text

    @staticmethod
    def format_special_environments(text: str) -> str:
        """Formate les environnements mathématiques spéciaux"""
        # Gère les environnements align, gather, etc.
        envs = ['align', 'gather', 'equation', 'array', 'matrix']
        for env in envs:
            pattern = f'\\\\begin{{{env}}}(.*?)\\\\end{{{env}}}'
            text = re.sub(pattern, f'\n\n$$\\\\begin{{{env}}}\\1\\\\end{{{env}}}$$\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('\\)', '$')
        
        # Application des différentes améliorations
        text = LatexFormatter.cleanup_latex_fractions(text)
        text = LatexFormatter.format_special_environments(text)
        text = LatexFormatter.format_inline_math(text)
        text = LatexFormatter.format_display_math(text)
        
        # Assure un bon espacement des équations
        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 + " Utilise la notation LaTeX appropriée avec $...$ pour les expressions en ligne et $$...$$ pour les équations importantes. Pour les fractions, utilise \\frac{num}{den}."
                ]
            )
            return response
        except Exception as e:
            logger.error(f"Erreur lors de l'analyse de l'image: {e}")
            raise


def setup_latex_display():
    """Configure l'affichage LaTeX dans Streamlit"""
    st.markdown("""
        <script>
        window.MathJax = {
            tex: {
                inlineMath: [['$', '$'], ['\\\\(', '\\\\)']],
                displayMath: [['$$', '$$'], ['\\\\[', '\\\\]']],
                processEscapes: true,
                macros: {
                    R: "{\\\\mathbb{R}}",
                    N: "{\\\\mathbb{N}}",
                    Z: "{\\\\mathbb{Z}}",
                    vecv: ["\\\\begin{pmatrix}#1\\\\\\\\#2\\\\\\\\#3\\\\end{pmatrix}", 3]
                }
            }
        };
        </script>
        <style>
        /* Styles de base pour 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 des fractions */
        .katex .frac-line {
            border-bottom-width: 0.08em;
        }
        .katex .mfrac .frac-line {
            margin: 0.1em 0;
        }
        /* Style des matrices */
        .katex .mord.matrix {
            margin: 0.2em 0;
        }
        /* Amélioration des indices et exposants */
        .katex .msupsub {
            font-size: 0.9em;
        }
        /* Meilleure lisibilité des symboles */
        .katex .mathdefault {
            color: inherit;
        }
        /* Adaptation mobile */
        @media (max-width: 768px) {
            .katex { 
                font-size: 1.1em !important; 
            }
            .katex-display { 
                padding: 0.3em;
                margin: 1em 0 !important;
                font-size: 0.9em !important;
            }
        }
        /* Conteneur de réponse */
        .response-container {
            margin: 1em 0;
            padding: 1em;
            border-radius: 8px;
            background: rgba(255, 255, 255, 0.05);
        }
        /* Mise en évidence des équations importantes */
        .important-equation {
            border-left: 3px solid #4CAF50;
            padding-left: 1em;
        }
        </style>
    """, unsafe_allow_html=True)

def stream_response(container, response: Generator) -> None:
    """Gère le streaming de la réponse avec support LaTeX amélioré"""
    mode = 'starting'
    thinking_placeholder = None
    answer_placeholder = None
    thinking_text = ""
    answer_text = ""
    
    setup_latex_display()
    
    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 = LatexFormatter.enhance_latex_display(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. 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",
        initial_sidebar_state="collapsed"
    )
    
    st.title("Mariam M-0")
    setup_latex_display()
    
    # 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 contenant des mathématiques",
        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()