test: fix tests

Author: shah-siddd

examples/tracing/litellm/test_local_litellm.py

@@ -342,19 +342,19 @@ def get_model_parameters(kwargs: Dict[str, Any]) -> Dict[str, Any]:
 
 def create_trace_args(
     end_time: float,
-    inputs: Dict,
+    inputs: Dict[str, Any],
     output: str,
     latency: float,
     tokens: int,
     prompt_tokens: int,
     completion_tokens: int,
     model: str,
-    model_parameters: Optional[Dict] = None,
-    metadata: Optional[Dict] = None,
+    model_parameters: Optional[Dict[str, Any]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
     raw_output: Optional[str] = None,
     id: Optional[str] = None,
     cost: Optional[float] = None,
-) -> Dict:
+) -> Dict[str, Any]:
     """Returns a dictionary with the trace arguments."""
     trace_args = {
         "end_time": end_time,

src/openlayer/lib/integrations/litellm_tracer.py

@@ -1,6 +1,7 @@
 """Test LiteLLM integration."""
 
 import builtins
+from typing import Any, Dict
 from unittest.mock import Mock, patch
 
 import pytest
@@ -30,7 +31,7 @@ def test_trace_litellm_raises_import_error_without_dependency(self):
 
     @patch('openlayer.lib.integrations.litellm_tracer.HAVE_LITELLM', True)
     @patch('openlayer.lib.integrations.litellm_tracer.litellm')
-    def test_trace_litellm_patches_completion(self, mock_litellm):
+    def test_trace_litellm_patches_completion(self, mock_litellm: Mock) -> None:
         """Test that trace_litellm successfully patches litellm.completion."""
         from openlayer.lib.integrations.litellm_tracer import trace_litellm
 
@@ -177,7 +178,7 @@ def test_create_trace_args(self):
         """Test trace arguments creation."""
         from openlayer.lib.integrations.litellm_tracer import create_trace_args
 
-        args = create_trace_args(
+        args: Dict[str, Any] = create_trace_args(
             end_time=1234567890.0,
             inputs={"prompt": "test"},
             output="response",
@@ -218,7 +219,7 @@ def test_lib_init_trace_litellm_import_error(self):
 
         # Mock import to fail for litellm specifically
         original_import = builtins.__import__
-        def mock_import(name, *args, **kwargs):
+        def mock_import(name: str, *args: Any, **kwargs: Any) -> Any:
             if name == 'litellm':
                 raise ImportError("No module named 'litellm'")
             return original_import(name, *args, **kwargs)
@@ -261,7 +262,7 @@ def test_extract_litellm_metadata(self):
 
     @patch('openlayer.lib.integrations.litellm_tracer.HAVE_LITELLM', True)
     @patch('openlayer.lib.integrations.litellm_tracer.litellm')
-    def test_detect_provider_with_litellm_method(self, mock_litellm):
+    def test_detect_provider_with_litellm_method(self, mock_litellm: Mock) -> None:
         """Test provider detection using LiteLLM's get_llm_provider method."""
         from openlayer.lib.integrations.litellm_tracer import detect_provider_from_response