Merge pull request #400 from NVIDIA/fix/autorization-header-leak-for-…

…nemollm Fix Authorization header leak for NeMo LLM.
NVIDIA · Mar 13, 2024 · 28f59c5 · 28f59c5
2 parents 89a5007 + 1138710
commit 28f59c5
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/nemoguardrails/llm/providers/nemollm.py b/nemoguardrails/llm/providers/nemollm.py
@@ -195,7 +195,13 @@ def _generate(
             generations=generations,
             llm_output={
                 "url": self._get_request_url(),
-                "headers": self._get_request_headers(),
+                "headers": {
+                    k: v
+                    for k, v in self._get_request_headers().items()
+                    # We make sure the Authorization header is not returned as this
+                    # can lean the authorization key.
+                    if k != "Authorization"
+                },
                 "model_name": self.model,
             },
         )
@@ -293,7 +299,13 @@ async def _agenerate(
             generations=generations,
             llm_output={
                 "url": self._get_request_url(),
-                "headers": self._get_request_headers(),
+                "headers": {
+                    k: v
+                    for k, v in self._get_request_headers().items()
+                    # We make sure the Authorization header is not returned as this
+                    # can lean the authorization key.
+                    if k != "Authorization"
+                },
                 "model_name": self.model,
             },
         )