chore: lint

5war00p · Dec 24, 2023 · f98dc0b · f98dc0b
1 parent b5ad45e
commit f98dc0b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/...entelemetry-instrumentation-replicate/opentelemetry/instrumentation/replicate/__init__.py b/...entelemetry-instrumentation-replicate/opentelemetry/instrumentation/replicate/__init__.py
@@ -57,6 +57,7 @@ def _set_span_attribute(span, name, value):
             span.set_attribute(name, value)
     return
 
+
 input_attribute_map = {
     "prompt": f"{SpanAttributes.LLM_PROMPTS}.0.user",
     "temperature": SpanAttributes.LLM_TEMPERATURE,
@@ -137,6 +138,7 @@ def _handle_request(span, args, kwargs):
             "Failed to set input attributes for replicate span, error: %s", str(ex)
         )
 
+
 def _handle_response(span, response):
     try:
         if span.is_recording():
@@ -150,6 +152,7 @@ def _handle_response(span, response):
     if span.is_recording():
         span.set_status(Status(StatusCode.OK))
 
+
 def _with_tracer_wrapper(func):
     """Helper for providing tracer for wrapper functions."""
 

diff --git a/packages/traceloop-sdk/tests/test_replicate_instrumentaion.py b/packages/traceloop-sdk/tests/test_replicate_instrumentaion.py
@@ -5,12 +5,11 @@
 def test_replicate_image_generation(exporter):
     @task(name="image_generation")
     def generate_image():
-        image = replicate.run(
+        replicate.run(
             "stability-ai/stable-diffusion:27b93a2413e7f36cd83da926f3656280b2931564ff050bf9575f1fdf9bcd7478",
-            input={"prompt": "robots"}
+            input={"prompt": "robots"},
         )
 
-
     @workflow(name="robot_image_generator")
     def image_workflow():
         generate_image()
@@ -24,6 +23,7 @@ def image_workflow():
         "robot_image_generator.workflow",
     ]
 
+
 def test_replicate_image_generation_stream(exporter):
     @task(name="image_generation_stream")
     def generate_image_stream():
@@ -49,16 +49,15 @@ def image_workflow():
         "robot_image_generator_stream.workflow",
     ]
 
+
 def test_replicate_image_generation_predictions(exporter):
     @task(name="image_generation_predictions")
     def generate_image_predictions():
         model = replicate.models.get("kvfrans/clipdraw")
-        version = model.versions.get("5797a99edc939ea0e9242d5e8c9cb3bc7d125b1eac21bda852e5cb79ede2cd9b")
-        prediction = replicate.predictions.create(
-            version,
-            input={"prompt": "robots"}
+        version = model.versions.get(
+            "5797a99edc939ea0e9242d5e8c9cb3bc7d125b1eac21bda852e5cb79ede2cd9b"
         )
-
+        replicate.predictions.create(version, input={"prompt": "robots"})
 
     @workflow(name="robot_image_generator_predictions")
     def image_workflow():