mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-18 21:09:00 +00:00
Refactor: use SecretStr for GPTRouter chat-model (#15101)
This commit is contained in:
parent
f506b4cfd2
commit
b6952d41e5
@ -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,11 +173,13 @@ 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(
|
||||||
|
get_from_dict_or_env(
|
||||||
values,
|
values,
|
||||||
"gpt_router_api_key",
|
"gpt_router_api_key",
|
||||||
"GPT_ROUTER_API_KEY",
|
"GPT_ROUTER_API_KEY",
|
||||||
)
|
)
|
||||||
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from gpt_router.client import GPTRouterClient
|
from gpt_router.client import GPTRouterClient
|
||||||
@ -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:
|
||||||
|
@ -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(
|
||||||
|
Loading…
Reference in New Issue
Block a user