Skip to content

Commit

Permalink
Fix issue #6273: [Feature]: Disable LitLLM Print Message (#6274)
Browse files Browse the repository at this point in the history
Co-authored-by: Xingyao Wang <[email protected]>
Co-authored-by: Xingyao Wang <[email protected]>
  • Loading branch information
3 people authored Jan 16, 2025
1 parent 0c961bf commit 0661c69
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 0 deletions.
21 changes: 21 additions & 0 deletions openhands/core/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ reportlab = "*"
[tool.coverage.run]
concurrency = ["gevent"]


[tool.poetry.group.runtime.dependencies]
jupyterlab = "*"
notebook = "*"
Expand Down Expand Up @@ -129,6 +130,7 @@ ignore = ["D1"]
[tool.ruff.lint.pydocstyle]
convention = "google"


[tool.poetry.group.evaluation.dependencies]
streamlit = "*"
whatthepatch = "*"
Expand Down
55 changes: 55 additions & 0 deletions tests/unit/test_logger_litellm.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 0661c69

Please sign in to comment.