diff --git a/libs/partners/ai21/poetry.lock b/libs/partners/ai21/poetry.lock
index 0766cf4238..67769e0e29 100644
--- a/libs/partners/ai21/poetry.lock
+++ b/libs/partners/ai21/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "ai21"
@@ -903,7 +903,6 @@ files = [
     {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
-    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
     {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
     {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
     {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@@ -1336,4 +1335,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.1,<4.0"
-content-hash = "8af2236a46bf24cbf18cfb15e7ef4967858e91764cc2f5bb260e59f225fe0047"
+content-hash = "3bc24f83ee4dd4fffd569b59ffb59fd85598e58aafbef01ad5ecc594baccda0d"
diff --git a/libs/partners/ai21/pyproject.toml b/libs/partners/ai21/pyproject.toml
index f9a8026f5e..baed0a4e42 100644
--- a/libs/partners/ai21/pyproject.toml
+++ b/libs/partners/ai21/pyproject.toml
@@ -14,7 +14,7 @@ license = "MIT"
 python = ">=3.8.1,<4.0"
 langchain-core = ">=0.1.48,<0.3"
 langchain-text-splitters = "^0.2.0"
-ai21 = "^2.2.5"
+ai21 = "^2.4.1"
 
 [tool.poetry.group.test]
 optional = true
diff --git a/libs/partners/ai21/tests/unit_tests/conftest.py b/libs/partners/ai21/tests/unit_tests/conftest.py
index c135181af7..9d95584daa 100644
--- a/libs/partners/ai21/tests/unit_tests/conftest.py
+++ b/libs/partners/ai21/tests/unit_tests/conftest.py
@@ -4,7 +4,7 @@ from typing import Generator
 from unittest.mock import Mock
 
 import pytest
-from ai21 import AI21Client, AI21EnvConfig
+from ai21 import AI21Client
 from ai21.models import (
     AnswerResponse,
     ChatOutput,
@@ -167,13 +167,10 @@ def temporarily_unset_api_key() -> Generator:
     """
     Unset and set environment key for testing purpose for when an API KEY is not set
     """
-    api_key = AI21EnvConfig.api_key
-    AI21EnvConfig.api_key = None
-    os.environ.pop("AI21_API_KEY", None)
+    api_key = os.environ.pop("AI21_API_KEY", None)
     yield
 
     if api_key is not None:
-        AI21EnvConfig.api_key = api_key
         os.environ["AI21_API_KEY"] = api_key