Add SecretStr for Symbl.ai Nebula API (#12896)

Description: This PR masks API key secrets for the Nebula model from
Symbl.ai
Issue: #12165 
Maintainer: @eyurtsev

---------

Co-authored-by: Praveen Venkateswaran <praveen.venkateswaran@ibm.com>
pull/12950/head
Praveen Venkateswaran 11 months ago committed by GitHub
parent 59d0bd2150
commit 8e0dcb37d2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -15,8 +15,9 @@ from tenacity import (
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.utils import get_from_dict_or_env
from langchain.pydantic_v1 import Extra, SecretStr, root_validator
from langchain.utils import convert_to_secret_str
from langchain.utils.env import get_from_dict_or_env
DEFAULT_NEBULA_SERVICE_URL = "https://api-nebula.symbl.ai"
DEFAULT_NEBULA_SERVICE_PATH = "/v1/model/generate"
@ -50,7 +51,7 @@ class Nebula(LLM):
nebula_service_url: Optional[str] = None
nebula_service_path: Optional[str] = None
nebula_api_key: Optional[str] = None
nebula_api_key: Optional[SecretStr] = None
model: Optional[str] = None
max_new_tokens: Optional[int] = 128
temperature: Optional[float] = 0.6
@ -81,8 +82,8 @@ class Nebula(LLM):
"NEBULA_SERVICE_PATH",
DEFAULT_NEBULA_SERVICE_PATH,
)
nebula_api_key = get_from_dict_or_env(
values, "nebula_api_key", "NEBULA_API_KEY", None
nebula_api_key = convert_to_secret_str(
get_from_dict_or_env(values, "nebula_api_key", "NEBULA_API_KEY", None)
)
if nebula_service_url.endswith("/"):
@ -187,9 +188,12 @@ def make_request(
) -> Any:
"""Generate text from the model."""
params = params or {}
api_key = None
if self.nebula_api_key is not None:
api_key = self.nebula_api_key.get_secret_value()
headers = {
"Content-Type": "application/json",
"ApiKey": f"{self.nebula_api_key}",
"ApiKey": f"{api_key}",
}
body = {

@ -0,0 +1,29 @@
"""Test the Nebula model by Symbl.ai"""
from pytest import CaptureFixture, MonkeyPatch
from langchain.llms.symblai_nebula import Nebula
from langchain.pydantic_v1 import SecretStr
def test_api_key_is_secret_string() -> None:
llm = Nebula(nebula_api_key="secret-api-key")
assert isinstance(llm.nebula_api_key, SecretStr)
assert llm.nebula_api_key.get_secret_value() == "secret-api-key"
def test_api_key_masked_when_passed_from_env(
monkeypatch: MonkeyPatch, capsys: CaptureFixture
) -> None:
monkeypatch.setenv("NEBULA_API_KEY", "secret-api-key")
llm = Nebula()
print(llm.nebula_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
def test_api_key_masked_when_passed_via_constructor(capsys: CaptureFixture) -> None:
llm = Nebula(nebula_api_key="secret-api-key")
print(llm.nebula_api_key, end="")
captured = capsys.readouterr()
assert captured.out == "**********"
Loading…
Cancel
Save