phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #4757 from BerriAI/litellm_langsmith_log_user_id

Browse source 
[Feat] Proxy + Langsmith - Log user_api_key_user_id, user_api_key_team_alias
This commit is contained in:
Ishaan Jaff 2024-07-17 16:40:42 -07:00 • committed by GitHub
commit bf793598b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 29 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/integrations/langsmith.py
View file

@ -27,6 +27,11 @@ class LangsmithInputs(BaseModel):
original_response: Optional[str] = None original_response: Optional[str] = None
response_cost: Optional[float] = 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): def is_serializable(value):
non_serializable_types = ( non_serializable_types = (
@ -57,6 +62,13 @@ class LangsmithLogger:
kwargs.get("litellm_params", {}).get("metadata", {}) or {} kwargs.get("litellm_params", {}).get("metadata", {}) or {}
) # if metadata is None ) # 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 # set project name and run_name for langsmith logging
# users can pass project_name and run name to litellm.completion() # 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"}) # Example: litellm.completion(model, messages, metadata={"project_name": "my-litellm-project", "run_name": "my-langsmith-run"})

18
litellm/tests/test_langsmith.py
View file

@ -26,7 +26,18 @@ def test_langsmith_logging():
messages=[{"role": "user", "content": "what llm are u"}], messages=[{"role": "user", "content": "what llm are u"}],
max_tokens=10, max_tokens=10,
temperature=0.2, 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) print(response)
time.sleep(3) 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 input_fields_on_langsmith
assert "api_key" not in extra_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: except Exception as e:
print(e) print(e)