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
2 changed files with 35 additions and 11 deletions

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(