github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-20 05:43:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: use SecretStr for VolcEngineMaas llms (#15117)

Browse Source
This commit is contained in:
chyroc 2023-12-27 04:59:08 +08:00 committed by GitHub
parent 3cc1da2b38
commit 674fde87d2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 43 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
libs/community/langchain_community/llms/volcengine_maas.py
View File

@ -5,8 +5,8 @@ from typing import Any, Dict, Iterator, List, Optional
from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunk from langchain_core.outputs import GenerationChunk
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, 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
class VolcEngineMaasBase(BaseModel): class VolcEngineMaasBase(BaseModel):
@ -14,9 +14,9 @@ class VolcEngineMaasBase(BaseModel):
client: Any client: Any
volc_engine_maas_ak: Optional[str] = None volc_engine_maas_ak: Optional[SecretStr] = None
"""access key for volc engine""" """access key for volc engine"""
volc_engine_maas_sk: Optional[str] = None volc_engine_maas_sk: Optional[SecretStr] = None
"""secret key for volc engine""" """secret key for volc engine"""
endpoint: Optional[str] = "maas-api.ml-platform-cn-beijing.volces.com" endpoint: Optional[str] = "maas-api.ml-platform-cn-beijing.volces.com"
@ -54,8 +54,12 @@ class VolcEngineMaasBase(BaseModel):
@root_validator() @root_validator()
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
ak = get_from_dict_or_env(values, "volc_engine_maas_ak", "VOLC_ACCESSKEY") volc_engine_maas_ak = convert_to_secret_str(
sk = get_from_dict_or_env(values, "volc_engine_maas_sk", "VOLC_SECRETKEY") get_from_dict_or_env(values, "volc_engine_maas_ak", "VOLC_ACCESSKEY")
)
volc_engine_maas_sk = convert_to_secret_str(
get_from_dict_or_env(values, "volc_engine_maas_sk", "VOLC_SECRETKEY")
)
endpoint = values["endpoint"] endpoint = values["endpoint"]
if values["endpoint"] is not None and values["endpoint"] != "": if values["endpoint"] is not None and values["endpoint"] != "":
endpoint = values["endpoint"] endpoint = values["endpoint"]
@ -68,10 +72,11 @@ class VolcEngineMaasBase(BaseModel):
connection_timeout=values["connect_timeout"], connection_timeout=values["connect_timeout"],
socket_timeout=values["read_timeout"], socket_timeout=values["read_timeout"],
) )
maas.set_ak(ak) maas.set_ak(volc_engine_maas_ak.get_secret_value())
values["volc_engine_maas_ak"] = ak maas.set_sk(volc_engine_maas_sk.get_secret_value())
values["volc_engine_maas_sk"] = sk
maas.set_sk(sk) values["volc_engine_maas_ak"] = volc_engine_maas_ak
values["volc_engine_maas_sk"] = volc_engine_maas_sk
values["client"] = maas values["client"] = maas
except ImportError: except ImportError:
raise ImportError( raise ImportError(

29
libs/community/tests/integration_tests/llms/test_volcengine_maas.py
View File

@ -3,8 +3,35 @@
from typing import Generator from typing import Generator
from langchain_core.outputs import LLMResult from langchain_core.outputs import LLMResult
from langchain_core.pydantic_v1 import SecretStr
from pytest import CaptureFixture
from langchain_community.llms.volcengine_maas import VolcEngineMaasLLM from langchain_community.llms.volcengine_maas import (
VolcEngineMaasBase,
VolcEngineMaasLLM,
)
def test_api_key_is_string() -> None:
llm = VolcEngineMaasBase(
volc_engine_maas_ak="secret-volc-ak",
volc_engine_maas_sk="secret-volc-sk",
)
assert isinstance(llm.volc_engine_maas_ak, SecretStr)
assert isinstance(llm.volc_engine_maas_sk, SecretStr)
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
llm = VolcEngineMaasBase(
volc_engine_maas_ak="secret-volc-ak",
volc_engine_maas_sk="secret-volc-sk",
)
print(llm.volc_engine_maas_ak, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
def test_default_call() -> None: def test_default_call() -> None: