Mask API key for Aleph Alpha LLM

libs/langchain/langchain/llms/aleph_alpha.py

@@ -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.
 
@@ -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"],

libs/langchain/pyproject.toml

@@ -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",

libs/langchain/tests/unit_tests/llms/test_aleph_alpha.py

@@ -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 == "**********"