diff --git a/docs/my-website/docs/observability/gcs_bucket_integration.md b/docs/my-website/docs/observability/gcs_bucket_integration.md new file mode 100644 index 00000000000..4894853e3ac --- /dev/null +++ b/docs/my-website/docs/observability/gcs_bucket_integration.md @@ -0,0 +1,111 @@ +import Image from '@theme/IdealImage'; + +# 🪣 Google Cloud Storage Buckets - Logging LLM Input/Output + +Log LLM Logs to [Google Cloud Storage Buckets](https://cloud.google.com/storage?hl=en) + +:::info + +✨ This is an Enterprise only feature [Get Started with Enterprise here](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat) + +::: + + +### Usage + +1. Add `gcs_bucket` to LiteLLM Config.yaml +```yaml +model_list: +- litellm_params: + api_base: https://openai-function-calling-workers.tasslexyz.workers.dev/ + api_key: my-fake-key + model: openai/my-fake-model + model_name: fake-openai-endpoint + +litellm_settings: + callbacks: ["gcs_bucket"] # 👈 KEY CHANGE # 👈 KEY CHANGE +``` + +2. Set required env variables + +```shell +GCS_BUCKET_NAME="" +GCS_PATH_SERVICE_ACCOUNT="/Users/ishaanjaffer/Downloads/adroit-crow-413218-a956eef1a2a8.json" # Add path to service account.json +``` + +3. Start Proxy + +``` +litellm --config /path/to/config.yaml +``` + +4. Test it! + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ +--header 'Content-Type: application/json' \ +--data ' { + "model": "fake-openai-endpoint", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ], + } +' +``` + + +## Expected Logs on GCS Buckets + + + +### Fields Logged on GCS Buckets + +Example payload of a `/chat/completion` request logged on GCS +```json +{ + "request_id": "chatcmpl-3946ddc2-bcfe-43f6-9b8e-2427951de85c", + "call_type": "acompletion", + "api_key": "", + "cache_hit": "None", + "startTime": "2024-08-01T14:27:12.563246", + "endTime": "2024-08-01T14:27:12.572709", + "completionStartTime": "2024-08-01T14:27:12.572709", + "model": "gpt-3.5-turbo", + "user": "", + "team_id": "", + "metadata": "{}", + "cache_key": "Cache OFF", + "spend": 0.000054999999999999995, + "total_tokens": 30, + "prompt_tokens": 10, + "completion_tokens": 20, + "request_tags": "[]", + "end_user": "ishaan-2", + "api_base": "", + "model_group": "", + "model_id": "", + "requester_ip_address": null, + "output": [ + "{\"finish_reason\":\"stop\",\"index\":0,\"message\":{\"content\":\"Hi!\",\"role\":\"assistant\",\"tool_calls\":null,\"function_call\":null}}" + ] +} +``` + +## Getting `service_account.json` from Google Cloud Console + +1. Go to [Google Cloud Console](https://console.cloud.google.com/) +2. Search for IAM & Admin +3. Click on Service Accounts +4. Select a Service Account +5. Click on 'Keys' -> Add Key -> Create New Key -> JSON +6. Save the JSON file and add the path to `GCS_PATH_SERVICE_ACCOUNT` + +## Support & Talk to Founders + +- [Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version) +- [Community Discord 💭](https://discord.gg/wuPM9dRgDw) +- Our numbers 📞 +1 (770) 8783-106 / ‭+1 (412) 618-6238‬ +- Our emails ✉️ ishaan@berri.ai / krrish@berri.ai diff --git a/docs/my-website/docs/proxy/bucket.md b/docs/my-website/docs/proxy/bucket.md new file mode 100644 index 00000000000..963b8a4f6ca --- /dev/null +++ b/docs/my-website/docs/proxy/bucket.md @@ -0,0 +1,129 @@ + +# 🪣 Logging GCS, s3 Buckets + +LiteLLM Supports Logging to the following Cloud Buckets +- (Enterprise) ✨ [Google Cloud Storage Buckets](#logging-proxy-inputoutput-to-google-cloud-storage-buckets) +- (Free OSS) [Amazon s3 Buckets](#logging-proxy-inputoutput---s3-buckets) + +## Logging Proxy Input/Output to Google Cloud Storage Buckets + +Log LLM Logs to [Google Cloud Storage Buckets](https://cloud.google.com/storage?hl=en) + +:::info + +✨ This is an Enterprise only feature [Get Started with Enterprise here](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat) + +::: + + +### Usage + +1. Add `gcs_bucket` to LiteLLM Config.yaml +```yaml +model_list: +- litellm_params: + api_base: https://openai-function-calling-workers.tasslexyz.workers.dev/ + api_key: my-fake-key + model: openai/my-fake-model + model_name: fake-openai-endpoint + +litellm_settings: + callbacks: ["gcs_bucket"] # 👈 KEY CHANGE # 👈 KEY CHANGE +``` + +2. Set required env variables + +```shell +GCS_BUCKET_NAME="" +GCS_PATH_SERVICE_ACCOUNT="/Users/ishaanjaffer/Downloads/adroit-crow-413218-a956eef1a2a8.json" # Add path to service account.json +``` + +3. Start Proxy + +``` +litellm --config /path/to/config.yaml +``` + +4. Test it! + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ +--header 'Content-Type: application/json' \ +--data ' { + "model": "fake-openai-endpoint", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ], + } +' +``` + + +### Expected Logs on GCS Buckets + + + + + +## Logging Proxy Input/Output - s3 Buckets + +We will use the `--config` to set + +- `litellm.success_callback = ["s3"]` + +This will log all successfull LLM calls to s3 Bucket + +**Step 1** Set AWS Credentials in .env + +```shell +AWS_ACCESS_KEY_ID = "" +AWS_SECRET_ACCESS_KEY = "" +AWS_REGION_NAME = "" +``` + +**Step 2**: Create a `config.yaml` file and set `litellm_settings`: `success_callback` + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo +litellm_settings: + success_callback: ["s3"] + s3_callback_params: + s3_bucket_name: logs-bucket-litellm # AWS Bucket Name for S3 + s3_region_name: us-west-2 # AWS Region Name for S3 + s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 + s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 + s3_path: my-test-path # [OPTIONAL] set path in bucket you want to write logs to + s3_endpoint_url: https://s3.amazonaws.com # [OPTIONAL] S3 endpoint URL, if you want to use Backblaze/cloudflare s3 buckets +``` + +**Step 3**: Start the proxy, make a test request + +Start proxy + +```shell +litellm --config config.yaml --debug +``` + +Test Request + +```shell +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --data ' { + "model": "Azure OpenAI GPT-4 East", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] + }' +``` + +Your logs should be available on the specified s3 Bucket diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 81d6c8767a8..c4098bf75f2 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -8,7 +8,6 @@ Log Proxy input, output, and exceptions using: - Langsmith - DataDog - DynamoDB -- s3 Bucket - etc. import Image from '@theme/IdealImage'; @@ -1379,66 +1378,6 @@ Expected output on Datadog -## Logging Proxy Input/Output - s3 Buckets - -We will use the `--config` to set - -- `litellm.success_callback = ["s3"]` - -This will log all successfull LLM calls to s3 Bucket - -**Step 1** Set AWS Credentials in .env - -```shell -AWS_ACCESS_KEY_ID = "" -AWS_SECRET_ACCESS_KEY = "" -AWS_REGION_NAME = "" -``` - -**Step 2**: Create a `config.yaml` file and set `litellm_settings`: `success_callback` - -```yaml -model_list: - - model_name: gpt-3.5-turbo - litellm_params: - model: gpt-3.5-turbo -litellm_settings: - success_callback: ["s3"] - s3_callback_params: - s3_bucket_name: logs-bucket-litellm # AWS Bucket Name for S3 - s3_region_name: us-west-2 # AWS Region Name for S3 - s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 - s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 - s3_path: my-test-path # [OPTIONAL] set path in bucket you want to write logs to - s3_endpoint_url: https://s3.amazonaws.com # [OPTIONAL] S3 endpoint URL, if you want to use Backblaze/cloudflare s3 buckets -``` - -**Step 3**: Start the proxy, make a test request - -Start proxy - -```shell -litellm --config config.yaml --debug -``` - -Test Request - -```shell -curl --location 'http://0.0.0.0:4000/chat/completions' \ - --header 'Content-Type: application/json' \ - --data ' { - "model": "Azure OpenAI GPT-4 East", - "messages": [ - { - "role": "user", - "content": "what llm are you" - } - ] - }' -``` - -Your logs should be available on the specified s3 Bucket - ## Logging Proxy Input/Output - DynamoDB We will use the `--config` to set diff --git a/docs/my-website/img/gcs_bucket.png b/docs/my-website/img/gcs_bucket.png new file mode 100644 index 00000000000..034053da65a Binary files /dev/null and b/docs/my-website/img/gcs_bucket.png differ diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 42a1bfe6d1e..61aca116980 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -47,7 +47,7 @@ const sidebars = { { type: "category", label: "🪢 Logging", - items: ["proxy/logging", "proxy/streaming_logging"], + items: ["proxy/logging", "proxy/bucket", "proxy/streaming_logging"], }, "proxy/team_logging", "proxy/guardrails", @@ -203,6 +203,7 @@ const sidebars = { items: [ "observability/langfuse_integration", "observability/logfire_integration", + "observability/gcs_bucket_integration", "observability/langsmith_integration", "observability/arize_integration", "debugging/local_debugging", diff --git a/litellm/__init__.py b/litellm/__init__.py index cc03867cff2..0bbccf82719 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -46,6 +46,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "galileo", "braintrust", "arize", + "gcs_bucket", ] _known_custom_logger_compatible_callbacks: List = list( get_args(_custom_logger_compatible_callbacks_literal) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py new file mode 100644 index 00000000000..88c885ebdb2 --- /dev/null +++ b/litellm/integrations/gcs_bucket.py @@ -0,0 +1,225 @@ +import json +import os +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +import httpx +from pydantic import BaseModel, Field + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy._types import CommonProxyErrors, SpendLogsPayload + + +class GCSBucketPayload(SpendLogsPayload): + messages: Optional[List] + output: Optional[Union[Dict, str, List]] + + +class GCSBucketLogger(CustomLogger): + def __init__(self) -> None: + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True: + raise ValueError( + f"GCS Bucket logging is a premium feature. Please upgrade to use it. {CommonProxyErrors.not_premium_user.value}" + ) + + self.async_httpx_client = AsyncHTTPHandler( + timeout=httpx.Timeout(timeout=600.0, connect=5.0) + ) + self.path_service_account_json = os.getenv("GCS_PATH_SERVICE_ACCOUNT", None) + self.BUCKET_NAME = os.getenv("GCS_BUCKET_NAME", None) + + if self.BUCKET_NAME is None: + raise ValueError( + "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." + ) + + if self.path_service_account_json is None: + raise ValueError( + "GCS_PATH_SERVICE_ACCOUNT is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_PATH_SERVICE_ACCOUNT' in the environment." + ) + pass + + #### ASYNC #### + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True: + raise ValueError( + f"GCS Bucket logging is a premium feature. Please upgrade to use it. {CommonProxyErrors.not_premium_user.value}" + ) + try: + verbose_logger.debug( + "GCS Logger: async_log_success_event logging kwargs: %s, response_obj: %s", + kwargs, + response_obj, + ) + headers = await self.construct_request_headers() + logging_payload: GCSBucketPayload = await self.get_gcs_payload( + kwargs, response_obj, start_time, end_time + ) + + object_name = logging_payload["request_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, + ) + + if response.status_code != 200: + verbose_logger.error("GCS Bucket logging error: %s", str(response.text)) + + verbose_logger.debug("GCS Bucket response %s", response) + verbose_logger.debug("GCS Bucket status code %s", response.status_code) + verbose_logger.debug("GCS Bucket response.text %s", response.text) + except Exception as e: + verbose_logger.error("GCS Bucket logging error: %s", str(e)) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + pass + + async def construct_request_headers(self) -> Dict[str, str]: + from litellm import vertex_chat_completion + + auth_header, _ = vertex_chat_completion._get_token_and_url( + model="gcs-bucket", + vertex_credentials=self.path_service_account_json, + vertex_project=None, + vertex_location=None, + gemini_api_key=None, + stream=None, + custom_llm_provider="vertex_ai", + api_base=None, + ) + verbose_logger.debug("constructed auth_header %s", auth_header) + headers = { + "Authorization": f"Bearer {auth_header}", # auth_header + "Content-Type": "application/json", + } + + return headers + + 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, + ) + + spend_logs_payload: SpendLogsPayload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + end_user_id=kwargs.get("user"), + ) + + gcs_payload: GCSBucketPayload = GCSBucketPayload( + **spend_logs_payload, messages=None, output=None + ) + gcs_payload["messages"] = kwargs.get("messages", None) + gcs_payload["startTime"] = start_time.isoformat() + gcs_payload["endTime"] = end_time.isoformat() + + if gcs_payload["completionStartTime"] is not None: + gcs_payload["completionStartTime"] = gcs_payload[ # type: ignore + "completionStartTime" # type: ignore + ].isoformat() + + output = None + if response_obj is not None and ( + kwargs.get("call_type", None) == "embedding" + or isinstance(response_obj, litellm.EmbeddingResponse) + ): + output = None + elif response_obj is not None and isinstance( + response_obj, litellm.ModelResponse + ): + output_list = [] + for choice in response_obj.choices: + output_list.append(choice.json()) + output = output_list + elif response_obj is not None and isinstance( + response_obj, litellm.TextCompletionResponse + ): + output_list = [] + for choice in response_obj.choices: + output_list.append(choice.json()) + output = output_list + elif response_obj is not None and isinstance( + response_obj, litellm.ImageResponse + ): + output = response_obj["data"] + elif response_obj is not None and isinstance( + response_obj, litellm.TranscriptionResponse + ): + output = response_obj["text"] + + gcs_payload["output"] = output + return gcs_payload + + async def download_gcs_object(self, object_name): + """ + Download an object from GCS. + + https://cloud.google.com/storage/docs/downloading-objects#download-object-json + """ + try: + headers = await self.construct_request_headers() + url = f"https://storage.googleapis.com/storage/v1/b/{self.BUCKET_NAME}/o/{object_name}?alt=media" + + # Send the GET request to download the object + response = await self.async_httpx_client.get(url=url, headers=headers) + + if response.status_code != 200: + verbose_logger.error( + "GCS object download error: %s", str(response.text) + ) + return None + + verbose_logger.debug( + "GCS object download response status code: %s", response.status_code + ) + + # Return the content of the downloaded object + return response.content + + except Exception as e: + verbose_logger.error("GCS object download error: %s", str(e)) + return None + + async def delete_gcs_object(self, object_name): + """ + Delete an object from GCS. + """ + try: + headers = await self.construct_request_headers() + url = f"https://storage.googleapis.com/storage/v1/b/{self.BUCKET_NAME}/o/{object_name}" + + # Send the DELETE request to delete the object + response = await self.async_httpx_client.delete(url=url, headers=headers) + + if (response.status_code != 200) or (response.status_code != 204): + verbose_logger.error( + "GCS object delete error: %s, status code: %s", + str(response.text), + response.status_code, + ) + return None + + verbose_logger.debug( + "GCS object delete response status code: %s, response: %s", + response.status_code, + response.text, + ) + + # Return the content of the downloaded object + return response.text + + except Exception as e: + verbose_logger.error("GCS object download error: %s", str(e)) + return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 83a88410ea8..619b91c3a91 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -59,6 +59,7 @@ from ..integrations.custom_logger import CustomLogger from ..integrations.datadog import DataDogLogger from ..integrations.dynamodb import DyanmoDBLogger from ..integrations.galileo import GalileoObserve +from ..integrations.gcs_bucket import GCSBucketLogger from ..integrations.greenscale import GreenscaleLogger from ..integrations.helicone import HeliconeLogger from ..integrations.lago import LagoLogger @@ -1987,6 +1988,14 @@ def _init_custom_logger_compatible_class( _langsmith_logger = LangsmithLogger() _in_memory_loggers.append(_langsmith_logger) return _langsmith_logger # type: ignore + elif logging_integration == "gcs_bucket": + for callback in _in_memory_loggers: + if isinstance(callback, GCSBucketLogger): + return callback # type: ignore + + _gcs_bucket_logger = GCSBucketLogger() + _in_memory_loggers.append(_gcs_bucket_logger) + return _gcs_bucket_logger # type: ignore elif logging_integration == "arize": if "ARIZE_SPACE_KEY" not in os.environ: raise ValueError("ARIZE_SPACE_KEY not found in environment variables") @@ -2101,6 +2110,10 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, LangsmithLogger): return callback + elif logging_integration == "gcs_bucket": + for callback in _in_memory_loggers: + if isinstance(callback, GCSBucketLogger): + return callback elif logging_integration == "otel": from litellm.integrations.opentelemetry import OpenTelemetry diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e3a0a4f1c5a..1828a92d2eb 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -129,6 +129,50 @@ class AsyncHTTPHandler: except Exception as e: raise e + async def delete( + self, + url: str, + data: Optional[Union[dict, str]] = None, # type: ignore + json: Optional[dict] = None, + params: Optional[dict] = None, + headers: Optional[dict] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + stream: bool = False, + ): + try: + if timeout is None: + timeout = self.timeout + req = self.client.build_request( + "DELETE", url, data=data, json=json, params=params, headers=headers, timeout=timeout # type: ignore + ) + response = await self.client.send(req, stream=stream) + response.raise_for_status() + return response + except (httpx.RemoteProtocolError, httpx.ConnectError): + # Retry the request with a new session if there is a connection error + new_client = self.create_client(timeout=timeout, concurrent_limit=1) + try: + return await self.single_connection_post_request( + url=url, + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + ) + finally: + await new_client.aclose() + except httpx.HTTPStatusError as e: + setattr(e, "status_code", e.response.status_code) + if stream is True: + setattr(e, "message", await e.response.aread()) + else: + setattr(e, "message", e.response.text) + raise e + except Exception as e: + raise e + async def single_connection_post_request( self, url: str, diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index aec92579bdd..9995373f32e 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -790,6 +790,16 @@ class VertexLLM(BaseLLM): if credentials is not None and isinstance(credentials, str): import google.oauth2.service_account + verbose_logger.debug( + "Vertex: Loading vertex credentials from %s", credentials + ) + verbose_logger.debug( + "Vertex: checking if credentials is a valid path, os.path.exists(%s)=%s, current dir %s", + credentials, + os.path.exists(credentials), + os.getcwd(), + ) + if os.path.exists(credentials): json_obj = json.load(open(credentials)) else: diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index b78832a10aa..4bc9a87692b 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -55,4 +55,4 @@ general_settings: max_response_size_mb: 10 litellm_settings: - success_callback: ["langfuse"] \ No newline at end of file + callbacks: ["gcs_bucket"] # 👈 KEY CHANGE \ No newline at end of file diff --git a/litellm/tests/adroit-crow-413218-bc47f303efc9.json b/litellm/tests/adroit-crow-413218-bc47f303efc9.json new file mode 100644 index 00000000000..e2fd8512b19 --- /dev/null +++ b/litellm/tests/adroit-crow-413218-bc47f303efc9.json @@ -0,0 +1,13 @@ +{ + "type": "service_account", + "project_id": "adroit-crow-413218", + "private_key_id": "", + "private_key": "", + "client_email": "test-adroit-crow@adroit-crow-413218.iam.gserviceaccount.com", + "client_id": "104886546564708740969", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-adroit-crow%40adroit-crow-413218.iam.gserviceaccount.com", + "universe_domain": "googleapis.com" +} diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py new file mode 100644 index 00000000000..53ef96fda68 --- /dev/null +++ b/litellm/tests/test_gcs_bucket.py @@ -0,0 +1,96 @@ +import io +import os +import sys + +sys.path.insert(0, os.path.abspath("../..")) + +import asyncio +import json +import logging +import tempfile +import uuid + +import pytest + +import litellm +from litellm import completion +from litellm._logging import verbose_logger +from litellm.integrations.gcs_bucket import GCSBucketLogger + +verbose_logger.setLevel(logging.DEBUG) + + +def load_vertex_ai_credentials(): + # Define the path to the vertex_key.json file + print("loading vertex ai credentials") + filepath = os.path.dirname(os.path.abspath(__file__)) + vertex_key_path = filepath + "/adroit-crow-413218-bc47f303efc9.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("GCS_PRIVATE_KEY_ID", "") + private_key = os.environ.get("GCS_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"]) + + +@pytest.mark.asyncio +async def test_basic_gcs_logger(): + load_vertex_ai_credentials() + gcs_logger = GCSBucketLogger() + print("GCSBucketLogger", gcs_logger) + + 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!", + ) + + print("response", response) + + await asyncio.sleep(5) + + # Check if object landed on GCS + object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) + # convert object_from_gcs from bytes to DICT + object_from_gcs = json.loads(object_from_gcs) + print("object_from_gcs", object_from_gcs) + + assert object_from_gcs["request_id"] == response.id + assert object_from_gcs["call_type"] == "acompletion" + assert object_from_gcs["model"] == "gpt-3.5-turbo" + + # Delete Object from GCS + print("deleting object from GCS") + await gcs_logger.delete_gcs_object(object_name=response.id)