assessment_agent.py

# Copyright 2025 - EWAAST Project
# Adapted from Google's appoint-ready architecture
# Licensed under the Apache License, Version 2.0

"""
EWAAST Assessment Agent

Core assessment logic with streaming output.
Now powered by Google Gemini 1.5 Pro (Real AI) via gemini_client.
"""

import json
import time
import os
import base64
import io
from typing import Generator
from PIL import Image

# Import Real AI Client
from gemini_client import get_gemini_client

# ===== MST VISUAL GUIDANCE =====

MST_GUIDANCE = {
    1: {"category": "Light", "visual": "Look for non-blanchable erythema (redness), warmth"},
    2: {"category": "Light", "visual": "Look for red to pink coloration, visible redness"},
    3: {"category": "Light", "visual": "Look for pink to light red, subtle redness"},
    4: {"category": "Medium", "visual": "Look for subtle darkening, warmth may be more reliable than color"},
    5: {"category": "Medium", "visual": "Look for color change rather than redness, warmth, shiny skin"},
    6: {"category": "Medium", "visual": "Look for darkening of skin, purple or brownish discoloration"},
    7: {"category": "Deep", "visual": "CRITICAL: Do NOT rely on redness. Look for purple/blue discoloration, warmth, induration"},
    8: {"category": "Deep", "visual": "CRITICAL: Do NOT rely on redness. Look for purple, maroon, or deep burgundy; may appear black"},
    9: {"category": "Deep", "visual": "CRITICAL: Do NOT rely on redness. Look for deep purple, indigo, or dark discoloration; warmth, texture changes"},
    10: {"category": "Deep", "visual": "CRITICAL: Do NOT rely on redness. Look for subtle purple/blue undertones; induration, warmth, localized heat"}
}

def _b64_to_pil(image_b64: str) -> Image.Image:
    """Convert base64 string to PIL Image."""
    try:
        if "," in image_b64:
            image_b64 = image_b64.split(",")[1]
        image_data = base64.b64decode(image_b64)
        return Image.open(io.BytesIO(image_data))
    except Exception as e:
        print(f"Error converting image: {e}")
        return Image.new('RGB', (100, 100), color='gray')


def classify_skin_tone(image_b64: str, context: str = "") -> dict:
    """
    Classify the Monk Skin Tone from an image using Gemini.
    """
    client = get_gemini_client()
    
    if client.is_available():
        try:
            image = _b64_to_pil(image_b64)
            result = client.analyze_mst(image)
            
            mst_value = result.get("mst_value", 5)
            # Normalize constraints
            mst_value = max(1, min(10, int(mst_value)))
            
            guidance = MST_GUIDANCE.get(mst_value, MST_GUIDANCE[5])
            
            return {
                "mst_value": mst_value,
                "category": guidance["category"],
                "visual_guidance": guidance["visual"],
                "confidence": result.get("confidence", 0.85)
            }
        except Exception as e:
            print(f"AI Classification failed: {e}")
    
    # Fallback / Demo Logic
    print("Falling back to heuristic MST classification")
    
    # Check if this is a "valid" fallback (user provided context) or just missing API
    if not context and not client.is_available():
        return {
            "mst_value": "Unknown",
            "category": "Unknown",
            "visual_guidance": "⚠️ AI Unavailable: Please add GOOGLE_API_KEY to .env or Space Secrets",
            "confidence": 0.0
        }

    # Rest of heuristic logic
    import re
    mst_match = re.search(r'MST[:\s]*(\d+)', context, re.IGNORECASE)
    if mst_match:
        mst_value = int(mst_match.group(1))
    elif 'deep' in context.lower() or 'dark' in context.lower():
        mst_value = 9
    elif 'medium' in context.lower():
        mst_value = 5
    else:
         # Default safe fallback - do not assume light skin if unknown
        mst_value = 5
        
    guidance = MST_GUIDANCE.get(mst_value, MST_GUIDANCE[5])
    return {
        "mst_value": mst_value,
        "category": guidance["category"],
        "visual_guidance": guidance["visual"] + " (Estimated)",
        "confidence": 0.5
    }

def generate_assessment_report(image_b64: str, mst_result: dict, context: str) -> dict:
    """
    Generate the clinical assessment report using Gemini.
    """
    client = get_gemini_client()
    
    if client.is_available():
        try:
            image = _b64_to_pil(image_b64)
            return client.assess_wound(image, mst_result, context)
        except Exception as e:
            print(f"AI Assessment failed: {e}")
            error_msg = str(e)

    # Fallback Logic
    return {
        "stage": "Assessment Failed",
        "rationale": f"AI Error: {error_msg}" if 'error_msg' in locals() else "AI Assessment Unavailable. Please check your GOOGLE_API_KEY configuration.",
        "care_plan": "1. Verify API Key in Settings\n2. Reload Application\n3. Check Server Logs for Details",
        "urgency": "N/A"
    }


def generate_preceptor_feedback(mst_value: int, wound_truth: dict, student_diagnosis: dict) -> dict:
    """
    Generate educational feedback comparing student diagnosis to truth.
    Acts as a "Clinical Nurse Educator" using EPUAP guidelines.
    """
    client = get_gemini_client()
    
    # Construct the scenario context for the AI
    prompt = f"""
    You are an expert Clinical Nurse Educator specializing in Wound Care and Health Equity.
    A student nurse has just assessed a patient. Provide feedback on their diagnosis.
    
    PATIENT CONTEXT:
    - Monk Skin Tone (MST): {mst_value} (Scale 1-10)
    
    GOLD STANDARD (CORRECT) DIAGNOSIS:
    - Wound Stage: {wound_truth.get('id')} ({wound_truth.get('name')})
    - Description: {wound_truth.get('description')}
    
    STUDENT NURSE DIAGNOSIS:
    - Stage: {student_diagnosis.get('stage')}
    - Tissue Type: {student_diagnosis.get('tissueType')}
    - Priority Action: {student_diagnosis.get('priorityAction')}
    
    YOUR TASK:
    Compare the student's diagnosis to the Gold Standard.
    Provide constructive feedback.
    
    CRITICAL EPUAP GUIDELINE RULES TO ENFORCE:
    1. If MST >= 7 (Deep Skin Tone) and the student relied on "Redness" or missed a DTI/Stage 1:
       - Educate them that "Redness is not a reliable indicator in dark skin".
       - Remind them to check for: Temperature (Warmth/Coolness), Induration (Firmness), and Texture changes.
    2. If the student missed a DTI (Deep Tissue Injury):
       - Explain that DTIs can look like a bruise or be "boggy" on palpation.
    
    OUTPUT FORMAT (JSON):
    {{
        "is_correct": boolean,
        "title": "Short title (e.g., 'Correct Diagnosis' or 'Missed Stage 1')",
        "feedback": "2-3 sentences of direct feedback to the student.",
        "learning_point": "State the relevant EPUAP guideline or clinical pearl."
    }}
    """
    
    if client.is_available():
        try:
            from google.genai import types
            response = client.client.models.generate_content(
                model=client.model_id,
                contents=prompt,
                config=types.GenerateContentConfig(
                    response_mime_type="application/json"
                )
            )
            return json.loads(response.text)
        except Exception as e:
            print(f"AI Preceptor Error: {e}")
            
    # Fallback Logic if AI fails
    is_correct = student_diagnosis.get('stage') == wound_truth.get('id')
    return {
        "is_correct": is_correct,
        "title": "Diagnosis Result" if is_correct else "Incorrect Diagnosis",
        "feedback": "AI Preceptor unavailable. Please compare your answer to the Gold Standard manually.",
        "learning_point": "EPUAP Guidelines emphasize checking for temperature and induration differences on dark skin."
    }


def stream_assessment(image_b64: str, context: str) -> Generator[str, None, None]:
    """
    Stream the wound assessment process.
    """
    
    # Step 1: Acknowledge
    yield json.dumps({
        "step": "received",
        "thinking": "Image received. validating...",
        "data": {}
    })
    time.sleep(0.5)
    
    # Step 2: Validate Image (New Step with Gemini)
    client = get_gemini_client()
    if client.is_available():
        try:
            img = _b64_to_pil(image_b64)
            val_result = client.validate_wound_image(img)
            if not val_result.get("is_valid", True):
                yield json.dumps({
                    "step": "error",
                    "thinking": f"Image Rejected: {val_result.get('reason')}",
                    "data": {"error": val_result.get('reason')}
                })
                yield json.dumps({"event": "end"})
                return
        except Exception:
            pass

    # Step 3: Classify skin tone
    yield json.dumps({
        "step": "classifying_mst",
        "thinking": "Analyzing skin tone using Monk Skin Tone (MST) scale with Gemini Vision...",
        "data": {}
    })
    
    mst_result = classify_skin_tone(image_b64, context)
    
    yield json.dumps({
        "step": "mst_complete",
        "thinking": f"Identified MST {mst_result['mst_value']} ({mst_result['category']}). Adjusted visual guidance: {mst_result['visual_guidance']}",
        "data": mst_result
    })
    
    # Step 4: Analyze wound
    yield json.dumps({
        "step": "analyzing_wound",
        "thinking": f"Gemini is analyzing wound features. Looking for MST-specific signs (e.g., discoloration vs redness)...",
        "data": {}
    })
    
    # Step 5: Generate report
    report = generate_assessment_report(image_b64, mst_result, context)
    
    yield json.dumps({
        "step": "complete",
        "thinking": "Assessment complete. Clinical rationale generated.",
        "data": {
            "mst": mst_result,
            "report": report
        }
    })
    
    yield json.dumps({"event": "end"})