Use env credentials instead of passing them through

Author: bvandermoon

src/maxtext/configs/pyconfig.py

@@ -324,6 +324,11 @@ def initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConfi
 
   # 2. Get overrides from CLI and kwargs
   cli_cfg = omegaconf.OmegaConf.from_cli(cli_args)
+  if "hf_access_token" in cli_cfg:
+    logger.warning(
+        "WARNING: Passing 'hf_access_token' via command-line arguments is deprecated and insecure because it makes "
+        "your token visible in 'ps' and shell history. Please set the 'HF_TOKEN' environment variable instead."
+    )
   kwargs_cfg = omegaconf.OmegaConf.create(kwargs)
   overrides_cfg = omegaconf.OmegaConf.merge(cli_cfg, kwargs_cfg)
 

src/maxtext/experimental/agent/ckpt_conversion_agent/README.md

@@ -4,8 +4,11 @@ The agent is used to automate the model-specific mappings of checkpoint conversi
 ## Quick starts
 To begin, you'll need:
 
-1. A Google atccount.
-2. An API key (create one in [Google AI Studio](https://aistudio.google.com/app/apikey)).
+1. A Google account.
+2. An API key (create one in [Google AI Studio](https://aistudio.google.com/app/apikey)), and set it as an environment variable:
+```bash
+export GEMINI_API_KEY="<Your-API-KEY>"
+```
 3. Install the Google Generative AI Python library:
 ```
 pip install -q -U "google-genai>=1.0.0"
@@ -30,7 +33,7 @@ After it, you can get two `*.json` files in `config.base_output_directory` folde
 
 ```bash
 python3 -m maxtext.experimental.agent.ckpt_conversion_agent.step1 --target_model=<model_name> \
-  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent --api_key=<Your-API-KEY>
+  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent
 ```
 
 Our engineer should check the `src/maxtext/experimental/agent/ckpt_conversion_agent/outputs/proposed_dsl.txt` for potential new DSL and assess if it's needed. Then we need to add this ops into `src/maxtext/experimental/agent/ckpt_conversion_agent/context/dsl.txt`.
@@ -39,7 +42,7 @@ Our engineer should check the `src/maxtext/experimental/agent/ckpt_conversion_ag
 
 ```bash
 python3 -m maxtext.experimental.agent.ckpt_conversion_agent.step2 --target_model=<model_name> \
-  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent --api_key=<Your-API-KEY>
+  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent
 ```
 
 ## Evaluation and Debugging
@@ -53,7 +56,7 @@ You can automatically verify the output by comparing the generated code against
 
 ```bash
 python3 -m maxtext.experimental.agent.ckpt_conversion_agent.evaluation --files ground_truth/<model>.py \
-  outputs/hook_fn.py --api_key=<Your-API-KEY> --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent
+  outputs/hook_fn.py --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent
 ```
 
 ### Manual Debugging (No Ground-Truth Code)
@@ -121,5 +124,5 @@ Run the [One-shot agent Jyputer notebook](./baselines/one-shot-agent.ipynb)
 ### Prompt-chain Agent:
 ```bash
 python3 -m maxtext.experimental.agent.ckpt_conversion_agent.prompt_chain --target_model=<model_name> \
-  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent --api_key=<Your-API-KEY>
+  --dir_path=src/maxtext/experimental/agent/ckpt_conversion_agent
 ```

src/maxtext/experimental/agent/ckpt_conversion_agent/analysis.py

@@ -27,15 +27,15 @@ class AnalysisAgent(BaseAgent):
   conversion script, with verification that every parameter is mapped.
   """
 
-  def __init__(self, api_key, dir_path, target_model="gemma3", max_retries=3):
+  def __init__(self, dir_path, target_model="gemma3", max_retries=3):
     """
     Initializes the PlanAgent.
 
     Args:
         target_model (str): The target model for conversion.
         max_retries (int): The maximum number of retries for generation.
     """
-    super().__init__(api_key)
+    super().__init__()
 
     self.target_model = target_model
     self.max_retries = max_retries

src/maxtext/experimental/agent/ckpt_conversion_agent/base.py

@@ -16,6 +16,8 @@
 A base class for other agents
 """
 
+import os
+
 from google import genai
 from google.genai.types import GenerateContentConfig
 
@@ -25,16 +27,14 @@
 class BaseAgent:
   """A base class for agents that provides text generation capabilities."""
 
-  def __init__(self, api_key, model_id=MODEL_ID):
+  def __init__(self, model_id=MODEL_ID):
     """
     Initializes the BaseAgent with a genai client.
-
-    Args:
-        client: An initialized genai.Client object.
     """
-    if not api_key:
-      raise ValueError("A valid api_key must be provided.")
-    client = genai.Client(api_key=api_key)
+    resolved_api_key = os.environ.get("GEMINI_API_KEY")
+    if not resolved_api_key:
+      raise ValueError("GEMINI_API_KEY environment variable is not set.")
+    client = genai.Client(api_key=resolved_api_key)
     self.client = client
     self.model_id = model_id
 

src/maxtext/experimental/agent/ckpt_conversion_agent/dsl.py

@@ -24,7 +24,7 @@
 class DSLAgent(BaseAgent):
   """DSL Agent"""
 
-  def __init__(self, api_key, dir_path, target_model="gemma3-4b", max_retries=3):
+  def __init__(self, dir_path, target_model="gemma3-4b", max_retries=3):
     """
     Initializes the DSLAgent.
 
@@ -33,7 +33,7 @@ def __init__(self, api_key, dir_path, target_model="gemma3-4b", max_retries=3):
         max_retries (int): The maximum number of retries for generation.
     """
     # Initialize the parent BaseAgent with the client
-    super().__init__(api_key)
+    super().__init__()
 
     self.target_model = target_model
     self.max_retries = max_retries
@@ -83,7 +83,6 @@ def verify_dsl(self):
   parser.add_argument(
       "--dir_path", type=str, required=True, help='The file path to the context directory (e.g., "context/gemma3").'
   )
-  parser.add_argument("--api_key", type=str, help="Optional API key for external services.")
   args = parser.parse_args()
-  agent = DSLAgent(api_key=args.api_key, dir_path=args.dir_path, target_model=TARGET_MODEL)
+  agent = DSLAgent(dir_path=args.dir_path, target_model=TARGET_MODEL)
   global_verification_dsl = agent.verify_dsl()

src/maxtext/experimental/agent/ckpt_conversion_agent/evaluation.py

@@ -66,12 +66,11 @@ def main():
       description="Gemini evaluate the agent code implementation against human-written ground truth code"
   )
   parser.add_argument("--files", nargs=2, help="Paths to code files to analyze.")
-  parser.add_argument("--api_key", type=str, help="API key.")
   parser.add_argument("--dir_path", type=str, help="Directory path.")
 
   args = parser.parse_args()
 
-  baseAgent = BaseAgent(api_key=args.api_key)
+  baseAgent = BaseAgent()
   dir_path = args.dir_path
 
   prompt_templates = {

src/maxtext/experimental/agent/ckpt_conversion_agent/mapping.py

@@ -28,15 +28,15 @@ class MappingAgent(BaseAgent):
   An agent that generates and verifies mapping functions for model conversion.
   """
 
-  def __init__(self, api_key, dir_path, target_model="gemma3-4b", max_retries=3):
+  def __init__(self, dir_path, target_model="gemma3-4b", max_retries=3):
     """
     Initializes the MappingAgent.
 
     Args:
         target_model (str): The target model for conversion.
         max_retries (int): The maximum number of retries for generation.
     """
-    super().__init__(api_key)
+    super().__init__()
 
     self.target_model = target_model
     self.max_retries = max_retries
@@ -171,9 +171,8 @@ def generate_shape_mapping(self):
   parser.add_argument(
       "--dir_path", type=str, required=True, help='The file path to the context directory (e.g., "context/gemma3").'
   )
-  parser.add_argument("--api_key", type=str, help="Optional API key for external services.")
   args = parser.parse_args()
-  agent = MappingAgent(api_key=args.api_key, dir_path=args.dir_path, target_model=TARGET_MODEL)
+  agent = MappingAgent(dir_path=args.dir_path, target_model=TARGET_MODEL)
   try:
     param_mapping_code = agent.generate_param_mapping()
     shape_mapping_code = agent.generate_shape_mapping()

src/maxtext/experimental/agent/ckpt_conversion_agent/plan.py

@@ -28,15 +28,15 @@ class PlanAgent(BaseAgent):
   conversion script, with verification that every parameter is mapped.
   """
 
-  def __init__(self, api_key, dir_path, target_model="gemma3", max_retries=3):
+  def __init__(self, dir_path, target_model="gemma3", max_retries=3):
     """
     Initializes the PlanAgent.
 
     Args:
         target_model (str): The target model for conversion.
         max_retries (int): The maximum number of retries for generation.
     """
-    super().__init__(api_key)
+    super().__init__()
 
     self.target_model = target_model
     self.max_retries = max_retries
@@ -109,7 +109,6 @@ def plan_conversion(self):
   parser.add_argument(
       "--dir_path", type=str, required=True, help='The file path to the context directory (e.g., "context/gemma3").'
   )
-  parser.add_argument("--api_key", type=str, help="Optional API key for external services.")
   args = parser.parse_args()
-  agent = PlanAgent(api_key=args.api_key, dir_path=args.dir_path, target_model=TARGET_MODEL)
+  agent = PlanAgent(dir_path=args.dir_path, target_model=TARGET_MODEL)
   agent.plan_conversion()

src/maxtext/experimental/agent/ckpt_conversion_agent/prompt_chain.py

@@ -29,9 +29,9 @@ class prompt_chaining_agent(BaseAgent):
   with verification that every parameter is actually mapped.
   """
 
-  def __init__(self, api_key, target_model="gemma3", max_retries=3, dir_path="context"):
+  def __init__(self, target_model="gemma3", max_retries=3, dir_path="context"):
     # Initialize the parent BaseAgent with the client
-    super().__init__(api_key)
+    super().__init__()
     self.target_model = target_model
     self.max_retries = max_retries
     self.dir_path = dir_path
@@ -178,7 +178,6 @@ def run_chain(self, max_retries=3):
 
   parser.add_argument("--target_model", type=str, required=True, help='The name of the target model (e.g., "GEMMA3").')
   parser.add_argument("--dir_path", type=str, required=True, help="The file path to the ckpt conversion agent directory.")
-  parser.add_argument("--api_key", type=str, required=True, help="Gemini API key.")
   args = parser.parse_args()
-  agent = prompt_chaining_agent(api_key=args.api_key, target_model=args.target_model, dir_path=args.dir_path)
+  agent = prompt_chaining_agent(target_model=args.target_model, dir_path=args.dir_path)
   agent.run_chain()

src/maxtext/experimental/agent/ckpt_conversion_agent/step1.py

@@ -28,17 +28,15 @@
   parser.add_argument(
       "--dir_path", type=str, required=True, help='The file path to the context directory (e.g., "context/gemma3").'
   )
-  parser.add_argument("--api_key", type=str, help="Optional API key for external services.")
   args = parser.parse_args()
 
   TARGET_MODEL = args.target_model
   dir_path = args.dir_path
-  api_key = args.api_key
 
-  analysisAgent = AnalysisAgent(api_key=api_key, dir_path=dir_path, target_model=TARGET_MODEL)
+  analysisAgent = AnalysisAgent(dir_path=dir_path, target_model=TARGET_MODEL)
   analysisAgent.analyze_model_structures()
 
-  dslAgent = DSLAgent(api_key=api_key, dir_path=dir_path, target_model=TARGET_MODEL)
+  dslAgent = DSLAgent(dir_path=dir_path, target_model=TARGET_MODEL)
   dslAgent.verify_dsl()
 
   # Human interaction needed,