Skip to content

Commit b22da81

Browse files
Mask API key for Aleph Alpha LLM (#12377)
- **Description:** Add masking of API Key for Aleph Alpha LLM when printed. - **Issue**: #12165 - **Dependencies:** None - **Tag maintainer:** @eyurtsev --------- Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
1 parent d6acb3e commit b22da81

File tree

4 files changed

+74
-8
lines changed

4 files changed

+74
-8
lines changed

libs/langchain/langchain/llms/aleph_alpha.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
1-
from typing import Any, Dict, List, Optional, Sequence
1+
from typing import Any, Dict, List, Optional, Sequence, Union
22

33
from langchain.callbacks.manager import CallbackManagerForLLMRun
44
from langchain.llms.base import LLM
55
from langchain.llms.utils import enforce_stop_tokens
6-
from langchain.pydantic_v1 import Extra, root_validator
6+
from langchain.pydantic_v1 import Extra, SecretStr, root_validator
77
from langchain.utils import get_from_dict_or_env
88

99

10+
def _to_secret(value: Union[SecretStr, str]) -> SecretStr:
11+
"""Convert a string to a SecretStr if needed."""
12+
if isinstance(value, SecretStr):
13+
return value
14+
return SecretStr(value)
15+
16+
1017
class AlephAlpha(LLM):
1118
"""Aleph Alpha large language models.
1219
@@ -169,14 +176,14 @@ class Config:
169176
@root_validator()
170177
def validate_environment(cls, values: Dict) -> Dict:
171178
"""Validate that api key and python package exists in environment."""
172-
aleph_alpha_api_key = get_from_dict_or_env(
173-
values, "aleph_alpha_api_key", "ALEPH_ALPHA_API_KEY"
179+
values["aleph_alpha_api_key"] = _to_secret(
180+
get_from_dict_or_env(values, "aleph_alpha_api_key", "ALEPH_ALPHA_API_KEY")
174181
)
175182
try:
176183
from aleph_alpha_client import Client
177184

178185
values["client"] = Client(
179-
token=aleph_alpha_api_key,
186+
token=values["aleph_alpha_api_key"].get_secret_value(),
180187
host=values["host"],
181188
hosting=values["hosting"],
182189
request_timeout_seconds=values["request_timeout_seconds"],

libs/langchain/poetry.lock

Lines changed: 25 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

libs/langchain/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,7 @@ cli = [
317317
# Please use new-line on formatting to make it easier to add new packages without
318318
# merge-conflicts
319319
extended_testing = [
320+
"aleph-alpha-client",
320321
"amazon-textract-caller",
321322
"aiosqlite",
322323
"assemblyai",
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
"""Test Aleph Alpha specific stuff."""
2+
3+
import pytest
4+
from pytest import CaptureFixture, MonkeyPatch
5+
6+
from langchain.llms.aleph_alpha import AlephAlpha
7+
from langchain.pydantic_v1 import SecretStr
8+
9+
10+
@pytest.mark.requires("aleph_alpha_client")
11+
def test_api_key_is_secret_string() -> None:
12+
llm = AlephAlpha(aleph_alpha_api_key="secret-api-key")
13+
assert isinstance(llm.aleph_alpha_api_key, SecretStr)
14+
15+
16+
@pytest.mark.requires("aleph_alpha_client")
17+
def test_api_key_masked_when_passed_via_constructor(
18+
capsys: CaptureFixture,
19+
) -> None:
20+
llm = AlephAlpha(aleph_alpha_api_key="secret-api-key")
21+
print(llm.aleph_alpha_api_key, end="")
22+
captured = capsys.readouterr()
23+
24+
assert captured.out == "**********"
25+
26+
27+
@pytest.mark.requires("aleph_alpha_client")
28+
def test_api_key_masked_when_passed_from_env(
29+
monkeypatch: MonkeyPatch, capsys: CaptureFixture
30+
) -> None:
31+
monkeypatch.setenv("ALEPH_ALPHA_API_KEY", "secret-api-key")
32+
llm = AlephAlpha()
33+
print(llm.aleph_alpha_api_key, end="")
34+
captured = capsys.readouterr()
35+
36+
assert captured.out == "**********"

0 commit comments

Comments
 (0)