github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-30 08:14:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: use SecretStr for PipelineAI llms (#15120)

Browse Source
This commit is contained in:
chyroc 2023-12-27 05:00:58 +08:00 committed by GitHub
parent d63ceb65b3
commit 0a9a73a9c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
libs/community/langchain_community/llms/pipelineai.py
View File

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

19
libs/community/tests/unit_tests/llms/test_pipelineai.py Normal file
View File

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