diff --git a/ai21/models/assistant/assistant.py b/ai21/models/assistant/assistant.py
index 44c3ed06..cdb39f36 100644
--- a/ai21/models/assistant/assistant.py
+++ b/ai21/models/assistant/assistant.py
@@ -3,10 +3,10 @@
 from typing_extensions import TypedDict
 
 Optimization = Literal["cost", "latency"]
-Tool = Literal["rag", "internet_research", "plan_approval"]
+Tool = Literal["rag", "web_search", "plan_approval"]
 
 
 class ToolResources(TypedDict, total=False):
     rag: Optional[dict]
-    internet_research: Optional[dict]
+    web_search: Optional[dict]
     plan_approval: Optional[dict]