diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index a74b943574b..33d9dd50a13 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -121,13 +121,18 @@ class AzureSentinelLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.log_queue: List[StandardLoggingPayload] = [] - # When True, string fields (messages, response) are truncated to the - # Azure Log Analytics column limit (256 KB / 262,144 chars). Azure - # silently truncates at this limit anyway; doing it ourselves lets us - # keep the tail (most recent content) and record metadata. - # Controlled by AZURE_SENTINEL_TRUNCATE_CONTENT env var (default: false). - truncate_env = os.getenv("AZURE_SENTINEL_TRUNCATE_CONTENT", "false") - self.truncate_content = truncate_env.lower() in ("true", "1", "yes") + # When set to a positive integer, string fields (messages, response) + # are truncated to that many characters. Azure Log Analytics silently + # truncates string columns when over the limits doing it here + # ourselves lets us keep the tail (most recent content) and record + # metadata. Set AZURE_SENTINEL_TRUNCATE_BYTES to the desired char + # limit (e.g. "262144"). When unset or "0", truncation is disabled. + truncate_bytes_env = os.getenv("AZURE_SENTINEL_TRUNCATE_BYTES", "0") + try: + self.truncate_max_chars = int(truncate_bytes_env) + except ValueError: + self.truncate_max_chars = 0 + self.truncate_content = self.truncate_max_chars > 0 async def _get_oauth_token(self) -> str: """ @@ -253,18 +258,13 @@ class AzureSentinelLogger(CustomBatchLogger): # We target a conservative threshold to stay safely under the limit. MAX_BATCH_SIZE_BYTES = 950_000 # ~950KB uncompressed target per batch - # Azure Log Analytics silently truncates string column values at 256 KB. - # We enforce this limit ourselves so we can keep the tail (most recent - # content) and record truncation metadata. - MAX_COLUMN_CHARS = 262_144 # 256 KB Azure Log Analytics column limit - def _enforce_column_limits( self, payload: StandardLoggingPayload ) -> StandardLoggingPayload: """ - Truncate messages/response string fields to the Azure Log Analytics - column limit (262,144 chars). Keeps the *tail* of each field so that - the most recent conversation turns and response text are preserved. + Truncate messages/response string fields to ``self.truncate_max_chars`` + characters. Keeps the *tail* of each field so that the most recent + conversation turns and response text are preserved. Only called when ``self.truncate_content`` is True. @@ -272,7 +272,7 @@ class AzureSentinelLogger(CustomBatchLogger): limit, otherwise returns a deep copy with truncated fields and truncation metadata added. """ - limit = self.MAX_COLUMN_CHARS + limit = self.truncate_max_chars messages = payload.get("messages") response = payload.get("response") msg_str = str(messages) if messages is not None else "" @@ -328,8 +328,8 @@ class AzureSentinelLogger(CustomBatchLogger): Splits payloads into gzip-compressed batches that stay under MAX_BATCH_SIZE_BYTES (uncompressed) per batch. - When truncate_content is enabled, enforces the Azure Log Analytics - column limit (256 KB) on messages/response fields before batching. + When truncate_content is enabled, enforces the configured character + limit on messages/response fields before batching. Returns a list of gzip-compressed byte strings, each representing a JSON array of log entries. diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 7bc0ff2f7c2..534cbf9de07 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -139,8 +139,7 @@ async def test_azure_sentinel_batch_splitting(): @pytest.mark.asyncio async def test_azure_sentinel_split_into_batches_single_oversized_entry(): """Test that a single entry larger than MAX_BATCH_SIZE_BYTES is sent alone""" - with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_CONTENT": "false"}): - logger = _make_logger() + logger = _make_logger() # One very large payload that exceeds the batch size on its own. # Use varied content so JSON serialization keeps it large. @@ -177,11 +176,11 @@ async def test_azure_sentinel_split_into_batches_single_oversized_entry(): def test_column_limit_truncates_large_fields(): - """Test that fields exceeding 262,144 chars are truncated (keeping the tail)""" - with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_CONTENT": "true"}): + """Test that fields exceeding the configured limit are truncated (keeping the tail)""" + with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_BYTES": "262144"}): logger = _make_logger() - # Content larger than 256 KB column limit + # Content larger than the configured limit big_messages = "A" * 300_000 big_response = "B" * 300_000 @@ -199,12 +198,12 @@ def test_column_limit_truncates_large_fields(): # Messages field should be truncated, keeping tail msg_str = str(result["messages"]) assert msg_str.startswith("[truncated by litellm]...") - assert len(msg_str) <= logger.MAX_COLUMN_CHARS + assert len(msg_str) <= logger.truncate_max_chars # Response field should be truncated, keeping tail resp_str = str(result["response"]) assert resp_str.startswith("[truncated by litellm]...") - assert len(resp_str) <= logger.MAX_COLUMN_CHARS + assert len(resp_str) <= logger.truncate_max_chars # Truncation metadata present metadata = result.get("metadata", {}) @@ -217,15 +216,15 @@ def test_column_limit_truncates_large_fields(): assert "messages" in trunc_info["truncated_fields"] assert "response" in trunc_info["truncated_fields"] assert trunc_info["original_messages_chars"] == len(str(payload["messages"])) - assert trunc_info["max_column_chars"] == 262_144 + assert trunc_info["max_column_chars"] == 262144 # Original payload not mutated - assert len(str(payload["messages"])) > 262_144 + assert len(str(payload["messages"])) > 262144 def test_column_limit_preserves_small_payloads(): - """Test that payloads under the column limit are returned unchanged""" - with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_CONTENT": "true"}): + """Test that payloads under the configured limit are returned unchanged""" + with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_BYTES": "262144"}): logger = _make_logger() payload = _make_payload( @@ -242,13 +241,15 @@ def test_column_limit_preserves_small_payloads(): assert "litellm_content_truncated" not in metadata -def test_truncate_disabled_via_env_var(): - """Test that truncation is skipped when AZURE_SENTINEL_TRUNCATE_CONTENT=false""" - with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_CONTENT": "false"}): +def test_truncate_disabled_when_env_unset(): + """Test that truncation is skipped when AZURE_SENTINEL_TRUNCATE_BYTES is not set""" + with patch.dict(os.environ, {}, clear=False): + # Ensure the env var is absent + os.environ.pop("AZURE_SENTINEL_TRUNCATE_BYTES", None) logger = _make_logger() assert logger.truncate_content is False - # Create payload with content exceeding 256 KB column limit + # Create payload with content exceeding 256 KB huge_content = "z" * 300_000 payloads = [ _make_payload( @@ -270,11 +271,12 @@ def test_truncate_disabled_via_env_var(): def test_truncate_enabled_in_split_batches(): """Test that _split_into_batches truncates large fields when enabled""" - with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_CONTENT": "true"}): + with patch.dict(os.environ, {"AZURE_SENTINEL_TRUNCATE_BYTES": "262144"}): logger = _make_logger() assert logger.truncate_content is True + assert logger.truncate_max_chars == 262144 - # Content exceeding 256 KB column limit + # Content exceeding the configured limit huge_content = "X" * 400_000 payloads = [ _make_payload(