Skip to content

Commit

Permalink
Mask API key for Aleph Alpha LLM (#12377)
Browse files Browse the repository at this point in the history
- **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>
  • Loading branch information
slangenbach and eyurtsev authored Oct 27, 2023
1 parent d6acb3e commit b22da81
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 8 deletions.
17 changes: 12 additions & 5 deletions libs/langchain/langchain/llms/aleph_alpha.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,19 @@
from typing import Any, Dict, List, Optional, Sequence
from typing import Any, Dict, List, Optional, Sequence, Union

from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from langchain.llms.utils import enforce_stop_tokens
from langchain.pydantic_v1 import Extra, root_validator
from langchain.pydantic_v1 import Extra, SecretStr, root_validator
from langchain.utils import get_from_dict_or_env


def _to_secret(value: Union[SecretStr, str]) -> SecretStr:
"""Convert a string to a SecretStr if needed."""
if isinstance(value, SecretStr):
return value
return SecretStr(value)


class AlephAlpha(LLM):
"""Aleph Alpha large language models.
Expand Down Expand Up @@ -169,14 +176,14 @@ class Config:
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment."""
aleph_alpha_api_key = get_from_dict_or_env(
values, "aleph_alpha_api_key", "ALEPH_ALPHA_API_KEY"
values["aleph_alpha_api_key"] = _to_secret(
get_from_dict_or_env(values, "aleph_alpha_api_key", "ALEPH_ALPHA_API_KEY")
)
try:
from aleph_alpha_client import Client

values["client"] = Client(
token=aleph_alpha_api_key,
token=values["aleph_alpha_api_key"].get_secret_value(),
host=values["host"],
hosting=values["hosting"],
request_timeout_seconds=values["request_timeout_seconds"],
Expand Down
28 changes: 25 additions & 3 deletions libs/langchain/poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions libs/langchain/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,7 @@ cli = [
# Please use new-line on formatting to make it easier to add new packages without
# merge-conflicts
extended_testing = [
"aleph-alpha-client",
"amazon-textract-caller",
"aiosqlite",
"assemblyai",
Expand Down
36 changes: 36 additions & 0 deletions libs/langchain/tests/unit_tests/llms/test_aleph_alpha.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
"""Test Aleph Alpha specific stuff."""

import pytest
from pytest import CaptureFixture, MonkeyPatch

from langchain.llms.aleph_alpha import AlephAlpha
from langchain.pydantic_v1 import SecretStr


@pytest.mark.requires("aleph_alpha_client")
def test_api_key_is_secret_string() -> None:
llm = AlephAlpha(aleph_alpha_api_key="secret-api-key")
assert isinstance(llm.aleph_alpha_api_key, SecretStr)


@pytest.mark.requires("aleph_alpha_client")
def test_api_key_masked_when_passed_via_constructor(
capsys: CaptureFixture,
) -> None:
llm = AlephAlpha(aleph_alpha_api_key="secret-api-key")
print(llm.aleph_alpha_api_key, end="")
captured = capsys.readouterr()

assert captured.out == "**********"


@pytest.mark.requires("aleph_alpha_client")
def test_api_key_masked_when_passed_from_env(
monkeypatch: MonkeyPatch, capsys: CaptureFixture
) -> None:
monkeypatch.setenv("ALEPH_ALPHA_API_KEY", "secret-api-key")
llm = AlephAlpha()
print(llm.aleph_alpha_api_key, end="")
captured = capsys.readouterr()

assert captured.out == "**********"

0 comments on commit b22da81

Please sign in to comment.