refactor prometheus to be a customLogger class

This commit is contained in:
Ishaan Jaff 2024-08-10 09:28:46 -07:00
parent 20ef72194c
commit ce9af1e1df
2 changed files with 13 additions and 3 deletions

View file

@ -43,6 +43,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"logfire",
"dynamic_rate_limiter",
"langsmith",
"prometheus",
"galileo",
"braintrust",
"arize",

View file

@ -1904,9 +1904,6 @@ def set_callbacks(callback_list, function_id=None):
openMeterLogger = OpenMeterLogger()
elif callback == "datadog":
dataDogLogger = DataDogLogger()
elif callback == "prometheus":
if prometheusLogger is None:
prometheusLogger = PrometheusLogger()
elif callback == "dynamodb":
dynamoLogger = DyanmoDBLogger()
elif callback == "s3":
@ -1980,6 +1977,14 @@ def _init_custom_logger_compatible_class(
_langsmith_logger = LangsmithLogger()
_in_memory_loggers.append(_langsmith_logger)
return _langsmith_logger # type: ignore
elif logging_integration == "prometheus":
for callback in _in_memory_loggers:
if isinstance(callback, PrometheusLogger):
return callback # type: ignore
_prometheus_logger = PrometheusLogger()
_in_memory_loggers.append(_prometheus_logger)
return _prometheus_logger # type: ignore
elif logging_integration == "gcs_bucket":
for callback in _in_memory_loggers:
if isinstance(callback, GCSBucketLogger):
@ -2102,6 +2107,10 @@ def get_custom_logger_compatible_class(
for callback in _in_memory_loggers:
if isinstance(callback, LangsmithLogger):
return callback
elif logging_integration == "prometheus":
for callback in _in_memory_loggers:
if isinstance(callback, PrometheusLogger):
return callback
elif logging_integration == "gcs_bucket":
for callback in _in_memory_loggers:
if isinstance(callback, GCSBucketLogger):