Use secretstr for api keys for javelin-ai-gateway (#13417)

- Make javelin_ai_gateway_api_key a SecretStr --------- Co-authored-by: Hiroshi Tashiro <[email protected]>
langchain-ai · Nov 15, 2023 · accadcc · accadcc
1 parent ba501b2
commit accadcc
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 31 deletions.
diff --git a/libs/langchain/langchain/chat_models/javelin_ai_gateway.py b/libs/langchain/langchain/chat_models/javelin_ai_gateway.py
@@ -1,12 +1,12 @@
 import logging
-from typing import Any, Dict, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional, cast
 
 from langchain.callbacks.manager import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain.chat_models.base import BaseChatModel
-from langchain.pydantic_v1 import BaseModel, Extra
+from langchain.pydantic_v1 import BaseModel, Extra, SecretStr
 from langchain.schema import (
     ChatGeneration,
     ChatResult,
@@ -65,7 +65,7 @@ class ChatJavelinAIGateway(BaseChatModel):
     client: Any
     """javelin client."""
 
-    javelin_api_key: Optional[str] = None
+    javelin_api_key: Optional[SecretStr] = None
     """The API key for the Javelin AI Gateway."""
 
     def __init__(self, **kwargs: Any):
@@ -84,7 +84,8 @@ def __init__(self, **kwargs: Any):
         if self.gateway_uri:
             try:
                 self.client = JavelinClient(
-                    base_url=self.gateway_uri, api_key=self.javelin_api_key
+                    base_url=self.gateway_uri,
+                    api_key=cast(SecretStr, self.javelin_api_key).get_secret_value(),
                 )
             except UnauthorizedError as e:
                 raise ValueError("Javelin: Incorrect API Key.") from e
@@ -93,7 +94,7 @@ def __init__(self, **kwargs: Any):
     def _default_params(self) -> Dict[str, Any]:
         params: Dict[str, Any] = {
             "gateway_uri": self.gateway_uri,
-            "javelin_api_key": self.javelin_api_key,
+            "javelin_api_key": cast(SecretStr, self.javelin_api_key).get_secret_value(),
             "route": self.route,
             **(self.params.dict() if self.params else {}),
         }

diff --git a/libs/langchain/poetry.lock b/libs/langchain/poetry.lock
diff --git a/libs/langchain/pyproject.toml b/libs/langchain/pyproject.toml
@@ -140,6 +140,7 @@ rspace_client = {version = "^2.5.0", optional = true}
 upstash-redis = {version = "^0.15.0", optional = true}
 google-cloud-documentai = {version = "^2.20.1", optional = true}
 fireworks-ai = {version = "^0.6.0", optional = true, python = ">=3.9,<4.0"}
+javelin-sdk = {version = "^0.1.8", optional = true}
 
 
 [tool.poetry.group.test.dependencies]
@@ -373,6 +374,7 @@ extended_testing = [
  "upstash-redis",
  "rspace_client",
  "fireworks-ai",
+ "javelin-sdk",
 ]
 
 [tool.ruff]

diff --git a/libs/langchain/tests/unit_tests/chat_models/test_javelin_ai_gateway.py b/libs/langchain/tests/unit_tests/chat_models/test_javelin_ai_gateway.py
@@ -0,0 +1,32 @@
+"""Test `Javelin AI Gateway` chat models"""
+
+import pytest
+
+from langchain.chat_models import ChatJavelinAIGateway
+from langchain.pydantic_v1 import SecretStr
+
+
+@pytest.mark.requires("javelin_sdk")
+def test_api_key_is_secret_string() -> None:
+    llm = ChatJavelinAIGateway(
+        gateway_uri="<javelin-ai-gateway-uri>",
+        route="<javelin-ai-gateway-chat-route>",
+        javelin_api_key="secret-api-key",
+        params={"temperature": 0.1},
+    )
+    assert isinstance(llm.javelin_api_key, SecretStr)
+    assert llm.javelin_api_key.get_secret_value() == "secret-api-key"
+
+
+@pytest.mark.requires("javelin_sdk")
+def test_api_key_masked_when_passed_via_constructor() -> None:
+    llm = ChatJavelinAIGateway(
+        gateway_uri="<javelin-ai-gateway-uri>",
+        route="<javelin-ai-gateway-chat-route>",
+        javelin_api_key="secret-api-key",
+        params={"temperature": 0.1},
+    )
+
+    assert str(llm.javelin_api_key) == "**********"
+    assert "secret-api-key" not in repr(llm.javelin_api_key)
+    assert "secret-api-key" not in repr(llm)