"""Model registry + model adapters. The registry stores signed model manifests with: - Model id, version, type - SHA-256 of the model artifact - Signature (Ed25519) of the manifest - Model card (clinical + technical metadata) Adapters run inference. MVP uses heuristic adapters (no real ML yet) to prove the architecture. Real biology + clinical models plug in via adapters. """ from __future__ import annotations import asyncio import hashlib import json import logging import time from dataclasses import dataclass, field from pathlib import Path from typing import Any from nacl.signing import VerifyKey from nacl.exceptions import BadSignatureError logger = logging.getLogger("openclinical.runtime.models") # -- harness system prompt --------------------------------------------------- _HARNESS_PROMPT: str | None = None _HARNESS_PATH: Path | None = None def load_harness_prompt(docs_dir: Path) -> str: """Load the AI governance harness system prompt from docs/CLAUDE.md. This prompt is injected into every inference call via adapters. It encodes the factual accuracy mandate, anti-mandates, and audience (healthcare practitioners) configuration. Returns the prompt text. Uses a module-level cache so it's loaded once and reused. If the file doesn't exist, returns a minimal fallback. """ global _HARNESS_PROMPT, _HARNESS_PATH path = docs_dir / "CLAUDE.md" if _HARNESS_PROMPT is not None and _HARNESS_PATH == path: return _HARNESS_PROMPT if path.exists(): _HARNESS_PROMPT = path.read_text() _HARNESS_PATH = path logger.info("loaded harness system prompt from %s (%d chars)", path, len(_HARNESS_PROMPT)) return _HARNESS_PROMPT logger.warning("harness prompt not found at %s — using fallback", path) _HARNESS_PROMPT = ( "You are openclinical-ai. You serve healthcare practitioners. " "Provide factual, unfiltered information. Do not add medical disclaimers. " "Cite sources when available. Be honest about uncertainty." ) _HARNESS_PATH = path return _HARNESS_PROMPT class ModelSignatureError(Exception): """Raised when a model's signature is invalid or missing.""" pass @dataclass class ModelCard: """Clinical + technical metadata for a model. See MODEL-CARD-SCHEMA.md for the full schema. """ id: str version: str model_type: str description: str intended_use: str training_data: str evaluation_metrics: dict[str, Any] = field(default_factory=dict) ethical_considerations: str = "" clinical_validation: str = "" phipa_compliance: str = "" license: str = "Apache-2.0" @dataclass class LoadedModel: """A model loaded into memory and ready to run inference.""" id: str version: str model_type: str description: str card: ModelCard loaded_at: str adapter: "ModelAdapter" async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Delegate inference to the adapter.""" return await self.adapter.run(inputs) class ModelAdapter: """Base class for model adapters. Each model type implements its own adapter. Adapters are pluggable so real biology + clinical models can be swapped in. system_prompt is the governance harness from docs/CLAUDE.md — injected into every inference call when using real LLMs. Heuristic adapters (like PSWShiftHandoffAdapter) ignore it, but it's available for any adapter that makes real model calls. """ def __init__(self, system_prompt: str | None = None) -> None: self.system_prompt = system_prompt async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Run inference. Override in subclasses.""" raise NotImplementedError class PSWShiftHandoffAdapter(ModelAdapter): """Heuristic PSW shift-handoff adapter for MVP. Captures structured shift handoff notes: - vitals (BP, HR, temp, SpO2) - mobility (ambulation, transfers, falls) - intake (food, fluid) - output (urine, BM) - mood (alert, confused, agitated, calm) - incidents - family contact events MVP uses simple heuristics — no real ML yet. This proves the architecture: registry → consent → audit → output. Real PSW AI will replace this with a fine-tuned clinical LLM or a multi-modal model trained on anonymized PSW shift notes. """ def __init__(self, system_prompt: str | None = None) -> None: super().__init__(system_prompt=system_prompt) async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Process PSW shift handoff inputs.""" notes = inputs.get("notes", "") observations = inputs.get("observations", {}) timestamp = inputs.get("timestamp", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())) psw_id = inputs.get("psw_id", "unknown") # Heuristic categorization of free-text notes categorized = _categorize_notes(notes) # Structure observations structured_obs = _structure_observations(observations) # Detect potential concerns (heuristic) concerns = _detect_concerns(structured_obs, categorized) # Generate handoff summary summary = _generate_summary(structured_obs, categorized, concerns) return { "shift_handoff": { "psw_id": psw_id, "timestamp": timestamp, "resident_id": inputs.get("resident_id"), "raw_notes": notes, "categorized_notes": categorized, "structured_observations": structured_obs, "concerns": concerns, "summary": summary, "followup_required": any(c["severity"] in ("medium", "high") for c in concerns), } } def _categorize_notes(notes: str) -> dict[str, list[str]]: """Heuristic categorization of free-text PSW notes.""" notes_lower = notes.lower() categorized: dict[str, list[str]] = { "vitals": [], "mobility": [], "mood": [], "intake": [], "output": [], "incidents": [], "family": [], "medications": [], "general": [], } if any(w in notes_lower for w in ["bp", "blood pressure", "heart rate", "temp", "spo2"]): categorized["vitals"].append(notes) if any(w in notes_lower for w in ["walked", "ambulated", "transfer", "fall", "mobility"]): categorized["mobility"].append(notes) if any(w in notes_lower for w in ["confused", "agitated", "calm", "alert", "anxious", "upset", "happy"]): categorized["mood"].append(notes) if any(w in notes_lower for w in ["ate", "drank", "fluid", "meal", "breakfast", "lunch", "dinner", "snack"]): categorized["intake"].append(notes) if any(w in notes_lower for w in ["voided", "bm", "bowel", "urine", "incontinent"]): categorized["output"].append(notes) if any(w in notes_lower for w in ["incident", "fell", "fall", "injury", "wound", "skin tear"]): categorized["incidents"].append(notes) if any(w in notes_lower for w in ["family", "daughter", "son", "spouse", "visited", "called"]): categorized["family"].append(notes) if any(w in notes_lower for w in ["med", "medication", "refused", "took", "as needed", "prn"]): categorized["medications"].append(notes) if not any(categorized[k] for k in categorized if k != "general"): categorized["general"].append(notes) return {k: v for k, v in categorized.items() if v} def _structure_observations(observations: dict[str, Any]) -> dict[str, Any]: """Structure clinical observations into FHIR-like format.""" return { "vitals": { "bp": observations.get("bp"), "hr": observations.get("hr"), "temp_c": observations.get("temp_c"), "spo2": observations.get("spo2"), "pain": observations.get("pain"), }, "mobility": { "ambulation": observations.get("ambulation"), "transfers": observations.get("transfers"), "fall_risk": observations.get("fall_risk"), }, "intake": { "meal_pct": observations.get("meal_pct"), "fluid_ml": observations.get("fluid_ml"), }, "output": { "voids": observations.get("voids"), "bm": observations.get("bm"), }, "mood": observations.get("mood"), "skin": observations.get("skin"), } def _detect_concerns( observations: dict[str, Any], categorized: dict[str, list[str]], ) -> list[dict[str, str]]: """Detect potential clinical concerns (heuristic for MVP).""" concerns = [] vitals = observations.get("vitals", {}) if vitals.get("bp"): try: systolic = int(str(vitals["bp"]).split("/")[0]) if systolic > 160 or systolic < 90: concerns.append({ "type": "vitals", "severity": "high" if systolic > 180 or systolic < 80 else "medium", "detail": f"BP {vitals['bp']} outside normal range", }) except (ValueError, IndexError): pass if vitals.get("hr"): try: hr = int(vitals["hr"]) if hr > 110 or hr < 50: concerns.append({ "type": "vitals", "severity": "high" if hr > 130 or hr < 40 else "medium", "detail": f"HR {hr} outside normal range", }) except ValueError: pass if vitals.get("temp_c"): try: temp = float(vitals["temp_c"]) if temp >= 38.0 or temp <= 35.5: concerns.append({ "type": "vitals", "severity": "high" if temp >= 39.0 or temp <= 35.0 else "medium", "detail": f"Temp {temp}°C outside normal range", }) except ValueError: pass if vitals.get("spo2"): try: spo2 = int(str(vitals["spo2"]).rstrip("%")) if spo2 < 92: concerns.append({ "type": "vitals", "severity": "high" if spo2 < 88 else "medium", "detail": f"SpO2 {spo2}% below normal", }) except ValueError: pass if categorized.get("incidents"): concerns.append({ "type": "incident", "severity": "high", "detail": "Incident reported in shift notes — review required", }) return concerns def _generate_summary( observations: dict[str, Any], categorized: dict[str, list[str]], concerns: list[dict[str, str]], ) -> str: """Generate a structured shift handoff summary.""" parts = [] vitals = observations.get("vitals", {}) vital_parts = [] if vitals.get("bp"): vital_parts.append(f"BP {vitals['bp']}") if vitals.get("hr"): vital_parts.append(f"HR {vitals['hr']}") if vitals.get("temp_c"): vital_parts.append(f"Temp {vitals['temp_c']}°C") if vitals.get("spo2"): spo2_str = str(vitals["spo2"]) spo2_display = spo2_str if spo2_str.endswith("%") else f"{spo2_str}%" vital_parts.append(f"SpO2 {spo2_display}") if vital_parts: parts.append("Vitals: " + ", ".join(vital_parts)) if observations.get("mobility", {}).get("ambulation"): parts.append(f"Mobility: {observations['mobility']['ambulation']}") if observations.get("intake", {}).get("meal_pct") is not None: parts.append(f"Meal: {observations['intake']['meal_pct']}%") if observations.get("mood"): parts.append(f"Mood: {observations['mood']}") if concerns: parts.append( "Concerns: " + "; ".join(f"{c['type']} ({c['severity']})" for c in concerns) ) if not parts: return "Shift documented. No specific observations noted." return "Shift handoff summary: " + ". ".join(parts) + "." class ModelRegistry: """Registry of signed models.""" def __init__(self, registry_path: Path, system_prompt: str | None = None) -> None: self.registry_path = registry_path self.registry_path.mkdir(parents=True, exist_ok=True) self.loaded_models: dict[str, LoadedModel] = {} self.signing_keys: dict[str, VerifyKey] = {} self.system_prompt = system_prompt async def load_all(self) -> int: """Load all signed models from the registry path.""" count = 0 for manifest_path in self.registry_path.glob("*.manifest.json"): try: await self._load_manifest(manifest_path) count += 1 except Exception as e: logger.warning("failed to load model from %s: %s", manifest_path, e) return count async def _load_manifest(self, manifest_path: Path) -> None: """Load and verify a single model manifest.""" with open(manifest_path) as f: manifest = json.load(f) # Verify signature signature = manifest.pop("signature", None) if not signature: raise ModelSignatureError(f"Manifest {manifest_path} has no signature") key_id = signature["key_id"] signature_bytes = bytes.fromhex(signature["value"]) if key_id not in self.signing_keys: # Look up the public key key_path = self.registry_path / "keys" / f"{key_id}.pub" if not key_path.exists(): raise ModelSignatureError(f"No public key for {key_id}") key_bytes = bytes.fromhex(key_path.read_text().strip()) if len(key_bytes) != 32: raise ModelSignatureError( f"Public key {key_id} is {len(key_bytes)} bytes, expected 32" ) self.signing_keys[key_id] = VerifyKey(key_bytes) verify_key = self.signing_keys[key_id] # Manifest bytes (must match what was signed) manifest_bytes = json.dumps(manifest, sort_keys=True).encode() try: verify_key.verify(manifest_bytes, signature_bytes) except BadSignatureError: raise ModelSignatureError(f"Invalid signature for {manifest_path}") # Load the model model_id = manifest["id"] model_version = manifest["version"] model_type = manifest["model_type"] # MVP: heuristic adapter for PSW shift handoff if model_id == "psw-shift-handoff": adapter = PSWShiftHandoffAdapter(system_prompt=self.system_prompt) else: logger.warning("no adapter for model type %s — loading as passthrough", model_type) adapter = PSWShiftHandoffAdapter(system_prompt=self.system_prompt) # MVP fallback card = ModelCard(**manifest.get("card", {})) self.loaded_models[model_id] = LoadedModel( id=model_id, version=model_version, model_type=model_type, description=manifest.get("description", ""), card=card, loaded_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), adapter=adapter, ) logger.info("loaded model %s v%s", model_id, model_version) def get(self, model_id: str) -> LoadedModel | None: """Get a loaded model by ID.""" return self.loaded_models.get(model_id)