From f77882948da2739e66c33f7424dd79ccb997ef01 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 24 Jan 2025 21:27:22 -0800 Subject: [PATCH] test_init_custom_logger_compatible_class_as_callback --- .../logging_callback_tests/test_unit_tests_init_callbacks.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 53ad2f7196..7d77e26aaf 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -28,6 +28,7 @@ from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger from litellm.integrations.opik.opik import OpikLogger from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.mlflow import MlflowLogger @@ -70,6 +71,7 @@ callback_class_str_to_classType = { "langfuse": LangfusePromptManagement, "otel": OpenTelemetry, "pagerduty": PagerDutyAlerting, + "gcs_pubsub": GcsPubSubLogger, } expected_env_vars = { @@ -90,6 +92,8 @@ expected_env_vars = { "ARIZE_API_KEY": "arize_api_key", "ARGILLA_API_KEY": "argilla_api_key", "PAGERDUTY_API_KEY": "pagerduty_api_key", + "GCS_PUBSUB_TOPIC_ID": "gcs_pubsub_topic_id", + "GCS_PUBSUB_PROJECT_ID": "gcs_pubsub_project_id", }