From 0f4e796eedb852a6eb62c99bf6e71e4e9091cc08 Mon Sep 17 00:00:00 2001 From: Samad Koita <> Date: Sun, 29 Oct 2023 01:59:15 +0530 Subject: [PATCH 1/4] add SecretStr --- libs/langchain/langchain/llms/gooseai.py | 16 ++++++--- .../tests/unit_tests/llms/test_gooseai.py | 33 +++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) create mode 100644 libs/langchain/tests/unit_tests/llms/test_gooseai.py diff --git a/libs/langchain/langchain/llms/gooseai.py b/libs/langchain/langchain/llms/gooseai.py index 831cff9a54ab0..910108e0a4a74 100644 --- a/libs/langchain/langchain/llms/gooseai.py +++ b/libs/langchain/langchain/llms/gooseai.py @@ -1,13 +1,18 @@ import logging -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, List, Mapping, Optional, Union from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.llms.base import LLM -from langchain.pydantic_v1 import Extra, Field, root_validator +from langchain.pydantic_v1 import Extra, Field, root_validator, SecretStr from langchain.utils import get_from_dict_or_env logger = logging.getLogger(__name__) +def _to_secret(value: Union[SecretStr, str]) -> SecretStr: + """Convert a string to a SecretStr if needed.""" + if isinstance(value, SecretStr): + return value + return SecretStr(value) class GooseAI(LLM): """GooseAI large language models. @@ -89,13 +94,14 @@ 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.""" - gooseai_api_key = get_from_dict_or_env( - values, "gooseai_api_key", "GOOSEAI_API_KEY" + gooseai_api_key = _to_secret( + get_from_dict_or_env(values, "gooseai_api_key", "GOOSEAI_API_KEY") ) + values["gooseai_api_key"] = gooseai_api_key try: import openai - openai.api_key = gooseai_api_key + openai.api_key = gooseai_api_key.get_secret_value() openai.api_base = "https://api.goose.ai/v1" values["client"] = openai.Completion except ImportError: diff --git a/libs/langchain/tests/unit_tests/llms/test_gooseai.py b/libs/langchain/tests/unit_tests/llms/test_gooseai.py new file mode 100644 index 0000000000000..44b7bd32e0d25 --- /dev/null +++ b/libs/langchain/tests/unit_tests/llms/test_gooseai.py @@ -0,0 +1,33 @@ +"""Test GooseAI""" + +import pytest +from pytest import MonkeyPatch + + +from langchain.llms.gooseai import GooseAI +from langchain.pydantic_v1 import SecretStr + +@pytest.mark.requires("openai") +def test_api_key_is_secret_string() -> None: + llm = GooseAI(gooseai_api_key="secret-api-key") + assert isinstance(llm.gooseai_api_key, SecretStr) + assert llm.gooseai_api_key.get_secret_value() == "secret-api-key" + + +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_via_constructor() -> None: + llm = GooseAI(gooseai_api_key="secret-api-key") + assert str(llm.gooseai_api_key)== "**********" + assert "secret-api-key" not in repr(llm.gooseai_api_key) + assert "secret-api-key" not in repr(llm) + + +@pytest.mark.requires("openai") +def test_api_key_masked_when_passed_from_env() -> None: + with MonkeyPatch.context() as mp: + mp.setenv("GOOSEAI_API_KEY", "secret-api-key") + llm = GooseAI() + assert str(llm.gooseai_api_key)== "**********" + assert "secret-api-key" not in repr(llm.gooseai_api_key) + assert "secret-api-key" not in repr(llm) + From 44260e5815f5cc4284042e759381902f5448e267 Mon Sep 17 00:00:00 2001 From: Samad Koita <> Date: Sun, 29 Oct 2023 03:03:40 +0530 Subject: [PATCH 2/4] fix lint --- libs/langchain/langchain/llms/gooseai.py | 4 +++- libs/langchain/tests/unit_tests/llms/test_gooseai.py | 7 +++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/libs/langchain/langchain/llms/gooseai.py b/libs/langchain/langchain/llms/gooseai.py index 910108e0a4a74..0fd129fcdc83a 100644 --- a/libs/langchain/langchain/llms/gooseai.py +++ b/libs/langchain/langchain/llms/gooseai.py @@ -3,17 +3,19 @@ from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.llms.base import LLM -from langchain.pydantic_v1 import Extra, Field, root_validator, SecretStr +from langchain.pydantic_v1 import Extra, Field, SecretStr, root_validator from langchain.utils import get_from_dict_or_env logger = logging.getLogger(__name__) + def _to_secret(value: Union[SecretStr, str]) -> SecretStr: """Convert a string to a SecretStr if needed.""" if isinstance(value, SecretStr): return value return SecretStr(value) + class GooseAI(LLM): """GooseAI large language models. diff --git a/libs/langchain/tests/unit_tests/llms/test_gooseai.py b/libs/langchain/tests/unit_tests/llms/test_gooseai.py index 44b7bd32e0d25..db7a25ee9f92f 100644 --- a/libs/langchain/tests/unit_tests/llms/test_gooseai.py +++ b/libs/langchain/tests/unit_tests/llms/test_gooseai.py @@ -3,10 +3,10 @@ import pytest from pytest import MonkeyPatch - from langchain.llms.gooseai import GooseAI from langchain.pydantic_v1 import SecretStr + @pytest.mark.requires("openai") def test_api_key_is_secret_string() -> None: llm = GooseAI(gooseai_api_key="secret-api-key") @@ -17,7 +17,7 @@ def test_api_key_is_secret_string() -> None: @pytest.mark.requires("openai") def test_api_key_masked_when_passed_via_constructor() -> None: llm = GooseAI(gooseai_api_key="secret-api-key") - assert str(llm.gooseai_api_key)== "**********" + assert str(llm.gooseai_api_key) == "**********" assert "secret-api-key" not in repr(llm.gooseai_api_key) assert "secret-api-key" not in repr(llm) @@ -27,7 +27,6 @@ def test_api_key_masked_when_passed_from_env() -> None: with MonkeyPatch.context() as mp: mp.setenv("GOOSEAI_API_KEY", "secret-api-key") llm = GooseAI() - assert str(llm.gooseai_api_key)== "**********" + assert str(llm.gooseai_api_key) == "**********" assert "secret-api-key" not in repr(llm.gooseai_api_key) assert "secret-api-key" not in repr(llm) - From e300629574e44d9ac8e080c55745403d39bd3cf0 Mon Sep 17 00:00:00 2001 From: Samad Koita <> Date: Sun, 29 Oct 2023 03:10:46 +0530 Subject: [PATCH 3/4] add typing --- libs/langchain/langchain/llms/gooseai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langchain/langchain/llms/gooseai.py b/libs/langchain/langchain/llms/gooseai.py index 0fd129fcdc83a..5c7a86483a148 100644 --- a/libs/langchain/langchain/llms/gooseai.py +++ b/libs/langchain/langchain/llms/gooseai.py @@ -67,7 +67,7 @@ class GooseAI(LLM): logit_bias: Optional[Dict[str, float]] = Field(default_factory=dict) """Adjust the probability of specific tokens being generated.""" - gooseai_api_key: Optional[str] = None + gooseai_api_key: Optional[Union[str, SecretStr]] = None class Config: """Configuration for this pydantic config.""" From b058954a8afaddb220105ab244d5e3a1e7571757 Mon Sep 17 00:00:00 2001 From: Samad Koita <> Date: Sun, 29 Oct 2023 15:11:45 +0530 Subject: [PATCH 4/4] fix typing --- libs/langchain/langchain/llms/gooseai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langchain/langchain/llms/gooseai.py b/libs/langchain/langchain/llms/gooseai.py index 5c7a86483a148..30947ddb1b48f 100644 --- a/libs/langchain/langchain/llms/gooseai.py +++ b/libs/langchain/langchain/llms/gooseai.py @@ -67,7 +67,7 @@ class GooseAI(LLM): logit_bias: Optional[Dict[str, float]] = Field(default_factory=dict) """Adjust the probability of specific tokens being generated.""" - gooseai_api_key: Optional[Union[str, SecretStr]] = None + gooseai_api_key: Optional[SecretStr] = None class Config: """Configuration for this pydantic config."""