agents/bio_analysis.py

import json
from langchain_core.messages import HumanMessage
from llm_factory import get_llm
import base64

# Initialize both models
text_llm = get_llm(model_type="text", temperature=0.1)
vision_llm = get_llm(model_type="vision", temperature=0.1)

import io
from pypdf import PdfReader
from services.rag_service import rag_service

async def run_bio_analysis_agent(file_data: bytes, mime_type: str) -> dict:
    try:
        print(f"Processing file with mime_type: {mime_type}")
        
        prompt_text = """
        Tu es un Expert Biologiste et Pharmacien.
        
        TA MISSION :
        Analyser ce rapport de laboratoire et extraire les données structurées.
        
        FORMAT DE SORTIE ATTENDU (JSON uniquement) :
        {
          "patient_name": "Nom du patient ou 'Inconnu'",
          "date": "Date du rapport ou 'Inconnue'",
          "critical_alerts": ["Liste des valeurs critiques nécessitant une attention immédiate"],
          "results": [
            {
              "test_name": "Nom de l'analyse (ex: Hémoglobine)",
              "value": "Valeur numérique (ex: 10.5)",
              "unit": "Unité (ex: g/dL)",
              "ref_range": "Plage de référence (ex: 12-16)",
              "status": "LOW" | "HIGH" | "NORMAL",
              "interpretation": "Brève interprétation clinique"
            }
          ],
          "global_interpretation": "Synthèse globale de l'état du patient basés sur les résultats.",
          "recommendations": ["Actions recommandées"]
        }
        
        RÉPOND UNIQUEMENT AVEC LE JSON. PAS DE MARKDOWN.
        """

        message_content = []
        selected_llm = text_llm # Default to text

        if mime_type == "application/pdf":
            # Extract text from PDF
            try:
                pdf_file = io.BytesIO(file_data)
                reader = PdfReader(pdf_file)
                extracted_text = ""
                for page in reader.pages:
                    extracted_text += page.extract_text() + "\n"
                
                print(f"Extracted {len(extracted_text)} chars from PDF")
                message_content.append({"type": "text", "text": prompt_text + "\n\nVOICI LE CONTENU DU RAPPORT :\n" + extracted_text})
                selected_llm = text_llm # Use DeepSeek for text
            except Exception as pdf_err:
                print(f"PDF Extraction failed: {pdf_err}")
                return {"error": "PDF parsing failed", "details": str(pdf_err)}

        elif mime_type.startswith("image/"):
            # Use Vision for images
            b64_data = base64.b64encode(file_data).decode("utf-8")
            message_content.append({"type": "text", "text": prompt_text})
            message_content.append({
                "type": "image_url",
                "image_url": {"url": f"data:{mime_type};base64,{b64_data}"}
            })
            selected_llm = vision_llm # Use Gemini for Vision
        else:
            return {"error": "Unsupported file type", "details": mime_type}
        
        message = HumanMessage(content=message_content)
        
        response = await selected_llm.ainvoke([message])
        
        # Clean up response to ensure valid JSON
        content = response.content.strip()
        if content.startswith("```json"):
            content = content[7:-3]
        elif content.startswith("```"):
            content = content[3:-3]
            
        try:
            # SAVE TO RAG (Mocking 'default' patient ID for this POC)
            rag_service.add_document("default", f"Résultat Analyse Labo:\n{content}", "lab_result")
            
            return json.loads(content)
        except json.JSONDecodeError as e:
            print(f"JSON Parse Error: {e}")
            print(f"Raw Content: {content}")
            
            # Even on JSON error, save raw text so context is kept
            rag_service.add_document("default", f"Tentative Analyse Labo (Brut):\n{content}", "lab_result_raw")
            
            return {
                "error": "Failed to parse JSON",
                "raw_content": content
            }
            
    except Exception as e:
        print(f"CRITICAL ERROR in run_bio_analysis_agent: {str(e)}")
        import traceback
        traceback.print_exc()
        return {
            "error": "Analysis failed",
            "details": str(e)
        }