scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(presidio): resolve runtime error by handling asyncio loops in bac… (#19714)

Browse source 
* fix(presidio): resolve runtime error by handling asyncio loops in background threads

* add test case for thread safety
This commit is contained in:
Harshit Jain 2026-01-25 05:06:49 +05:30 • committed by GitHub
parent 53d3868ff2
commit 05fdd099ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 188 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

132
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -9,8 +9,10 @@
import asyncio
import threading
import json
from datetime import datetime
from contextlib import asynccontextmanager
from typing import (
TYPE_CHECKING,
Any,
@ -107,6 +109,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
# Lock to prevent race conditions when creating session under concurrent load
# Note: asyncio.Lock() can be created without an event loop; it only needs one when awaited
self._session_lock: asyncio.Lock = asyncio.Lock()
# Track main thread ID to safely identity when we are running in main loop vs background thread
self._main_thread_id = threading.get_ident()
if mock_testing is True: # for testing purposes only
return
@ -172,46 +179,47 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"http://" + self.presidio_anonymizer_api_base
)
async def _get_http_session(self) -> aiohttp.ClientSession:
@asynccontextmanager
async def _get_session_iterator(
self,
) -> AsyncGenerator[aiohttp.ClientSession, None]:
"""
Get or create the shared HTTP session for Presidio API calls.
Async context manager for yielding an HTTP session.
Fixes memory leak (issue #14540) where every guardrail check created
a new aiohttp.ClientSession that was never properly closed.
Thread-safe: Uses asyncio.Lock to prevent race conditions when
multiple concurrent requests try to create the session simultaneously.
Logic:
1. If running in the main thread (where the object was initialized/destined to live normally),
use the shared `self._http_session` (protected by a lock).
2. If running in a background thread (e.g. logging hook), yield a NEW ephemeral session
and ensure it is closed after use.
"""
async with self._session_lock:
if self._http_session is None or self._http_session.closed:
self._http_session = aiohttp.ClientSession()
return self._http_session
# Check if we are in the stored main thread
if threading.get_ident() == self._main_thread_id:
# Main thread -> use shared session
async with self._session_lock:
if self._http_session is None or self._http_session.closed:
self._http_session = aiohttp.ClientSession()
yield self._http_session
else:
# Background thread -> create ephemeral session
# This avoids "attached to a different loop" or "no running event loop" errors
# when accessing the shared session created in the main loop
session = aiohttp.ClientSession()
try:
yield session
finally:
if not session.closed:
await session.close()
async def _close_http_session(self) -> None:
"""Close the HTTP session if it exists."""
"""Close the shared HTTP session if it exists."""
if self._http_session is not None and not self._http_session.closed:
await self._http_session.close()
self._http_session = None
def __del__(self):
"""Cleanup: close HTTP session on instance destruction."""
if self._http_session is not None and not self._http_session.closed:
try:
# Try to close the session, but don't fail if event loop is gone
import asyncio
try:
loop = asyncio.get_event_loop()
if loop.is_running():
# Schedule cleanup, don't block __del__
asyncio.create_task(self._close_http_session())
else:
loop.run_until_complete(self._close_http_session())
except RuntimeError:
# Event loop is closed, can't clean up - not ideal but better than crashing
pass
except Exception:
# Suppress all exceptions in __del__ to avoid issues during shutdown
pass
"""Cleanup: we try to close, but doing async cleanup in __del__ is risky."""
pass
def _get_presidio_analyze_request_payload(
self,
@ -273,28 +281,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return self.mock_redacted_text
# Use shared session to prevent memory leak (issue #14540)
session = await self._get_http_session()
async with self._get_session_iterator() as session:
# Make the request to /analyze
analyze_url = f"{self.presidio_analyzer_api_base}analyze"
# Make the request to /analyze
analyze_url = f"{self.presidio_analyzer_api_base}analyze"
analyze_payload: PresidioAnalyzeRequest = (
self._get_presidio_analyze_request_payload(
text=text,
presidio_config=presidio_config,
request_data=request_data,
analyze_payload: PresidioAnalyzeRequest = (
self._get_presidio_analyze_request_payload(
text=text,
presidio_config=presidio_config,
request_data=request_data,
)
)
)
verbose_proxy_logger.debug(
"Making request to: %s with payload: %s",
analyze_url,
analyze_payload,
)
verbose_proxy_logger.debug(
"Making request to: %s with payload: %s",
analyze_url,
analyze_payload,
)
async with session.post(analyze_url, json=analyze_payload) as response:
analyze_results = await response.json()
verbose_proxy_logger.debug("analyze_results: %s", analyze_results)
async with session.post(analyze_url, json=analyze_payload) as response:
analyze_results = await response.json()
verbose_proxy_logger.debug("analyze_results: %s", analyze_results)
# Handle error responses from Presidio (e.g., {'error': 'No text provided'})
# Presidio may return a dict instead of a list when errors occur
@ -302,7 +309,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
if "error" in analyze_results:
verbose_proxy_logger.warning(
"Presidio analyzer returned error: %s, returning empty list",
analyze_results.get("error")
analyze_results.get("error"),
)
return []
# If it's a dict but not an error, try to process it as a single item
@ -314,7 +321,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
except Exception as e:
verbose_proxy_logger.warning(
"Failed to parse Presidio dict response: %s, returning empty list",
e
e,
)
return []
@ -351,20 +358,19 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return text
# Use shared session to prevent memory leak (issue #14540)
session = await self._get_http_session()
async with self._get_session_iterator() as session:
# Make the request to /anonymize
anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize"
verbose_proxy_logger.debug("Making request to: %s", anonymize_url)
anonymize_payload = {
"text": text,
"analyzer_results": analyze_results,
}
# Make the request to /anonymize
anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize"
verbose_proxy_logger.debug("Making request to: %s", anonymize_url)
anonymize_payload = {
"text": text,
"analyzer_results": analyze_results,
}
async with session.post(
anonymize_url, json=anonymize_payload
) as response:
redacted_text = await response.json()
async with session.post(
anonymize_url, json=anonymize_payload
) as response:
redacted_text = await response.json()
new_text = text
if redacted_text is not None:

141
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -4,10 +4,9 @@ Tests PII detection and masking for different message formats
"""
import asyncio
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import MagicMock, patch
import pytest
@ -625,7 +624,7 @@ async def test_request_data_flows_to_apply_guardrail():
return text
with patch.object(presidio, "check_pii", mock_check_pii):
result = await presidio.apply_guardrail(
await presidio.apply_guardrail(
inputs={"texts": ["Test message"]},
request_data=request_data,
input_type="request",
@ -709,16 +708,21 @@ async def test_presidio_filter_scope_initializer(monkeypatch):
monkeypatch.setattr(litellm, "logging_callback_manager", mgr, raising=False)
import litellm.proxy.guardrails.guardrail_initializers as gi
import litellm.proxy.guardrails.guardrail_hooks.presidio as presidio_mod
monkeypatch.setattr(
presidio_mod, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False
)
monkeypatch.setattr(gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False)
monkeypatch.setattr(
gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False
)
# input-only
created.clear()
from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio
params_input = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="input")
params_input = LitellmParams(
guardrail="presidio", mode="pre_call", presidio_filter_scope="input"
)
guardrail_dict = {"guardrail_name": "g1"}
cb = initialize_presidio(params_input, guardrail_dict)
assert cb is created[0]
@ -726,14 +730,18 @@ async def test_presidio_filter_scope_initializer(monkeypatch):
# output-only
created.clear()
params_output = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="output")
params_output = LitellmParams(
guardrail="presidio", mode="pre_call", presidio_filter_scope="output"
)
cb = initialize_presidio(params_output, guardrail_dict)
assert len(created) == 1
assert created[0].apply_to_output is True
# both -> expect two callbacks (input + output)
created.clear()
params_both = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="both")
params_both = LitellmParams(
guardrail="presidio", mode="pre_call", presidio_filter_scope="both"
)
cb = initialize_presidio(params_both, guardrail_dict)
assert len(created) == 2
assert any(not c.apply_to_output for c in created)
@ -741,13 +749,15 @@ async def test_presidio_filter_scope_initializer(monkeypatch):
@pytest.mark.asyncio
async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, mock_cache):
async def test_empty_content_handling(
presidio_guardrail, mock_user_api_key, mock_cache
):
"""
Test that Presidio handles empty content gracefully.
This is common in tool/function calling where assistant messages have
empty content but include tool_calls.
Bug fix: Previously crashed with:
TypeError: argument after ** must be a mapping, not str
"""
@ -761,7 +771,10 @@ async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, moc
{
"id": "call_123",
"type": "function",
"function": {"name": "calculator", "arguments": '{"a":2,"b":2}'},
"function": {
"name": "calculator",
"arguments": '{"a":2,"b":2}',
},
}
],
},
@ -794,10 +807,12 @@ async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, moc
@pytest.mark.asyncio
async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mock_cache):
async def test_whitespace_only_content(
presidio_guardrail, mock_user_api_key, mock_cache
):
"""
Test that Presidio handles whitespace-only content gracefully.
Whitespace-only content should be treated the same as empty content.
"""
test_data = {
@ -832,7 +847,7 @@ async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mo
async def test_analyze_text_with_empty_string():
"""
Test analyze_text method directly with empty string.
Should return empty list without making API call to Presidio.
"""
presidio = _OPTIONAL_PresidioPIIMasking(
@ -864,7 +879,7 @@ async def test_analyze_text_with_empty_string():
async def test_analyze_text_error_dict_handling():
"""
Test that analyze_text handles error dict responses from Presidio API.
When Presidio returns {'error': 'No text provided'}, should handle gracefully
instead of crashing with TypeError.
"""
@ -878,16 +893,20 @@ async def test_analyze_text_error_dict_handling():
class MockResponse:
async def json(self):
return {"error": "No text provided"}
async def __aenter__(self):
return self
async def __aexit__(self, *args):
pass
class MockSession:
def post(self, *args, **kwargs):
return MockResponse()
async def __aenter__(self):
return self
async def __aexit__(self, *args):
pass
@ -904,10 +923,12 @@ async def test_analyze_text_error_dict_handling():
@pytest.mark.asyncio
async def test_tool_calling_complete_scenario(presidio_guardrail, mock_user_api_key, mock_cache):
async def test_tool_calling_complete_scenario(
presidio_guardrail, mock_user_api_key, mock_cache
):
"""
Test complete tool calling scenario with PII in user message.
This tests the real-world scenario where:
1. User provides a query with PII
2. Assistant responds with empty content + tool_calls
@ -1002,7 +1023,12 @@ def test_no_thresholds_returns_all():
guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True)
analyze_results = [
{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.1, "start": 0, "end": 4},
{"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.2, "start": 5, "end": 9},
{
"entity_type": PiiEntityType.EMAIL_ADDRESS,
"score": 0.2,
"start": 5,
"end": 9,
},
]
filtered = guardrail.filter_analyze_results_by_score(analyze_results)
@ -1019,7 +1045,12 @@ def test_entity_specific_threshold_only_applies_to_that_entity():
)
analyze_results = [
{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4},
{"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.1, "start": 5, "end": 9},
{
"entity_type": PiiEntityType.EMAIL_ADDRESS,
"score": 0.1,
"start": 5,
"end": 9,
},
]
filtered = guardrail.filter_analyze_results_by_score(analyze_results)
@ -1038,7 +1069,12 @@ def test_filter_uses_default_all_threshold():
)
analyze_results = [
{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4},
{"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.8, "start": 5, "end": 9},
{
"entity_type": PiiEntityType.EMAIL_ADDRESS,
"score": 0.8,
"start": 5,
"end": 9,
},
]
filtered = guardrail.filter_analyze_results_by_score(analyze_results)
@ -1059,7 +1095,12 @@ def test_entity_specific_overrides_default_threshold():
)
analyze_results = [
{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.65, "start": 0, "end": 4},
{"entity_type": PiiEntityType.EMAIL_ADDRESS, "score": 0.75, "start": 5, "end": 9},
{
"entity_type": PiiEntityType.EMAIL_ADDRESS,
"score": 0.75,
"start": 5,
"end": 9,
},
]
filtered = guardrail.filter_analyze_results_by_score(analyze_results)
@ -1134,3 +1175,59 @@ def test_update_in_memory_applies_score_thresholds():
guardrail.update_in_memory_litellm_params(params)
assert guardrail.presidio_score_thresholds == {PiiEntityType.CREDIT_CARD: 0.85}
@pytest.mark.asyncio
async def test_get_session_iterator_thread_safety(presidio_guardrail):
"""
Test that _get_session_iterator yields:
1. The shared session when in the main thread.
2. A new session when in a background thread.
"""
import threading
import aiohttp
# 1. Main Thread Case
# We are in the "main thread" relative to the guardrail initialization
async with presidio_guardrail._get_session_iterator() as session:
assert isinstance(session, aiohttp.ClientSession)
assert session is presidio_guardrail._http_session
shared_session_id = id(session)
# 2. Background Thread Case
# Define a helper function to run in a thread
def thread_target(loop, result_future):
async def run_in_loop():
# This runs in the thread's loop
async with presidio_guardrail._get_session_iterator() as session:
return session, id(session)
try:
# Create a new loop for this thread to run async code
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
session_obj, session_id = new_loop.run_until_complete(run_in_loop())
result_future.set_result((session_obj, session_id))
new_loop.close()
except Exception as e:
result_future.set_exception(e)
# Run the background thread test
bg_future = asyncio.Future()
t = threading.Thread(
target=thread_target, args=(asyncio.get_running_loop(), bg_future)
)
t.start()
t.join()
bg_session, bg_session_id = await bg_future
# Assertions
# The background session should be DIFFERENT from the shared session
assert bg_session_id != shared_session_id
# The shared session should still be open (not closed by the background thread)
assert not presidio_guardrail._http_session.closed
# The background session should be closed (handled by the context manager in the thread)
assert bg_session.closed
print("✓ Session iterator thread safety test passed")