mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-21 14:18:52 +00:00
Mask api key for Together LLM (#13981)
- **Description:** Add unit tests and mask api key for Together LLM - **Issue:** the issue https://github.com/langchain-ai/langchain/issues/12165 , - **Dependencies:** N/A - **Tag maintainer:** ?, - **Twitter handle:** N/A --------- Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
This commit is contained in:
parent
5f5c701f2c
commit
a578076aea
@ -3,7 +3,7 @@ import logging
|
|||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from aiohttp import ClientSession
|
from aiohttp import ClientSession
|
||||||
from langchain_core.pydantic_v1 import Extra, root_validator
|
from langchain_core.pydantic_v1 import Extra, SecretStr, root_validator
|
||||||
|
|
||||||
from langchain.callbacks.manager import (
|
from langchain.callbacks.manager import (
|
||||||
AsyncCallbackManagerForLLMRun,
|
AsyncCallbackManagerForLLMRun,
|
||||||
@ -11,7 +11,7 @@ from langchain.callbacks.manager import (
|
|||||||
)
|
)
|
||||||
from langchain.llms.base import LLM
|
from langchain.llms.base import LLM
|
||||||
from langchain.utilities.requests import Requests
|
from langchain.utilities.requests import Requests
|
||||||
from langchain.utils import get_from_dict_or_env
|
from langchain.utils import convert_to_secret_str, get_from_dict_or_env
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@ -28,7 +28,7 @@ class Together(LLM):
|
|||||||
|
|
||||||
base_url: str = "https://api.together.xyz/inference"
|
base_url: str = "https://api.together.xyz/inference"
|
||||||
"""Base inference API URL."""
|
"""Base inference API URL."""
|
||||||
together_api_key: str
|
together_api_key: SecretStr
|
||||||
"""Together AI API key. Get it here: https://api.together.xyz/settings/api-keys"""
|
"""Together AI API key. Get it here: https://api.together.xyz/settings/api-keys"""
|
||||||
model: str
|
model: str
|
||||||
"""Model name. Available models listed here:
|
"""Model name. Available models listed here:
|
||||||
@ -69,8 +69,8 @@ class Together(LLM):
|
|||||||
@root_validator(pre=True)
|
@root_validator(pre=True)
|
||||||
def validate_environment(cls, values: Dict) -> Dict:
|
def validate_environment(cls, values: Dict) -> Dict:
|
||||||
"""Validate that api key exists in environment."""
|
"""Validate that api key exists in environment."""
|
||||||
values["together_api_key"] = get_from_dict_or_env(
|
values["together_api_key"] = convert_to_secret_str(
|
||||||
values, "together_api_key", "TOGETHER_API_KEY"
|
get_from_dict_or_env(values, "together_api_key", "TOGETHER_API_KEY")
|
||||||
)
|
)
|
||||||
return values
|
return values
|
||||||
|
|
||||||
@ -116,7 +116,7 @@ class Together(LLM):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
headers = {
|
headers = {
|
||||||
"Authorization": f"Bearer {self.together_api_key}",
|
"Authorization": f"Bearer {self.together_api_key.get_secret_value()}",
|
||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
}
|
}
|
||||||
stop_to_use = stop[0] if stop and len(stop) == 1 else stop
|
stop_to_use = stop[0] if stop and len(stop) == 1 else stop
|
||||||
@ -167,7 +167,7 @@ class Together(LLM):
|
|||||||
The string generated by the model.
|
The string generated by the model.
|
||||||
"""
|
"""
|
||||||
headers = {
|
headers = {
|
||||||
"Authorization": f"Bearer {self.together_api_key}",
|
"Authorization": f"Bearer {self.together_api_key.get_secret_value()}",
|
||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
}
|
}
|
||||||
stop_to_use = stop[0] if stop and len(stop) == 1 else stop
|
stop_to_use = stop[0] if stop and len(stop) == 1 else stop
|
||||||
|
61
libs/langchain/tests/unit_tests/llms/test_together.py
Normal file
61
libs/langchain/tests/unit_tests/llms/test_together.py
Normal file
@ -0,0 +1,61 @@
|
|||||||
|
"""Test Together LLM"""
|
||||||
|
from typing import cast
|
||||||
|
|
||||||
|
from langchain_core.pydantic_v1 import SecretStr
|
||||||
|
from pytest import CaptureFixture, MonkeyPatch
|
||||||
|
|
||||||
|
from langchain.llms.together import Together
|
||||||
|
|
||||||
|
|
||||||
|
def test_together_api_key_is_secret_string() -> None:
|
||||||
|
"""Test that the API key is stored as a SecretStr."""
|
||||||
|
llm = Together(
|
||||||
|
together_api_key="secret-api-key",
|
||||||
|
model="togethercomputer/RedPajama-INCITE-7B-Base",
|
||||||
|
temperature=0.2,
|
||||||
|
max_tokens=250,
|
||||||
|
)
|
||||||
|
assert isinstance(llm.together_api_key, SecretStr)
|
||||||
|
|
||||||
|
|
||||||
|
def test_together_api_key_masked_when_passed_from_env(
|
||||||
|
monkeypatch: MonkeyPatch, capsys: CaptureFixture
|
||||||
|
) -> None:
|
||||||
|
"""Test that the API key is masked when passed from an environment variable."""
|
||||||
|
monkeypatch.setenv("TOGETHER_API_KEY", "secret-api-key")
|
||||||
|
llm = Together(
|
||||||
|
model="togethercomputer/RedPajama-INCITE-7B-Base",
|
||||||
|
temperature=0.2,
|
||||||
|
max_tokens=250,
|
||||||
|
)
|
||||||
|
print(llm.together_api_key, end="")
|
||||||
|
captured = capsys.readouterr()
|
||||||
|
|
||||||
|
assert captured.out == "**********"
|
||||||
|
|
||||||
|
|
||||||
|
def test_together_api_key_masked_when_passed_via_constructor(
|
||||||
|
capsys: CaptureFixture,
|
||||||
|
) -> None:
|
||||||
|
"""Test that the API key is masked when passed via the constructor."""
|
||||||
|
llm = Together(
|
||||||
|
together_api_key="secret-api-key",
|
||||||
|
model="togethercomputer/RedPajama-INCITE-7B-Base",
|
||||||
|
temperature=0.2,
|
||||||
|
max_tokens=250,
|
||||||
|
)
|
||||||
|
print(llm.together_api_key, end="")
|
||||||
|
captured = capsys.readouterr()
|
||||||
|
|
||||||
|
assert captured.out == "**********"
|
||||||
|
|
||||||
|
|
||||||
|
def test_together_uses_actual_secret_value_from_secretstr() -> None:
|
||||||
|
"""Test that the actual secret value is correctly retrieved."""
|
||||||
|
llm = Together(
|
||||||
|
together_api_key="secret-api-key",
|
||||||
|
model="togethercomputer/RedPajama-INCITE-7B-Base",
|
||||||
|
temperature=0.2,
|
||||||
|
max_tokens=250,
|
||||||
|
)
|
||||||
|
assert cast(SecretStr, llm.together_api_key).get_secret_value() == "secret-api-key"
|
Loading…
Reference in New Issue
Block a user