mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-22 20:43:08 +00:00
Refactor: use SecretStr for PipelineAI llms (#15120)
This commit is contained in:
parent
d63ceb65b3
commit
0a9a73a9c9
@ -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
19
libs/community/tests/unit_tests/llms/test_pipelineai.py
Normal 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 == "**********"
|
Loading…
Reference in New Issue
Block a user