Create agent.py

agent.py

@@ -0,0 +1,263 @@
+import random
+import string
+import re
+import datetime
+import base64
+import pandas as pd
+from pathlib import Path
+import os
+import logging
+from langchain_core.tools import tool
+from langchain_groq import ChatGroq
+from langgraph.prebuilt import create_react_agent
+from langgraph.checkpoint.memory import MemorySaver
+
+import cv2
+import numpy as np
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+# Global variables
+_SESSIONS = None
+_CURRENT_THREAD_ID = None
+_API_CLIENT = None
+_RETRIEVER = None
+
+def set_session_storage(sessions_dict: dict) -> None:
+    global _SESSIONS
+    _SESSIONS = sessions_dict
+
+def set_current_thread_id(thread_id: str) -> None:
+    global _CURRENT_THREAD_ID
+    _CURRENT_THREAD_ID = thread_id
+
+def set_api_client(api_client) -> None:
+    global _API_CLIENT
+    _API_CLIENT = api_client
+
+def get_latest_image(thread_id: str) -> dict | None:
+    """Get latest image from external API"""
+    if _API_CLIENT is None:
+        return None
+    
+    result = _API_CLIENT.get_patient_image(thread_id)
+    if result.get('ok') and result.get('base64'):
+        # Decode base64 image
+        try:
+            img_data = base64.b64decode(result['base64'])
+            image = Image.open(io.BytesIO(img_data))
+            return {"image": image}
+        except Exception as e:
+            logger.error(f"Error decoding image: {e}")
+            return None
+    return None
+
+def _store_image_blobs(blobs: dict) -> None:
+    """Store image blobs in session"""
+    if not _SESSIONS or not _CURRENT_THREAD_ID:
+        return
+    if _CURRENT_THREAD_ID not in _SESSIONS:
+        _SESSIONS[_CURRENT_THREAD_ID] = {}
+    existing = _SESSIONS[_CURRENT_THREAD_ID].get("_image_blobs", {})
+    existing.update(blobs)
+    _SESSIONS[_CURRENT_THREAD_ID]["_image_blobs"] = existing
+    names = list(_SESSIONS[_CURRENT_THREAD_ID].get("_session_image_names", []))
+    for n in blobs.keys():
+        if n not in names:
+            names.append(n)
+    _SESSIONS[_CURRENT_THREAD_ID]["_session_image_names"] = names
+
+# System Prompt
+SYSTEM_PROMPT = """You are DermaScan AI, a dermatology assistant on a dataset of 1000 skin cancer patients. Built by Eng. Youssef Bastawisy.
+
+IMAGES: handled by tools only, on the image already uploaded this session. NEVER say you can't see/read images. On any image request, call the matching tool immediately:
+- segment_my_image (segment/mask/overlay), reanalyze_my_image (analyze/re-analyze), get_my_uploaded_image (show/view photo)
+No image is ever downloaded from the internet or from a patient_id — matching an uploaded image to a dataset record is handled internally by the tool. NEVER mention how the matching works internally (no "feature vectors", "global_features", "similarity", "threshold", "columns", etc.) — just state plainly whether a matching record was found or not.
+If a tool says no matching record was found (new patient), you MUST say clearly this is a new patient not in the database, and STOP there for anything hereditary/genetic — never mention a mutation, family history, or risk score for them, and never reuse a different (previously matched) patient's name or data for this new image.
+After a tool returns [ANALYSIS_RESULT:{data}], discuss only that data (label, confidence, affected area) — never say "shown"/"displayed".
+
+LANGUAGE: reply in the same language/dialect the user just used (Arabic or English). Never mix in Russian/Thai/Chinese/Japanese/Korean script — if unsure of a term, describe simply instead.
+
+PATIENT MODE STYLE: no jargon (no "asymmetry index", "ABCDE"). Plain everyday words, max 4-6 short sentences/bullets. Never mix English medical terms into Arabic sentences.
+
+ROLE ACCESS:
+- You operate in PATIENT mode only. Never reveal other patients' data, only their own uploaded result.
+- If the result is Malignant: stress urgency + next steps + suggest the "📅 Book appointment" button (create a referral ticket only if explicitly asked).
+- If the result is Benign: prevention + monitoring tips.
+
+MATCHED PATIENT RECORD: if this block appears, the uploaded image matched an existing patient (same person uploading). Use it to answer diagnosis/family-history questions directly and specifically — summarize relevant parts, don't dump verbatim.
+
+STRICT RAG RULE: your knowledge is limited to (1) patient records via matched block, (2) knowledge-base via search_knowledge_base, (3) AI image results, (4) this conversation. Never invent symptoms, history, numbers. If info isn't available say exactly: "This information is not available in the current patient record." (or Arabic equivalent).
+
+TOOLS:
+- search_knowledge_base is the default first call for general clinical/knowledge questions (cite the source filename).
+- Use segment_my_image for segmentation/mask/overlay requests.
+- Use reanalyze_my_image for classification + segmentation analysis.
+- STOP after one tool call for a simple lookup, and at most two tool calls for a combined request.
+- Never call the same tool twice.
+- Never call a second different tool just because the first one returned a valid answer.
+- Keep replies SHORT by default (2–4 sentences) unless the user explicitly asks for more detail.
+
+QUICK ROUTING (pick one, don't overthink):
+- General symptom/skin question, no specific image → search_knowledge_base only.
+- "my/this image", "analyze it", "the result" → segment_my_image/reanalyze_my_image/get_my_uploaded_image.
+- Own history/family/risk → answer from the [MATCHED PATIENT RECORD] block if present; if absent, say the info isn't available yet.
+
+Booking is via the "📅 Book appointment" button. Never print booking IDs, tickets, or debug/system text yourself.
+
+ANALYSIS MARKER: if a tool response starts with `[ANALYSIS_RESULT:{...}]`, keep that exact marker at the very start of your reply, unmodified. After it, briefly cover: classification + confidence, affected area, what it means in plain terms, relevant family/hereditary context if present, and next steps. Keep it short — a routine benign result needs only 2-3 sentences; a malignant one with family history can go a bit longer. Remind patients once (not every message) that this doesn't replace a professional diagnosis.
+
+If a user gives a local file path, tell them to use the 📎 upload button instead.
+"""
+
+# Tools
+@tool
+def search_knowledge_base(query: str) -> str:
+    """Search the dermatology knowledge base for clinical information."""
+    if _RETRIEVER is None:
+        return "Retriever not initialized."
+    docs = _RETRIEVER.invoke(query)
+    if not docs:
+        return "No relevant information found."
+    formatted = []
+    for i, d in enumerate(docs, 1):
+        source = d.metadata.get("source", "unknown").replace("\\", "/").split("/")[-1]
+        formatted.append(f"[Source {i}: {source}]\n{d.page_content}")
+    return "\n\n---\n\n".join(formatted)
+
+@tool
+def get_my_uploaded_image() -> str:
+    """[PATIENT MODE] Tell patient to use the Show My Image button."""
+    return "Your image is stored for this session. Use the 📷 Show My Image button to view it."
+
+@tool
+def segment_my_image() -> str:
+    """[PATIENT MODE] Run U-Net segmentation on the patient's uploaded image."""
+    if not _SESSIONS or not _CURRENT_THREAD_ID:
+        return "No session found."
+    
+    if _API_CLIENT is None:
+        return "API client not initialized."
+    
+    try:
+        result = _API_CLIENT.segment_patient_image(_CURRENT_THREAD_ID)
+        
+        if not result.get('ok'):
+            return f"Error: {result.get('error', 'Segmentation failed')}"
+        
+        # Store images
+        if 'images' in result:
+            _store_image_blobs(result['images'])
+        
+        import json as _json
+        result_data = {
+            "label": "",
+            "confidence_pct": 0,
+            "infection_pct": result.get('infection_pct', 0)
+        }
+        marker = f"[ANALYSIS_RESULT:{_json.dumps(result_data)}]"
+        
+        return f"{marker}\n\nSEGMENTATION OF YOUR IMAGE\nAffected Area: {result['infection_pct']}%"
+        
+    except Exception as e:
+        logger.error(f"Error in segment_my_image: {e}", exc_info=True)
+        return f"Error segmenting image: {str(e)}"
+
+@tool
+def reanalyze_my_image() -> str:
+    """[PATIENT MODE] Full AI analysis (classification + segmentation) on the patient's
+    most recently uploaded image."""
+    if not _SESSIONS or not _CURRENT_THREAD_ID:
+        return "No session found."
+    
+    if _API_CLIENT is None:
+        return "API client not initialized."
+    
+    try:
+        result = _API_CLIENT.reanalyze_patient_image(_CURRENT_THREAD_ID)
+        
+        if not result.get('ok'):
+            return f"Error: {result.get('error', 'Analysis failed')}"
+        
+        # Store images
+        if 'images' in result:
+            _store_image_blobs(result['images'])
+        
+        import json as _json
+        result_data = {
+            "label": result.get("label", ""),
+            "confidence_pct": result.get("confidence_pct", 0),
+            "infection_pct": result.get("infection_pct", 0),
+        }
+        marker = f"[ANALYSIS_RESULT:{_json.dumps(result_data)}]"
+        
+        lines = [
+            marker,
+            "",
+            "ANALYSIS OF YOUR IMAGE",
+            f"Classification: {result['label']} ({result['confidence_pct']}%)",
+            f"Affected Area: {result['infection_pct']}%"
+        ]
+        
+        return "\n".join(lines)
+        
+    except Exception as e:
+        logger.error(f"Error in reanalyze_my_image: {e}", exc_info=True)
+        return f"Error analyzing image: {str(e)}"
+
+@tool
+def create_referral_ticket(issue_summary: str, urgency: str = "routine") -> str:
+    """Create a dermatologist referral ticket."""
+    ticket_id = "DS-" + "".join(random.choices(string.digits, k=6))
+    times = {
+        "routine":   "within 2–4 weeks",
+        "urgent":    "within 48–72 hours",
+        "emergency": "within 24 hours — go to nearest dermatology clinic",
+    }
+    timeframe = times.get(urgency, "within 2–4 weeks")
+    
+    DEFAULT_CLINIC = {
+        "name": "مركز القاهرة للأمراض الجلدية",
+        "address": "شارع التحرير، الدقي، الجيزة",
+        "phone": "01011112222",
+    }
+    
+    return (
+        f"Referral created. Ticket: {ticket_id} | Urgency: {urgency.upper()} | "
+        f"Timeframe: {timeframe} | Summary: {issue_summary} | "
+        f"Clinic: {DEFAULT_CLINIC['name']} — {DEFAULT_CLINIC['address']}"
+    )
+
+def build_agent(retriever, model_name: str = "llama-3.1-8b-instant", temperature: float = 0.25):
+    global _RETRIEVER
+    _RETRIEVER = retriever
+
+    groq_key = os.environ.get("GROQ_API_KEY")
+    if not groq_key:
+        raise ValueError("GROQ_API_KEY not set in environment variables.")
+
+    llm = ChatGroq(
+        groq_api_key=groq_key,
+        model_name=model_name,
+        temperature=temperature,
+        max_tokens=512,
+        request_timeout=120,
+    )
+
+    tools = [
+        search_knowledge_base,
+        get_my_uploaded_image,
+        segment_my_image,
+        reanalyze_my_image,
+        create_referral_ticket,
+    ]
+    
+    memory = MemorySaver()
+
+    return create_react_agent(
+        model=llm,
+        tools=tools,
+        prompt=SYSTEM_PROMPT,
+        checkpointer=memory,
+    )