github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-20 13:54:48 +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, ChatGenerationChunk,
ChatResult, ChatResult,
) )
from langchain_core.pydantic_v1 import Field, root_validator from langchain_core.pydantic_v1 import 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
from requests.exceptions import HTTPError from requests.exceptions import HTTPError
from tenacity import ( from tenacity import (
before_sleep_log, before_sleep_log,
@ -153,7 +153,7 @@ class ChatTongyi(BaseChatModel):
top_p: float = 0.8 top_p: float = 0.8
"""Total probability mass of tokens to consider at each step.""" """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.""" """Dashscope api key provide by Alibaba Cloud."""
streaming: bool = False streaming: bool = False
@ -170,8 +170,8 @@ class ChatTongyi(BaseChatModel):
@root_validator() @root_validator()
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment.""" """Validate that api key and python package exists in environment."""
values["dashscope_api_key"] = get_from_dict_or_env( values["dashscope_api_key"] = convert_to_secret_str(
values, "dashscope_api_key", "DASHSCOPE_API_KEY" get_from_dict_or_env(values, "dashscope_api_key", "DASHSCOPE_API_KEY")
) )
try: try:
import dashscope import dashscope
@ -197,7 +197,7 @@ class ChatTongyi(BaseChatModel):
return { return {
"model": self.model_name, "model": self.model_name,
"top_p": self.top_p, "top_p": self.top_p,
"api_key": self.dashscope_api_key, "api_key": self.dashscope_api_key.get_secret_value(),
"result_format": "message", "result_format": "message",
**self.model_kwargs, **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.callbacks import CallbackManager
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, LLMResult 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 langchain_community.chat_models.tongyi import ChatTongyi
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler 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: def test_default_call() -> None:
"""Test default model call.""" """Test default model call."""
chat = ChatTongyi() chat = ChatTongyi()