diff --git a/libs/community/langchain_community/llms/pipelineai.py b/libs/community/langchain_community/llms/pipelineai.py index 91182d99760..d309c9cab55 100644 --- a/libs/community/langchain_community/llms/pipelineai.py +++ b/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} diff --git a/libs/community/tests/unit_tests/llms/test_pipelineai.py b/libs/community/tests/unit_tests/llms/test_pipelineai.py new file mode 100644 index 00000000000..63f750d0b22 --- /dev/null +++ b/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 == "**********"