From 0661c69bd3b287e124c6f22ca6e69064f5b97561 Mon Sep 17 00:00:00 2001
From: OpenHands <opendevin@all-hands.dev>
Date: Fri, 17 Jan 2025 03:43:55 +0900
Subject: [PATCH] Fix issue #6273: [Feature]: Disable LitLLM Print Message
 (#6274)

Co-authored-by: Xingyao Wang <xingyao6@illinois.edu>
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>
---
 openhands/core/logger.py          | 21 ++++++++++++
 pyproject.toml                    |  2 ++
 tests/unit/test_logger_litellm.py | 55 +++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+)
 create mode 100644 tests/unit/test_logger_litellm.py

diff --git a/openhands/core/logger.py b/openhands/core/logger.py
index d6bd0ffd9253..7b7fd89a97a5 100644
--- a/openhands/core/logger.py
+++ b/openhands/core/logger.py
@@ -8,10 +8,31 @@
 from types import TracebackType
 from typing import Any, Literal, Mapping
 
+import litellm
 from termcolor import colored
 
 LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO').upper()
 DEBUG = os.getenv('DEBUG', 'False').lower() in ['true', '1', 'yes']
+DEBUG_LLM = os.getenv('DEBUG_LLM', 'False').lower() in ['true', '1', 'yes']
+
+# Configure litellm logging based on DEBUG_LLM
+if DEBUG_LLM:
+    confirmation = input(
+        '\n⚠️ WARNING: You are enabling DEBUG_LLM which may expose sensitive information like API keys.\n'
+        'This should NEVER be enabled in production.\n'
+        "Type 'y' to confirm you understand the risks: "
+    )
+    if confirmation.lower() == 'y':
+        litellm.suppress_debug_info = False
+        litellm.set_verbose = True
+    else:
+        print('DEBUG_LLM disabled due to lack of confirmation')
+        litellm.suppress_debug_info = True
+        litellm.set_verbose = False
+else:
+    litellm.suppress_debug_info = True
+    litellm.set_verbose = False
+
 if DEBUG:
     LOG_LEVEL = 'DEBUG'
 
diff --git a/pyproject.toml b/pyproject.toml
index fb538d5aa3cc..8375b7b7b47d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -101,6 +101,7 @@ reportlab = "*"
 [tool.coverage.run]
 concurrency = ["gevent"]
 
+
 [tool.poetry.group.runtime.dependencies]
 jupyterlab = "*"
 notebook = "*"
@@ -129,6 +130,7 @@ ignore = ["D1"]
 [tool.ruff.lint.pydocstyle]
 convention = "google"
 
+
 [tool.poetry.group.evaluation.dependencies]
 streamlit = "*"
 whatthepatch = "*"
diff --git a/tests/unit/test_logger_litellm.py b/tests/unit/test_logger_litellm.py
new file mode 100644
index 000000000000..916f7a058351
--- /dev/null
+++ b/tests/unit/test_logger_litellm.py
@@ -0,0 +1,55 @@
+import importlib
+import os
+import sys
+from unittest import mock
+
+import litellm
+import pytest
+
+
+@pytest.fixture
+def reset_litellm():
+    """Reset litellm settings and logger module after each test."""
+    yield
+    litellm.suppress_debug_info = False
+    litellm.set_verbose = False
+    # Remove logger module from sys.modules to force reload
+    if 'openhands.core.logger' in sys.modules:
+        del sys.modules['openhands.core.logger']
+
+
+def test_litellm_settings_debug_llm_disabled(reset_litellm):
+    """Test that litellm settings are properly configured when DEBUG_LLM is disabled."""
+    with mock.patch.dict(os.environ, {'DEBUG_LLM': 'false'}):
+        import openhands.core.logger  # noqa: F401
+
+        importlib.reload(openhands.core.logger)
+
+        assert litellm.suppress_debug_info is True
+        assert litellm.set_verbose is False
+
+
+def test_litellm_settings_debug_llm_enabled(reset_litellm):
+    """Test that litellm settings are properly configured when DEBUG_LLM is enabled and confirmed."""
+    with mock.patch.dict(os.environ, {'DEBUG_LLM': 'true'}), mock.patch(
+        'builtins.input', return_value='y'
+    ):
+        import openhands.core.logger  # noqa: F401
+
+        importlib.reload(openhands.core.logger)
+
+        assert litellm.suppress_debug_info is False
+        assert litellm.set_verbose is True
+
+
+def test_litellm_settings_debug_llm_enabled_but_declined(reset_litellm):
+    """Test that litellm settings remain disabled when DEBUG_LLM is enabled but user declines."""
+    with mock.patch.dict(os.environ, {'DEBUG_LLM': 'true'}), mock.patch(
+        'builtins.input', return_value='n'
+    ):
+        import openhands.core.logger  # noqa: F401
+
+        importlib.reload(openhands.core.logger)
+
+        assert litellm.suppress_debug_info is True
+        assert litellm.set_verbose is False