From 96641b9b131b3838287cb40ff1012c8bb2c0d6a3 Mon Sep 17 00:00:00 2001 From: hsolpark Date: Tue, 5 Dec 2023 16:15:10 +0900 Subject: [PATCH 1/3] Mask API key for ChatOpenAI based chat_models --- .../langchain/chat_models/anyscale.py | 14 ++--- .../langchain/chat_models/azure_openai.py | 9 ++-- .../langchain/chat_models/everlyai.py | 12 +++-- .../langchain/langchain/chat_models/openai.py | 19 +++---- .../chat_models/test_azureopenai.py | 51 ++++++++++++++++++- .../unit_tests/chat_models/test_openai.py | 29 +++++++++++ 6 files changed, 108 insertions(+), 26 deletions(-) diff --git a/libs/langchain/langchain/chat_models/anyscale.py b/libs/langchain/langchain/chat_models/anyscale.py index 27efc4a1abfb1..7314fe78ab5f5 100644 --- a/libs/langchain/langchain/chat_models/anyscale.py +++ b/libs/langchain/langchain/chat_models/anyscale.py @@ -102,10 +102,12 @@ def get_available_models( @root_validator(pre=True) def validate_environment_override(cls, values: dict) -> dict: """Validate that api key and python package exists in environment.""" - values["openai_api_key"] = get_from_dict_or_env( - values, - "anyscale_api_key", - "ANYSCALE_API_KEY", + values["openai_api_key"] = convert_to_secret_str( + get_from_dict_or_env( + values, + "anyscale_api_key", + "ANYSCALE_API_KEY", + ) ) values["anyscale_api_key"] = convert_to_secret_str( get_from_dict_or_env( @@ -137,7 +139,7 @@ def validate_environment_override(cls, values: dict) -> dict: try: if is_openai_v1(): client_params = { - "api_key": values["openai_api_key"], + "api_key": values["openai_api_key"].get_secret_value(), "base_url": values["openai_api_base"], # To do: future support # "organization": values["openai_organization"], @@ -163,7 +165,7 @@ def validate_environment_override(cls, values: dict) -> dict: model_name = values["model_name"] available_models = cls.get_available_models( - values["openai_api_key"], + values["openai_api_key"].get_secret_value(), values["openai_api_base"], ) diff --git a/libs/langchain/langchain/chat_models/azure_openai.py b/libs/langchain/langchain/chat_models/azure_openai.py index 2b9cb1b5f6da9..34526f92c697d 100644 --- a/libs/langchain/langchain/chat_models/azure_openai.py +++ b/libs/langchain/langchain/chat_models/azure_openai.py @@ -7,10 +7,10 @@ from typing import Any, Callable, Dict, Union from langchain_core.outputs import ChatResult -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator +from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator from langchain.chat_models.openai import ChatOpenAI -from langchain.utils import get_from_dict_or_env +from langchain.utils import convert_to_secret_str, get_from_dict_or_env from langchain.utils.openai import is_openai_v1 logger = logging.getLogger(__name__) @@ -70,7 +70,7 @@ class AzureChatOpenAI(ChatOpenAI): """ openai_api_version: str = Field(default="", alias="api_version") """Automatically inferred from env var `OPENAI_API_VERSION` if not provided.""" - openai_api_key: Union[str, None] = Field(default=None, alias="api_key") + openai_api_key: Union[SecretStr, None] = Field(default=None, alias="api_key") """Automatically inferred from env var `AZURE_OPENAI_API_KEY` if not provided.""" azure_ad_token: Union[str, None] = None """Your Azure Active Directory token. @@ -110,6 +110,7 @@ def validate_environment(cls, values: Dict) -> Dict: or os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY") ) + values["openai_api_key"] = convert_to_secret_str(values["openai_api_key"]) values["openai_api_base"] = values["openai_api_base"] or os.getenv( "OPENAI_API_BASE" ) @@ -184,7 +185,7 @@ def validate_environment(cls, values: Dict) -> Dict: "api_version": values["openai_api_version"], "azure_endpoint": values["azure_endpoint"], "azure_deployment": values["deployment_name"], - "api_key": values["openai_api_key"], + "api_key": values["openai_api_key"].get_secret_value(), "azure_ad_token": values["azure_ad_token"], "azure_ad_token_provider": values["azure_ad_token_provider"], "organization": values["openai_organization"], diff --git a/libs/langchain/langchain/chat_models/everlyai.py b/libs/langchain/langchain/chat_models/everlyai.py index 95bddaba5b06d..c5d8503e69b6a 100644 --- a/libs/langchain/langchain/chat_models/everlyai.py +++ b/libs/langchain/langchain/chat_models/everlyai.py @@ -13,7 +13,7 @@ ChatOpenAI, _import_tiktoken, ) -from langchain.utils import get_from_dict_or_env +from langchain.utils import convert_to_secret_str, get_from_dict_or_env if TYPE_CHECKING: import tiktoken @@ -74,10 +74,12 @@ def get_available_models() -> Set[str]: @root_validator(pre=True) def validate_environment_override(cls, values: dict) -> dict: """Validate that api key and python package exists in environment.""" - values["openai_api_key"] = get_from_dict_or_env( - values, - "everlyai_api_key", - "EVERLYAI_API_KEY", + values["openai_api_key"] = convert_to_secret_str( + get_from_dict_or_env( + values, + "everlyai_api_key", + "EVERLYAI_API_KEY", + ) ) values["openai_api_base"] = DEFAULT_API_BASE diff --git a/libs/langchain/langchain/chat_models/openai.py b/libs/langchain/langchain/chat_models/openai.py index 74c24066cb8d2..3fbea2857a88f 100644 --- a/libs/langchain/langchain/chat_models/openai.py +++ b/libs/langchain/langchain/chat_models/openai.py @@ -18,6 +18,7 @@ Tuple, Type, Union, + cast, ) from langchain_core.language_models import LanguageModelInput @@ -32,9 +33,10 @@ ToolMessageChunk, ) from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator +from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator from langchain_core.runnables import Runnable from langchain_core.utils import ( + convert_to_secret_str, get_pydantic_field_names, ) @@ -188,10 +190,9 @@ def is_lc_serializable(cls) -> bool: """What sampling temperature to use.""" model_kwargs: Dict[str, Any] = Field(default_factory=dict) """Holds any model parameters valid for `create` call not explicitly specified.""" - # When updating this to use a SecretStr - # Check for classes that derive from this class (as some of them - # may assume openai_api_key is a str) - openai_api_key: Optional[str] = Field(default=None, alias="api_key") + # All classes(i.e. AzureOpenAI) that derive from this class should + # set openai_api_key as SecretStr + openai_api_key: Optional[SecretStr] = Field(default=None, alias="api_key") """Automatically inferred from env var `OPENAI_API_KEY` if not provided.""" openai_api_base: Optional[str] = Field(default=None, alias="base_url") """Base URL path for API requests, leave blank if not using a proxy or service @@ -269,8 +270,8 @@ def validate_environment(cls, values: Dict) -> Dict: if values["n"] > 1 and values["streaming"]: raise ValueError("n must be 1 when streaming.") - values["openai_api_key"] = get_from_dict_or_env( - values, "openai_api_key", "OPENAI_API_KEY" + values["openai_api_key"] = convert_to_secret_str( + get_from_dict_or_env(values, "openai_api_key", "OPENAI_API_KEY") ) # Check OPENAI_ORGANIZATION for backwards compatibility. values["openai_organization"] = ( @@ -298,7 +299,7 @@ def validate_environment(cls, values: Dict) -> Dict: if is_openai_v1(): client_params = { - "api_key": values["openai_api_key"], + "api_key": values["openai_api_key"].get_secret_value(), "organization": values["openai_organization"], "base_url": values["openai_api_base"], "timeout": values["request_timeout"], @@ -530,7 +531,7 @@ def _client_params(self) -> Dict[str, Any]: if not is_openai_v1(): openai_creds.update( { - "api_key": self.openai_api_key, + "api_key": cast(SecretStr, self.openai_api_key).get_secret_value(), "api_base": self.openai_api_base, "organization": self.openai_organization, } diff --git a/libs/langchain/tests/unit_tests/chat_models/test_azureopenai.py b/libs/langchain/tests/unit_tests/chat_models/test_azureopenai.py index fd1ec775b00c9..2e75fd3589993 100644 --- a/libs/langchain/tests/unit_tests/chat_models/test_azureopenai.py +++ b/libs/langchain/tests/unit_tests/chat_models/test_azureopenai.py @@ -5,13 +5,14 @@ import pytest from langchain.chat_models.azure_openai import AzureChatOpenAI +from langchain.pydantic_v1 import SecretStr @mock.patch.dict( os.environ, { - "OPENAI_API_KEY": "test", - "OPENAI_API_BASE": "https://oai.azure.com/", + "AZURE_OPENAI_API_KEY": "test", + "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", "OPENAI_API_VERSION": "2023-05-01", }, ) @@ -53,3 +54,49 @@ def test_model_name_set_on_chat_result_when_present_in_response( chat_result.llm_output is not None and chat_result.llm_output["model_name"] == model_name ) + + +@mock.patch.dict( + os.environ, + { + "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", + "OPENAI_API_VERSION": "2023-05-01", + }, +) +@pytest.mark.requires("openai") +def test_api_key_is_secret_string_and_matches_input() -> None: + llm = AzureChatOpenAI(openai_api_key="secret-api-key") + assert isinstance(llm.openai_api_key, SecretStr) + assert llm.openai_api_key.get_secret_value() == "secret-api-key" + + +@mock.patch.dict( + os.environ, + { + "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", + "OPENAI_API_VERSION": "2023-05-01", + }, +) +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_via_constructor() -> None: + llm = AzureChatOpenAI(openai_api_key="secret-api-key") + assert str(llm.openai_api_key) == "**********" + assert "secret-api-key" not in repr(llm.openai_api_key) + assert "secret-api-key" not in repr(llm) + + +@mock.patch.dict( + os.environ, + { + "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", + "OPENAI_API_VERSION": "2023-05-01", + }, +) +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_via_env() -> None: + with pytest.MonkeyPatch.context() as mp: + mp.setenv("OPENAI_API_KEY", "secret-api-key") + llm = AzureChatOpenAI() + assert str(llm.openai_api_key) == "**********" + assert "secret-api-key" not in repr(llm.openai_api_key) + assert "secret-api-key" not in repr(llm) diff --git a/libs/langchain/tests/unit_tests/chat_models/test_openai.py b/libs/langchain/tests/unit_tests/chat_models/test_openai.py index 54b74e2341d0a..27970e02e2e5e 100644 --- a/libs/langchain/tests/unit_tests/chat_models/test_openai.py +++ b/libs/langchain/tests/unit_tests/chat_models/test_openai.py @@ -1,5 +1,6 @@ """Test OpenAI Chat API wrapper.""" import json +import os from typing import Any from unittest.mock import MagicMock, patch @@ -13,6 +14,9 @@ from langchain.adapters.openai import convert_dict_to_message from langchain.chat_models.openai import ChatOpenAI +from langchain.pydantic_v1 import SecretStr + +os.environ["OPENAI_API_KEY"] = "foo" @pytest.mark.requires("openai") @@ -121,3 +125,28 @@ def mock_create(*args: Any, **kwargs: Any) -> Any: res = llm.predict("bar") assert res == "Bar Baz" assert completed + + +@pytest.mark.requires("openai") +def test_api_key_is_secret_string_and_matches_input() -> None: + llm = ChatOpenAI(openai_api_key="secret-api-key") + assert isinstance(llm.openai_api_key, SecretStr) + assert llm.openai_api_key.get_secret_value() == "secret-api-key" + + +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_via_constructor() -> None: + llm = ChatOpenAI(openai_api_key="secret-api-key") + assert str(llm.openai_api_key) == "**********" + assert "secret-api-key" not in repr(llm.openai_api_key) + assert "secret-api-key" not in repr(llm) + + +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_via_env() -> None: + with pytest.MonkeyPatch.context() as mp: + mp.setenv("OPENAI_API_KEY", "secret-api-key") + llm = ChatOpenAI() + assert str(llm.openai_api_key) == "**********" + assert "secret-api-key" not in repr(llm.openai_api_key) + assert "secret-api-key" not in repr(llm) From 34064a1368f2aa67ad55bf43405d4340ad1d7e8d Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Thu, 18 Jul 2024 18:25:54 -0400 Subject: [PATCH 2/3] merge --- .../chat_models/test_azureopenai.py | 50 +------------------ .../unit_tests/chat_models/test_openai.py | 25 ---------- 2 files changed, 2 insertions(+), 73 deletions(-) diff --git a/libs/community/tests/unit_tests/chat_models/test_azureopenai.py b/libs/community/tests/unit_tests/chat_models/test_azureopenai.py index de6e41f9cd4ac..49419ed7dc327 100644 --- a/libs/community/tests/unit_tests/chat_models/test_azureopenai.py +++ b/libs/community/tests/unit_tests/chat_models/test_azureopenai.py @@ -10,8 +10,8 @@ @mock.patch.dict( os.environ, { - "AZURE_OPENAI_API_KEY": "test", - "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", + "OPENAI_API_KEY": "test", + "OPENAI_API_BASE": "https://oai.azure.com/", "OPENAI_API_VERSION": "2023-05-01", }, ) @@ -53,49 +53,3 @@ def test_model_name_set_on_chat_result_when_present_in_response( chat_result.llm_output is not None and chat_result.llm_output["model_name"] == model_name ) - - -@mock.patch.dict( - os.environ, - { - "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", - "OPENAI_API_VERSION": "2023-05-01", - }, -) -@pytest.mark.requires("openai") -def test_api_key_is_secret_string_and_matches_input() -> None: - llm = AzureChatOpenAI(openai_api_key="secret-api-key") - assert isinstance(llm.openai_api_key, SecretStr) - assert llm.openai_api_key.get_secret_value() == "secret-api-key" - - -@mock.patch.dict( - os.environ, - { - "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", - "OPENAI_API_VERSION": "2023-05-01", - }, -) -@pytest.mark.requires("openai") -def test_api_key_masked_when_passed_via_constructor() -> None: - llm = AzureChatOpenAI(openai_api_key="secret-api-key") - assert str(llm.openai_api_key) == "**********" - assert "secret-api-key" not in repr(llm.openai_api_key) - assert "secret-api-key" not in repr(llm) - - -@mock.patch.dict( - os.environ, - { - "AZURE_OPENAI_ENDPOINT": "https://oai.azure.com/", - "OPENAI_API_VERSION": "2023-05-01", - }, -) -@pytest.mark.requires("openai") -def test_api_key_masked_when_passed_via_env() -> None: - with pytest.MonkeyPatch.context() as mp: - mp.setenv("OPENAI_API_KEY", "secret-api-key") - llm = AzureChatOpenAI() - assert str(llm.openai_api_key) == "**********" - assert "secret-api-key" not in repr(llm.openai_api_key) - assert "secret-api-key" not in repr(llm) diff --git a/libs/community/tests/unit_tests/chat_models/test_openai.py b/libs/community/tests/unit_tests/chat_models/test_openai.py index 38aa080f26004..4cb3efc6d8e74 100644 --- a/libs/community/tests/unit_tests/chat_models/test_openai.py +++ b/libs/community/tests/unit_tests/chat_models/test_openai.py @@ -131,28 +131,3 @@ async def mock_create(*args: Any, **kwargs: Any) -> Any: res = await llm.apredict("bar") assert res == "Bar Baz" assert completed - - -@pytest.mark.requires("openai") -def test_api_key_is_secret_string_and_matches_input() -> None: - llm = ChatOpenAI(openai_api_key="secret-api-key") - assert isinstance(llm.openai_api_key, SecretStr) - assert llm.openai_api_key.get_secret_value() == "secret-api-key" - - -@pytest.mark.requires("openai") -def test_api_key_masked_when_passed_via_constructor() -> None: - llm = ChatOpenAI(openai_api_key="secret-api-key") - assert str(llm.openai_api_key) == "**********" - assert "secret-api-key" not in repr(llm.openai_api_key) - assert "secret-api-key" not in repr(llm) - - -@pytest.mark.requires("openai") -def test_api_key_masked_when_passed_via_env() -> None: - with pytest.MonkeyPatch.context() as mp: - mp.setenv("OPENAI_API_KEY", "secret-api-key") - llm = ChatOpenAI() - assert str(llm.openai_api_key) == "**********" - assert "secret-api-key" not in repr(llm.openai_api_key) - assert "secret-api-key" not in repr(llm) From 177ef55110837a1ec66b6170998d028ff9249439 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Thu, 18 Jul 2024 18:30:18 -0400 Subject: [PATCH 3/3] update everlyai --- .../langchain_community/chat_models/everlyai.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/libs/community/langchain_community/chat_models/everlyai.py b/libs/community/langchain_community/chat_models/everlyai.py index c85d4c9e9a0f8..4122dae280cb7 100644 --- a/libs/community/langchain_community/chat_models/everlyai.py +++ b/libs/community/langchain_community/chat_models/everlyai.py @@ -8,7 +8,7 @@ from langchain_core.messages import BaseMessage from langchain_core.pydantic_v1 import Field, root_validator -from langchain_core.utils import get_from_dict_or_env +from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env from langchain_community.adapters.openai import convert_message_to_dict from langchain_community.chat_models.openai import ( @@ -79,10 +79,12 @@ def get_available_models() -> Set[str]: @root_validator(pre=True) def validate_environment_override(cls, values: dict) -> dict: """Validate that api key and python package exists in environment.""" - values["openai_api_key"] = get_from_dict_or_env( - values, - "everlyai_api_key", - "EVERLYAI_API_KEY", + values["openai_api_key"] = convert_to_secret_str( + get_from_dict_or_env( + values, + "everlyai_api_key", + "EVERLYAI_API_KEY", + ) ) values["openai_api_base"] = DEFAULT_API_BASE