github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-18 12:58:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

23
libs/community/tests/integration_tests/chat_models/test_gpt_router.py
View File

@ -7,11 +7,34 @@ from langchain_core.callbacks import (
)
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.gpt_router import GPTRouter, GPTRouterModel
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:
"""Test valid call to GPTRouter."""
anthropic_claude = GPTRouterModel(