Add UI analytics, quality gates, and security docs

adit24dhaya · adit24dhaya · commit ef4b1a1e01b0 · 2026-04-26T23:33:05.000-07:00
Introduce risk trend/history visualizations, add lint+coverage CI and pre-commit hooks, and document architecture plus PII/retention hardening guidance for recruiter-ready project quality.

Made-with: Cursor
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -26,5 +26,11 @@ jobs:
       - name: Compile
         run: python -m compileall agents api models tools memory ui main.py scripts tests
 
-      - name: Run tests
-        run: pytest -q
+      - name: Lint
+        run: ruff check .
+
+      - name: Format check
+        run: ruff format --check .
+
+      - name: Run tests with coverage
+        run: pytest -q --cov=. --cov-report=term-missing --cov-fail-under=60
diff --git a/.gitignore b/.gitignore
@@ -3,3 +3,4 @@ __pycache__/
 *.pyc
 .DS_Store
 logs/
+.coverage
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.11
+    hooks:
+      - id: ruff-check
+        args: [--fix]
+      - id: ruff-format
diff --git a/README.md b/README.md
@@ -29,6 +29,8 @@ User Input
    -> Final Output (dashboard + API/chat-ready response)
 ```
 
+Detailed diagram: [`docs/architecture.md`](docs/architecture.md)
+
 ## Key Features
 
 - **Risk prediction** from patient vitals and demographic inputs
@@ -112,6 +114,13 @@ source .venv/bin/activate
 pip install -r requirements.txt
 ```
 
+### 2.1) Enable local pre-commit checks (recommended)
+
+```bash
+pre-commit install
+pre-commit run --all-files
+```
+
 ### 3) Configure environment variables
 
 Create a `.env` file (or export directly):
@@ -252,6 +261,12 @@ Always consult qualified healthcare professionals for real medical advice.
 - Added structured decision audit logging (`logs/decisions.jsonl`)
 - Added Dockerfile + Docker Compose for reproducible local deployment
 - Added CI workflow (compile + test on push/PR)
+- Added lint + coverage quality gates in CI
+- Added local pre-commit hooks for code quality
+
+## Security Notes
+
+See [`docs/security.md`](docs/security.md) for PII handling guidance, retention policy, and hardening recommendations.
 
 ## Contributing
 
diff --git a/agents/explainer.py b/agents/explainer.py
@@ -10,7 +10,9 @@ def __init__(self, model_name="gpt-4o-mini"):
 
     def explain(self, patient_data, prob, feature_explanation=None, similar_cases=None):
         if self.client is None:
-            return self._fallback_explanation(patient_data, prob, feature_explanation, similar_cases)
+            return self._fallback_explanation(
+                patient_data, prob, feature_explanation, similar_cases
+            )
 
         prompt = f"""
         Patient data: {patient_data}
@@ -23,14 +25,15 @@ def explain(self, patient_data, prob, feature_explanation=None, similar_cases=No
 
         try:
             response = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[{"role": "user", "content": prompt}]
+                model=self.model_name, messages=[{"role": "user", "content": prompt}]
             )
             return response.choices[0].message.content
         except Exception as exc:
             return f"{self._fallback_explanation(patient_data, prob, feature_explanation, similar_cases)} LLM explanation unavailable: {exc}"
 
-    def _fallback_explanation(self, patient_data, prob, feature_explanation=None, similar_cases=None):
+    def _fallback_explanation(
+        self, patient_data, prob, feature_explanation=None, similar_cases=None
+    ):
         age = patient_data["age"]
         bmi = patient_data["bmi"]
         bp = patient_data["bp"]
@@ -55,7 +58,9 @@ def _fallback_explanation(self, patient_data, prob, feature_explanation=None, si
         top_impacts = self._top_feature_text(feature_explanation)
         memory_text = ""
         if similar_cases:
-            memory_text = f" I found {len(similar_cases)} similar prior case(s) in memory for comparison."
+            memory_text = (
+                f" I found {len(similar_cases)} similar prior case(s) in memory for comparison."
+            )
 
         return (
             f"The model estimated a {prob:.1%} risk using a calculated BMI of {bmi:.1f}. "
@@ -67,9 +72,6 @@ def _top_feature_text(self, feature_explanation):
             return ""
 
         top_features = feature_explanation["features"][:3]
-        parts = [
-            f"{item['feature']} {item['direction']}"
-            for item in top_features
-        ]
+        parts = [f"{item['feature']} {item['direction']}" for item in top_features]
         method = feature_explanation.get("method", "model")
         return f"The {method.upper()} explanation highlights: {', '.join(parts)}. "
diff --git a/agents/orchestrator.py b/agents/orchestrator.py
@@ -71,5 +71,5 @@ def run(self, patient_data):
             "similar_cases": similar_cases,
             "safety": safety,
             "explanation": explanation,
-            "recommendation": recommendation
+            "recommendation": recommendation,
         }
diff --git a/agents/recommender.py b/agents/recommender.py
@@ -8,9 +8,13 @@ def __init__(self, model_name="gpt-4o-mini"):
         self.model_name = model_name
         self.client = OpenAI() if os.getenv("OPENAI_API_KEY") else None
 
-    def recommend(self, patient_data, risk, retrieved_context=None, similar_cases=None, safety=None):
+    def recommend(
+        self, patient_data, risk, retrieved_context=None, similar_cases=None, safety=None
+    ):
         if self.client is None:
-            return self._fallback_recommendation(patient_data, risk, retrieved_context, similar_cases, safety)
+            return self._fallback_recommendation(
+                patient_data, risk, retrieved_context, similar_cases, safety
+            )
 
         prompt = f"""
         Patient data: {patient_data}
@@ -24,14 +28,15 @@ def recommend(self, patient_data, risk, retrieved_context=None, similar_cases=No
 
         try:
             response = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[{"role": "user", "content": prompt}]
+                model=self.model_name, messages=[{"role": "user", "content": prompt}]
             )
             return response.choices[0].message.content
         except Exception as exc:
             return f"{self._fallback_recommendation(patient_data, risk, retrieved_context, similar_cases, safety)} LLM recommendation unavailable: {exc}"
 
-    def _fallback_recommendation(self, patient_data, risk, retrieved_context=None, similar_cases=None, safety=None):
+    def _fallback_recommendation(
+        self, patient_data, risk, retrieved_context=None, similar_cases=None, safety=None
+    ):
         glucose_note = ""
         if not patient_data.get("glucose_measured", True):
             glucose_note = " Because glucose was not measured, get an A1C or fasting blood glucose test before making clinical decisions."
diff --git a/agents/retriever.py b/agents/retriever.py
@@ -31,11 +31,13 @@ def retrieve(self, patient_data, risk, feature_explanation, n_results=3):
 
         contexts = []
         for document, metadata, distance in zip(documents, metadatas, distances):
-            contexts.append({
-                "title": metadata.get("title", "Medical note"),
-                "text": document,
-                "distance": float(distance),
-            })
+            contexts.append(
+                {
+                    "title": metadata.get("title", "Medical note"),
+                    "text": document,
+                    "distance": float(distance),
+                }
+            )
         return self._rerank(contexts, patient_data, risk)[:n_results]
 
     def _seed_knowledge(self):
@@ -88,7 +90,10 @@ def score(item):
             if title == "High risk clinical review":
                 value += 2.0 if risk == "High" else -3.0
             if title == "Glucose testing":
-                if not patient_data.get("glucose_measured", True) or patient_data.get("glucose", 0) >= 140:
+                if (
+                    not patient_data.get("glucose_measured", True)
+                    or patient_data.get("glucose", 0) >= 140
+                ):
                     value += 1.5
             if title == "Blood pressure follow-up" and patient_data.get("bp", 0) >= 130:
                 value += 1.25
diff --git a/api/app.py b/api/app.py
@@ -6,7 +6,6 @@
 from fastapi import Depends, FastAPI, Header, HTTPException
 from pydantic import BaseModel, Field
 
-
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
@@ -15,7 +14,6 @@
 from main import agent
 from tools.audit_logger import AuditLogger
 
-
 app = FastAPI(title="Healthcare Agent API")
 settings = get_settings()
 audit_logger = AuditLogger(settings.log_dir)
diff --git a/docs/architecture.md b/docs/architecture.md
@@ -0,0 +1,36 @@
+# Architecture Diagram
+
+This diagram is designed for recruiter/interviewer walkthroughs.
+
+```mermaid
+flowchart TD
+    A[User Input<br/>UI or API] --> B[RiskModel<br/>Logistic Regression]
+    B --> C[Orchestrator Agent]
+
+    C --> D[Explanation Agent<br/>LLM + fallback rules]
+    C --> E[Recommendation Agent<br/>LLM + safety context]
+    C --> F[Retrieval Agent<br/>Medical knowledge RAG]
+    C --> G[Memory Agent<br/>ChromaDB patient memory]
+    C --> H[Feature Explainer<br/>SHAP attributions]
+    C --> I[Safety Guard<br/>alerts + escalation + confidence]
+
+    F --> J[data/medical_knowledge.jsonl]
+    G --> K[data/chroma]
+    C --> L[Audit Logger]
+    L --> M[logs/decisions.jsonl]
+
+    C --> N[Final Response]
+    N --> O[Streamlit Dashboard]
+    N --> P[FastAPI /predict + /v1/predict]
+```
+
+## Request Lifecycle (High Level)
+
+1. Validate and normalize input (including BMI calculation).
+2. Predict probability and classify risk.
+3. Compute SHAP feature impacts.
+4. Retrieve similar past cases + relevant medical context.
+5. Generate explanation and recommendation.
+6. Apply safety assessment (alerts, escalation, confidence).
+7. Persist memory and write audit trail.
+8. Return structured response to UI/API client.
diff --git a/docs/security.md b/docs/security.md
@@ -0,0 +1,39 @@
+# Security Hardening Notes
+
+This project is an educational prototype. The points below describe a practical hardening baseline for handling sensitive health-adjacent data.
+
+## PII/PHI Handling Policy
+
+- **Data minimization**: only collect fields required for risk scoring and explainability.
+- **No raw identifiers by default**: avoid name, email, phone, address, national IDs in request payloads.
+- **Redaction in logs**: audit logs should include only normalized clinical features and decision metadata.
+- **Need-to-know access**: restrict memory/history endpoints behind token auth in non-local environments.
+- **Environment separation**: keep development/test data isolated from demo/production-like data.
+
+## Data Retention Policy
+
+- **Default retention (prototype)**: keep decision/audit history for 30 days.
+- **Memory retention**: rotate or archive old ChromaDB patient vectors on a schedule.
+- **Right to delete**: support deleting patient-associated memory records on request.
+- **Backups**: encrypt backups and define retention windows for snapshots.
+
+## Recommended Production Controls
+
+- Enforce API auth (`REQUIRE_API_TOKEN=true`) and rotate tokens regularly.
+- Move from shared token auth to per-user auth (JWT/OAuth) for multi-user environments.
+- Use HTTPS termination and secure secrets storage (vault/KMS, not plaintext env files).
+- Add rate limiting and request size limits at API gateway/reverse proxy level.
+- Add structured security events and alerting for auth failures and unusual access patterns.
+
+## Storage and Encryption Guidance
+
+- Encrypt data at rest for vector store, logs, and backups.
+- Encrypt data in transit between UI, API, and storage.
+- Use separate credentials/roles for read-only analytics vs write paths.
+- Prevent accidental commits of logs/databases by keeping them in `.gitignore`.
+
+## Compliance Notes
+
+- Treat this as non-diagnostic decision support.
+- For regulated use, complete threat modeling, DPIA/HIPAA review, and clinical validation.
+- Maintain auditability of model version, feature inputs, and recommendation provenance.
diff --git a/main.py b/main.py
@@ -10,7 +10,6 @@
 from tools.explainability import FeatureExplainer
 from tools.safety import SafetyGuard
 
-
 BASE_DIR = Path(__file__).resolve().parent
 DATA_PATH = BASE_DIR / "data" / "patients.csv"
 KNOWLEDGE_PATH = BASE_DIR / "data" / "medical_knowledge.jsonl"
diff --git a/memory/store.py b/memory/store.py
@@ -34,11 +34,13 @@ def get_all(self, limit=25):
             results.get("documents", []),
             results.get("metadatas", []),
         ):
-            records.append({
-                "id": record_id,
-                "summary": document,
-                "metadata": metadata,
-            })
+            records.append(
+                {
+                    "id": record_id,
+                    "summary": document,
+                    "metadata": metadata,
+                }
+            )
 
         records.sort(key=lambda item: item["metadata"].get("timestamp", ""), reverse=True)
         return records[:limit]
@@ -61,11 +63,13 @@ def find_similar(self, patient_data, n_results=3):
 
         matches = []
         for document, metadata, distance in zip(documents, metadatas, distances):
-            matches.append({
-                "summary": document,
-                "metadata": metadata,
-                "distance": float(distance),
-            })
+            matches.append(
+                {
+                    "summary": document,
+                    "metadata": metadata,
+                    "distance": float(distance),
+                }
+            )
         return matches
 
     def _patient_summary(self, patient_data, prob, risk):
diff --git a/models/risk_model.py b/models/risk_model.py
@@ -17,14 +17,13 @@ def calculate_bmi(height_cm, weight_kg):
         height_m = height_cm / 100
         if height_m <= 0:
             raise ValueError("Height must be greater than 0.")
-        return weight_kg / (height_m ** 2)
+        return weight_kg / (height_m**2)
 
     def train(self, df):
         df = df.copy()
         if "bmi" not in df.columns:
             df["bmi"] = df.apply(
-                lambda row: self.calculate_bmi(row["height_cm"], row["weight_kg"]),
-                axis=1
+                lambda row: self.calculate_bmi(row["height_cm"], row["weight_kg"]), axis=1
             )
 
         self.default_glucose = float(df["glucose"].median())
@@ -86,9 +85,14 @@ def prepare_patient(self, patient_data):
         }
 
     def feature_frame(self, patient):
-        return pd.DataFrame([{
-            "age": patient["age"],
-            "bmi": patient["bmi"],
-            "bp": patient["bp"],
-            "glucose": patient["glucose"],
-        }], columns=self.FEATURE_NAMES)
+        return pd.DataFrame(
+            [
+                {
+                    "age": patient["age"],
+                    "bmi": patient["bmi"],
+                    "bp": patient["bp"],
+                    "glucose": patient["glucose"],
+                }
+            ],
+            columns=self.FEATURE_NAMES,
+        )
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,8 @@
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+exclude = [".venv", "data/chroma", "logs"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+ignore = ["E402", "E501"]
diff --git a/requirements.txt b/requirements.txt
@@ -8,3 +8,6 @@ langchain
 chromadb
 shap
 pytest
+pytest-cov
+ruff
+pre-commit
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -2,7 +2,6 @@
 
 from api.app import app
 
-
 client = TestClient(app)
 
 
diff --git a/tools/explainability.py b/tools/explainability.py
diff --git a/tools/safety.py b/tools/safety.py
diff --git a/ui/app.py b/ui/app.py

Original file line number	Diff line number	Diff line change
`@@ -71,5 +71,5 @@ def run(self, patient_data):`
`71`	`71`	`"similar_cases": similar_cases,`
`72`	`72`	`"safety": safety,`
`73`	`73`	`"explanation": explanation,`
`74`		`- "recommendation": recommendation`
	`74`	`+ "recommendation": recommendation,`
`75`	`75`	`}`