diff --git a/src/agentlab/agents/generic_agent/agent_configs.py b/src/agentlab/agents/generic_agent/agent_configs.py
index 2d4f9bc39..a5db8c906 100644
--- a/src/agentlab/agents/generic_agent/agent_configs.py
+++ b/src/agentlab/agents/generic_agent/agent_configs.py
@@ -1,4 +1,5 @@
 import bgym
+
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.experiments import args
 from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
@@ -41,7 +42,7 @@
     use_abstract_example=True,
     use_hints=True,
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
 )
@@ -89,7 +90,7 @@
     use_abstract_example=True,  # useful
     use_hints=True,  # useful
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
 )
@@ -136,7 +137,7 @@
     use_abstract_example=True,
     use_hints=True,
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
     add_missparsed_messages=True,
@@ -186,7 +187,7 @@
     use_abstract_example=True,
     use_hints=True,
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
     add_missparsed_messages=True,
@@ -234,7 +235,7 @@
     use_abstract_example=True,
     use_hints=True,
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
 )
@@ -298,7 +299,7 @@
     use_hints=args.Choice([True, False], p=[0.7, 0.3]),
     be_cautious=args.Choice([True, False]),
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     extra_instructions=None,
 )
 
diff --git a/src/agentlab/agents/generic_agent/tmlr_config.py b/src/agentlab/agents/generic_agent/tmlr_config.py
index 11860e691..48a28c682 100644
--- a/src/agentlab/agents/generic_agent/tmlr_config.py
+++ b/src/agentlab/agents/generic_agent/tmlr_config.py
@@ -40,7 +40,7 @@
     use_abstract_example=True,
     use_hints=True,
     enable_chat=False,
-    max_prompt_tokens=None,
+    max_prompt_tokens=40_000,
     be_cautious=True,
     extra_instructions=None,
 )
diff --git a/src/agentlab/llm/llm_configs.py b/src/agentlab/llm/llm_configs.py
index 8376b5c25..feb4d1c84 100644
--- a/src/agentlab/llm/llm_configs.py
+++ b/src/agentlab/llm/llm_configs.py
@@ -20,28 +20,28 @@
     "openai/gpt-4o-mini-2024-07-18": OpenAIModelArgs(
         model_name="gpt-4o-mini-2024-07-18",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         vision_support=True,
     ),
     "openai/gpt-4-1106-preview": OpenAIModelArgs(
         model_name="gpt-4-1106-preview",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,  # make sure we don't bust budget
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
     ),
     "openai/gpt-4-vision-preview": OpenAIModelArgs(
         model_name="gpt-4-vision-preview",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,  # make sure we don't bust budget
-        max_new_tokens=4000,  # I think this model has very small default value if we don't set max_new_tokens
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,  # I think this model has very small default value if we don't set max_new_tokens
         vision_support=True,
     ),
     "openai/gpt-4o-2024-05-13": OpenAIModelArgs(
         model_name="gpt-4o-2024-05-13",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,  # make sure we don't bust budget
-        max_new_tokens=4000,  # I think this model has very small default value if we don't set max_new_tokens
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,  # I think this model has very small default value if we don't set max_new_tokens
         vision_support=True,
     ),
     "openai/gpt-3.5-turbo-0125": OpenAIModelArgs(
@@ -67,22 +67,25 @@
         model_name="gpt-4o",
         deployment_name="gpt-4o-2024-05-13",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4_000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
+        vision_support=True,
     ),
     "azure/gpt-4o-2024-08-06": AzureModelArgs(
         model_name="gpt-4o",
         deployment_name="gpt-4o-2024-08-06",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4_000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
+        vision_support=True,
     ),
     "azure/gpt-4o-mini-2024-07-18": AzureModelArgs(
         model_name="gpt-4o-mini",
         deployment_name="gpt-4o-mini-2024-07-18",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4_000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
+        vision_support=True,
     ),
     # ---------------- OSS LLMs ----------------#
     "meta-llama/Meta-Llama-3-70B-Instruct": SelfHostedModelArgs(
@@ -113,43 +116,43 @@
     "openrouter/meta-llama/llama-3.1-405b-instruct": OpenRouterModelArgs(
         model_name="meta-llama/llama-3.1-405b-instruct",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
     "openrouter/meta-llama/llama-3.1-70b-instruct": OpenRouterModelArgs(
         model_name="meta-llama/llama-3.1-70b-instruct",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
     "openrouter/meta-llama/llama-3-70b-instruct": OpenRouterModelArgs(
         model_name="meta-llama/llama-3-70b-instruct",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
     "openrouter/meta-llama/llama-3.1-8b-instruct:free": OpenRouterModelArgs(
         model_name="meta-llama/llama-3.1-8b-instruct:free",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
     "openrouter/meta-llama/llama-3.1-8b-instruct": OpenRouterModelArgs(
         model_name="meta-llama/llama-3.1-8b-instruct",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
     "openrouter/anthropic/claude-3.5-sonnet:beta": OpenRouterModelArgs(
         model_name="anthropic/claude-3.5-sonnet:beta",
         max_total_tokens=200_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=160_000,
+        max_new_tokens=40_000,
         temperature=1e-1,
         vision_support=True,
     ),
@@ -163,8 +166,8 @@
     "openrouter/openai/o1-mini-2024-09-12": OpenRouterModelArgs(
         model_name="openai/o1-mini-2024-09-12",
         max_total_tokens=128_000,
-        max_input_tokens=40_000,
-        max_new_tokens=4000,
+        max_input_tokens=100_000,
+        max_new_tokens=28_000,
         temperature=1e-1,
     ),
 }