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

Refactor: use SecretStr for Predibase llms (#15119)

Browse Source
This commit is contained in:
chyroc 2023-12-27 05:01:42 +08:00 committed by GitHub
parent 0a9a73a9c9
commit 1abcf441ae
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libs/community/langchain_community/llms/predibase.py
View File

@ -2,7 +2,7 @@ from typing import Any, Dict, List, Mapping, Optional
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Field
from langchain_core.pydantic_v1 import Field, SecretStr
class Predibase(LLM):
@ -13,7 +13,7 @@ class Predibase(LLM):
"""
model: str
predibase_api_key: str
predibase_api_key: SecretStr
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
@property
@ -30,7 +30,7 @@ class Predibase(LLM):
try:
from predibase import PredibaseClient
pc = PredibaseClient(token=self.predibase_api_key)
pc = PredibaseClient(token=self.predibase_api_key.get_secret_value())
except ImportError as e:
raise ImportError(
"Could not import Predibase Python package. "

19
libs/community/tests/integration_tests/llms/test_predibase.py Normal file
View File

@ -0,0 +1,19 @@
from langchain_core.pydantic_v1 import SecretStr
from pytest import CaptureFixture
from langchain_community.llms.predibase import Predibase
def test_api_key_is_string() -> None:
llm = Predibase(predibase_api_key="secret-api-key")
assert isinstance(llm.predibase_api_key, SecretStr)
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
llm = Predibase(predibase_api_key="secret-api-key")
print(llm.predibase_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"