fix: Set LITELLM_MODE to PRODUCTION before importing LiteLLM

LiteLLM defaults to DEV mode, which automatically loads environment variables from a local `.env` file. This change sets LITELLM_MODE to PRODUCTION to prevent LiteLLM from implicitly loading `.env` files when used within ADK.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 858723362

Author: GWeale

src/google/adk/models/lite_llm.py

@@ -16,6 +16,7 @@
 
 import base64
 import copy
+import importlib.util
 import json
 import logging
 import mimetypes
@@ -32,27 +33,20 @@
 from typing import Literal
 from typing import Optional
 from typing import Tuple
+from typing import TYPE_CHECKING
 from typing import TypedDict
 from typing import Union
 from urllib.parse import urlparse
 import uuid
 import warnings
 
 from google.genai import types
-import litellm
-from litellm import acompletion
-from litellm import ChatCompletionAssistantMessage
-from litellm import ChatCompletionAssistantToolCall
-from litellm import ChatCompletionMessageToolCall
-from litellm import ChatCompletionSystemMessage
-from litellm import ChatCompletionToolMessage
-from litellm import ChatCompletionUserMessage
-from litellm import completion
-from litellm import CustomStreamWrapper
-from litellm import Function
-from litellm import Message
-from litellm import ModelResponse
-from litellm import OpenAIMessageContent
+
+if not TYPE_CHECKING and importlib.util.find_spec("litellm") is None:
+  raise ImportError(
+      "LiteLLM support requires: pip install google-adk[extensions]"
+  )
+
 from pydantic import BaseModel
 from pydantic import Field
 from typing_extensions import override
@@ -61,8 +55,36 @@
 from .llm_request import LlmRequest
 from .llm_response import LlmResponse
 
-# This will add functions to prompts if functions are provided.
-litellm.add_function_to_prompt = True
+if TYPE_CHECKING:
+  import litellm
+  from litellm import acompletion
+  from litellm import ChatCompletionAssistantMessage
+  from litellm import ChatCompletionAssistantToolCall
+  from litellm import ChatCompletionMessageToolCall
+  from litellm import ChatCompletionSystemMessage
+  from litellm import ChatCompletionToolMessage
+  from litellm import ChatCompletionUserMessage
+  from litellm import completion
+  from litellm import CustomStreamWrapper
+  from litellm import Function
+  from litellm import Message
+  from litellm import ModelResponse
+  from litellm import OpenAIMessageContent
+else:
+  litellm = None
+  acompletion = None
+  ChatCompletionAssistantMessage = None
+  ChatCompletionAssistantToolCall = None
+  ChatCompletionMessageToolCall = None
+  ChatCompletionSystemMessage = None
+  ChatCompletionToolMessage = None
+  ChatCompletionUserMessage = None
+  completion = None
+  CustomStreamWrapper = None
+  Function = None
+  Message = None
+  ModelResponse = None
+  OpenAIMessageContent = None
 
 logger = logging.getLogger("google_adk." + __name__)
 
@@ -109,6 +131,50 @@
     "before a response was recorded)."
 )
 
+_LITELLM_IMPORTED = False
+_LITELLM_GLOBAL_SYMBOLS = (
+    "ChatCompletionAssistantMessage",
+    "ChatCompletionAssistantToolCall",
+    "ChatCompletionMessageToolCall",
+    "ChatCompletionSystemMessage",
+    "ChatCompletionToolMessage",
+    "ChatCompletionUserMessage",
+    "CustomStreamWrapper",
+    "Function",
+    "Message",
+    "ModelResponse",
+    "OpenAIMessageContent",
+    "acompletion",
+    "completion",
+)
+
+
+def _ensure_litellm_imported() -> None:
+  """Imports LiteLLM with safe defaults.
+
+  LiteLLM defaults to DEV mode, which auto-loads a local `.env` at import time.
+  ADK should not implicitly load `.env` just because LiteLLM is installed.
+
+  Users can opt into LiteLLM's default behavior by setting LITELLM_MODE=DEV.
+  """
+  global _LITELLM_IMPORTED
+  if _LITELLM_IMPORTED:
+    return
+
+  # https://github.com/BerriAI/litellm/blob/main/litellm/__init__.py#L80-L82
+  os.environ.setdefault("LITELLM_MODE", "PRODUCTION")
+
+  import litellm as litellm_module
+
+  litellm_module.add_function_to_prompt = True
+
+  globals()["litellm"] = litellm_module
+  for symbol in _LITELLM_GLOBAL_SYMBOLS:
+    globals()[symbol] = getattr(litellm_module, symbol)
+
+  _redirect_litellm_loggers_to_stdout()
+  _LITELLM_IMPORTED = True
+
 
 def _map_finish_reason(
     finish_reason: Any,
@@ -344,6 +410,7 @@ async def acompletion(
     Returns:
       The model response as a message.
     """
+    _ensure_litellm_imported()
 
     return await acompletion(
         model=model,
@@ -367,6 +434,7 @@ def completion(
     Returns:
       The response from the model.
     """
+    _ensure_litellm_imported()
 
     return completion(
         model=model,
@@ -513,6 +581,7 @@ async def _content_to_message_param(
   Returns:
     A litellm Message, a list of litellm Messages.
   """
+  _ensure_litellm_imported()
 
   tool_messages: list[Message] = []
   non_tool_parts: list[types.Part] = []
@@ -622,6 +691,8 @@ def _ensure_tool_results(messages: List[Message]) -> List[Message]:
   if not messages:
     return messages
 
+  _ensure_litellm_imported()
+
   healed_messages: List[Message] = []
   pending_tool_call_ids: List[str] = []
 
@@ -691,6 +762,7 @@ async def _get_content(
   Returns:
     The litellm content.
   """
+  _ensure_litellm_imported()
 
   parts_list = list(parts)
   if len(parts_list) == 1:
@@ -925,6 +997,7 @@ def _build_tool_call_from_json_dict(
     candidate: Any, *, index: int
 ) -> Optional[ChatCompletionMessageToolCall]:
   """Creates a tool call object from JSON content embedded in text."""
+  _ensure_litellm_imported()
 
   if not isinstance(candidate, dict):
     return None
@@ -972,11 +1045,12 @@ def _parse_tool_calls_from_text(
     text_block: str,
 ) -> tuple[list[ChatCompletionMessageToolCall], Optional[str]]:
   """Extracts inline JSON tool calls from LiteLLM text responses."""
-
   tool_calls = []
   if not text_block:
     return tool_calls, None
 
+  _ensure_litellm_imported()
+
   remainder_segments = []
   cursor = 0
   text_length = len(text_block)
@@ -1014,7 +1088,6 @@ def _split_message_content_and_tool_calls(
     message: Message,
 ) -> tuple[Optional[OpenAIMessageContent], list[ChatCompletionMessageToolCall]]:
   """Returns message content and tool calls, parsing inline JSON when needed."""
-
   existing_tool_calls = message.get("tool_calls") or []
   normalized_tool_calls = (
       list(existing_tool_calls) if existing_tool_calls else []
@@ -1180,6 +1253,7 @@ def _model_response_to_chunk(
   Yields:
     A tuple of text or function or usage metadata chunk and finish reason.
   """
+  _ensure_litellm_imported()
 
   message = None
   if response.get("choices", None):
@@ -1255,6 +1329,7 @@ def _model_response_to_generate_content_response(
   Returns:
     The LlmResponse.
   """
+  _ensure_litellm_imported()
 
   message = None
   finish_reason = None
@@ -1313,6 +1388,7 @@ def _message_to_generate_content_response(
   Returns:
     The LlmResponse.
   """
+  _ensure_litellm_imported()
 
   parts: List[types.Part] = []
   if not thought_parts:
@@ -1440,6 +1516,8 @@ async def _get_completion_inputs(
     The litellm inputs (message list, tool dictionary, response format and
     generation params).
   """
+  _ensure_litellm_imported()
+
   # Determine provider for file handling
   provider = _get_provider_from_model(model)
 
@@ -1665,11 +1743,6 @@ def _redirect_litellm_loggers_to_stdout() -> None:
         handler.stream = sys.stdout
 
 
-# Redirect LiteLLM loggers to stdout immediately after import to ensure
-# INFO-level logs are not incorrectly treated as errors in cloud environments.
-_redirect_litellm_loggers_to_stdout()
-
-
 class LiteLlm(BaseLlm):
   """Wrapper around litellm.
 
@@ -1732,6 +1805,7 @@ async def generate_content_async(
     Yields:
       LlmResponse: The model response.
     """
+    _ensure_litellm_imported()
 
     self._maybe_append_user_content(llm_request)
     _append_fallback_user_content_if_missing(llm_request)

tests/unittests/models/test_litellm_import.py

@@ -0,0 +1,108 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import os
+import subprocess
+import sys
+
+import pytest
+
+
+def _subprocess_env() -> dict[str, str]:
+  env = dict(os.environ)
+  src_path = os.path.join(os.getcwd(), "src")
+  pythonpath = env.get("PYTHONPATH", "")
+  env["PYTHONPATH"] = (
+      f"{src_path}{os.pathsep}{pythonpath}" if pythonpath else src_path
+  )
+  return env
+
+
+def test_importing_models_does_not_import_litellm_or_set_mode():
+  env = _subprocess_env()
+  env.pop("LITELLM_MODE", None)
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os, sys\n"
+              "import google.adk.models\n"
+              "print('litellm' in sys.modules)\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  stdout_lines = result.stdout.strip().splitlines()
+  assert stdout_lines == ["False", "None"]
+
+
+def test_ensure_litellm_imported_defaults_to_production():
+  if importlib.util.find_spec("litellm") is None:
+    pytest.skip("litellm is not installed")
+
+  env = _subprocess_env()
+  env.pop("LITELLM_MODE", None)
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os\n"
+              "from google.adk.models.lite_llm import"
+              " _ensure_litellm_imported\n"
+              "_ensure_litellm_imported()\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  assert result.stdout.strip() == "PRODUCTION"
+
+
+def test_ensure_litellm_imported_does_not_override():
+  if importlib.util.find_spec("litellm") is None:
+    pytest.skip("litellm is not installed")
+
+  env = _subprocess_env()
+  env["LITELLM_MODE"] = "DEV"
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os\n"
+              "from google.adk.models.lite_llm import"
+              " _ensure_litellm_imported\n"
+              "_ensure_litellm_imported()\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  assert result.stdout.strip() == "DEV"