Fix mypy issues

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:
         """

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(

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)

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

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"""

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):
         """

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(

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:

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)

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: