diff --git a/promptops/promptops/guard.py b/promptops/promptops/guard.py index da7e03b8c..986d5adfa 100644 --- a/promptops/promptops/guard.py +++ b/promptops/promptops/guard.py @@ -55,8 +55,61 @@ def wrapper(*args, **kwargs) -> Any: try: response_text = extract_text(response) except ValueError as e: - logger.warning(str(e)) - return response + import hmac + import hashlib + import json + import os + import uuid + from datetime import datetime + from promptops.engine import redact_sensitive_data, get_workspace_audit_dir, get_signing_key + + timestamp = datetime.now().isoformat() + audit_dir = get_workspace_audit_dir() + key = get_signing_key() + + if mode == "fail_fast": + state = { + "prompt_id": prompt_id, + "status": "extraction_failure", + "mode": "fail_fast", + "error": str(e), + "timestamp": timestamp + } + else: + stringified_response = str(response) + redacted_text = redact_sensitive_data(stringified_response) + state = { + "prompt_id": prompt_id, + "status": "extraction_failure", + "mode": "warning", + "error": str(e), + "timestamp": timestamp, + "fallback_content": redacted_text + } + + state_json = json.dumps(state, sort_keys=True) + payload_to_sign = f"guard_failure|||{prompt_id}|||{timestamp}|||{state_json}" + signature = hmac.new(key, payload_to_sign.encode('utf-8'), hashlib.sha256).hexdigest() + + failure_id = str(uuid.uuid4()) + checkpoint_file = os.path.join(audit_dir, "guard_failures", f"{failure_id}.json") + sig_file = os.path.join(audit_dir, "guard_failures", f"{failure_id}.sig") + + os.makedirs(os.path.dirname(checkpoint_file), exist_ok=True) + with open(checkpoint_file, 'w', encoding='utf-8') as f: + f.write(state_json) + with open(sig_file, 'w', encoding='utf-8') as f: + json.dump({ + "timestamp": timestamp, + "signature": signature, + "algorithm": "HMAC-SHA256" + }, f) + + if mode == "fail_fast": + raise ProomptsValidationError(f"Extraction failed: {e}") + else: + logger.warning(f"Extraction failed: {e}. Falling back to redacted warning mode output.") + return redacted_text # 3. Load prompt schema from file prompt_path = ROOT / "prompts" / f"{prompt_id}.prompt.yaml" diff --git a/tests/test_guard.py b/tests/test_guard.py index 46714cad2..fd97f78ee 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -90,5 +90,154 @@ def mock_llm_call(): mock_llm_call() self.assertIn("must contain hello", str(ctx.exception)) + @patch("promptops.guard.load_yaml") + @patch("pathlib.Path.exists") + def test_extraction_failure_fail_fast(self, mock_exists, mock_load_yaml): + import tempfile + import os + import json + import hmac + import hashlib + from promptops.engine import get_signing_key + + mock_exists.return_value = True + mock_load_yaml.return_value = { + "name": "Test", + "description": "Test", + "model": "gpt-4", + "modelParameters": {"temperature": 0.0}, + "metadata": {"domain": "test", "complexity": "low", "tags": ["skill"]}, + "messages": [{"role": "system", "content": "hello"}], + "testData": [], + "evaluators": [] + } + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch.dict(os.environ, {"PROMPTOPS_WORKSPACE_AUDIT": tmp_dir}): + @guard(prompt_id="test_prompt_fail_fast", mode="fail_fast") + def mock_llm_call(): + # Return an unrecognized format + return {"unrecognized": "format", "patient_ssn": "123-45-6789"} + + with self.assertRaises(ProomptsValidationError) as ctx: + mock_llm_call() + self.assertIn("Extraction failed", str(ctx.exception)) + + # Check that signed audit was written + failures_dir = os.path.join(tmp_dir, "guard_failures") + self.assertTrue(os.path.exists(failures_dir)) + files = os.listdir(failures_dir) + json_files = [f for f in files if f.endswith(".json")] + sig_files = [f for f in files if f.endswith(".sig")] + self.assertEqual(len(json_files), 1) + self.assertEqual(len(sig_files), 1) + + # Verify contents and signature + json_path = os.path.join(failures_dir, json_files[0]) + sig_path = os.path.join(failures_dir, sig_files[0]) + + with open(json_path, "r", encoding="utf-8") as f: + state = json.load(f) + + self.assertEqual(state["prompt_id"], "test_prompt_fail_fast") + self.assertEqual(state["status"], "extraction_failure") + self.assertEqual(state["mode"], "fail_fast") + self.assertIn("Could not extract text", state["error"]) + + with open(sig_path, "r", encoding="utf-8") as f: + sig_meta = json.load(f) + + self.assertEqual(sig_meta["algorithm"], "HMAC-SHA256") + + # Recalculate signature + key = get_signing_key() + timestamp = state["timestamp"] + state_json = json.dumps(state, sort_keys=True) + payload_to_sign = f"guard_failure|||test_prompt_fail_fast|||{timestamp}|||{state_json}" + expected_sig = hmac.new(key, payload_to_sign.encode("utf-8"), hashlib.sha256).hexdigest() + self.assertEqual(sig_meta["signature"], expected_sig) + + @patch("promptops.guard.load_yaml") + @patch("pathlib.Path.exists") + def test_extraction_failure_warning(self, mock_exists, mock_load_yaml): + import tempfile + import os + import json + import hmac + import hashlib + from promptops.engine import get_signing_key + + mock_exists.return_value = True + mock_load_yaml.return_value = { + "name": "Test", + "description": "Test", + "model": "gpt-4", + "modelParameters": {"temperature": 0.0}, + "metadata": {"domain": "test", "complexity": "low", "tags": ["skill"]}, + "messages": [{"role": "system", "content": "hello"}], + "testData": [], + "evaluators": [] + } + + with tempfile.TemporaryDirectory() as tmp_dir: + with patch.dict(os.environ, {"PROMPTOPS_WORKSPACE_AUDIT": tmp_dir}): + @guard(prompt_id="test_prompt_warning", mode="warning") + def mock_llm_call(): + # Return an unrecognized format with patient data + return { + "ssn": "000-12-3456", + "email": "john.doe@hospital.org", + "phone": "+1 (555) 019-2834", + "dob": "1978/11/05" + } + + result = mock_llm_call() + + # Check output redaction + self.assertIsInstance(result, str) + self.assertNotIn("000-12-3456", result) + self.assertNotIn("john.doe@hospital.org", result) + self.assertNotIn("555", result) + self.assertNotIn("1978", result) + + self.assertIn("[REDACTED_SSN]", result) + self.assertIn("[REDACTED_EMAIL]", result) + self.assertIn("[REDACTED_PHONE]", result) + self.assertIn("[REDACTED_DATE]", result) + + # Check that signed audit was written + failures_dir = os.path.join(tmp_dir, "guard_failures") + self.assertTrue(os.path.exists(failures_dir)) + files = os.listdir(failures_dir) + json_files = [f for f in files if f.endswith(".json")] + sig_files = [f for f in files if f.endswith(".sig")] + self.assertEqual(len(json_files), 1) + self.assertEqual(len(sig_files), 1) + + # Verify contents and signature + json_path = os.path.join(failures_dir, json_files[0]) + sig_path = os.path.join(failures_dir, sig_files[0]) + + with open(json_path, "r", encoding="utf-8") as f: + state = json.load(f) + + self.assertEqual(state["prompt_id"], "test_prompt_warning") + self.assertEqual(state["status"], "extraction_failure") + self.assertEqual(state["mode"], "warning") + self.assertIn("[REDACTED_SSN]", state["fallback_content"]) + + with open(sig_path, "r", encoding="utf-8") as f: + sig_meta = json.load(f) + + self.assertEqual(sig_meta["algorithm"], "HMAC-SHA256") + + # Recalculate signature + key = get_signing_key() + timestamp = state["timestamp"] + state_json = json.dumps(state, sort_keys=True) + payload_to_sign = f"guard_failure|||test_prompt_warning|||{timestamp}|||{state_json}" + expected_sig = hmac.new(key, payload_to_sign.encode("utf-8"), hashlib.sha256).hexdigest() + self.assertEqual(sig_meta["signature"], expected_sig) + if __name__ == "__main__": unittest.main()