github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-19 13:23:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: use SecretStr for tongyi chat-model (#15102)

Browse Source
This commit is contained in:
chyroc 2024-01-03 07:45:23 +08:00 committed by GitHub
parent e1c2cd7a28
commit 37ad6ec248
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
libs/community/langchain_community/chat_models/tongyi.py
View File

@ -37,8 +37,8 @@ from langchain_core.outputs import (
ChatGenerationChunk,
ChatResult,
)
from langchain_core.pydantic_v1 import Field, root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from requests.exceptions import HTTPError
from tenacity import (
before_sleep_log,
@ -153,7 +153,7 @@ class ChatTongyi(BaseChatModel):
top_p: float = 0.8
"""Total probability mass of tokens to consider at each step."""
dashscope_api_key: Optional[str] = None
dashscope_api_key: Optional[SecretStr] = None
"""Dashscope api key provide by Alibaba Cloud."""
streaming: bool = False
@ -170,8 +170,8 @@ class ChatTongyi(BaseChatModel):
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment."""
values["dashscope_api_key"] = get_from_dict_or_env(
values, "dashscope_api_key", "DASHSCOPE_API_KEY"
values["dashscope_api_key"] = convert_to_secret_str(
get_from_dict_or_env(values, "dashscope_api_key", "DASHSCOPE_API_KEY")
)
try:
import dashscope
@ -197,7 +197,7 @@ class ChatTongyi(BaseChatModel):
return {
"model": self.model_name,
"top_p": self.top_p,
"api_key": self.dashscope_api_key,
"api_key": self.dashscope_api_key.get_secret_value(),
"result_format": "message",
**self.model_kwargs,
}

17
libs/community/tests/integration_tests/chat_models/test_tongyi.py
View File

@ -3,11 +3,28 @@
from langchain_core.callbacks import CallbackManager
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, LLMResult
from langchain_core.pydantic_v1 import SecretStr
from pytest import CaptureFixture
from langchain_community.chat_models.tongyi import ChatTongyi
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
def test_api_key_is_string() -> None:
llm = ChatTongyi(dashscope_api_key="secret-api-key")
assert isinstance(llm.dashscope_api_key, SecretStr)
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
llm = ChatTongyi(dashscope_api_key="secret-api-key")
print(llm.dashscope_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
def test_default_call() -> None:
"""Test default model call."""
chat = ChatTongyi()