From 88f0f8e36091fd4fb34b722ba0c0e71112e87d9a Mon Sep 17 00:00:00 2001 From: Kamol Treewatchararat Date: Wed, 27 Dec 2023 17:16:24 +0700 Subject: [PATCH 1/5] use secretStr to mask api key for banana llm use convert_to_secret_str() util create unit tests --- .../langchain_community/llms/bananadev.py | 6 +-- .../tests/unit_tests/llms/test_bananadev.py | 37 +++++++++++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) create mode 100644 libs/community/tests/unit_tests/llms/test_bananadev.py diff --git a/libs/community/langchain_community/llms/bananadev.py b/libs/community/langchain_community/llms/bananadev.py index 88ab7f5e58c1f..4ed1583a907d5 100644 --- a/libs/community/langchain_community/llms/bananadev.py +++ b/libs/community/langchain_community/llms/bananadev.py @@ -4,7 +4,7 @@ from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.language_models.llms import LLM from langchain_core.pydantic_v1 import Extra, 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.llms.utils import enforce_stop_tokens @@ -66,8 +66,8 @@ def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]: @root_validator() def validate_environment(cls, values: Dict) -> Dict: """Validate that api key and python package exists in environment.""" - banana_api_key = get_from_dict_or_env( - values, "banana_api_key", "BANANA_API_KEY" + banana_api_key = convert_to_secret_str( + get_from_dict_or_env(values, "banana_api_key", "BANANA_API_KEY") ) values["banana_api_key"] = banana_api_key return values diff --git a/libs/community/tests/unit_tests/llms/test_bananadev.py b/libs/community/tests/unit_tests/llms/test_bananadev.py new file mode 100644 index 0000000000000..e83db8c0f28c8 --- /dev/null +++ b/libs/community/tests/unit_tests/llms/test_bananadev.py @@ -0,0 +1,37 @@ +"""Test Banana llm""" +from typing import cast + +from langchain_core.pydantic_v1 import SecretStr +from pytest import CaptureFixture, MonkeyPatch + +from langchain_community.llms.bananadev import Banana + +def test_api_key_is_secret_string() -> None: + llm = Banana(banana_api_key="secret-api-key") + assert isinstance(llm.banana_api_key, SecretStr) + +def test_api_key_masked_when_passed_from_env( + monkeypatch: MonkeyPatch, capsys: CaptureFixture +) -> None: + """Test initialization with an API key provided via an env variable""" + monkeypatch.setenv("BANANA_API_KEY", "secret-api-key") + llm = Banana() + print(llm.banana_api_key, end="") + captured = capsys.readouterr() + + assert captured.out == "**********" + +def test_api_key_masked_when_passed_via_constructor( + capsys: CaptureFixture, +) -> None: + """Test initialization with an API key provided via the initializer""" + llm = Banana(banana_api_key="secret-api-key") + print(llm.banana_api_key, end="") + captured = capsys.readouterr() + + assert captured.out == "**********" + +def test_uses_actual_secret_value_from_secretstr() -> None: + """Test that actual secret is retrieved using `.get_secret_value()`.""" + llm = Banana(banana_api_key="secret-api-key") + assert cast(SecretStr, llm.banana_api_key).get_secret_value() == "secret-api-key" From 04bb3562599e37622949265afccddc37fb5dfe46 Mon Sep 17 00:00:00 2001 From: Kamol Treewatchararat Date: Wed, 27 Dec 2023 17:23:55 +0700 Subject: [PATCH 2/5] lint --- libs/community/tests/unit_tests/llms/test_bananadev.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/community/tests/unit_tests/llms/test_bananadev.py b/libs/community/tests/unit_tests/llms/test_bananadev.py index e83db8c0f28c8..0f7f58fa2d3a5 100644 --- a/libs/community/tests/unit_tests/llms/test_bananadev.py +++ b/libs/community/tests/unit_tests/llms/test_bananadev.py @@ -6,6 +6,7 @@ from langchain_community.llms.bananadev import Banana + def test_api_key_is_secret_string() -> None: llm = Banana(banana_api_key="secret-api-key") assert isinstance(llm.banana_api_key, SecretStr) From 176be72d39eff343eaa1eeaf46a3dc1f1a8a4c4c Mon Sep 17 00:00:00 2001 From: Kamol Treewatchararat Date: Wed, 27 Dec 2023 17:25:19 +0700 Subject: [PATCH 3/5] lint --- libs/community/tests/unit_tests/llms/test_bananadev.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libs/community/tests/unit_tests/llms/test_bananadev.py b/libs/community/tests/unit_tests/llms/test_bananadev.py index 0f7f58fa2d3a5..d6cfaa4cdeea5 100644 --- a/libs/community/tests/unit_tests/llms/test_bananadev.py +++ b/libs/community/tests/unit_tests/llms/test_bananadev.py @@ -11,6 +11,7 @@ def test_api_key_is_secret_string() -> None: llm = Banana(banana_api_key="secret-api-key") assert isinstance(llm.banana_api_key, SecretStr) + def test_api_key_masked_when_passed_from_env( monkeypatch: MonkeyPatch, capsys: CaptureFixture ) -> None: @@ -22,6 +23,7 @@ def test_api_key_masked_when_passed_from_env( assert captured.out == "**********" + def test_api_key_masked_when_passed_via_constructor( capsys: CaptureFixture, ) -> None: @@ -32,6 +34,7 @@ def test_api_key_masked_when_passed_via_constructor( assert captured.out == "**********" + def test_uses_actual_secret_value_from_secretstr() -> None: """Test that actual secret is retrieved using `.get_secret_value()`.""" llm = Banana(banana_api_key="secret-api-key") From 31066e2c72ce31ef103b3764095c59a49d7e4609 Mon Sep 17 00:00:00 2001 From: Kamol Treewatchararat Date: Wed, 27 Dec 2023 17:37:20 +0700 Subject: [PATCH 4/5] Update banana_api_key type --- libs/community/langchain_community/llms/bananadev.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/libs/community/langchain_community/llms/bananadev.py b/libs/community/langchain_community/llms/bananadev.py index 4ed1583a907d5..43ee44b83dfd6 100644 --- a/libs/community/langchain_community/llms/bananadev.py +++ b/libs/community/langchain_community/llms/bananadev.py @@ -1,9 +1,9 @@ import logging -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, List, Mapping, Optional, cast from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.language_models.llms import LLM -from langchain_core.pydantic_v1 import Extra, Field, root_validator +from langchain_core.pydantic_v1 import Extra, Field, SecretStr, root_validator from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env from langchain_community.llms.utils import enforce_stop_tokens @@ -38,7 +38,7 @@ class Banana(LLM): """Holds any model parameters valid for `create` call not explicitly specified.""" - banana_api_key: Optional[str] = None + banana_api_key: Optional[SecretStr] = None class Config: """Configuration for this pydantic config.""" @@ -103,7 +103,7 @@ def _call( ) params = self.model_kwargs or {} params = {**params, **kwargs} - api_key = self.banana_api_key + api_key = cast(SecretStr, self.banana_api_key) model_key = self.model_key model_url_slug = self.model_url_slug model_inputs = { @@ -113,7 +113,7 @@ def _call( } model = Client( # Found in main dashboard - api_key=api_key, + api_key=api_key.get_secret_value(), # Both found in model details page model_key=model_key, url=f"https://{model_url_slug}.run.banana.dev", From 103ba21e004ff65ba453c34b5900249df0821305 Mon Sep 17 00:00:00 2001 From: Bagatur Date: Fri, 29 Mar 2024 17:49:20 -0700 Subject: [PATCH 5/5] fmt --- libs/community/tests/unit_tests/llms/test_bananadev.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/community/tests/unit_tests/llms/test_bananadev.py b/libs/community/tests/unit_tests/llms/test_bananadev.py index d6cfaa4cdeea5..ea219cc815e15 100644 --- a/libs/community/tests/unit_tests/llms/test_bananadev.py +++ b/libs/community/tests/unit_tests/llms/test_bananadev.py @@ -18,7 +18,7 @@ def test_api_key_masked_when_passed_from_env( """Test initialization with an API key provided via an env variable""" monkeypatch.setenv("BANANA_API_KEY", "secret-api-key") llm = Banana() - print(llm.banana_api_key, end="") + print(llm.banana_api_key, end="") # noqa: T201 captured = capsys.readouterr() assert captured.out == "**********" @@ -29,7 +29,7 @@ def test_api_key_masked_when_passed_via_constructor( ) -> None: """Test initialization with an API key provided via the initializer""" llm = Banana(banana_api_key="secret-api-key") - print(llm.banana_api_key, end="") + print(llm.banana_api_key, end="") # noqa: T201 captured = capsys.readouterr() assert captured.out == "**********"