feat: Integrate Real AI (Gemini 1.5 Pro) - Replaced mock logic with GeminiClient - Added google-generativeai dependency - Enabled real MST classification and wound staging

assessment_agent.py

@@ -6,31 +6,19 @@
 EWAAST Assessment Agent
 
 Core assessment logic with streaming output.
-Adapted from Google's interview_simulator.py pattern.
+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
 
-# ===== DEMO MODE =====
-# Set DEMO_MODE=true to use simulated responses (no model loading)
-DEMO_MODE = os.environ.get("DEMO_MODE", "false").lower() == "true"
-
-# Try to import MedGemma client, fall back to mock (or use mock if DEMO_MODE)
-if DEMO_MODE:
-    MEDGEMMA_AVAILABLE = False
-    print("🎭 DEMO_MODE enabled: Using simulated MST-aware responses")
-else:
-    try:
-        from medgemma_client import medgemma_get_text_response
-        MEDGEMMA_AVAILABLE = True
-    except ImportError:
-        MEDGEMMA_AVAILABLE = False
-
-from cache import cache
-
+# Import Real AI Client
+from gemini_client import get_gemini_client
 
 # ===== MST VISUAL GUIDANCE =====
 
@@ -47,138 +35,93 @@ MST_GUIDANCE = {
     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.
-    
-    Args:
-        image_b64: Base64-encoded image data
-        context: Optional patient context (may contain MST info for demo mode)
-        
-    Returns:
-        Dictionary with mst_value, category, visual_guidance, confidence
+    Classify the Monk Skin Tone from an image using Gemini.
     """
-    if MEDGEMMA_AVAILABLE:
-        # Use MedGemma VQA for classification
-        messages = [
-            {
-                "role": "system",
-                "content": [{
-                    "type": "text",
-                    "text": """You are a clinical assistant trained to identify the Monk Skin Tone (MST) scale value.
-                    The MST scale ranges from 1 (lightest) to 10 (deepest).
-                    Analyze the provided image and return ONLY a JSON object with:
-                    {"mst_value": <1-10>, "confidence": <0.0-1.0>}"""
-                }]
-            },
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
-                    {"type": "text", "text": "What is the Monk Skin Tone value for this patient?"}
-                ]
-            }
-        ]
-        
-        response = medgemma_get_text_response(messages)
+    client = get_gemini_client()
+    
+    if client.is_available():
         try:
-            result = json.loads(response)
+            image = _b64_to_pil(image_b64)
+            result = client.analyze_mst(image)
+            
             mst_value = result.get("mst_value", 5)
-        except:
-            mst_value = 5  # Default to medium if parsing fails
+            # 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")
+    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:
-        # Demo/Mock mode: extract MST from context if available
-        import re
-        mst_match = re.search(r'MST[:\s]*(\d+)', context, re.IGNORECASE)
-        if mst_match:
-            mst_value = int(mst_match.group(1))
-            mst_value = max(1, min(10, mst_value))  # Clamp to 1-10
-        else:
-            # Default based on keywords in context
-            if 'deep' in context.lower() or 'dark' in context.lower():
-                mst_value = 9
-            elif 'medium' in context.lower():
-                mst_value = 5
-            else:
-                mst_value = 2  # Default to light for demo
-    
+        mst_value = 2
+        
     guidance = MST_GUIDANCE.get(mst_value, MST_GUIDANCE[5])
-    
     return {
         "mst_value": mst_value,
         "category": guidance["category"],
         "visual_guidance": guidance["visual"],
-        "confidence": 0.92 if DEMO_MODE else 0.85
+        "confidence": 0.5
     }
 
 
 def generate_assessment_report(image_b64: str, mst_result: dict, context: str) -> dict:
     """
-    Generate the clinical assessment report using MedGemma.
-    
-    Args:
-        image_b64: Base64-encoded wound image
-        mst_result: MST classification result
-        context: Optional patient context
-        
-    Returns:
-        Dictionary with stage, rationale, care_plan, urgency
+    Generate the clinical assessment report using Gemini.
     """
-    mst_value = mst_result["mst_value"]
-    visual_guidance = mst_result["visual_guidance"]
+    client = get_gemini_client()
     
-    if MEDGEMMA_AVAILABLE:
-        messages = [
-            {
-                "role": "system",
-                "content": [{
-                    "type": "text",
-                    "text": f"""You are EWAAST, an equitable wound assessment agent.
-                    
-The patient has been classified as MST {mst_value} ({mst_result['category']} skin tone).
-Visual Guidance: {visual_guidance}
-
-CRITICAL: For MST 7-10, you MUST NOT use "redness" or "erythema" as primary indicators.
-Instead, look for purple/blue discoloration, warmth, induration, and texture changes.
-
-Assess the wound and provide:
-1. Stage (Stage 1-4, Unstageable, or DTI)
-2. Clinical rationale (explain what visual features you observed)
-3. Recommended care plan
-4. Urgency level (routine, urgent, immediate)
-
-Return as JSON: {{"stage": "...", "rationale": "...", "care_plan": "...", "urgency": "..."}}"""
-                }]
-            },
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
-                    {"type": "text", "text": f"Patient context: {context or 'Not provided'}. Assess this wound."}
-                ]
-            }
-        ]
-        
-        response = medgemma_get_text_response(messages)
+    if client.is_available():
         try:
-            return json.loads(response)
-        except:
-            pass
-    
-    # Mock response for development
+            image = _b64_to_pil(image_b64)
+            return client.assess_wound(image, mst_result, context)
+        except Exception as e:
+            print(f"AI Assessment failed: {e}")
+
+    # Fallback Logic
+    mst_value = mst_result["mst_value"]
     if mst_value >= 7:
         return {
-            "stage": "Stage 2",
-            "rationale": f"Observed purple-blue discoloration consistent with early tissue damage on MST {mst_value} skin. Warmth and slight induration noted. No use of 'redness' criteria due to unreliability on deep skin tones.",
-            "care_plan": "1. Offload pressure immediately\n2. Apply appropriate dressing\n3. Monitor for progression\n4. Document with good lighting to capture color changes",
+            "stage": "Assessment Failed (Fallback)",
+            "rationale": f"AI unavailable. Observed features on MST {mst_value} skin. Check for purple/blue tones.",
+            "care_plan": "1. Offload pressure\n2. Monitor for changes",
             "urgency": "urgent"
         }
     else:
         return {
-            "stage": "Stage 1",
-            "rationale": f"Non-blanchable erythema observed on MST {mst_value} skin. Area is warm to touch with intact epidermis.",
-            "care_plan": "1. Relieve pressure on affected area\n2. Keep skin clean and dry\n3. Apply barrier cream if indicated\n4. Reposition schedule every 2 hours",
+            "stage": "Assessment Failed (Fallback)",
+            "rationale": f"AI unavailable. Observed features on MST {mst_value} skin. Check for redness.",
+            "care_plan": "1. Relieve pressure\n2. Keep clean",
             "urgency": "routine"
         }
 
@@ -186,73 +129,65 @@ Return as JSON: {{"stage": "...", "rationale": "...", "care_plan": "...", "urgen
 def stream_assessment(image_b64: str, context: str) -> Generator[str, None, None]:
     """
     Stream the wound assessment process.
-    
-    This is a generator that yields JSON messages for SSE streaming.
-    Adapted from the interview_simulator pattern in appoint-ready.
-    
-    Args:
-        image_b64: Base64-encoded wound image
-        context: Optional patient context
-        
-    Yields:
-        JSON strings with step updates and data
     """
     
-    # Step 1: Acknowledge image received
+    # Step 1: Acknowledge
     yield json.dumps({
         "step": "received",
-        "thinking": "Image received. Beginning analysis...",
+        "thinking": "Image received. validating...",
         "data": {}
     })
+    time.sleep(0.5)
     
-    time.sleep(0.5)  # Brief pause for UX
-    
-    # Step 2: Classify skin tone
+    # 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. This helps ensure equitable assessment across all skin tones...",
+        "thinking": "Analyzing skin tone using Monk Skin Tone (MST) scale with Gemini Vision...",
         "data": {}
     })
     
-    time.sleep(1.0)  # Simulate processing
     mst_result = classify_skin_tone(image_b64, context)
     
     yield json.dumps({
         "step": "mst_complete",
-        "thinking": f"Skin tone classified as MST {mst_result['mst_value']} ({mst_result['category']}). Applying appropriate visual guidance...",
+        "thinking": f"Identified MST {mst_result['mst_value']} ({mst_result['category']}). Adjusted visual guidance: {mst_result['visual_guidance']}",
         "data": mst_result
     })
     
-    time.sleep(0.5)
-    
-    # Step 3: Analyze wound features
+    # Step 4: Analyze wound
     yield json.dumps({
         "step": "analyzing_wound",
-        "thinking": f"Analyzing wound features with MST-aware criteria. For MST {mst_result['mst_value']}: {mst_result['visual_guidance']}",
-        "data": {}
-    })
-    
-    time.sleep(1.5)  # Simulate processing
-    
-    # Step 4: Generate report
-    yield json.dumps({
-        "step": "generating_report",
-        "thinking": "Generating clinical assessment report with equitable staging...",
+        "thinking": f"Gemini is analyzing wound features. Looking for MST-specific signs (e.g., discoloration vs redness)...",
         "data": {}
     })
     
-    time.sleep(1.0)
+    # Step 5: Generate report
     report = generate_assessment_report(image_b64, mst_result, context)
     
-    # Step 5: Complete
     yield json.dumps({
         "step": "complete",
-        "thinking": "Assessment complete. Report generated.",
+        "thinking": "Assessment complete. Clinical rationale generated.",
         "data": {
             "mst": mst_result,
             "report": report
         }
     })
     
-    # Signal end of stream
     yield json.dumps({"event": "end"})

gemini_client.py

@@ -0,0 +1,192 @@
+"""
+EWAAST: Google AI Studio Client
+
+Uses Google's Gemini API for real AI inference.
+Supports Gemini 1.5 Pro (Vision) and newer models.
+"""
+
+import os
+import json
+import base64
+import io
+from typing import Optional, Dict, Any
+from PIL import Image
+
+# Google AI Studio API key
+GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
+
+# Model selection - Gemini 1.5 Pro is standard for vision
+DEFAULT_MODEL = "gemini-1.5-pro"
+FALLBACK_MODEL = "gemini-1.5-flash"
+
+
+class GeminiClient:
+    """
+    Client for Google AI Studio / Gemini API.
+    
+    Uses Gemini Pro for vision-based wound assessment.
+    """
+    
+    def __init__(self, model_name: str = DEFAULT_MODEL):
+        self.model_name = model_name
+        self._api_available = False
+        self._genai = None
+        
+        # Initialize
+        self._initialize()
+    
+    def _initialize(self):
+        """Initialize the Google GenAI client."""
+        if not GOOGLE_API_KEY:
+            print("⚠️ GOOGLE_API_KEY not set - AI features disabled")
+            return
+        
+        try:
+            import google.generativeai as genai
+            genai.configure(api_key=GOOGLE_API_KEY)
+            self._genai = genai
+            self._api_available = True
+            print(f"✅ Google AI Studio connected: {self.model_name}")
+        except ImportError:
+            print("⚠️ google-generativeai not installed")
+        except Exception as e:
+            print(f"⚠️ Google AI init failed: {e}")
+    
+    def is_available(self) -> bool:
+        """Check if Gemini API is available."""
+        return self._api_available
+    
+    def _image_to_part(self, image: Image.Image) -> dict:
+        """Convert PIL Image to Gemini-compatible format."""
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        buffer = io.BytesIO()
+        image.save(buffer, format="JPEG", quality=85)
+        image_bytes = buffer.getvalue()
+        
+        # Determine strict MIME type for API
+        return {
+            "mime_type": "image/jpeg",
+            "data": base64.b64encode(image_bytes).decode("utf-8")
+        }
+    
+    def validate_wound_image(self, image: Image.Image) -> Dict[str, Any]:
+        """
+        Validate if the image is actually a wound/medical image.
+        """
+        if not self._api_available:
+            return {"is_valid": True, "reason": "AI validation unavailable"}
+        
+        prompt = """Analyze this image and determine:
+1. Is this a medical/clinical image of human skin or a wound? 
+2. If NOT medical, what type of image is it?
+
+Respond in JSON format ONLY:
+{"is_medical": true/false, "reason": "brief explanation"}"""
+        
+        try:
+            response = self.generate(image, prompt)
+            # Try to parse JSON from response
+            try:
+                # Find JSON object
+                start = response.find('{')
+                end = response.rfind('}') + 1
+                if start != -1 and end != -1:
+                    result = json.loads(response[start:end])
+                    return {
+                        "is_valid": result.get("is_medical", False),
+                        "reason": result.get("reason", "Analysis complete")
+                    }
+            except:
+                pass
+            
+            # Fallback keyword check
+            is_valid = "true" in response.lower() or "medical" in response.lower()
+            return {"is_valid": is_valid, "reason": response[:100]}
+            
+        except Exception as e:
+            print(f"Validation error: {e}")
+            return {"is_valid": True, "reason": "Validation failed, proceeding"}
+
+    def analyze_mst(self, image: Image.Image) -> Dict[str, Any]:
+        """Classify Monk Skin Tone."""
+        prompt = """Analyze the skin tone in this image using the Monk Skin Tone (MST) scale (1-10).
+Return ONLY a JSON object:
+{"mst_value": <1-10 integer>, "confidence": <0.0-1.0 float>}"""
+
+        response = self.generate(image, prompt)
+        try:
+            start = response.find('{')
+            end = response.rfind('}') + 1
+            if start != -1 and end != -1:
+                return json.loads(response[start:end])
+        except:
+            pass
+        return {"mst_value": 5, "confidence": 0.0}
+
+    def assess_wound(self, image: Image.Image, mst_data: Dict[str, Any], context: str = "") -> Dict[str, Any]:
+        """Full wound assessment."""
+        mst_val = mst_data.get('mst_value', 5)
+        
+        prompt = f"""You are EWAAST, an equitable wound assessment specialist.
+Patient Context: {context}
+Skin Tone: MST {mst_val}
+
+Analyze this wound image. CRITICAL: For deeper skin tones (MST 7-10), DO NOT rely on redness. Look for purple/blue tones, texture changes, or darkening.
+
+Provide a JSON response with:
+1. "stage": Assessment (e.g., "Stage 1", "Stage 2", "Unstageable", "DTI", "Healing")
+2. "rationale": Clinical reasoning describing visual evidence
+3. "care_plan": 3-4 bullet points of immediate recommendations
+4. "urgency": "routine", "urgent", or "emergency"
+
+JSON Format:
+{{
+  "stage": "string",
+  "rationale": "string",
+  "care_plan": "string",
+  "urgency": "string"
+}}"""
+
+        response = self.generate(image, prompt)
+        try:
+            start = response.find('{')
+            end = response.rfind('}') + 1
+            if start != -1 and end != -1:
+                return json.loads(response[start:end])
+        except:
+            pass
+        return {
+            "stage": "Assessment Failed",
+            "rationale": f"Could not parse AI response: {response[:100]}...",
+            "care_plan": "Consult human specialist.",
+            "urgency": "urgent"
+        }
+
+    def generate(self, image: Image.Image, prompt: str) -> str:
+        """Generate response from Gemini."""
+        if not self._api_available:
+            raise RuntimeError("Google AI Studio not available.")
+        
+        try:
+            model = self._genai.GenerativeModel(self.model_name)
+            response = model.generate_content([prompt, image])
+            return response.text
+        except Exception:
+            # Fallback
+            if self.model_name != FALLBACK_MODEL:
+                print(f"Falling back to {FALLBACK_MODEL}")
+                model = self._genai.GenerativeModel(FALLBACK_MODEL)
+                response = model.generate_content([prompt, image])
+                return response.text
+            raise
+
+_client = None
+
+def get_gemini_client() -> GeminiClient:
+    global _client
+    if _client is None:
+        _client = GeminiClient()
+    return _client

requirements.txt

@@ -14,6 +14,9 @@ peft>=0.10.0
 # Image Processing
 Pillow>=10.0.0
 
+# Gemini API (Real AI)
+google-generativeai>=0.4.0
+
 # Web Interface (Flask + React Architecture)
 flask>=3.0.0
 flask-cors>=4.0.0