move import litellm to header

Author: yaozheng-fang

veadk/agent.py

@@ -17,10 +17,18 @@
 import os
 from typing import Optional, Union
 
+# If user didn't set LITELLM_LOCAL_MODEL_COST_MAP, set it to True
+# to enable local model cost map.
+# This value is `false` by default, which brings heavy performance burden,
+# for instance, importing `Litellm` needs about 10s latency.
+if not os.getenv("LITELLM_LOCAL_MODEL_COST_MAP"):
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+
 from google.adk.agents import LlmAgent, RunConfig
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.llm_agent import InstructionProvider, ToolUnion
 from google.adk.agents.run_config import StreamingMode
+from google.adk.models.lite_llm import LiteLlm
 from google.adk.runners import Runner
 from google.genai import types
 from pydantic import ConfigDict, Field
@@ -141,14 +149,6 @@ def model_post_init(self, __context: Any) -> None:
         logger.info(f"Model extra config: {self.model_extra_config}")
 
         if not self.model:
-            # If user didn't set LITELLM_LOCAL_MODEL_COST_MAP, set it to True
-            # to enable local model cost map.
-            # This value is `false` by default, which brings heavy performance burden,
-            # for instance, about 10s latency.
-            if not os.getenv("LITELLM_LOCAL_MODEL_COST_MAP"):
-                os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-            from google.adk.models.lite_llm import LiteLlm
-
             self.model = LiteLlm(
                 model=f"{self.model_provider}/{self.model_name}",
                 api_key=self.model_api_key,