ai21: fix ai21 unittests (#22526)

Co-authored-by: Asaf Gardin <asafg@ai21.com>
Co-authored-by: Erick Friis <erick@langchain.dev>
pull/22544/head
Asaf Joseph Gardin 2 months ago committed by GitHub
parent 58192d617f
commit 75cba742e5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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"

@ -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

@ -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

Loading…
Cancel
Save