From df446c52e470c24a3332e2b01c644b5aed905167 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 17 Jul 2024 10:56:33 -0700
Subject: [PATCH] test api_key not logged on langsmith

---
 litellm/integrations/langsmith.py | 16 ++++++++++++++++
 litellm/tests/test_langsmith.py   | 26 ++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py
index 817c0357e2a..44e1dd888eb 100644
--- a/litellm/integrations/langsmith.py
+++ b/litellm/integrations/langsmith.py
@@ -46,6 +46,9 @@ class LangsmithLogger:
         self.langsmith_default_run_name = os.getenv(
             "LANGSMITH_DEFAULT_RUN_NAME", "LLMRun"
         )
+        self.langsmith_base_url = os.getenv(
+            "LANGSMITH_BASE_URL", "https://api.smith.langchain.com"
+        )
 
     def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
         # Method definition
@@ -60,6 +63,7 @@ class LangsmithLogger:
         # if not set litellm will fallback to the environment variable LANGSMITH_PROJECT, then to the default project_name = litellm-completion, run_name = LLMRun
         project_name = metadata.get("project_name", self.langsmith_project)
         run_name = metadata.get("run_name", self.langsmith_default_run_name)
+        run_id = metadata.get("id", None)
         print_verbose(
             f"Langsmith Logging - project_name: {project_name}, run_name {run_name}"
         )
@@ -111,6 +115,7 @@ class LangsmithLogger:
                 "session_name": project_name,
                 "start_time": start_time,
                 "end_time": end_time,
+                "id": run_id,
             }
 
             url = f"{langsmith_base_url}/runs"
@@ -128,6 +133,17 @@ class LangsmithLogger:
             print_verbose(
                 f"Langsmith Layer Logging - final response object: {response_obj}. Response text from langsmith={response.text}"
             )
+            return
         except:
             print_verbose(f"Langsmith Layer Error - {traceback.format_exc()}")
             pass
+
+    def get_run_by_id(self, run_id):
+
+        url = f"{self.langsmith_base_url}/runs/{run_id}"
+        response = requests.get(
+            url=url,
+            headers={"x-api-key": self.langsmith_api_key},
+        )
+
+        return response.json()
diff --git a/litellm/tests/test_langsmith.py b/litellm/tests/test_langsmith.py
index 5f3837b264c..a4025c851f5 100644
--- a/litellm/tests/test_langsmith.py
+++ b/litellm/tests/test_langsmith.py
@@ -6,13 +6,19 @@ sys.path.insert(0, os.path.abspath("../.."))
 
 import litellm
 from litellm import completion
+from litellm.integrations.langsmith import LangsmithLogger
 
 litellm.set_verbose = True
 import time
 
+test_langsmith_logger = LangsmithLogger()
+
 
 def test_langsmith_logging():
     try:
+        import uuid
+
+        run_id = str(uuid.uuid4())
         litellm.set_verbose = True
         litellm.success_callback = ["langsmith"]
         response = completion(
@@ -20,9 +26,29 @@ def test_langsmith_logging():
             messages=[{"role": "user", "content": "what llm are u"}],
             max_tokens=10,
             temperature=0.2,
+            metadata={"id": run_id},
         )
         print(response)
         time.sleep(3)
+
+        print("run_id", run_id)
+        logged_run_on_langsmith = test_langsmith_logger.get_run_by_id(run_id=run_id)
+
+        print("logged_run_on_langsmith", logged_run_on_langsmith)
+
+        print("fields in logged_run_on_langsmith", logged_run_on_langsmith.keys())
+
+        input_fields_on_langsmith = logged_run_on_langsmith.get("inputs")
+        extra_fields_on_langsmith = logged_run_on_langsmith.get("extra")
+
+        print("\nLogged INPUT ON LANGSMITH", input_fields_on_langsmith)
+
+        print("\nextra fields on langsmith", extra_fields_on_langsmith)
+
+        assert input_fields_on_langsmith is not None
+        assert "api_key" not in input_fields_on_langsmith
+        assert "api_key" not in extra_fields_on_langsmith
+
     except Exception as e:
         print(e)