diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 4aeb9d4d64..7423e55b62 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -228,9 +228,11 @@ def safe_deep_copy(data): """ Safe Deep Copy - The LiteLLM Request has some object that can-not be pickled / deep copied - - Use this function to safely deep copy the LiteLLM Request + The LiteLLM request may contain objects that cannot be pickled/deep-copied + (e.g., tracing spans, locks, clients). + + This helper deep-copies each top-level key independently; on failure keeps + original ref """ import copy @@ -255,9 +257,22 @@ def safe_deep_copy(data): "litellm_parent_otel_span" ) data["litellm_metadata"]["litellm_parent_otel_span"] = "placeholder" - new_data = copy.deepcopy(data) - # Step 2: re-add the litellm_parent_otel_span after doing a deep copy + # Step 2: Per-key deepcopy with fallback + if isinstance(data, dict): + new_data = {} + for k, v in data.items(): + try: + new_data[k] = copy.deepcopy(v) + except Exception: + new_data[k] = v + else: + try: + new_data = copy.deepcopy(data) + except Exception: + new_data = data + + # Step 3: re-add the litellm_parent_otel_span after doing a deep copy if isinstance(data, dict) and litellm_parent_otel_span is not None: if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]: data["metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span @@ -268,4 +283,4 @@ def safe_deep_copy(data): data["litellm_metadata"][ "litellm_parent_otel_span" ] = litellm_parent_otel_span - return new_data + return new_data \ No newline at end of file diff --git a/litellm/main.py b/litellm/main.py index 786a0196e5..e7130caa49 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -116,6 +116,7 @@ from litellm.utils import ( from ._logging import verbose_logger from .caching.caching import disable_cache, enable_cache, update_cache +from .litellm_core_utils.core_helpers import safe_deep_copy from .litellm_core_utils.fallback_utils import ( async_completion_with_fallbacks, completion_with_fallbacks, @@ -2772,8 +2773,7 @@ def completion( # type: ignore # noqa: PLR0915 ) api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE") - - new_params = deepcopy(optional_params) + new_params = safe_deep_copy(optional_params or {}) response = vertex_chat_completion.completion( # type: ignore model=model, messages=messages, @@ -2817,7 +2817,7 @@ def completion( # type: ignore # noqa: PLR0915 api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") - new_params = deepcopy(optional_params) + new_params = safe_deep_copy(optional_params or {}) if vertex_partner_models_chat_completion.is_vertex_partner_model(model): model_response = vertex_partner_models_chat_completion.completion( model=model, diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 32f3ad3f55..89cc11c40d 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -9,7 +9,11 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide +from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + safe_divide, + safe_deep_copy +) def test_get_litellm_metadata_from_kwargs(): @@ -127,3 +131,43 @@ def test_safe_divide_weight_scenario(): expected_zero = [0, 0, 0] assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}" + + +def test_safe_deep_copy_with_non_pickleables_and_span(): + """ + Verify safe_deep_copy: + - does not crash when non-pickleables are present, + - preserves structure/keys, + - deep-copies JSON-y payloads (e.g., messages), + - keeps non-pickleables by reference, + - redacts OTEL span in the copy and restores it in the original. + """ + import threading + rlock = threading.RLock() + data = { + "metadata": {"litellm_parent_otel_span": rlock, "x": 1}, + "messages": [{"role": "user", "content": "hi"}], + "optional_params": {"handle": rlock}, + "ok": True, + } + + copied = safe_deep_copy(data) + + # Structure preserved + assert set(copied.keys()) == set(data.keys()) + + # Messages are deep-copied (new object, same content) + assert copied["messages"] is not data["messages"] + assert copied["messages"][0] == data["messages"][0] + + # Non-pickleable subtree kept by reference (no crash) + assert copied["optional_params"] is data["optional_params"] + assert copied["optional_params"]["handle"] is rlock + + # OTEL span: redacted in the copy, restored in original + assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder" + assert data["metadata"]["litellm_parent_otel_span"] is rlock + + # Other simple fields unchanged + assert copied["ok"] is True + assert copied["metadata"]["x"] == 1