fmt

DataDog · May 1, 2024 · 54c0620 · 54c0620
1 parent db30e50
commit 54c0620
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/tests/llmobs/test_llmobs_decorators.py b/tests/llmobs/test_llmobs_decorators.py
@@ -20,7 +20,10 @@ def mock_logs():
 
 def test_llm_decorator_with_llmobs_disabled_logs_warning(LLMObs, mock_logs):
     for decorator_name, decorator in (("llm", llm), ("embedding", embedding)):
-        @decorator(model_name="test_model", model_provider="test_provider", name="test_function", session_id="test_session_id")
+
+        @decorator(
+            model_name="test_model", model_provider="test_provider", name="test_function", session_id="test_session_id"
+        )
         def f():
             pass
 
@@ -32,7 +35,11 @@ def f():
 
 def test_non_llm_decorator_with_llmobs_disabled_logs_warning(LLMObs, mock_logs):
     for decorator_name, decorator in (
-        ("task", task), ("workflow", workflow), ("tool", tool), ("agent", agent), ("retrieval", retrieval),
+        ("task", task),
+        ("workflow", workflow),
+        ("tool", tool),
+        ("agent", agent),
+        ("retrieval", retrieval),
     ):
 
         @decorator(name="test_function", session_id="test_session_id")

diff --git a/tests/llmobs/test_llmobs_service.py b/tests/llmobs/test_llmobs_service.py
@@ -665,8 +665,9 @@ def test_ml_app_override(LLMObs, mock_llmobs_span_writer):
         pass
     mock_llmobs_span_writer.enqueue.assert_called_with(
         _expected_llmobs_non_llm_span_event(span, "retrieval", tags={"ml_app": "test_app"})
+    )
+
 
-
 def test_export_span_llmobs_not_enabled_raises_warning(LLMObs, mock_logs):
     LLMObs.disable()
     LLMObs.export_span()