diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260415120000_health_check_latest_per_model_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260415120000_health_check_latest_per_model_index/migration.sql new file mode 100644 index 0000000000..189191b901 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260415120000_health_check_latest_per_model_index/migration.sql @@ -0,0 +1,12 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is +-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration +-- without IF NOT EXISTS if you must support older versions). +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx" ON "LiteLLM_HealthCheckTable"("model_id", "model_name", "checked_at" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b5..1b3db52eaa 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1045,6 +1045,7 @@ model LiteLLM_HealthCheckTable { @@index([model_name]) @@index([checked_at]) @@index([status]) + @@index([model_id, model_name, checked_at(sort: Desc)], map: "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx") } // Search Tools table for storing search tool configurations diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b4b8e68113..ed34ec6682 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -83,7 +83,12 @@ def _redact_pii_matches(response_json: dict) -> dict: redacted_response = copy.deepcopy(response_json) # Get assessments from the response - assessments = redacted_response.get("assessments", []) + # NOTE: We use `.get("key") or []` instead of `.get("key", [])` because + # the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null). + # In Python, dict.get("key", []) returns None (not []) when the key exists + # with a None/null value. The `or []` ensures we always get an iterable, + # preventing "TypeError: 'NoneType' object is not iterable". + assessments = redacted_response.get("assessments") or [] if not assessments: return redacted_response @@ -91,13 +96,13 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact PII entities in sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] for pii_entity in pii_entities: if "match" in pii_entity: pii_entity["match"] = "[REDACTED]" # Redact regex matches - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] for regex_match in regexes: if "match" in regex_match: regex_match["match"] = "[REDACTED]" @@ -105,12 +110,12 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact custom word matches in word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if "match" in custom_word: custom_word["match"] = "[REDACTED]" - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if "match" in managed_word: managed_word["match"] = "[REDACTED]" @@ -825,7 +830,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return False # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED) - assessments = response.get("assessments", []) + # NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API. + # See _redact_pii_matches() for detailed explanation of the null safety pattern. + assessments = response.get("assessments") or [] if not assessments: return False @@ -833,7 +840,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check topic policy topic_policy = assessment.get("topicPolicy") if topic_policy: - topics = topic_policy.get("topics", []) + topics = topic_policy.get("topics") or [] for topic in topics: if topic.get("action") == "BLOCKED": return True @@ -841,7 +848,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check content policy content_policy = assessment.get("contentPolicy") if content_policy: - filters = content_policy.get("filters", []) + filters = content_policy.get("filters") or [] for filter_item in filters: if filter_item.get("action") == "BLOCKED": return True @@ -849,11 +856,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if custom_word.get("action") == "BLOCKED": return True - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if managed_word.get("action") == "BLOCKED": return True @@ -861,12 +868,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] if pii_entities: for pii_entity in pii_entities: if pii_entity.get("action") == "BLOCKED": return True - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] if regexes: for regex in regexes: if regex.get("action") == "BLOCKED": @@ -875,7 +882,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check contextual grounding policy contextual_grounding_policy = assessment.get("contextualGroundingPolicy") if contextual_grounding_policy: - grounding_filters = contextual_grounding_policy.get("filters", []) + grounding_filters = contextual_grounding_policy.get("filters") or [] for grounding_filter in grounding_filters: if grounding_filter.get("action") == "BLOCKED": return True @@ -1534,7 +1541,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Raises: Exception: If content is blocked by Bedrock guardrail """ - texts = inputs.get("texts", []) + # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. + # dict.get("texts", []) would return None if the key exists with a None value. + texts = inputs.get("texts") or [] try: verbose_proxy_logger.debug( f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index 2ecee5095b..5b8370fece 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -84,7 +84,7 @@ class SharedHealthCheckManager: "Pod %s failed to acquire health check lock", self.pod_id ) - return acquired + return bool(acquired) except Exception as e: verbose_proxy_logger.error("Error acquiring health check lock: %s", str(e)) return False diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fce95465b5..1b3db52eaa 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1045,6 +1045,7 @@ model LiteLLM_HealthCheckTable { @@index([model_name]) @@index([checked_at]) @@index([status]) + @@index([model_id, model_name, checked_at(sort: Desc)], map: "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx") } // Search Tools table for storing search tool configurations diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a6f81986a6..d5cc7d2562 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4525,29 +4525,20 @@ class PrismaClient: async def get_all_latest_health_checks(self): """ - Get the latest health check for each model + Get the latest health check for each model. + + Uses DB-level DISTINCT ON (model_id, model_name) with ORDER BY checked_at DESC + (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. """ try: - # Get all unique model names first - all_checks = await self.db.litellm_healthchecktable.find_many( - order={"checked_at": "desc"} + return await self.db.litellm_healthchecktable.find_many( + distinct=["model_id", "model_name"], + order=[ + {"model_id": "asc"}, + {"model_name": "asc"}, + {"checked_at": "desc"}, + ], ) - - # Group by model_name and get the latest for each - latest_checks = {} - for check in all_checks: - # Create a unique key: prefer model_id if available, otherwise use model_name - # This ensures we get the latest check for each unique model - if check.model_id: - key = (check.model_id, check.model_name) - else: - key = (None, check.model_name) - - # Only add if we haven't seen this key yet (since checks are ordered by checked_at desc) - if key not in latest_checks: - latest_checks[key] = check - - return list(latest_checks.values()) except Exception as e: verbose_proxy_logger.error(f"Error getting all latest health checks: {e}") return [] diff --git a/schema.prisma b/schema.prisma index fce95465b5..1b3db52eaa 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1045,6 +1045,7 @@ model LiteLLM_HealthCheckTable { @@index([model_name]) @@index([checked_at]) @@index([status]) + @@index([model_id, model_name, checked_at(sort: Desc)], map: "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx") } // Search Tools table for storing search tool configurations diff --git a/scripts/health_check/benchmark_get_all_latest_health_checks.py b/scripts/health_check/benchmark_get_all_latest_health_checks.py new file mode 100644 index 0000000000..9161861883 --- /dev/null +++ b/scripts/health_check/benchmark_get_all_latest_health_checks.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +Bench LiteLLM_HealthCheckTable + PrismaClient + - set DATABASE_URL to your Postgres + - Run ```prisma generate``` to install prisma client before running test ) + - This test writes to the default "public" database. Make sure to run cleanup after testing + +""" + +from __future__ import annotations + +import argparse +import asyncio +import gc +import os +import sys +import time +import tracemalloc +from datetime import datetime, timedelta, timezone +from typing import Any, List + +SEED_MARKER = "benchmark_get_all_latest_health_checks.py" # Utility Marker for cleanup process. + + +def _rss_kb_linux() -> int: + try: + with open("/proc/self/status", encoding="utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + return int(line.split()[1]) + except OSError: + pass + return 0 + + +def _fmt_kb(kb: int) -> str: + if kb <= 0: + return "n/a" + return f"{kb} KiB (~{kb / 1024.0:.1f} MiB)" + + +def _build_batch( + *, + batch_index: int, + batch_size: int, + num_models: int, + base_time: datetime, +) -> List[dict[str, Any]]: + rows: List[dict[str, Any]] = [] + for i in range(batch_size): + global_i = batch_index * batch_size + i + model_idx = global_i % max(num_models, 1) + model_name = f"bench-model-{model_idx}" + model_id = f"bench-mid-{model_idx}" if model_idx % 2 == 0 else None + checked_at = base_time - timedelta(seconds=global_i) + rows.append( + { + "model_name": model_name, + "model_id": model_id, + "status": "healthy" if global_i % 3 else "unhealthy", + "healthy_count": 1, + "unhealthy_count": 0, + "checked_by": SEED_MARKER, + "checked_at": checked_at, + } + ) + return rows + + +async def _seed( + prisma: Any, + *, + total_rows: int, + batch_size: int, + num_models: int, +) -> None: + db = prisma.db + base_time = datetime.now(timezone.utc) + inserted = 0 + batch_idx = 0 + while inserted < total_rows: + n = min(batch_size, total_rows - inserted) + await db.litellm_healthchecktable.create_many( + data=_build_batch( + batch_index=batch_idx, + batch_size=n, + num_models=num_models, + base_time=base_time, + ) + ) + inserted += n + batch_idx += 1 + if batch_idx % 10 == 0: + print(f" {inserted}/{total_rows}", flush=True) + print(f"Seeded {inserted} rows ({SEED_MARKER}).") + + +async def _cleanup(prisma: Any) -> None: + result = await prisma.db.litellm_healthchecktable.delete_many( + where={"checked_by": SEED_MARKER}, + ) + n = getattr(result, "count", result) + print(f"Deleted {n} rows.") + + +async def _bench(prisma: Any) -> None: + gc.collect() + rss0 = _rss_kb_linux() + print(f"RSS (after gc): {_fmt_kb(rss0)}") + + tracemalloc.start() + t0 = time.perf_counter() + try: + rows = await prisma.get_all_latest_health_checks() + finally: + elapsed = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + gc.collect() + rss1 = _rss_kb_linux() + print(f"get_all_latest_health_checks: {len(rows)} rows in {elapsed:.2f}s") + print(f"tracemalloc peak: {peak / 1e6:.2f} MiB") + print(f"RSS after: {_fmt_kb(rss1)}") + + +async def _amain() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("action", choices=("seed", "bench", "cleanup")) + p.add_argument("--rows", type=int, default=10_000) + p.add_argument("--batch-size", type=int, default=1000) + p.add_argument("--num-models", type=int, default=50) + args = p.parse_args() + + database_url = os.getenv("DATABASE_URL") + if not database_url: + print("Set DATABASE_URL.", file=sys.stderr) + return 1 + + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + if repo_root not in sys.path: + sys.path.insert(0, repo_root) + + from litellm.caching.caching import DualCache + from litellm.proxy.proxy_cli import append_query_params + from litellm.proxy.utils import PrismaClient, ProxyLogging + + db_url = append_query_params( + database_url, {"connection_limit": 100, "pool_timeout": 60} + ) + prisma = PrismaClient( + database_url=db_url, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + try: + await prisma.connect() + except Exception as e: + print(f"Connect failed: {e}", file=sys.stderr) + return 1 + + try: + if args.action == "seed": + await _seed( + prisma, + total_rows=args.rows, + batch_size=args.batch_size, + num_models=args.num_models, + ) + elif args.action == "bench": + await _bench(prisma) + else: + await _cleanup(prisma) + finally: + try: + await prisma.disconnect() + except Exception: + pass + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(_amain())) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index cb594c221c..146d16c242 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -5,7 +5,10 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) import litellm -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail +from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + _redact_pii_matches, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache from unittest.mock import MagicMock, AsyncMock, patch @@ -1601,3 +1604,128 @@ async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): # If no error is raised and result is None, then the test passes assert result is None print("✅ No output text in response test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_null_list_fields(): + """Test that explicit null values from Bedrock API are handled correctly. + + The Bedrock API can return explicit JSON null for list fields like + piiEntities, regexes, customWords, managedWordLists. This would cause + TypeError: 'NoneType' object is not iterable if not handled. + """ + # Test 1: null piiEntities and regexes + response_with_null_pii = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + } + } + ], + } + redacted = _redact_pii_matches(response_with_null_pii) + assert redacted is not None + assert redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] is None + assert redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] is None + + # Test 2: null customWords and managedWordLists + response_with_null_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + } + } + ], + } + redacted = _redact_pii_matches(response_with_null_words) + assert redacted is not None + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + # Test 3: null assessments at top level + response_with_null_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, + } + redacted = _redact_pii_matches(response_with_null_assessments) + assert redacted is not None + + +@pytest.mark.asyncio +async def test__redact_pii_matches_malformed_response(): + """Test _redact_pii_matches with malformed response (should not crash)""" + + # Test with completely malformed response + malformed_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": "not_a_list", + } + redacted_response = _redact_pii_matches(malformed_response) + assert redacted_response == malformed_response + + # Test with missing keys + missing_keys_response = { + "action": "GUARDRAIL_INTERVENED", + } + redacted_response = _redact_pii_matches(missing_keys_response) + assert redacted_response == missing_keys_response + + +@pytest.mark.asyncio +async def test_should_raise_guardrail_blocked_exception_null_fields(): + """Test that _should_raise_guardrail_blocked_exception handles null list fields. + + Validates the or [] null-safety pattern works for all policy fields + in _should_raise_guardrail_blocked_exception. + """ + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Test with null assessments + response_null_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_assessments) is False + + # Test with null topics in topicPolicy + response_null_topics = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"topicPolicy": {"topics": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_topics) is False + + # Test with null filters in contentPolicy + response_null_filters = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"contentPolicy": {"filters": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_filters) is False + + # Test with null customWords and managedWordLists in wordPolicy + response_null_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"wordPolicy": {"customWords": None, "managedWordLists": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_words) is False + + # Test with null piiEntities and regexes in sensitiveInformationPolicy + response_null_pii = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": None, "regexes": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_pii) is False + + # Test with null filters in contextualGroundingPolicy + response_null_grounding = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"contextualGroundingPolicy": {"filters": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_grounding) is False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 010ead425c..0472cc565e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1190,6 +1190,483 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): print("✅ BLOCKED content with masking enabled raises exception correctly") +# ────────────────────────────────────────────────────────────────────────────── +# Null-safety tests for Bedrock guardrail responses +# +# The Bedrock ApplyGuardrail API can return explicit null/None for list fields +# such as "regexes", "piiEntities", "topics", "filters", "customWords", and +# "managedWordLists" when a particular policy category is present in the +# assessment but has no matches. +# +# Python's dict.get("key", []) returns None (NOT []) when the key exists with +# a None value. The `or []` fallback ensures we always iterate over a list. +# +# Without the fix, iterating over None raises: +# TypeError: 'NoneType' object is not iterable +# which surfaces to callers as: +# openai.InternalServerError: Error code: 500 +# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} +# ────────────────────────────────────────────────────────────────────────────── + + +class TestRedactPiiMatchesNullSafety: + """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" + + @pytest.mark.asyncio + async def test_should_handle_null_regexes_in_sensitive_info_policy(self): + """Bedrock can return regexes: null while piiEntities has data. + + Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no + custom regex patterns configured, so the API returns regexes: null. + """ + response = { + "action": "NONE", + "actionReason": "No action.", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "NONE", + "detected": True, + "match": "joebloggs@gmail.com", + "type": "EMAIL", + } + ], + "regexes": None, # Explicit null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError: 'NoneType' object is not iterable + redacted = _redact_pii_matches(response) + + # PII match should be redacted + pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert pii[0]["match"] == "[REDACTED]" + assert pii[0]["type"] == "EMAIL" + + @pytest.mark.asyncio + async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): + """Bedrock can return piiEntities: null while regexes has data.""" + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, # null from Bedrock API + "regexes": [ + { + "name": "CUSTOM_PATTERN", + "match": "secret-abc-123", + "action": "BLOCKED", + } + ], + }, + } + ], + } + + redacted = _redact_pii_matches(response) + + regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_should_handle_null_custom_words_and_managed_words(self): + """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" + response = { + "action": "NONE", + "assessments": [ + { + "wordPolicy": { + "customWords": None, # null from Bedrock API + "managedWordLists": None, # null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + + # Values should remain None (no crash) + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """Bedrock can return assessments: null.""" + response = { + "action": "NONE", + "assessments": None, # null from Bedrock API + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + assert redacted["assessments"] is None + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists_together(self): + """All sub-list fields are null at the same time — worst-case scenario.""" + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + }, + "topicPolicy": None, + "contentPolicy": None, + "contextualGroundingPolicy": None, + } + ], + } + + # Should not raise any exception + redacted = _redact_pii_matches(response) + assert redacted is not None + + +class TestShouldRaiseGuardrailBlockedExceptionNullSafety: + """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" + + def _create_guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists(self): + """All policy sub-lists are null — should not crash, should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null from Bedrock API + }, + "contentPolicy": { + "filters": None, # null + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": { + "filters": None, # null + }, + } + ], + } + + # No BLOCKED actions found (all lists null) → should return False + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_detect_blocked_despite_other_null_lists(self): + """A mix of null lists and a real BLOCKED action — should still detect it.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null — should not crash + }, + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": None, # entire policy is null + } + ], + } + + # Should return True because contentPolicy has a BLOCKED filter + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """assessments itself is null — should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, # null from Bedrock API + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_handle_null_topics_with_blocked_word_policy(self): + """topics is null but wordPolicy has a BLOCKED customWord.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, + }, + "wordPolicy": { + "customWords": [ + {"match": "badword", "action": "BLOCKED"} + ], + "managedWordLists": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_pii_with_blocked_regex(self): + """piiEntities is null but regexes has a BLOCKED match.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": [ + {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} + ], + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_grounding_filters(self): + """contextualGroundingPolicy.filters is null — should not crash.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contextualGroundingPolicy": { + "filters": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_not_crash_when_action_is_not_intervened(self): + """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" + guardrail = self._create_guardrail() + + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + +class TestApplyGuardrailNullSafety: + """Tests for apply_guardrail handling of null/None texts input.""" + + @pytest.mark.asyncio + async def test_should_handle_none_texts_in_inputs(self): + """inputs[\"texts\"] is explicitly None — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {"texts": None} # Explicit None + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + # With empty texts (from None → []), no Bedrock API call should be made + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + # Should return empty texts without crashing + assert result.get("texts") == [] + # No Bedrock API call should be made for empty input + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_should_handle_missing_texts_key(self): + """inputs has no \"texts\" key at all — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {} # No "texts" key + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result.get("texts") == [] + mock_post.assert_not_called() + + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): + """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Test 1: ANONYMIZED action should NOT raise exception + anonymized_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + anonymized_response + ) + assert should_raise is False, "ANONYMIZED actions should not raise exceptions" + + # Test 2: BLOCKED action should raise exception + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) + assert should_raise is True, "BLOCKED actions should raise exceptions" + + # Test 3: Mixed actions - should raise if ANY action is BLOCKED + mixed_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + }, + "topicPolicy": { + "topics": [ + {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} + ] + }, + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) + assert ( + should_raise is True + ), "Mixed actions with any BLOCKED should raise exceptions" + + # Test 4: NONE action should not raise exception + none_response = { + "action": "NONE", + "outputs": [], + "assessments": [], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) + assert should_raise is False, "NONE action should not raise exceptions" + + print("\u2705 BLOCKED vs ANONYMIZED actions test passed") + + # --------------------------------------------------------------------------- # L3: _extract_blocked_assessments + _get_http_exception_for_blocked_guardrail # Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error. @@ -1338,4 +1815,3 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): assert isinstance(exc, HTTPException) assert "assessments" not in exc.detail assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" - diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 13d2131efa..8b23e526b6 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -406,12 +406,6 @@ async def test_save_background_health_checks_to_db_exception_handling(): @pytest.mark.asyncio async def test_get_all_latest_health_checks_with_model_id(mock_prisma): """Test get_all_latest_health_checks properly groups by model_id""" - # Create mock checks with same model_name but different model_id - mock_check1 = MagicMock() - mock_check1.model_id = "model-123" - mock_check1.model_name = "gpt-3.5-turbo" - mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10) - mock_check2 = MagicMock() mock_check2.model_id = "model-456" mock_check2.model_name = "gpt-3.5-turbo" @@ -424,7 +418,7 @@ async def test_get_all_latest_health_checks_with_model_id(mock_prisma): # Order by checked_at desc mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check3, mock_check2, mock_check1] + return_value=[mock_check3, mock_check2] ) result = await mock_prisma.get_all_latest_health_checks() @@ -445,18 +439,13 @@ async def test_get_all_latest_health_checks_with_model_id(mock_prisma): @pytest.mark.asyncio async def test_get_all_latest_health_checks_without_model_id(mock_prisma): """Test get_all_latest_health_checks groups by model_name when model_id is None""" - mock_check1 = MagicMock() - mock_check1.model_id = None - mock_check1.model_name = "gpt-3.5-turbo" - mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10) - mock_check2 = MagicMock() mock_check2.model_id = None mock_check2.model_name = "gpt-3.5-turbo" mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check2, mock_check1] + return_value=[mock_check2] ) result = await mock_prisma.get_all_latest_health_checks() @@ -467,6 +456,41 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma): assert result[0].checked_at == mock_check2.checked_at # Latest +@pytest.mark.asyncio +async def test_get_all_latest_health_checks_same_name_with_and_without_model_id(mock_prisma): + """ + Same model_name can appear twice after DISTINCT ON: once keyed by (model_id, name) + and once by (NULL, name) — different Postgres groups than a single row with id. + """ + now = datetime.now(timezone.utc) + with_id = MagicMock() + with_id.model_id = "deployment-abc" + with_id.model_name = "gpt-4" + with_id.checked_at = now - timedelta(minutes=2) + + without_id = MagicMock() + without_id.model_id = None + without_id.model_name = "gpt-4" + without_id.checked_at = now - timedelta(minutes=1) + + mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( + return_value=[without_id, with_id] + ) + + result = await mock_prisma.get_all_latest_health_checks() + + assert len(result) == 2 + names = {r.model_name for r in result} + assert names == {"gpt-4"} + ids = {r.model_id for r in result} + assert "deployment-abc" in ids + assert None in ids + + by_key = {(r.model_id, r.model_name): r for r in result} + assert by_key[("deployment-abc", "gpt-4")].checked_at == with_id.checked_at + assert by_key[(None, "gpt-4")].checked_at == without_id.checked_at + + @pytest.mark.asyncio async def test_perform_health_check_and_save_passes_model_id_to_perform_health_check(): """Test that _perform_health_check_and_save passes model_id to perform_health_check so health checks run by model id.""" diff --git a/uv.lock b/uv.lock index 04224dc537..d9c4375d80 100644 --- a/uv.lock +++ b/uv.lock @@ -11,7 +11,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-08T16:01:27.663665Z" +exclude-newer = "2026-04-13T01:16:14.00322Z" exclude-newer-span = "P3D" [manifest] @@ -3602,7 +3602,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.83.6" +version = "1.83.8" source = { editable = "." } dependencies = [ { name = "aiohttp" },