From b982d7c9f55e2130f563a91ca06088a499ca4b6b Mon Sep 17 00:00:00 2001 From: raspawar Date: Thu, 6 Jun 2024 15:55:57 +0530 Subject: [PATCH 1/6] ensure api_key is not present in client obj --- .../langchain_nvidia_ai_endpoints/_common.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index 3ba8d3c1..dfb6f8ed 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -96,13 +96,11 @@ class Config: { "call": { "Accept": "application/json", - "Authorization": "Bearer {api_key}", "User-Agent": "langchain-nvidia-ai-endpoints", }, "stream": { "Accept": "text/event-stream", "Content-Type": "application/json", - "Authorization": "Bearer {api_key}", "User-Agent": "langchain-nvidia-ai-endpoints", }, }, @@ -121,13 +119,7 @@ def lc_secrets(self) -> Dict[str, str]: @property def headers(self) -> dict: """Return headers with API key injected""" - headers_ = self.headers_tmpl.copy() - for header in headers_.values(): - if "{api_key}" in header["Authorization"] and self.api_key: - header["Authorization"] = header["Authorization"].format( - api_key=self.api_key.get_secret_value(), - ) - return headers_ + return self.headers_tmpl.copy() @validator("base_url") def _validate_base_url(cls, v: str) -> str: @@ -195,12 +187,20 @@ def _post( """Method for posting to the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, - "headers": self.headers["call"], + # "headers": self.headers["call"], "json": self.payload_fn(payload), "stream": False, } + headers = { + "Authorization": f"Bearer {self.api_key.get_secret_value()}" + if self.api_key + else None, + **self.headers["call"], + } session = self.get_session_fn() - self.last_response = response = session.post(**self.last_inputs) + self.last_response = response = session.post( + headers=headers, **self.last_inputs + ) self._try_raise(response) return response, session @@ -212,13 +212,19 @@ def _get( """Method for getting from the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, - "headers": self.headers["call"], + # "headers": self.headers["call"], "stream": False, } if payload: self.last_inputs["json"] = self.payload_fn(payload) + headers = { + "Authorization": f"Bearer {self.api_key.get_secret_value()}" + if self.api_key + else None, + **self.headers["call"], + } session = self.get_session_fn() - self.last_response = response = session.get(**self.last_inputs) + self.last_response = response = session.get(headers=headers, **self.last_inputs) self._try_raise(response) return response, session From c97a31e0733e0e61cc5d289b5dc09ccab72d9c10 Mon Sep 17 00:00:00 2001 From: raspawar Date: Thu, 6 Jun 2024 16:27:49 +0530 Subject: [PATCH 2/6] api_key fix for stream requests --- .../langchain_nvidia_ai_endpoints/_common.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index dfb6f8ed..5ef43a7b 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -442,11 +442,17 @@ def get_req_stream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, - "headers": self.headers["stream"], + # "headers": self.headers["stream"], "json": self.payload_fn(payload), "stream": True, } - response = self.get_session_fn().post(**self.last_inputs) + headers = { + "Authorization": f"Bearer {self.api_key.get_secret_value()}" + if self.api_key + else None, + **self.headers["stream"], + } + response = self.get_session_fn().post(headers=headers, **self.last_inputs) self._try_raise(response) call = self.copy() @@ -477,11 +483,17 @@ async def get_req_astream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, - "headers": self.headers["stream"], + # "headers": self.headers["stream"], "json": self.payload_fn(payload), } + headers = { + "Authorization": f"Bearer {self.api_key.get_secret_value()}" + if self.api_key + else None, + **self.headers["stream"], + } async with self.get_asession_fn() as session: - async with session.post(**self.last_inputs) as response: + async with session.post(headers=headers, **self.last_inputs) as response: self._try_raise(response) async for line in response.content.iter_any(): if line and line.strip() != b"data: [DONE]": From 5eba3a69e50422744e4b8a0480a48564e35ad958 Mon Sep 17 00:00:00 2001 From: raspawar Date: Fri, 7 Jun 2024 21:44:54 +0530 Subject: [PATCH 3/6] remove duplicate code --- .../langchain_nvidia_ai_endpoints/_common.py | 46 ++++++++----------- 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index 5ef43a7b..68055f77 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -143,6 +143,13 @@ def _validate_model(cls, values: Dict[str, Any]) -> Dict[str, Any]: ) return values + def _add_authorization(self) -> Dict: + return { + "Authorization": f"Bearer {self.api_key.get_secret_value()}" + if self.api_key + else None, + } + @property def available_models(self) -> list[Model]: """List the available models that can be invoked.""" @@ -187,16 +194,11 @@ def _post( """Method for posting to the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, - # "headers": self.headers["call"], "json": self.payload_fn(payload), "stream": False, } - headers = { - "Authorization": f"Bearer {self.api_key.get_secret_value()}" - if self.api_key - else None, - **self.headers["call"], - } + headers = self._add_authorization() + headers.update(**self.headers["call"]) session = self.get_session_fn() self.last_response = response = session.post( headers=headers, **self.last_inputs @@ -212,17 +214,13 @@ def _get( """Method for getting from the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, - # "headers": self.headers["call"], "stream": False, } if payload: self.last_inputs["json"] = self.payload_fn(payload) - headers = { - "Authorization": f"Bearer {self.api_key.get_secret_value()}" - if self.api_key - else None, - **self.headers["call"], - } + + headers = self._add_authorization() + headers.update(**self.headers["call"]) session = self.get_session_fn() self.last_response = response = session.get(headers=headers, **self.last_inputs) self._try_raise(response) @@ -442,16 +440,12 @@ def get_req_stream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, - # "headers": self.headers["stream"], "json": self.payload_fn(payload), "stream": True, } - headers = { - "Authorization": f"Bearer {self.api_key.get_secret_value()}" - if self.api_key - else None, - **self.headers["stream"], - } + + headers = self._add_authorization() + headers.update(**self.headers["stream"]) response = self.get_session_fn().post(headers=headers, **self.last_inputs) self._try_raise(response) call = self.copy() @@ -483,15 +477,11 @@ async def get_req_astream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, - # "headers": self.headers["stream"], "json": self.payload_fn(payload), } - headers = { - "Authorization": f"Bearer {self.api_key.get_secret_value()}" - if self.api_key - else None, - **self.headers["stream"], - } + + headers = self._add_authorization() + headers.update(**self.headers["stream"]) async with self.get_asession_fn() as session: async with session.post(headers=headers, **self.last_inputs) as response: self._try_raise(response) From c8724cee7d00a9154e51e99004ee3717be5901fe Mon Sep 17 00:00:00 2001 From: raspawar Date: Fri, 7 Jun 2024 21:55:01 +0530 Subject: [PATCH 4/6] test case for api_key --- .../langchain_nvidia_ai_endpoints/_common.py | 10 +++++----- libs/ai-endpoints/tests/unit_tests/test_api_key.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index 68055f77..dcfb4e8f 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -143,7 +143,7 @@ def _validate_model(cls, values: Dict[str, Any]) -> Dict[str, Any]: ) return values - def _add_authorization(self) -> Dict: + def __add_authorization(self) -> Dict: return { "Authorization": f"Bearer {self.api_key.get_secret_value()}" if self.api_key @@ -197,7 +197,7 @@ def _post( "json": self.payload_fn(payload), "stream": False, } - headers = self._add_authorization() + headers = self.__add_authorization() headers.update(**self.headers["call"]) session = self.get_session_fn() self.last_response = response = session.post( @@ -219,7 +219,7 @@ def _get( if payload: self.last_inputs["json"] = self.payload_fn(payload) - headers = self._add_authorization() + headers = self.__add_authorization() headers.update(**self.headers["call"]) session = self.get_session_fn() self.last_response = response = session.get(headers=headers, **self.last_inputs) @@ -444,7 +444,7 @@ def get_req_stream( "stream": True, } - headers = self._add_authorization() + headers = self.__add_authorization() headers.update(**self.headers["stream"]) response = self.get_session_fn().post(headers=headers, **self.last_inputs) self._try_raise(response) @@ -480,7 +480,7 @@ async def get_req_astream( "json": self.payload_fn(payload), } - headers = self._add_authorization() + headers = self.__add_authorization() headers.update(**self.headers["stream"]) async with self.get_asession_fn() as session: async with session.post(headers=headers, **self.last_inputs) as response: diff --git a/libs/ai-endpoints/tests/unit_tests/test_api_key.py b/libs/ai-endpoints/tests/unit_tests/test_api_key.py index 2400d5c9..c1d23324 100644 --- a/libs/ai-endpoints/tests/unit_tests/test_api_key.py +++ b/libs/ai-endpoints/tests/unit_tests/test_api_key.py @@ -3,6 +3,7 @@ from typing import Any, Generator import pytest +from langchain_core.pydantic_v1 import SecretStr @contextmanager @@ -43,3 +44,13 @@ def get_api_key(instance: Any) -> str: assert get_api_key(public_class(nvidia_api_key="PARAM")) == "PARAM" assert get_api_key(public_class(api_key="PARAM")) == "PARAM" assert get_api_key(public_class(api_key="LOW", nvidia_api_key="HIGH")) == "HIGH" + + +def test_api_key_type(public_class: type) -> None: + # Test case to make sure the api_key is SecretStr and not str + def get_api_key(instance: Any) -> str: + return instance._client.client.api_key + + with no_env_var("NVIDIA_API_KEY"): + os.environ["NVIDIA_API_KEY"] = "ENV" + assert type(get_api_key(public_class())) == SecretStr From d22978053d6684b380df0ae09e2e8bcdb8df4d7f Mon Sep 17 00:00:00 2001 From: raspawar Date: Mon, 10 Jun 2024 09:36:39 +0530 Subject: [PATCH 5/6] add deepcopy impl for last_inputs --- .../langchain_nvidia_ai_endpoints/_common.py | 49 +++++++++++-------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index dcfb4e8f..9e4a6f11 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -5,6 +5,7 @@ import os import time import warnings +from copy import deepcopy from typing import ( Any, AsyncIterator, @@ -96,11 +97,13 @@ class Config: { "call": { "Accept": "application/json", + "Authorization": "Bearer {api_key}", "User-Agent": "langchain-nvidia-ai-endpoints", }, "stream": { "Accept": "text/event-stream", "Content-Type": "application/json", + "Authorization": "Bearer {api_key}", "User-Agent": "langchain-nvidia-ai-endpoints", }, }, @@ -119,7 +122,13 @@ def lc_secrets(self) -> Dict[str, str]: @property def headers(self) -> dict: """Return headers with API key injected""" - return self.headers_tmpl.copy() + headers_ = self.headers_tmpl.copy() + for header in headers_.values(): + if "{api_key}" in header["Authorization"] and self.api_key: + header["Authorization"] = header["Authorization"].format( + api_key=self.api_key, + ) + return headers_ @validator("base_url") def _validate_base_url(cls, v: str) -> str: @@ -143,12 +152,12 @@ def _validate_model(cls, values: Dict[str, Any]) -> Dict[str, Any]: ) return values - def __add_authorization(self) -> Dict: - return { - "Authorization": f"Bearer {self.api_key.get_secret_value()}" - if self.api_key - else None, - } + def __add_authorization(self, payload: dict) -> dict: + if self.api_key: + payload["headers"].update( + {"Authorization": f"Bearer {self.api_key.get_secret_value()}"} + ) + return payload @property def available_models(self) -> list[Model]: @@ -194,15 +203,13 @@ def _post( """Method for posting to the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, + "headers": self.headers["call"], "json": self.payload_fn(payload), "stream": False, } - headers = self.__add_authorization() - headers.update(**self.headers["call"]) + payload = self.__add_authorization(deepcopy(self.last_inputs)) session = self.get_session_fn() - self.last_response = response = session.post( - headers=headers, **self.last_inputs - ) + self.last_response = response = session.post(**payload) self._try_raise(response) return response, session @@ -214,15 +221,15 @@ def _get( """Method for getting from the AI Foundation Model Function API.""" self.last_inputs = { "url": invoke_url, + "headers": self.headers["call"], "stream": False, } if payload: self.last_inputs["json"] = self.payload_fn(payload) - headers = self.__add_authorization() - headers.update(**self.headers["call"]) + payload = self.__add_authorization(deepcopy(self.last_inputs)) session = self.get_session_fn() - self.last_response = response = session.get(headers=headers, **self.last_inputs) + self.last_response = response = session.get(**payload) self._try_raise(response) return response, session @@ -440,13 +447,13 @@ def get_req_stream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, + "headers": self.headers["stream"], "json": self.payload_fn(payload), "stream": True, } - headers = self.__add_authorization() - headers.update(**self.headers["stream"]) - response = self.get_session_fn().post(headers=headers, **self.last_inputs) + payload = self.__add_authorization(deepcopy(self.last_inputs)) + response = self.get_session_fn().post(**payload) self._try_raise(response) call = self.copy() @@ -477,13 +484,13 @@ async def get_req_astream( payload = {**payload, "stream": True} self.last_inputs = { "url": invoke_url, + "headers": self.headers["stream"], "json": self.payload_fn(payload), } - headers = self.__add_authorization() - headers.update(**self.headers["stream"]) + payload = self.__add_authorization(deepcopy(self.last_inputs)) async with self.get_asession_fn() as session: - async with session.post(headers=headers, **self.last_inputs) as response: + async with session.post(**payload) as response: self._try_raise(response) async for line in response.content.iter_any(): if line and line.strip() != b"data: [DONE]": From c484b4c4707f7c9ebd729d99939f23cb9c384f45 Mon Sep 17 00:00:00 2001 From: raspawar Date: Mon, 10 Jun 2024 20:54:45 +0530 Subject: [PATCH 6/6] test case for validation of key --- .../langchain_nvidia_ai_endpoints/_common.py | 20 +++++++++++-------- .../tests/integration_tests/test_api_key.py | 20 +++++++++++++++++++ 2 files changed, 32 insertions(+), 8 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py index 9e4a6f11..5a560f54 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py @@ -207,9 +207,10 @@ def _post( "json": self.payload_fn(payload), "stream": False, } - payload = self.__add_authorization(deepcopy(self.last_inputs)) session = self.get_session_fn() - self.last_response = response = session.post(**payload) + self.last_response = response = session.post( + **self.__add_authorization(deepcopy(self.last_inputs)) + ) self._try_raise(response) return response, session @@ -227,9 +228,10 @@ def _get( if payload: self.last_inputs["json"] = self.payload_fn(payload) - payload = self.__add_authorization(deepcopy(self.last_inputs)) session = self.get_session_fn() - self.last_response = response = session.get(**payload) + self.last_response = response = session.get( + **self.__add_authorization(deepcopy(self.last_inputs)) + ) self._try_raise(response) return response, session @@ -452,8 +454,9 @@ def get_req_stream( "stream": True, } - payload = self.__add_authorization(deepcopy(self.last_inputs)) - response = self.get_session_fn().post(**payload) + response = self.get_session_fn().post( + **self.__add_authorization(deepcopy(self.last_inputs)) + ) self._try_raise(response) call = self.copy() @@ -488,9 +491,10 @@ async def get_req_astream( "json": self.payload_fn(payload), } - payload = self.__add_authorization(deepcopy(self.last_inputs)) async with self.get_asession_fn() as session: - async with session.post(**payload) as response: + async with session.post( + **self.__add_authorization(deepcopy(self.last_inputs)) + ) as response: self._try_raise(response) async for line in response.content.iter_any(): if line and line.strip() != b"data: [DONE]": diff --git a/libs/ai-endpoints/tests/integration_tests/test_api_key.py b/libs/ai-endpoints/tests/integration_tests/test_api_key.py index 335124ea..b78ad1c3 100644 --- a/libs/ai-endpoints/tests/integration_tests/test_api_key.py +++ b/libs/ai-endpoints/tests/integration_tests/test_api_key.py @@ -3,6 +3,7 @@ import pytest from langchain_core.documents import Document +from langchain_core.messages import HumanMessage from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings, NVIDIARerank @@ -49,3 +50,22 @@ def test_api_key(public_class: type, param: str) -> None: with no_env_var("NVIDIA_API_KEY"): client = public_class(**{param: api_key}) contact_service(client) + + +def test_api_key_leakage(chat_model: str, mode: dict) -> None: + """Test ChatNVIDIA wrapper.""" + chat = ChatNVIDIA(model=chat_model, temperature=0.7, **mode) + message = HumanMessage(content="Hello") + chat.invoke([message]) + + # check last_input post request + last_inputs = chat._client.client.last_inputs + assert last_inputs + + authorization_header = last_inputs.get("headers", {}).get("Authorization") + + if authorization_header: + key = authorization_header.split("Bearer ")[1] + + assert not key.startswith("nvapi-") + assert key == "**********"