github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-10 21:35:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: use SecretStr for Petals llms (#15121)

Browse Source
This commit is contained in:
chyroc 2023-12-27 04:57:37 +08:00 committed by GitHub
parent 7ef25a3c1b
commit 3cc1da2b38
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
libs/community/langchain_community/llms/petals.py
View File

@ -3,8 +3,8 @@ 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 Extra, Field, root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_core.pydantic_v1 import Extra, Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from langchain_community.llms.utils import enforce_stop_tokens
@ -60,7 +60,7 @@ class Petals(LLM):
"""Holds any model parameters valid for `create` call
not explicitly specified."""
huggingface_api_key: Optional[str] = None
huggingface_api_key: Optional[SecretStr] = None
class Config:
"""Configuration for this pydantic config."""
@ -89,8 +89,8 @@ class Petals(LLM):
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment."""
huggingface_api_key = get_from_dict_or_env(
values, "huggingface_api_key", "HUGGINGFACE_API_KEY"
huggingface_api_key = convert_to_secret_str(
get_from_dict_or_env(values, "huggingface_api_key", "HUGGINGFACE_API_KEY")
)
try:
from petals import AutoDistributedModelForCausalLM
@ -101,7 +101,7 @@ class Petals(LLM):
values["client"] = AutoDistributedModelForCausalLM.from_pretrained(
model_name
)
values["huggingface_api_key"] = huggingface_api_key
values["huggingface_api_key"] = huggingface_api_key.get_secret_value()
except ImportError:
raise ImportError(

18
libs/community/tests/integration_tests/llms/test_petals.py
View File

@ -1,8 +1,26 @@
"""Test Petals API wrapper."""
from langchain_core.pydantic_v1 import SecretStr
from pytest import CaptureFixture
from langchain_community.llms.petals import Petals
def test_api_key_is_string() -> None:
llm = Petals(huggingface_api_key="secret-api-key")
assert isinstance(llm.huggingface_api_key, SecretStr)
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
llm = Petals(huggingface_api_key="secret-api-key")
print(llm.huggingface_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
def test_gooseai_call() -> None:
"""Test valid call to gooseai."""
llm = Petals(max_new_tokens=10)