community[patch]: use SecretStr for yandex model secrets (#15463)

This commit is contained in:
chyroc 2024-01-24 12:08:53 +08:00 committed by GitHub
parent d628a80a5d
commit 61da2ff24c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 26 additions and 20 deletions

View File

@ -5,8 +5,8 @@ import logging
from typing import Any, Callable, Dict, List from typing import Any, Callable, Dict, List
from langchain_core.embeddings import Embeddings from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, root_validator from langchain_core.pydantic_v1 import BaseModel, SecretStr, root_validator
from langchain_core.utils import get_from_dict_or_env from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from tenacity import ( from tenacity import (
before_sleep_log, before_sleep_log,
retry, retry,
@ -41,18 +41,16 @@ class YandexGPTEmbeddings(BaseModel, Embeddings):
embeddings = YandexGPTEmbeddings(iam_token="t1.9eu...", model_uri="emb://<folder-id>/text-search-query/latest") embeddings = YandexGPTEmbeddings(iam_token="t1.9eu...", model_uri="emb://<folder-id>/text-search-query/latest")
""" """
iam_token: str = "" iam_token: SecretStr = ""
"""Yandex Cloud IAM token for service account """Yandex Cloud IAM token for service account
with the `ai.languageModels.user` role""" with the `ai.languageModels.user` role"""
api_key: str = "" api_key: SecretStr = ""
"""Yandex Cloud Api Key for service account """Yandex Cloud Api Key for service account
with the `ai.languageModels.user` role""" with the `ai.languageModels.user` role"""
model_uri: str = "" model_uri: str = ""
"""Model uri to use.""" """Model uri to use."""
folder_id: str = "" folder_id: str = ""
"""Yandex Cloud folder ID""" """Yandex Cloud folder ID"""
model_uri: str = ""
"""Model uri to use."""
model_name: str = "text-search-query" model_name: str = "text-search-query"
"""Model name to use.""" """Model name to use."""
model_version: str = "latest" model_version: str = "latest"
@ -66,23 +64,27 @@ class YandexGPTEmbeddings(BaseModel, Embeddings):
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
"""Validate that iam token exists in environment.""" """Validate that iam token exists in environment."""
iam_token = get_from_dict_or_env(values, "iam_token", "YC_IAM_TOKEN", "") iam_token = convert_to_secret_str(
get_from_dict_or_env(values, "iam_token", "YC_IAM_TOKEN", "")
)
values["iam_token"] = iam_token values["iam_token"] = iam_token
api_key = get_from_dict_or_env(values, "api_key", "YC_API_KEY", "") api_key = convert_to_secret_str(
get_from_dict_or_env(values, "api_key", "YC_API_KEY", "")
)
values["api_key"] = api_key values["api_key"] = api_key
folder_id = get_from_dict_or_env(values, "folder_id", "YC_FOLDER_ID", "") folder_id = get_from_dict_or_env(values, "folder_id", "YC_FOLDER_ID", "")
values["folder_id"] = folder_id values["folder_id"] = folder_id
if api_key == "" and iam_token == "": if api_key.get_secret_value() == "" and iam_token.get_secret_value() == "":
raise ValueError("Either 'YC_API_KEY' or 'YC_IAM_TOKEN' must be provided.") raise ValueError("Either 'YC_API_KEY' or 'YC_IAM_TOKEN' must be provided.")
if values["iam_token"]: if values["iam_token"]:
values["_grpc_metadata"] = [ values["_grpc_metadata"] = [
("authorization", f"Bearer {values['iam_token']}") ("authorization", f"Bearer {values['iam_token'].get_secret_value()}")
] ]
if values["folder_id"]: if values["folder_id"]:
values["_grpc_metadata"].append(("x-folder-id", values["folder_id"])) values["_grpc_metadata"].append(("x-folder-id", values["folder_id"]))
else: else:
values["_grpc_metadata"] = ( values["_grpc_metadata"] = (
("authorization", f"Api-Key {values['api_key']}"), ("authorization", f"Api-Key {values['api_key'].get_secret_value()}"),
) )
if values["model_uri"] == "" and values["folder_id"] == "": if values["model_uri"] == "" and values["folder_id"] == "":
raise ValueError("Either 'model_uri' or 'folder_id' must be provided.") raise ValueError("Either 'model_uri' or 'folder_id' must be provided.")

View File

@ -9,8 +9,8 @@ from langchain_core.callbacks import (
) )
from langchain_core.language_models.llms import LLM from langchain_core.language_models.llms import LLM
from langchain_core.load.serializable import Serializable from langchain_core.load.serializable import Serializable
from langchain_core.pydantic_v1 import root_validator from langchain_core.pydantic_v1 import SecretStr, root_validator
from langchain_core.utils import get_from_dict_or_env from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from tenacity import ( from tenacity import (
before_sleep_log, before_sleep_log,
retry, retry,
@ -25,10 +25,10 @@ logger = logging.getLogger(__name__)
class _BaseYandexGPT(Serializable): class _BaseYandexGPT(Serializable):
iam_token: str = "" iam_token: SecretStr = ""
"""Yandex Cloud IAM token for service or user account """Yandex Cloud IAM token for service or user account
with the `ai.languageModels.user` role""" with the `ai.languageModels.user` role"""
api_key: str = "" api_key: SecretStr = ""
"""Yandex Cloud Api Key for service account """Yandex Cloud Api Key for service account
with the `ai.languageModels.user` role""" with the `ai.languageModels.user` role"""
folder_id: str = "" folder_id: str = ""
@ -72,24 +72,28 @@ class _BaseYandexGPT(Serializable):
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
"""Validate that iam token exists in environment.""" """Validate that iam token exists in environment."""
iam_token = get_from_dict_or_env(values, "iam_token", "YC_IAM_TOKEN", "") iam_token = convert_to_secret_str(
get_from_dict_or_env(values, "iam_token", "YC_IAM_TOKEN", "")
)
values["iam_token"] = iam_token values["iam_token"] = iam_token
api_key = get_from_dict_or_env(values, "api_key", "YC_API_KEY", "") api_key = convert_to_secret_str(
get_from_dict_or_env(values, "api_key", "YC_API_KEY", "")
)
values["api_key"] = api_key values["api_key"] = api_key
folder_id = get_from_dict_or_env(values, "folder_id", "YC_FOLDER_ID", "") folder_id = get_from_dict_or_env(values, "folder_id", "YC_FOLDER_ID", "")
values["folder_id"] = folder_id values["folder_id"] = folder_id
if api_key == "" and iam_token == "": if api_key.get_secret_value() == "" and iam_token.get_secret_value() == "":
raise ValueError("Either 'YC_API_KEY' or 'YC_IAM_TOKEN' must be provided.") raise ValueError("Either 'YC_API_KEY' or 'YC_IAM_TOKEN' must be provided.")
if values["iam_token"]: if values["iam_token"]:
values["_grpc_metadata"] = [ values["_grpc_metadata"] = [
("authorization", f"Bearer {values['iam_token']}") ("authorization", f"Bearer {values['iam_token'].get_secret_value()}")
] ]
if values["folder_id"]: if values["folder_id"]:
values["_grpc_metadata"].append(("x-folder-id", values["folder_id"])) values["_grpc_metadata"].append(("x-folder-id", values["folder_id"]))
else: else:
values["_grpc_metadata"] = ( values["_grpc_metadata"] = (
("authorization", f"Api-Key {values['api_key']}"), ("authorization", f"Api-Key {values['api_key'].get_secret_value()}"),
) )
if values["model_uri"] == "" and values["folder_id"] == "": if values["model_uri"] == "" and values["folder_id"] == "":
raise ValueError("Either 'model_uri' or 'folder_id' must be provided.") raise ValueError("Either 'model_uri' or 'folder_id' must be provided.")