fix custom logger

This commit is contained in:
Ishaan Jaff 2025-01-17 07:39:49 -08:00
parent 632ba92af1
commit 9d2707ecfe

View file

@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Union
import litellm import litellm
from litellm._logging import verbose_logger from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client, get_async_httpx_client,
httpxSpecialProvider, httpxSpecialProvider,
@ -24,7 +25,7 @@ from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import StandardLoggingPayload from litellm.types.utils import StandardLoggingPayload
class DataDogLLMObsLogger(CustomBatchLogger): class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
def __init__(self, **kwargs): def __init__(self, **kwargs):
try: try:
verbose_logger.debug("DataDogLLMObs: Initializing logger") verbose_logger.debug("DataDogLLMObs: Initializing logger")
@ -52,7 +53,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
asyncio.create_task(self.periodic_flush()) asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock() self.flush_lock = asyncio.Lock()
self.log_queue: List[LLMObsPayload] = [] self.log_queue: List[LLMObsPayload] = []
super().__init__(**kwargs, flush_lock=self.flush_lock) CustomBatchLogger.__init__(self, **kwargs, flush_lock=self.flush_lock)
except Exception as e: except Exception as e:
verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}") verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}")
raise e raise e
@ -89,11 +90,8 @@ class DataDogLLMObsLogger(CustomBatchLogger):
"data": DDIntakePayload( "data": DDIntakePayload(
type="span", type="span",
attributes=DDSpanAttributes( attributes=DDSpanAttributes(
ml_app="litellm", ml_app=self._get_datadog_service(),
tags=[ tags=[self._get_datadog_tags()],
"service:litellm",
f"env:{os.getenv('DD_ENV', 'production')}",
],
spans=self.log_queue, spans=self.log_queue,
), ),
), ),