Merge pull request #4754 from BerriAI/litellm_fix_langsmith_api_key_logged

[Fix] Langsmith - Don't Log Provider API Keys

litellm/integrations/langsmith.py

@@ -1,14 +1,33 @@
 #### What this does ####
 #    On success, logs events to Langsmith
-import dotenv, os  # type: ignore
-import requests  # type: ignore
-from datetime import datetime
-import traceback
 import asyncio
+import os
+import traceback
 import types
+from datetime import datetime
+from typing import Any, List, Optional
+
+import dotenv  # type: ignore
+import requests  # type: ignore
 from pydantic import BaseModel  # type: ignore
 
 
+class LangsmithInputs(BaseModel):
+    model: Optional[str] = None
+    messages: Optional[List[Any]] = None
+    stream: Optional[bool] = None
+    call_type: Optional[str] = None
+    litellm_call_id: Optional[str] = None
+    completion_start_time: Optional[datetime] = None
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    custom_llm_provider: Optional[str] = None
+    input: Optional[List[Any]] = None
+    log_event_type: Optional[str] = None
+    original_response: Optional[str] = None
+    response_cost: Optional[float] = None
+
+
 def is_serializable(value):
     non_serializable_types = (
         types.CoroutineType,
@@ -27,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
@@ -41,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}"
         )
@@ -52,10 +75,11 @@ class LangsmithLogger:
             print_verbose(
                 f"Langsmith Logging - Enters logging function for model {kwargs}"
             )
-            import requests
             import datetime
             from datetime import timezone
 
+            import requests
+
             try:
                 start_time = kwargs["start_time"].astimezone(timezone.utc).isoformat()
                 end_time = kwargs["end_time"].astimezone(timezone.utc).isoformat()
@@ -64,6 +88,9 @@ class LangsmithLogger:
                 end_time = datetime.datetime.utcnow().isoformat()
 
             # filter out kwargs to not include any dicts, langsmith throws an erros when trying to log kwargs
+            logged_kwargs = LangsmithInputs(**kwargs)
+            kwargs = logged_kwargs.model_dump()
+
             new_kwargs = {}
             for key in kwargs:
                 value = kwargs[key]
@@ -88,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"
@@ -103,8 +131,19 @@ class LangsmithLogger:
             else:
                 print_verbose("Run successfully created")
             print_verbose(
-                f"Langsmith Layer Logging - final response object: {response_obj}"
+                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()

litellm/proxy/proxy_config.yaml

@@ -17,9 +17,7 @@ model_list:
 
 general_settings: 
   master_key: sk-1234 
-  litellm_key_header_name: "X-Litellm-Key"
 
 litellm_settings:
-  cache: true
-  callbacks: ["otel"]
+  success_callback: ["langsmith"]
 

litellm/tests/test_langsmith.py

@@ -1,26 +1,54 @@
-import sys
-import os
 import io
+import os
+import sys
 
 sys.path.insert(0, os.path.abspath("../.."))
 
-from litellm import completion
 import litellm
+from litellm import completion
+from litellm.integrations.langsmith import LangsmithLogger
 
-litellm.success_callback = ["langsmith"]
 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(
             model="claude-instant-1.2",
             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)
 
@@ -30,17 +58,15 @@ def test_langsmith_logging():
 
 def test_langsmith_logging_with_metadata():
     try:
+        litellm.success_callback = ["langsmith"]
         response = completion(
             model="gpt-3.5-turbo",
             messages=[{"role": "user", "content": "what llm are u"}],
             max_tokens=10,
             temperature=0.2,
-            metadata={
-                "run_name": "litellmRUN",
-                "project_name": "litellm-completion",
-            },
         )
         print(response)
+        time.sleep(3)
     except Exception as e:
         print(e)
 
@@ -50,21 +76,19 @@ def test_langsmith_logging_with_metadata():
 
 def test_langsmith_logging_with_streaming_and_metadata():
     try:
+        litellm.success_callback = ["langsmith"]
         response = completion(
             model="gpt-3.5-turbo",
             messages=[{"role": "user", "content": "what llm are u"}],
             max_tokens=10,
             temperature=0.2,
-            metadata={
-                "run_name": "litellmRUN",
-                "project_name": "litellm-completion",
-            },
             stream=True,
         )
         for chunk in response:
             continue
+        time.sleep(3)
     except Exception as e:
         print(e)
 
 
-test_langsmith_logging_with_streaming_and_metadata()
+# test_langsmith_logging_with_streaming_and_metadata()