From 3714b83ee4e8087aa793c8895e9ec5361ffe732e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:06:10 -0700 Subject: [PATCH 01/18] feat gcs log user api key metadata --- litellm/integrations/gcs_bucket.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index 3fb778e2428..a16d952861e 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -13,7 +13,7 @@ from litellm.litellm_core_utils.logging_utils import ( convert_litellm_response_object_to_dict, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import CommonProxyErrors, SpendLogsPayload +from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload class RequestKwargs(TypedDict): @@ -27,6 +27,8 @@ class GCSBucketPayload(TypedDict): response_obj: Optional[Dict] start_time: str end_time: str + response_cost: Optional[float] + spend_log_metadata: str class GCSBucketLogger(CustomLogger): @@ -78,11 +80,12 @@ class GCSBucketLogger(CustomLogger): kwargs, response_obj, start_time_str, end_time_str ) + json_logged_payload = json.dumps(logging_payload) object_name = response_obj["id"] response = await self.async_httpx_client.post( headers=headers, url=f"https://storage.googleapis.com/upload/storage/v1/b/{self.BUCKET_NAME}/o?uploadType=media&name={object_name}", - json=logging_payload, + data=json_logged_payload, ) if response.status_code != 200: @@ -121,6 +124,10 @@ class GCSBucketLogger(CustomLogger): async def get_gcs_payload( self, kwargs, response_obj, start_time, end_time ) -> GCSBucketPayload: + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_logging_payload, + ) + request_kwargs = RequestKwargs( model=kwargs.get("model", None), messages=kwargs.get("messages", None), @@ -131,11 +138,21 @@ class GCSBucketLogger(CustomLogger): response_obj=response_obj ) + _spend_log_payload: SpendLogsPayload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + end_user_id=kwargs.get("end_user_id", None), + ) + gcs_payload: GCSBucketPayload = GCSBucketPayload( request_kwargs=request_kwargs, response_obj=response_dict, start_time=start_time, end_time=end_time, + spend_log_metadata=_spend_log_payload["metadata"], + response_cost=kwargs.get("response_cost", None), ) return gcs_payload From 91fb1ad019b4eadede97fb54b6787347a2487b59 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:07:08 -0700 Subject: [PATCH 02/18] test gcs logging payload --- litellm/tests/test_gcs_bucket.py | 59 +++++++++++++++++++++++++++++--- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index c5a6fb76acc..754b4993424 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -63,7 +63,7 @@ def load_vertex_ai_credentials(): @pytest.mark.asyncio async def test_basic_gcs_logger(): - load_vertex_ai_credentials() + # load_vertex_ai_credentials() gcs_logger = GCSBucketLogger() print("GCSBucketLogger", gcs_logger) @@ -75,6 +75,41 @@ async def test_basic_gcs_logger(): max_tokens=10, user="ishaan-2", mock_response="Hi!", + metadata={ + "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"], + "user_api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", + "user_api_key_alias": None, + "user_api_end_user_max_budget": None, + "litellm_api_version": "0.0.0", + "global_max_parallel_requests": None, + "user_api_key_user_id": "116544810872468347480", + "user_api_key_org_id": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + "user_api_key_metadata": {}, + "requester_ip_address": "127.0.0.1", + "spend_logs_metadata": {"hello": "world"}, + "headers": { + "content-type": "application/json", + "user-agent": "PostmanRuntime/7.32.3", + "accept": "*/*", + "postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4", + "host": "localhost:4000", + "accept-encoding": "gzip, deflate, br", + "connection": "keep-alive", + "content-length": "163", + }, + "endpoint": "http://localhost:4000/chat/completions", + "model_group": "gpt-3.5-turbo", + "deployment": "azure/chatgpt-v-2", + "model_info": { + "id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4", + "db_model": False, + }, + "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", + "caching_groups": None, + "raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-2', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n", + }, ) print("response", response) @@ -83,11 +118,14 @@ async def test_basic_gcs_logger(): # Check if object landed on GCS object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) + print("object from gcs=", object_from_gcs) # convert object_from_gcs from bytes to DICT - object_from_gcs = json.loads(object_from_gcs) - print("object_from_gcs", object_from_gcs) + parsed_data = json.loads(object_from_gcs) + print("object_from_gcs as dict", parsed_data) - gcs_payload = GCSBucketPayload(**object_from_gcs) + print("type of object_from_gcs", type(parsed_data)) + + gcs_payload = GCSBucketPayload(**parsed_data) print("gcs_payload", gcs_payload) @@ -97,6 +135,19 @@ async def test_basic_gcs_logger(): ] assert gcs_payload["response_obj"]["choices"][0]["message"]["content"] == "Hi!" + assert gcs_payload["response_cost"] > 0.0 + + gcs_payload["spend_log_metadata"] = json.loads(gcs_payload["spend_log_metadata"]) + + assert ( + gcs_payload["spend_log_metadata"]["user_api_key"] + == "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ) + assert ( + gcs_payload["spend_log_metadata"]["user_api_key_user_id"] + == "116544810872468347480" + ) + # Delete Object from GCS print("deleting object from GCS") await gcs_logger.delete_gcs_object(object_name=response.id) From ff80e90febd4d726a1bfb9be1baf979ef4867268 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:28:12 -0700 Subject: [PATCH 03/18] feat log responses in folders --- litellm/integrations/gcs_bucket.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index a16d952861e..46f55f8f01c 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -81,7 +81,12 @@ class GCSBucketLogger(CustomLogger): ) json_logged_payload = json.dumps(logging_payload) - object_name = response_obj["id"] + + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") + + # Modify the object_name to include the date-based folder + object_name = f"{current_date}/{response_obj['id']}" response = await self.async_httpx_client.post( headers=headers, url=f"https://storage.googleapis.com/upload/storage/v1/b/{self.BUCKET_NAME}/o?uploadType=media&name={object_name}", From 043c70063f204d9ec212b19351c45c5c5d82b55f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:33:35 -0700 Subject: [PATCH 04/18] tes logging to gcs buckets --- litellm/tests/test_gcs_bucket.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index 754b4993424..607599d903d 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -9,6 +9,7 @@ import json import logging import tempfile import uuid +from datetime import datetime import pytest @@ -116,8 +117,17 @@ async def test_basic_gcs_logger(): await asyncio.sleep(5) + # Get the current date + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") + + # Modify the object_name to include the date-based folder + object_name = f"{current_date}%2F{response.id}" + + print("object_name", object_name) + # Check if object landed on GCS - object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) + object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name) print("object from gcs=", object_from_gcs) # convert object_from_gcs from bytes to DICT parsed_data = json.loads(object_from_gcs) @@ -150,4 +160,4 @@ async def test_basic_gcs_logger(): # Delete Object from GCS print("deleting object from GCS") - await gcs_logger.delete_gcs_object(object_name=response.id) + # await gcs_logger.delete_gcs_object(object_name=response.id) From 04433254fe15e6136efd49c1805b8d863d3e5663 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:34:27 -0700 Subject: [PATCH 05/18] fix gcs test --- litellm/tests/test_gcs_bucket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index 607599d903d..b30978bad5d 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -160,4 +160,4 @@ async def test_basic_gcs_logger(): # Delete Object from GCS print("deleting object from GCS") - # await gcs_logger.delete_gcs_object(object_name=response.id) + await gcs_logger.delete_gcs_object(object_name=object_name) From c9ea8cdf416b660ae1d469ab6f4b43d029f008eb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Aug 2024 16:44:44 -0700 Subject: [PATCH 06/18] fix(cost_calculator.py): fix cost calc --- litellm/cost_calculator.py | 14 +++++++++++--- litellm/tests/test_custom_logger.py | 16 +++++++++++----- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6eec8d3cd5c..a3cb847a4f6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -490,10 +490,18 @@ def completion_cost( isinstance(completion_response, BaseModel) or isinstance(completion_response, dict) ): # tts returns a custom class - if isinstance(completion_response, BaseModel) and not isinstance( - completion_response, litellm.Usage + + usage_obj: Optional[Union[dict, litellm.Usage]] = completion_response.get( + "usage", {} + ) + if isinstance(usage_obj, BaseModel) and not isinstance( + usage_obj, litellm.Usage ): - completion_response = litellm.Usage(**completion_response.model_dump()) + setattr( + completion_response, + "usage", + litellm.Usage(**usage_obj.model_dump()), + ) # get input/output tokens from completion_response prompt_tokens = completion_response.get("usage", {}).get("prompt_tokens", 0) completion_tokens = completion_response.get("usage", {}).get( diff --git a/litellm/tests/test_custom_logger.py b/litellm/tests/test_custom_logger.py index e3407c9e119..465012bffb5 100644 --- a/litellm/tests/test_custom_logger.py +++ b/litellm/tests/test_custom_logger.py @@ -1,11 +1,17 @@ ### What this tests #### -import sys, os, time, inspect, asyncio, traceback +import asyncio +import inspect +import os +import sys +import time +import traceback + import pytest sys.path.insert(0, os.path.abspath("../..")) -from litellm import completion, embedding import litellm +from litellm import completion, embedding from litellm.integrations.custom_logger import CustomLogger @@ -201,7 +207,7 @@ def test_async_custom_handler_stream(): print("complete_streaming_response: ", complete_streaming_response) assert response_in_success_handler == complete_streaming_response except Exception as e: - pytest.fail(f"Error occurred: {e}") + pytest.fail(f"Error occurred: {e}\n{traceback.format_exc()}") # test_async_custom_handler_stream() @@ -457,11 +463,11 @@ async def test_cost_tracking_with_caching(): def test_redis_cache_completion_stream(): - from litellm import Cache - # Important Test - This tests if we can add to streaming cache, when custom callbacks are set import random + from litellm import Cache + try: print("\nrunning test_redis_cache_completion_stream") litellm.set_verbose = True From 5569e8790a516efebcc9333c950950a4b8ce7a70 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:06:10 -0700 Subject: [PATCH 07/18] feat gcs log user api key metadata --- litellm/integrations/gcs_bucket.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index 46f55f8f01c..3a76c6de230 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.logging_utils import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload +from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload class RequestKwargs(TypedDict): @@ -29,6 +30,8 @@ class GCSBucketPayload(TypedDict): end_time: str response_cost: Optional[float] spend_log_metadata: str + response_cost: Optional[float] + spend_log_metadata: str class GCSBucketLogger(CustomLogger): @@ -81,12 +84,7 @@ class GCSBucketLogger(CustomLogger): ) json_logged_payload = json.dumps(logging_payload) - - # Get the current date - current_date = datetime.now().strftime("%Y-%m-%d") - - # Modify the object_name to include the date-based folder - object_name = f"{current_date}/{response_obj['id']}" + object_name = response_obj["id"] response = await self.async_httpx_client.post( headers=headers, url=f"https://storage.googleapis.com/upload/storage/v1/b/{self.BUCKET_NAME}/o?uploadType=media&name={object_name}", @@ -133,6 +131,10 @@ class GCSBucketLogger(CustomLogger): get_logging_payload, ) + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_logging_payload, + ) + request_kwargs = RequestKwargs( model=kwargs.get("model", None), messages=kwargs.get("messages", None), @@ -151,6 +153,14 @@ class GCSBucketLogger(CustomLogger): end_user_id=kwargs.get("end_user_id", None), ) + _spend_log_payload: SpendLogsPayload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + end_user_id=kwargs.get("end_user_id", None), + ) + gcs_payload: GCSBucketPayload = GCSBucketPayload( request_kwargs=request_kwargs, response_obj=response_dict, @@ -158,6 +168,8 @@ class GCSBucketLogger(CustomLogger): end_time=end_time, spend_log_metadata=_spend_log_payload["metadata"], response_cost=kwargs.get("response_cost", None), + spend_log_metadata=_spend_log_payload["metadata"], + response_cost=kwargs.get("response_cost", None), ) return gcs_payload From de9bd4abe68789254cb554d2573cd488dad9fe20 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:07:08 -0700 Subject: [PATCH 08/18] test gcs logging payload --- litellm/tests/test_gcs_bucket.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index b30978bad5d..4fa9d8ef437 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -117,17 +117,8 @@ async def test_basic_gcs_logger(): await asyncio.sleep(5) - # Get the current date - # Get the current date - current_date = datetime.now().strftime("%Y-%m-%d") - - # Modify the object_name to include the date-based folder - object_name = f"{current_date}%2F{response.id}" - - print("object_name", object_name) - # Check if object landed on GCS - object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name) + object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) print("object from gcs=", object_from_gcs) # convert object_from_gcs from bytes to DICT parsed_data = json.loads(object_from_gcs) From 6ddd86abab308b6229bc5cef59db40eac42ef0ba Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:28:12 -0700 Subject: [PATCH 09/18] feat log responses in folders --- litellm/integrations/gcs_bucket.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index 3a76c6de230..c948668eb5c 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -84,7 +84,12 @@ class GCSBucketLogger(CustomLogger): ) json_logged_payload = json.dumps(logging_payload) - object_name = response_obj["id"] + + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") + + # Modify the object_name to include the date-based folder + object_name = f"{current_date}/{response_obj['id']}" response = await self.async_httpx_client.post( headers=headers, url=f"https://storage.googleapis.com/upload/storage/v1/b/{self.BUCKET_NAME}/o?uploadType=media&name={object_name}", From 3dfe2bfd85e6fb619cb4c5148cd91999e486408b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:33:35 -0700 Subject: [PATCH 10/18] tes logging to gcs buckets --- litellm/tests/test_gcs_bucket.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index 4fa9d8ef437..607599d903d 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -117,8 +117,17 @@ async def test_basic_gcs_logger(): await asyncio.sleep(5) + # Get the current date + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") + + # Modify the object_name to include the date-based folder + object_name = f"{current_date}%2F{response.id}" + + print("object_name", object_name) + # Check if object landed on GCS - object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) + object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name) print("object from gcs=", object_from_gcs) # convert object_from_gcs from bytes to DICT parsed_data = json.loads(object_from_gcs) @@ -151,4 +160,4 @@ async def test_basic_gcs_logger(): # Delete Object from GCS print("deleting object from GCS") - await gcs_logger.delete_gcs_object(object_name=object_name) + # await gcs_logger.delete_gcs_object(object_name=response.id) From 86978b7a9bc853a37befec371b946c99bf6238f2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:34:27 -0700 Subject: [PATCH 11/18] fix gcs test --- litellm/tests/test_gcs_bucket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index 607599d903d..b30978bad5d 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -160,4 +160,4 @@ async def test_basic_gcs_logger(): # Delete Object from GCS print("deleting object from GCS") - # await gcs_logger.delete_gcs_object(object_name=response.id) + await gcs_logger.delete_gcs_object(object_name=object_name) From 049f3e1e0c105816f7fe03408ec1095e3e2e44e8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 17:42:04 -0700 Subject: [PATCH 12/18] fix gcs logging test --- litellm/tests/test_gcs_bucket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index b30978bad5d..c21988c73dc 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -64,7 +64,7 @@ def load_vertex_ai_credentials(): @pytest.mark.asyncio async def test_basic_gcs_logger(): - # load_vertex_ai_credentials() + load_vertex_ai_credentials() gcs_logger = GCSBucketLogger() print("GCSBucketLogger", gcs_logger) From f322ffc413cfc252f825493470fabff7bb673b3c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Aug 2024 18:47:25 -0700 Subject: [PATCH 13/18] refactor(test_users.py): refactor test for user info to use mock endpoints --- .../internal_user_endpoints.py | 11 +++++- litellm/tests/test_proxy_server.py | 38 +++++++++++++++++++ tests/test_users.py | 7 ---- 3 files changed, 47 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8e2358c9920..a0e020b11fc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -312,7 +312,7 @@ async def user_info( try: if prisma_client is None: raise Exception( - f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) ## GET USER ROW ## if user_id is not None: @@ -365,7 +365,14 @@ async def user_info( getattr(caller_user_info, "user_role", None) == LitellmUserRoles.PROXY_ADMIN ): - teams_2 = await prisma_client.db.litellm_teamtable.find_many() + from litellm.proxy.management_endpoints.team_endpoints import list_team + + teams_2 = await list_team( + http_request=Request( + scope={"type": "http", "path": "/user/info"}, + ), + user_api_key_dict=user_api_key_dict, + ) else: teams_2 = await prisma_client.get_data( team_id_list=caller_user_info.teams, diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index dee20a273ce..757eef6d62d 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -928,3 +928,41 @@ async def test_create_team_member_add(prisma_client, new_member_method): mock_client.call_args.kwargs["data"]["create"]["budget_duration"] == litellm.internal_user_budget_duration ) + + +@pytest.mark.asyncio +async def test_user_info_team_list(prisma_client): + """Assert user_info for admin calls team_list function""" + from litellm.proxy._types import LiteLLM_UserTable + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.list_team", + new_callable=AsyncMock, + ) as mock_client: + + prisma_client.get_data = AsyncMock( + return_value=LiteLLM_UserTable( + user_role="proxy_admin", + user_id="default_user_id", + max_budget=None, + user_email="", + ) + ) + + try: + await user_info( + user_id=None, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", user_id="default_user_id" + ), + ) + except Exception: + pass + + mock_client.assert_called() diff --git a/tests/test_users.py b/tests/test_users.py index 632dd8f36c1..8113fd0801a 100644 --- a/tests/test_users.py +++ b/tests/test_users.py @@ -99,13 +99,6 @@ async def test_user_info(): ) assert status == 403 - ## check if returned teams as admin == all teams ## - admin_info = await get_user_info( - session=session, get_user="", call_user="sk-1234", view_all=True - ) - all_teams = await list_teams(session=session, i=0) - assert len(admin_info["teams"]) == len(all_teams) - @pytest.mark.asyncio async def test_user_update(): From 93a1335e46f2d69ac5a8ea46a99f050ab7c56e83 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Aug 2024 21:21:40 -0700 Subject: [PATCH 14/18] fix(litellm_pre_call_utils.py): support routing to logging project by api key --- litellm/integrations/gcs_bucket.py | 17 ----- litellm/integrations/langfuse.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 68 +++++++++++++++++-- litellm/tests/test_proxy_server.py | 89 +++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 25 deletions(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index c948668eb5c..46f55f8f01c 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -14,7 +14,6 @@ from litellm.litellm_core_utils.logging_utils import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload -from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload class RequestKwargs(TypedDict): @@ -30,8 +29,6 @@ class GCSBucketPayload(TypedDict): end_time: str response_cost: Optional[float] spend_log_metadata: str - response_cost: Optional[float] - spend_log_metadata: str class GCSBucketLogger(CustomLogger): @@ -136,10 +133,6 @@ class GCSBucketLogger(CustomLogger): get_logging_payload, ) - from litellm.proxy.spend_tracking.spend_tracking_utils import ( - get_logging_payload, - ) - request_kwargs = RequestKwargs( model=kwargs.get("model", None), messages=kwargs.get("messages", None), @@ -158,14 +151,6 @@ class GCSBucketLogger(CustomLogger): end_user_id=kwargs.get("end_user_id", None), ) - _spend_log_payload: SpendLogsPayload = get_logging_payload( - kwargs=kwargs, - response_obj=response_obj, - start_time=start_time, - end_time=end_time, - end_user_id=kwargs.get("end_user_id", None), - ) - gcs_payload: GCSBucketPayload = GCSBucketPayload( request_kwargs=request_kwargs, response_obj=response_dict, @@ -173,8 +158,6 @@ class GCSBucketLogger(CustomLogger): end_time=end_time, spend_log_metadata=_spend_log_payload["metadata"], response_cost=kwargs.get("response_cost", None), - spend_log_metadata=_spend_log_payload["metadata"], - response_cost=kwargs.get("response_cost", None), ) return gcs_payload diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index df4be3a5bcf..7a127f912b6 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -48,7 +48,7 @@ class LangFuseLogger: "secret_key": self.secret_key, "host": self.langfuse_host, "release": self.langfuse_release, - "debug": self.langfuse_debug, + "debug": True, "flush_interval": flush_interval, # flush interval in seconds } diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 13f9475c5c7..631f4769225 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -5,7 +5,12 @@ from fastapi import Request import litellm from litellm._logging import verbose_logger, verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, TeamCallbackMetadata, UserAPIKeyAuth +from litellm.proxy._types import ( + AddTeamCallback, + CommonProxyErrors, + TeamCallbackMetadata, + UserAPIKeyAuth, +) from litellm.types.utils import SupportedCacheControls if TYPE_CHECKING: @@ -59,6 +64,42 @@ def safe_add_api_version_from_query_params(data: dict, request: Request): verbose_logger.error("error checking api version in query params: %s", str(e)) +def convert_key_logging_metadata_to_callback( + data: AddTeamCallback, team_callback_settings_obj: Optional[TeamCallbackMetadata] +) -> TeamCallbackMetadata: + if team_callback_settings_obj is None: + team_callback_settings_obj = TeamCallbackMetadata() + if data.callback_type == "success": + if team_callback_settings_obj.success_callback is None: + team_callback_settings_obj.success_callback = [] + + if data.callback_name not in team_callback_settings_obj.success_callback: + team_callback_settings_obj.success_callback.append(data.callback_name) + elif data.callback_type == "failure": + if team_callback_settings_obj.failure_callback is None: + team_callback_settings_obj.failure_callback = [] + + if data.callback_name not in team_callback_settings_obj.failure_callback: + team_callback_settings_obj.failure_callback.append(data.callback_name) + elif data.callback_type == "success_and_failure": + if team_callback_settings_obj.success_callback is None: + team_callback_settings_obj.success_callback = [] + if team_callback_settings_obj.failure_callback is None: + team_callback_settings_obj.failure_callback = [] + if data.callback_name not in team_callback_settings_obj.success_callback: + team_callback_settings_obj.success_callback.append(data.callback_name) + + if data.callback_name in team_callback_settings_obj.failure_callback: + team_callback_settings_obj.failure_callback.append(data.callback_name) + + for var, value in data.callback_vars.items(): + if team_callback_settings_obj.callback_vars is None: + team_callback_settings_obj.callback_vars = {} + team_callback_settings_obj.callback_vars[var] = litellm.get_secret(value) + + return team_callback_settings_obj + + async def add_litellm_data_to_request( data: dict, request: Request, @@ -214,6 +255,7 @@ async def add_litellm_data_to_request( } # add the team-specific configs to the completion call # Team Callbacks controls + callback_settings_obj: Optional[TeamCallbackMetadata] = None if user_api_key_dict.team_metadata is not None: team_metadata = user_api_key_dict.team_metadata if "callback_settings" in team_metadata: @@ -231,13 +273,25 @@ async def add_litellm_data_to_request( } } """ - data["success_callback"] = callback_settings_obj.success_callback - data["failure_callback"] = callback_settings_obj.failure_callback + elif ( + user_api_key_dict.metadata is not None + and "logging" in user_api_key_dict.metadata + ): + for item in user_api_key_dict.metadata["logging"]: - if callback_settings_obj.callback_vars is not None: - # unpack callback_vars in data - for k, v in callback_settings_obj.callback_vars.items(): - data[k] = v + callback_settings_obj = convert_key_logging_metadata_to_callback( + data=AddTeamCallback(**item), + team_callback_settings_obj=callback_settings_obj, + ) + + if callback_settings_obj is not None: + data["success_callback"] = callback_settings_obj.success_callback + data["failure_callback"] = callback_settings_obj.failure_callback + + if callback_settings_obj.callback_vars is not None: + # unpack callback_vars in data + for k, v in callback_settings_obj.callback_vars.items(): + data[k] = v return data diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 757eef6d62d..890446e5663 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -966,3 +966,92 @@ async def test_user_info_team_list(prisma_client): pass mock_client.assert_called() + + +@pytest.mark.asyncio +async def test_add_callback_via_key(prisma_client): + """ + Test if callback specified in key, is used. + """ + global headers + import json + + from fastapi import HTTPException, Request, Response + from starlette.datastructures import URL + + from litellm.proxy.proxy_server import chat_completion + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + litellm.set_verbose = True + + try: + # Your test data + test_data = { + "model": "azure/chatgpt-v-2", + "messages": [ + {"role": "user", "content": "write 1 sentence poem"}, + ], + "max_tokens": 10, + "mock_response": "Hello world", + "api_key": "my-fake-key", + } + + request = Request(scope={"type": "http", "method": "POST", "headers": {}}) + request._url = URL(url="/chat/completions") + + json_bytes = json.dumps(test_data).encode("utf-8") + + request._body = json_bytes + + with patch.object( + litellm.litellm_core_utils.litellm_logging, + "LangFuseLogger", + new=MagicMock(), + ) as mock_client: + resp = await chat_completion( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth( + metadata={ + "logging": [ + { + "callback_name": "langfuse", # 'otel', 'langfuse', 'lunary' + "callback_type": "success", # set, if required by integration - future improvement, have logging tools work for success + failure by default + "callback_vars": { + "langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY", + "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY", + "langfuse_host": "https://us.cloud.langfuse.com", + }, + } + ] + } + ), + ) + print(resp) + mock_client.assert_called() + mock_client.return_value.log_event.assert_called() + args, kwargs = mock_client.return_value.log_event.call_args + print("KWARGS - {}".format(kwargs)) + kwargs = kwargs["kwargs"] + print(kwargs) + assert "user_api_key_metadata" in kwargs["litellm_params"]["metadata"] + assert ( + "logging" + in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"] + ) + checked_keys = False + for item in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"][ + "logging" + ]: + for k, v in item["callback_vars"].items(): + print("k={}, v={}".format(k, v)) + if "key" in k: + assert "os.environ" in v + checked_keys = True + + assert checked_keys + except Exception as e: + pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") From b1cf46faaa00c0662fdf6791722e87702fe12d27 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Aug 2024 23:20:43 -0700 Subject: [PATCH 15/18] fix(langfuse.py'): cleanup --- litellm/integrations/langfuse.py | 2 +- litellm/tests/test_proxy_server.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index 7a127f912b6..df4be3a5bcf 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -48,7 +48,7 @@ class LangFuseLogger: "secret_key": self.secret_key, "host": self.langfuse_host, "release": self.langfuse_release, - "debug": True, + "debug": self.langfuse_debug, "flush_interval": flush_interval, # flush interval in seconds } diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index b943096396f..00c58d1243a 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -1033,9 +1033,7 @@ async def test_add_callback_via_key(prisma_client): mock_client.assert_called() mock_client.return_value.log_event.assert_called() args, kwargs = mock_client.return_value.log_event.call_args - print("KWARGS - {}".format(kwargs)) kwargs = kwargs["kwargs"] - print(kwargs) assert "user_api_key_metadata" in kwargs["litellm_params"]["metadata"] assert ( "logging" From 0d0a793e200ae2c1ad756c74377a268d465b17d6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 13 Aug 2024 21:27:59 -0700 Subject: [PATCH 16/18] test(test_proxy_server.py): refactor test to work on ci/cd --- litellm/tests/test_proxy_server.py | 116 ++++++++++++++++++++++++++++- 1 file changed, 115 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 00c58d1243a..92202565714 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -967,6 +967,8 @@ async def test_user_info_team_list(prisma_client): mock_client.assert_called() + +# @pytest.mark.skip(reason="Local test") @pytest.mark.asyncio async def test_add_callback_via_key(prisma_client): """ @@ -1051,4 +1053,116 @@ async def test_add_callback_via_key(prisma_client): assert checked_keys except Exception as e: - pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") \ No newline at end of file + pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") + + +@pytest.mark.asyncio +async def test_add_callback_via_key_litellm_pre_call_utils(prisma_client): + import json + + from fastapi import HTTPException, Request, Response + from starlette.datastructures import URL + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") + + request = Request(scope={"type": "http", "method": "POST", "headers": {}}) + request._url = URL(url="/chat/completions") + + test_data = { + "model": "azure/chatgpt-v-2", + "messages": [ + {"role": "user", "content": "write 1 sentence poem"}, + ], + "max_tokens": 10, + "mock_response": "Hello world", + "api_key": "my-fake-key", + } + + json_bytes = json.dumps(test_data).encode("utf-8") + + request._body = json_bytes + + data = { + "data": { + "model": "azure/chatgpt-v-2", + "messages": [{"role": "user", "content": "write 1 sentence poem"}], + "max_tokens": 10, + "mock_response": "Hello world", + "api_key": "my-fake-key", + }, + "request": request, + "user_api_key_dict": UserAPIKeyAuth( + token=None, + key_name=None, + key_alias=None, + spend=0.0, + max_budget=None, + expires=None, + models=[], + aliases={}, + config={}, + user_id=None, + team_id=None, + max_parallel_requests=None, + metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY", + "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY", + "langfuse_host": "https://us.cloud.langfuse.com", + }, + } + ] + }, + tpm_limit=None, + rpm_limit=None, + budget_duration=None, + budget_reset_at=None, + allowed_cache_controls=[], + permissions={}, + model_spend={}, + model_max_budget={}, + soft_budget_cooldown=False, + litellm_budget_table=None, + org_id=None, + team_spend=None, + team_alias=None, + team_tpm_limit=None, + team_rpm_limit=None, + team_max_budget=None, + team_models=[], + team_blocked=False, + soft_budget=None, + team_model_aliases=None, + team_member_spend=None, + team_metadata=None, + end_user_id=None, + end_user_tpm_limit=None, + end_user_rpm_limit=None, + end_user_max_budget=None, + last_refreshed_at=None, + api_key=None, + user_role=None, + allowed_model_region=None, + parent_otel_span=None, + ), + "proxy_config": proxy_config, + "general_settings": {}, + "version": "0.0.0", + } + + new_data = await add_litellm_data_to_request(**data) + + assert "success_callback" in new_data + assert new_data["success_callback"] == ["langfuse"] + assert "langfuse_public_key" in new_data + assert "langfuse_secret_key" in new_data From 2b7a64ee288b221d24273bd312265e76ee0e6281 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 13 Aug 2024 21:36:16 -0700 Subject: [PATCH 17/18] test(test_proxy_server.py): skip local test --- litellm/tests/test_proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 92202565714..9a1c0912675 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -968,7 +968,7 @@ async def test_user_info_team_list(prisma_client): mock_client.assert_called() -# @pytest.mark.skip(reason="Local test") +@pytest.mark.skip(reason="Local test") @pytest.mark.asyncio async def test_add_callback_via_key(prisma_client): """ From edbe9e0741d00865f323ae129e51ce5816aa3b0b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 14 Aug 2024 09:59:13 -0700 Subject: [PATCH 18/18] test(test_function_call_parsing.py): fix test --- litellm/tests/test_function_call_parsing.py | 108 +++++++++++--------- 1 file changed, 57 insertions(+), 51 deletions(-) diff --git a/litellm/tests/test_function_call_parsing.py b/litellm/tests/test_function_call_parsing.py index d223a7c8f6b..fab9cf110c9 100644 --- a/litellm/tests/test_function_call_parsing.py +++ b/litellm/tests/test_function_call_parsing.py @@ -1,23 +1,27 @@ # What is this? ## Test to make sure function call response always works with json.loads() -> no extra parsing required. Relevant issue - https://github.com/BerriAI/litellm/issues/2654 -import sys, os +import os +import sys import traceback + from dotenv import load_dotenv load_dotenv() -import os, io +import io +import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import pytest -import litellm import json import warnings - -from litellm import completion from typing import List +import pytest + +import litellm +from litellm import completion + # Just a stub to keep the sample code simple class Trade: @@ -78,58 +82,60 @@ def trade(model_name: str) -> List[Trade]: }, } - response = completion( - model_name, - [ - { - "role": "system", - "content": """You are an expert asset manager, managing a portfolio. + try: + response = completion( + model_name, + [ + { + "role": "system", + "content": """You are an expert asset manager, managing a portfolio. - Always use the `trade` function. Make sure that you call it correctly. For example, the following is a valid call: + Always use the `trade` function. Make sure that you call it correctly. For example, the following is a valid call: + ``` + trade({ + "orders": [ + {"action": "buy", "asset": "BTC", "amount": 0.1}, + {"action": "sell", "asset": "ETH", "amount": 0.2} + ] + }) + ``` + + If there are no trades to make, call `trade` with an empty array: + ``` + trade({ "orders": [] }) + ``` + """, + }, + { + "role": "user", + "content": """Manage the portfolio. + + Don't jabber. + + This is the current market data: ``` - trade({ - "orders": [ - {"action": "buy", "asset": "BTC", "amount": 0.1}, - {"action": "sell", "asset": "ETH", "amount": 0.2} - ] - }) + {market_data} ``` - If there are no trades to make, call `trade` with an empty array: + Your portfolio is as follows: ``` - trade({ "orders": [] }) + {portfolio} ``` - """, + """.replace( + "{market_data}", "BTC: 64,000 USD\nETH: 3,500 USD" + ).replace( + "{portfolio}", "USD: 1000, BTC: 0.1, ETH: 0.2" + ), + }, + ], + tools=[tool_spec], + tool_choice={ + "type": "function", + "function": {"name": tool_spec["function"]["name"]}, # type: ignore }, - { - "role": "user", - "content": """Manage the portfolio. - - Don't jabber. - - This is the current market data: - ``` - {market_data} - ``` - - Your portfolio is as follows: - ``` - {portfolio} - ``` - """.replace( - "{market_data}", "BTC: 64,000 USD\nETH: 3,500 USD" - ).replace( - "{portfolio}", "USD: 1000, BTC: 0.1, ETH: 0.2" - ), - }, - ], - tools=[tool_spec], - tool_choice={ - "type": "function", - "function": {"name": tool_spec["function"]["name"]}, # type: ignore - }, - ) - + ) + except litellm.InternalServerError: + pass calls = response.choices[0].message.tool_calls trades = [trade for call in calls for trade in parse_call(call)] return trades