Refactor: use SecretStr for GPTRouter chat-model (#15101)

This commit is contained in:
chyroc 2024-01-02 07:20:26 +08:00 committed by GitHub
parent f506b4cfd2
commit b6952d41e5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 35 additions and 11 deletions

View File

@ -29,8 +29,8 @@ from langchain_core.language_models.chat_models import (
from langchain_core.language_models.llms import create_base_retry_decorator from langchain_core.language_models.llms import create_base_retry_decorator
from langchain_core.messages import AIMessageChunk, BaseMessage from langchain_core.messages import AIMessageChunk, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
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
from langchain_community.adapters.openai import ( from langchain_community.adapters.openai import (
convert_dict_to_message, convert_dict_to_message,
@ -150,7 +150,7 @@ class GPTRouter(BaseChatModel):
models_priority_list: List[GPTRouterModel] = Field(min_items=1) models_priority_list: List[GPTRouterModel] = Field(min_items=1)
gpt_router_api_base: str = Field(default=None) gpt_router_api_base: str = Field(default=None)
"""WriteSonic GPTRouter custom endpoint""" """WriteSonic GPTRouter custom endpoint"""
gpt_router_api_key: Optional[str] = None gpt_router_api_key: Optional[SecretStr] = None
"""WriteSonic GPTRouter API Key""" """WriteSonic GPTRouter API Key"""
temperature: float = 0.7 temperature: float = 0.7
"""What sampling temperature to use.""" """What sampling temperature to use."""
@ -173,10 +173,12 @@ class GPTRouter(BaseChatModel):
DEFAULT_API_BASE_URL, DEFAULT_API_BASE_URL,
) )
values["gpt_router_api_key"] = get_from_dict_or_env( values["gpt_router_api_key"] = convert_to_secret_str(
values, get_from_dict_or_env(
"gpt_router_api_key", values,
"GPT_ROUTER_API_KEY", "gpt_router_api_key",
"GPT_ROUTER_API_KEY",
)
) )
try: try:
@ -189,7 +191,8 @@ class GPTRouter(BaseChatModel):
) )
gpt_router_client = GPTRouterClient( gpt_router_client = GPTRouterClient(
values["gpt_router_api_base"], values["gpt_router_api_key"] values["gpt_router_api_base"],
values["gpt_router_api_key"].get_secret_value(),
) )
values["client"] = gpt_router_client values["client"] = gpt_router_client
@ -197,9 +200,7 @@ class GPTRouter(BaseChatModel):
@property @property
def lc_secrets(self) -> Dict[str, str]: def lc_secrets(self) -> Dict[str, str]:
return { return {"gpt_router_api_key": "GPT_ROUTER_API_KEY"}
"gpt_router_api_key": "GPT_ROUTER_API_KEY",
}
@property @property
def lc_serializable(self) -> bool: def lc_serializable(self) -> bool:

View File

@ -7,11 +7,34 @@ from langchain_core.callbacks import (
) )
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.gpt_router import GPTRouter, GPTRouterModel from langchain_community.chat_models.gpt_router import GPTRouter, GPTRouterModel
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:
gpt_router = GPTRouter(
gpt_router_api_base="https://example.com",
gpt_router_api_key="secret-api-key",
)
assert isinstance(gpt_router.gpt_router_api_key, SecretStr)
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
gpt_router = GPTRouter(
gpt_router_api_base="https://example.com",
gpt_router_api_key="secret-api-key",
)
print(gpt_router.gpt_router_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
def test_gpt_router_call() -> None: def test_gpt_router_call() -> None:
"""Test valid call to GPTRouter.""" """Test valid call to GPTRouter."""
anthropic_claude = GPTRouterModel( anthropic_claude = GPTRouterModel(