Refactor: use SecretStr for PipelineAI llms (#15120)

libs/community/langchain_community/llms/pipelineai.py

@@ -3,8 +3,14 @@ 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 BaseModel, Extra, Field, root_validator
-from langchain_core.utils import get_from_dict_or_env
+from langchain_core.pydantic_v1 import (
+    BaseModel,
+    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
 
@@ -34,7 +40,7 @@ class PipelineAI(LLM, BaseModel):
     """Holds any pipeline parameters valid for `create` call not
     explicitly specified."""
 
-    pipeline_api_key: Optional[str] = None
+    pipeline_api_key: Optional[SecretStr] = None
 
     class Config:
         """Configuration for this pydantic config."""
@@ -62,8 +68,8 @@ class PipelineAI(LLM, BaseModel):
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validate that api key and python package exists in environment."""
-        pipeline_api_key = get_from_dict_or_env(
-            values, "pipeline_api_key", "PIPELINE_API_KEY"
+        pipeline_api_key = convert_to_secret_str(
+            get_from_dict_or_env(values, "pipeline_api_key", "PIPELINE_API_KEY")
         )
         values["pipeline_api_key"] = pipeline_api_key
         return values
@@ -96,7 +102,7 @@ class PipelineAI(LLM, BaseModel):
                 "Could not import pipeline-ai python package. "
                 "Please install it with `pip install pipeline-ai`."
             )
-        client = PipelineCloud(token=self.pipeline_api_key)
+        client = PipelineCloud(token=self.pipeline_api_key.get_secret_value())
         params = self.pipeline_kwargs or {}
         params = {**params, **kwargs}
 

libs/community/tests/unit_tests/llms/test_pipelineai.py

@@ -0,0 +1,19 @@
+from langchain_core.pydantic_v1 import SecretStr
+from pytest import CaptureFixture
+
+from langchain_community.llms.pipelineai import PipelineAI
+
+
+def test_api_key_is_string() -> None:
+    llm = PipelineAI(pipeline_api_key="secret-api-key")
+    assert isinstance(llm.pipeline_api_key, SecretStr)
+
+
+def test_api_key_masked_when_passed_via_constructor(
+    capsys: CaptureFixture,
+) -> None:
+    llm = PipelineAI(pipeline_api_key="secret-api-key")
+    print(llm.pipeline_api_key, end="")
+    captured = capsys.readouterr()
+
+    assert captured.out == "**********"