From 698a9ddbcea85ae2bee7e30de6eef5744cbca012 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Fri, 30 Jan 2026 10:03:13 -0800 Subject: [PATCH] Fix mypy issues --- .../azure_sentinel/azure_sentinel.py | 3 +- .../azure_storage/azure_storage.py | 3 +- litellm/integrations/custom_batch_logger.py | 75 +++++++++++-------- .../integrations/datadog/datadog_llm_obs.py | 3 +- litellm/integrations/gcs_pubsub/pub_sub.py | 3 +- .../generic_api/generic_api_callback.py | 3 +- litellm/integrations/langsmith.py | 3 +- litellm/integrations/opik/opik.py | 4 +- litellm/integrations/s3_v2.py | 3 +- litellm/integrations/sqs.py | 3 +- 10 files changed, 53 insertions(+), 50 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 875432de876..9643dda5170 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -26,7 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.types.utils import StandardLoggingPayload -class AzureSentinelLogger(CustomBatchLogger): +class AzureSentinelLogger(CustomBatchLogger[StandardLoggingPayload]): """ Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API """ @@ -115,7 +115,6 @@ class AzureSentinelLogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) - self.log_queue: List[StandardLoggingPayload] = [] async def _get_oauth_token(self) -> str: """ diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 85f91199c1c..7980c5f5952 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -18,7 +18,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.types.utils import StandardLoggingPayload -class AzureBlobStorageLogger(CustomBatchLogger): +class AzureBlobStorageLogger(CustomBatchLogger[StandardLoggingPayload]): def __init__( self, **kwargs, @@ -63,7 +63,6 @@ class AzureBlobStorageLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: List[StandardLoggingPayload] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception( diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index b9b46f4db3f..5056af1f5c7 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -4,18 +4,30 @@ Custom Logger that handles batching logic Use this if you want your logs to be stored in memory and flushed periodically. """ +from __future__ import annotations + import asyncio import os import time -from typing import Any, Callable, List, Optional +from typing import Any, Callable, Generic, List, Optional, TypeVar, cast import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_QUEUE_SIZE, DEFAULT_QUEUE_OVERFLOW_STRATEGY from litellm.integrations.custom_logger import CustomLogger +# Type of items stored in the log queue. Subclasses use CustomBatchLogger[TheirPayloadType]. +T = TypeVar("T") + + +class CustomBatchLogger(CustomLogger, Generic[T]): + """ + Base class for batch loggers. Subclass with a type parameter for typed queues: + e.g. class LangsmithLogger(CustomBatchLogger[LangsmithQueueObject]) + """ + + log_queue: BoundedQueue[T] -class CustomBatchLogger(CustomLogger): def __init__( self, flush_lock: Optional[asyncio.Lock] = None, @@ -31,11 +43,14 @@ class CustomBatchLogger(CustomLogger): """ # Use BoundedQueue instead of plain list - all existing append() calls automatically get protection - # Pass flush callback for "force_flush" strategy - self.log_queue = BoundedQueue( - max_size=DEFAULT_MAX_QUEUE_SIZE, - overflow_strategy=DEFAULT_QUEUE_OVERFLOW_STRATEGY, - flush_callback=lambda: asyncio.create_task(self.flush_queue()) + # Pass flush callback for "force_flush" strategy. Cast for generic T (runtime is same class). + self.log_queue = cast( + BoundedQueue[T], + BoundedQueue( + max_size=DEFAULT_MAX_QUEUE_SIZE, + overflow_strategy=DEFAULT_QUEUE_OVERFLOW_STRATEGY, + flush_callback=lambda: asyncio.create_task(self.flush_queue()), + ), ) self.flush_interval = flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE @@ -69,15 +84,12 @@ class CustomBatchLogger(CustomLogger): pass -class BoundedQueue: +class BoundedQueue(Generic[T]): """ A list-like queue with optional size limit to prevent OOM. - - Wraps a list and overrides append() to enforce max_size limit. - This allows all existing code using log_queue.append() to automatically - get OOM protection without any code changes. + Generic over the item type so subclasses get typed append/iterate. """ - + def __init__( self, max_size: Optional[int] = None, @@ -94,7 +106,7 @@ class BoundedQueue: - "force_flush": Trigger flush callback to send logs, then add item - default flush_callback: Optional callback function to trigger flush (required for "force_flush" strategy) """ - self._items: List = [] + self._items: List[T] = [] self.max_size = max_size self.overflow_strategy = overflow_strategy self._dropped_count = 0 @@ -109,7 +121,7 @@ class BoundedQueue: } self._handler = self._strategy_handlers.get(overflow_strategy, self._handle_force_flush) - def _handle_drop_oldest(self, item): + def _handle_drop_oldest(self, item: T) -> None: """Handle overflow by removing oldest item (FIFO).""" self._items.pop(0) self._items.append(item) @@ -119,7 +131,7 @@ class BoundedQueue: f"CustomBatchLogger: Queue full ({self.max_size}), dropped {self._dropped_count} oldest logs" ) - def _handle_drop_newest(self, item): + def _handle_drop_newest(self, item: T) -> None: """Handle overflow by rejecting new item.""" self._dropped_count += 1 if self._dropped_count % 100 == 0: @@ -127,14 +139,14 @@ class BoundedQueue: f"CustomBatchLogger: Queue full ({self.max_size}), rejected {self._dropped_count} new logs" ) - def _handle_reject(self, item): + def _handle_reject(self, item: T) -> None: """Handle overflow by raising RuntimeError.""" raise RuntimeError( f"CustomBatchLogger: Queue full ({self.max_size}). " f"Cannot add more logs. Consider increasing max_queue_size or fixing flush issues." ) - def _handle_force_flush(self, item): + def _handle_force_flush(self, item: T) -> None: """Handle overflow by triggering flush callback, then adding item.""" if self._flush_callback is None: # Fallback to drop_oldest if no flush callback provided @@ -160,7 +172,7 @@ class BoundedQueue: # Add item after flush is triggered (flush is async, but we've scheduled it) self._items.append(item) - def append(self, item): + def append(self, item: T) -> None: """Append item with OOM protection.""" # No limit - allow unlimited growth (backward compatible) if self.max_size is None: @@ -190,11 +202,11 @@ class BoundedQueue: """Iterate over items.""" return iter(self._items) - def __getitem__(self, index): + def __getitem__(self, index: int) -> T: """Get item by index.""" return self._items[index] - - def __setitem__(self, index, value): + + def __setitem__(self, index: int, value: T) -> None: """Set item by index.""" self._items[index] = value @@ -210,24 +222,23 @@ class BoundedQueue: """String representation.""" return f"BoundedQueue(max_size={self.max_size}, size={len(self._items)}, dropped={self._dropped_count})" - # Delegate other list methods as needed - def extend(self, items): + def extend(self, items: List[T]) -> None: """Extend queue with items.""" for item in items: self.append(item) - - def pop(self, index: int = -1): + + def pop(self, index: int = -1) -> T: """Pop item from queue.""" return self._items.pop(index) - - def remove(self, item): + + def remove(self, item: T) -> None: """Remove item from queue.""" self._items.remove(item) - - def index(self, item): + + def index(self, item: T) -> int: """Get index of item.""" return self._items.index(item) - - def count(self, item): + + def count(self, item: T) -> int: """Count occurrences of item.""" return self._items.count(item) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 6ffdbc0a005..75003e38cd6 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -39,7 +39,7 @@ from litellm.types.utils import ( ) -class DataDogLLMObsLogger(CustomBatchLogger): +class DataDogLLMObsLogger(CustomBatchLogger[LLMObsPayload]): def __init__(self, **kwargs): try: verbose_logger.debug("DataDogLLMObs: Initializing logger") @@ -66,7 +66,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: List[LLMObsPayload] = [] ######################################################### # Handle datadog_llm_observability_params set as litellm.datadog_llm_observability_params diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index db7f9bb4d0b..579f4a7607a 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -28,7 +28,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) -class GcsPubSubLogger(CustomBatchLogger): +class GcsPubSubLogger(CustomBatchLogger[Union[SpendLogsPayload, StandardLoggingPayload]]): def __init__( self, project_id: Optional[str] = None, @@ -64,7 +64,6 @@ class GcsPubSubLogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) - self.log_queue: List[Union[SpendLogsPayload, StandardLoggingPayload]] = [] async def construct_request_headers(self) -> Dict[str, str]: """Construct authorization headers using Vertex AI auth""" diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 1c62ce9fcc3..5b00f2dd623 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -95,7 +95,7 @@ def substitute_env_variables(value: str) -> str: return re.sub(pattern, replace_env_var, value) -class GenericAPILogger(CustomBatchLogger): +class GenericAPILogger(CustomBatchLogger[Union[Dict, StandardLoggingPayload]]): def __init__( self, endpoint: Optional[str] = None, @@ -180,7 +180,6 @@ class GenericAPILogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) - self.log_queue: List[Union[Dict, StandardLoggingPayload]] = [] def _get_headers(self, headers: Optional[dict] = None): """ diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 5893f14105d..93e9a8b4b90 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -33,7 +33,7 @@ def is_serializable(value): return not isinstance(value, non_serializable_types) -class LangsmithLogger(CustomBatchLogger): +class LangsmithLogger(CustomBatchLogger[LangsmithQueueObject]): def __init__( self, langsmith_api_key: Optional[str] = None, @@ -70,7 +70,6 @@ class LangsmithLogger(CustomBatchLogger): if _batch_size: self.batch_size = int(_batch_size) - self.log_queue: List[LangsmithQueueObject] = [] asyncio.create_task(self.periodic_flush()) def get_credentials_from_env( diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index 7b687d34d1c..16ee94815ad 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -33,7 +33,7 @@ def _should_skip_event(kwargs: Dict[str, Any]) -> bool: return False -class OpikLogger(CustomBatchLogger): +class OpikLogger(CustomBatchLogger[Dict[str, Any]]): """ Opik Logger for logging events to an Opik Server """ @@ -294,7 +294,7 @@ class OpikLogger(CustomBatchLogger): return # Split the log_queue into traces and spans - traces, spans = utils.get_traces_and_spans_from_payload(self.log_queue) + traces, spans = utils.get_traces_and_spans_from_payload(list(self.log_queue)) # Send trace batch if len(traces) > 0: diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 534b85e4752..2a884bc0c21 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -27,7 +27,7 @@ from litellm.types.utils import StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger -class S3Logger(CustomBatchLogger, BaseAWSLLM): +class S3Logger(CustomBatchLogger[s3BatchLoggingElement], BaseAWSLLM): def __init__( self, s3_bucket_name: Optional[str] = None, @@ -105,7 +105,6 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): flush_interval=s3_flush_interval, batch_size=s3_batch_size, ) - self.log_queue: List[s3BatchLoggingElement] = [] # Call BaseAWSLLM's __init__ BaseAWSLLM.__init__(self) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 97a4c5723d8..0d19ff20d88 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -38,7 +38,7 @@ _BASE64_INLINE_PATTERN = re.compile( ) -class SQSLogger(CustomBatchLogger, BaseAWSLLM): +class SQSLogger(CustomBatchLogger[StandardLoggingPayload], BaseAWSLLM): """Batching logger that writes logs to an AWS SQS queue, optionally encrypting the payload.""" def __init__( @@ -114,7 +114,6 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): batch_size=sqs_batch_size, ) - self.log_queue: List[StandardLoggingPayload] = [] BaseAWSLLM.__init__(self) except Exception as e: