diff --git a/tests/local_testing/test_gcs_bucket.py b/tests/local_testing/test_gcs_bucket.py index 03321627bcb..9d72ff873ad 100644 --- a/tests/local_testing/test_gcs_bucket.py +++ b/tests/local_testing/test_gcs_bucket.py @@ -21,177 +21,68 @@ from litellm.integrations.gcs_bucket.gcs_bucket import ( StandardLoggingPayload, ) from litellm.types.utils import StandardCallbackDynamicParams -from unittest.mock import patch +from litellm.types.integrations.gcs_bucket import GCSLoggingConfig +from unittest.mock import patch, AsyncMock, MagicMock verbose_logger.setLevel(logging.DEBUG) -def load_vertex_ai_credentials(): - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - os.environ["GCS_FLUSH_INTERVAL"] = "1" - os.environ["GCS_USE_BATCHED_LOGGING"] = "false" - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GCS_PATH_SERVICE_ACCOUNT"] = os.path.abspath(temp_file.name) - print("created gcs path service account=", os.environ["GCS_PATH_SERVICE_ACCOUNT"]) +def _make_mock_gcs_logging_config(): + return GCSLoggingConfig( + bucket_name="test-bucket", + vertex_instance=MagicMock(), + path_service_account=None, + ) @pytest.mark.asyncio async def test_aaabasic_gcs_logger(): - load_vertex_ai_credentials() - gcs_logger = GCSBucketLogger() - print("GCSBucketLogger", gcs_logger) + os.environ["GCS_FLUSH_INTERVAL"] = "1" + os.environ["GCS_USE_BATCHED_LOGGING"] = "false" + os.environ["GCS_BUCKET_NAME"] = "test-bucket" - litellm.callbacks = [gcs_logger] - response = await litellm.acompletion( - model="gpt-3.5-turbo", - temperature=0.7, - messages=[{"role": "user", "content": "This is a test"}], - max_tokens=10, - user="ishaan-2", - mock_response="Hi!", - metadata={ - "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"], - "user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456", - "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", - "requester_metadata": {"foo": "bar"}, - "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", - "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-3', '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", - }, - ) + captured_payloads = [] - print("response", response) + async def mock_log_json_data_on_gcs( + self, headers, bucket_name, object_name, logging_payload + ): + captured_payloads.append( + { + "bucket_name": bucket_name, + "object_name": object_name, + "logging_payload": logging_payload, + } + ) + return {"kind": "storage#object", "name": object_name} - await asyncio.sleep(5) + with patch( + "litellm.proxy.proxy_server.premium_user", True + ), patch.object( + GCSBucketLogger, + "construct_request_headers", + new_callable=AsyncMock, + return_value={"Authorization": "Bearer mock_token"}, + ), patch.object( + GCSBucketLogger, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_make_mock_gcs_logging_config(), + ), patch.object( + GCSBucketLogger, + "_log_json_data_on_gcs", + mock_log_json_data_on_gcs, + ): + gcs_logger = GCSBucketLogger() - # 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) - print("object from gcs=", object_from_gcs) - # convert object_from_gcs from bytes to DICT - parsed_data = json.loads(object_from_gcs) - print("object_from_gcs as dict", parsed_data) - - print("type of object_from_gcs", type(parsed_data)) - - gcs_payload = StandardLoggingPayload(**parsed_data) - - print("gcs_payload", gcs_payload) - - assert gcs_payload["model"] == "gpt-3.5-turbo" - assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}] - - assert gcs_payload["response"]["choices"][0]["message"]["content"] == "Hi!" - - assert gcs_payload["response_cost"] > 0.0 - - assert gcs_payload["status"] == "success" - - assert ( - gcs_payload["metadata"]["user_api_key_hash"] - == "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456" - ) - assert gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480" - - assert gcs_payload["metadata"]["requester_metadata"] == {"foo": "bar"} - - # Delete Object from GCS - print("deleting object from GCS") - await gcs_logger.delete_gcs_object(object_name=object_name) - - -@pytest.mark.asyncio -async def test_basic_gcs_logger_failure(): - load_vertex_ai_credentials() - gcs_logger = GCSBucketLogger() - print("GCSBucketLogger", gcs_logger) - - gcs_log_id = f"failure-test-{uuid.uuid4().hex}" - - litellm.callbacks = [gcs_logger] - - try: + litellm.callbacks = [gcs_logger] response = await litellm.acompletion( model="gpt-3.5-turbo", temperature=0.7, messages=[{"role": "user", "content": "This is a test"}], max_tokens=10, user="ishaan-2", - mock_response=litellm.BadRequestError( - model="gpt-3.5-turbo", - message="Error: 400: Bad Request: Invalid API key, please check your API key and try again.", - llm_provider="openai", - ), + mock_response="Hi!", metadata={ - "gcs_log_id": gcs_log_id, "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"], "user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456", "user_api_key_alias": None, @@ -204,6 +95,7 @@ async def test_basic_gcs_logger_failure(): "user_api_key_team_alias": None, "user_api_key_metadata": {}, "requester_ip_address": "127.0.0.1", + "requester_metadata": {"foo": "bar"}, "spend_logs_metadata": {"hello": "world"}, "headers": { "content-type": "application/json", @@ -226,47 +118,153 @@ async def test_basic_gcs_logger_failure(): "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-3', '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", }, ) - except Exception: - pass - await asyncio.sleep(5) + print("response", response) - # Get the current date - # Get the current date - current_date = datetime.now().strftime("%Y-%m-%d") + await asyncio.sleep(3) - # Modify the object_name to include the date-based folder - object_name = gcs_log_id + assert len(captured_payloads) == 1, ( + f"Expected 1 GCS upload, got {len(captured_payloads)}" + ) - print("object_name", object_name) + gcs_payload = captured_payloads[0]["logging_payload"] - # Check if object landed on GCS - 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) - print("object_from_gcs as dict", parsed_data) + assert gcs_payload["model"] == "gpt-3.5-turbo" + assert gcs_payload["messages"] == [ + {"role": "user", "content": "This is a test"} + ] - print("type of object_from_gcs", type(parsed_data)) + assert gcs_payload["response"]["choices"][0]["message"]["content"] == "Hi!" - gcs_payload = StandardLoggingPayload(**parsed_data) + assert gcs_payload["response_cost"] > 0.0 - print("gcs_payload", gcs_payload) + assert gcs_payload["status"] == "success" - assert gcs_payload["model"] == "gpt-3.5-turbo" - assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}] + assert ( + gcs_payload["metadata"]["user_api_key_hash"] + == "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456" + ) + assert ( + gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480" + ) - assert gcs_payload["response_cost"] == 0 - assert gcs_payload["status"] == "failure" - - assert ( - gcs_payload["metadata"]["user_api_key_hash"] - == "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456" - ) - assert gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480" - - # Delete Object from GCS - print("deleting object from GCS") - await gcs_logger.delete_gcs_object(object_name=object_name) + assert gcs_payload["metadata"]["requester_metadata"] == {"foo": "bar"} +@pytest.mark.asyncio +async def test_basic_gcs_logger_failure(): + os.environ["GCS_FLUSH_INTERVAL"] = "1" + os.environ["GCS_USE_BATCHED_LOGGING"] = "false" + os.environ["GCS_BUCKET_NAME"] = "test-bucket" + + captured_payloads = [] + + async def mock_log_json_data_on_gcs( + self, headers, bucket_name, object_name, logging_payload + ): + captured_payloads.append( + { + "bucket_name": bucket_name, + "object_name": object_name, + "logging_payload": logging_payload, + } + ) + return {"kind": "storage#object", "name": object_name} + + gcs_log_id = f"failure-test-{uuid.uuid4().hex}" + + with patch( + "litellm.proxy.proxy_server.premium_user", True + ), patch.object( + GCSBucketLogger, + "construct_request_headers", + new_callable=AsyncMock, + return_value={"Authorization": "Bearer mock_token"}, + ), patch.object( + GCSBucketLogger, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_make_mock_gcs_logging_config(), + ), patch.object( + GCSBucketLogger, + "_log_json_data_on_gcs", + mock_log_json_data_on_gcs, + ): + gcs_logger = GCSBucketLogger() + + litellm.callbacks = [gcs_logger] + + try: + response = await litellm.acompletion( + model="gpt-3.5-turbo", + temperature=0.7, + messages=[{"role": "user", "content": "This is a test"}], + max_tokens=10, + user="ishaan-2", + mock_response=litellm.BadRequestError( + model="gpt-3.5-turbo", + message="Error: 400: Bad Request: Invalid API key, please check your API key and try again.", + llm_provider="openai", + ), + metadata={ + "gcs_log_id": gcs_log_id, + "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"], + "user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456", + "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", + "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-3', '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", + }, + ) + except Exception: + pass + + await asyncio.sleep(3) + + assert len(captured_payloads) == 1, ( + f"Expected 1 GCS upload, got {len(captured_payloads)}" + ) + + gcs_payload = captured_payloads[0]["logging_payload"] + + assert gcs_payload["model"] == "gpt-3.5-turbo" + assert gcs_payload["messages"] == [ + {"role": "user", "content": "This is a test"} + ] + + assert gcs_payload["response_cost"] == 0 + assert gcs_payload["status"] == "failure" + + assert ( + gcs_payload["metadata"]["user_api_key_hash"] + == "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456" + ) + assert ( + gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480" + )