diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 44e1dd888..d8b01c344 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -27,6 +27,11 @@ class LangsmithInputs(BaseModel): original_response: Optional[str] = None response_cost: Optional[float] = None + # LiteLLM Virtual Key specific fields + user_api_key: Optional[str] = None + user_api_key_user_id: Optional[str] = None + user_api_key_team_alias: Optional[str] = None + def is_serializable(value): non_serializable_types = ( @@ -57,6 +62,13 @@ class LangsmithLogger: kwargs.get("litellm_params", {}).get("metadata", {}) or {} ) # if metadata is None + # set user_api_key, user_team_id, user_api_key_user_id + kwargs["user_api_key"] = metadata.get("user_api_key", None) + kwargs["user_api_key_user_id"] = metadata.get("user_api_key_user_id", None) + kwargs["user_api_key_team_alias"] = metadata.get( + "user_api_key_team_alias", None + ) + # set project name and run_name for langsmith logging # users can pass project_name and run name to litellm.completion() # Example: litellm.completion(model, messages, metadata={"project_name": "my-litellm-project", "run_name": "my-langsmith-run"}) diff --git a/litellm/tests/test_langsmith.py b/litellm/tests/test_langsmith.py index a4025c851..8af0c9cbb 100644 --- a/litellm/tests/test_langsmith.py +++ b/litellm/tests/test_langsmith.py @@ -26,7 +26,18 @@ def test_langsmith_logging(): messages=[{"role": "user", "content": "what llm are u"}], max_tokens=10, temperature=0.2, - metadata={"id": run_id}, + metadata={ + "id": run_id, + "user_api_key": "6eb81e014497d89f3cc1aa9da7c2b37bda6b7fea68e4b710d33d94201e68970c", + "user_api_key_alias": "ishaans-langmsith-key", + "user_api_end_user_max_budget": None, + "litellm_api_version": "1.40.19", + "global_max_parallel_requests": None, + "user_api_key_user_id": "admin", + "user_api_key_org_id": None, + "user_api_key_team_id": "dbe2f686-a686-4896-864a-4c3924458709", + "user_api_key_team_alias": "testing-team", + }, ) print(response) time.sleep(3) @@ -49,6 +60,11 @@ def test_langsmith_logging(): assert "api_key" not in input_fields_on_langsmith assert "api_key" not in extra_fields_on_langsmith + # assert user_api_key in extra_fields_on_langsmith + assert "user_api_key" in extra_fields_on_langsmith + assert "user_api_key_user_id" in extra_fields_on_langsmith + assert "user_api_key_team_alias" in extra_fields_on_langsmith + except Exception as e: print(e)