mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-18 12:58:59 +00:00
Refactor: use SecretStr for StochasticAI llms (#15118)
This commit is contained in:
parent
674fde87d2
commit
d63ceb65b3
@ -5,8 +5,8 @@ from typing import Any, Dict, List, Mapping, Optional
|
||||
import requests
|
||||
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
|
||||
|
||||
@ -33,7 +33,7 @@ class StochasticAI(LLM):
|
||||
"""Holds any model parameters valid for `create` call not
|
||||
explicitly specified."""
|
||||
|
||||
stochasticai_api_key: Optional[str] = None
|
||||
stochasticai_api_key: Optional[SecretStr] = None
|
||||
|
||||
class Config:
|
||||
"""Configuration for this pydantic object."""
|
||||
@ -61,8 +61,8 @@ class StochasticAI(LLM):
|
||||
@root_validator()
|
||||
def validate_environment(cls, values: Dict) -> Dict:
|
||||
"""Validate that api key exists in environment."""
|
||||
stochasticai_api_key = get_from_dict_or_env(
|
||||
values, "stochasticai_api_key", "STOCHASTICAI_API_KEY"
|
||||
stochasticai_api_key = convert_to_secret_str(
|
||||
get_from_dict_or_env(values, "stochasticai_api_key", "STOCHASTICAI_API_KEY")
|
||||
)
|
||||
values["stochasticai_api_key"] = stochasticai_api_key
|
||||
return values
|
||||
@ -107,7 +107,7 @@ class StochasticAI(LLM):
|
||||
url=self.api_url,
|
||||
json={"prompt": prompt, "params": params},
|
||||
headers={
|
||||
"apiKey": f"{self.stochasticai_api_key}",
|
||||
"apiKey": f"{self.stochasticai_api_key.get_secret_value()}",
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
@ -119,7 +119,7 @@ class StochasticAI(LLM):
|
||||
response_get = requests.get(
|
||||
url=response_post_json["data"]["responseUrl"],
|
||||
headers={
|
||||
"apiKey": f"{self.stochasticai_api_key}",
|
||||
"apiKey": f"{self.stochasticai_api_key.get_secret_value()}",
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
|
19
libs/community/tests/unit_tests/llms/test_stochasticai.py
Normal file
19
libs/community/tests/unit_tests/llms/test_stochasticai.py
Normal file
@ -0,0 +1,19 @@
|
||||
from langchain_core.pydantic_v1 import SecretStr
|
||||
from pytest import CaptureFixture
|
||||
|
||||
from langchain_community.llms.stochasticai import StochasticAI
|
||||
|
||||
|
||||
def test_api_key_is_string() -> None:
|
||||
llm = StochasticAI(stochasticai_api_key="secret-api-key")
|
||||
assert isinstance(llm.stochasticai_api_key, SecretStr)
|
||||
|
||||
|
||||
def test_api_key_masked_when_passed_via_constructor(
|
||||
capsys: CaptureFixture,
|
||||
) -> None:
|
||||
llm = StochasticAI(stochasticai_api_key="secret-api-key")
|
||||
print(llm.stochasticai_api_key, end="")
|
||||
captured = capsys.readouterr()
|
||||
|
||||
assert captured.out == "**********"
|
Loading…
Reference in New Issue
Block a user