src/tools/patient_memory.py

import os
import uuid
import datetime
from supabase import create_client, Client
from langchain_core.tools import tool
from src.utils.logger import setup_logger

logger = setup_logger("FHIRPatientMemory")

def _get_client() -> Client:
    url = os.getenv("SUPABASE_URL")
    key = os.getenv("SUPABASE_KEY")
    return create_client(url, key)

@tool
def save_patient_memory(patient_id: str, glucose_history: str = None, medications: str = None, diet: str = None):
    """
    Save patient memory (narrative context) as strict FHIR Observation resources.
    """
    logger.info(f"Saving FHIR narrative memory for patient: {patient_id}")
    client = _get_client()
    
    results = []
    
    # Map narrative categories to FHIR-like coding
    categories = {
        "glucose_history": {"code": "narrative-glucose", "display": "Glucose History Narrative"},
        "medications": {"code": "narrative-meds", "display": "Medications Narrative"},
        "diet": {"code": "narrative-diet", "display": "Dietary Narrative"}
    }
    
    for key, value in [("glucose_history", glucose_history), ("medications", medications), ("diet", diet)]:
        if value is not None:
            obs_id = str(uuid.uuid4())
            fhir_obs = {
                "resourceType": "Observation",
                "id": obs_id,
                "status": "final",
                "code": {
                    "coding": [{
                        "system": "http://dm-chatbot.ai/codes",
                        "code": categories[key]["code"],
                        "display": categories[key]["display"]
                    }]
                },
                "subject": {"reference": f"Patient/{patient_id}"},
                "effectiveDateTime": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                "valueString": value # Using valueString for narrative text
            }
            
            data = {
                "id": obs_id,
                "patient_id": patient_id,
                "resource": fhir_obs,
                "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
            }
            client.table("observations").insert(data).execute()
            results.append(key)
            
    return f"Successfully saved FHIR narrative memory for: {', '.join(results)}"

@tool
def get_patient_memory(patient_id: str):
    """
    Retrieve patient narrative memory from FHIR Observation resources.
    """
    logger.info(f"Retrieving FHIR narrative memory for patient: {patient_id}")
    client = _get_client()
    
    try:
        # Fetch observations with narrative codes
        response = client.table("observations").select("resource").eq("patient_id", patient_id).execute()
        observations = [r["resource"] for r in response.data]
        
        # Filter for narrative codes
        memory = {}
        codes_map = {
            "narrative-glucose": "Glucose History",
            "narrative-meds": "Medications",
            "narrative-diet": "Diet"
        }
        
        for obs in observations:
            for coding in obs.get("code", {}).get("coding", []):
                code = coding.get("code")
                if code in codes_map:
                    # Keep only the latest one for each category
                    memory[codes_map[code]] = obs.get("valueString", "N/A")
        
        if not memory:
            return f"No FHIR narrative memory found for patient {patient_id}."
            
        output = f"Patient Narrative Memory (FHIR) for {patient_id}:\n"
        for cat, val in memory.items():
            output += f"- {cat}: {val}\n"
        return output
        
    except Exception as e:
        logger.error(f"Error retrieving FHIR memory: {e}")
        return f"Error retrieving FHIR memory: {str(e)}"