Run ruff

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/src/opentelemetry/instrumentation/vertexai_v2/__init__.py

@@ -18,26 +18,25 @@
 import os
 import types
 from typing import Collection
-from opentelemetry.instrumentation.vertexai.config import Config
-from opentelemetry.instrumentation.vertexai.utils import dont_throw
+
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import context as context_api
-from opentelemetry.trace import get_tracer, SpanKind
-from opentelemetry.trace.status import Status, StatusCode
-
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import (
     _SUPPRESS_INSTRUMENTATION_KEY,
     unwrap,
 )
-
+from opentelemetry.instrumentation.vertexai.config import Config
+from opentelemetry.instrumentation.vertexai.utils import dont_throw
+from opentelemetry.instrumentation.vertexai.version import __version__
 from opentelemetry.semconv_ai import (
     SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
-    SpanAttributes,
     LLMRequestTypeValues,
+    SpanAttributes,
 )
-from opentelemetry.instrumentation.vertexai.version import __version__
+from opentelemetry.trace import SpanKind, get_tracer
+from opentelemetry.trace.status import Status, StatusCode
 
 logger = logging.getLogger(__name__)
 

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/tests/conftest.py

@@ -1,11 +1,14 @@
 """Unit tests configuration module."""
 
 import pytest
+
 from opentelemetry import trace
+from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
 from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
-from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
 
 pytest_plugins = []
 

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/tests/disabled_test_bison.py

@@ -2,8 +2,13 @@
 
 import pytest
 import vertexai
+from vertexai.language_models import (
+    ChatModel,
+    InputOutputTextPair,
+    TextGenerationModel,
+)
+
 from opentelemetry.semconv_ai import SpanAttributes
-from vertexai.language_models import TextGenerationModel, ChatModel, InputOutputTextPair
 
 vertexai.init()
 
@@ -31,14 +36,17 @@ def test_vertexai_predict(exporter):
 
     vertexai_span = spans[0]
     assert (
-        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "text-bison@001"
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "text-bison@001"
     )
     assert (
         "Give me ten interview questions for the role of program manager."
         in vertexai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.user"]
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.8
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert (
         vertexai_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"]
@@ -74,14 +82,17 @@ async def async_predict_text() -> str:
 
     vertexai_span = spans[0]
     assert (
-        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "text-bison@001"
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "text-bison@001"
     )
     assert (
         "Give me ten interview questions for the role of program manager."
         in vertexai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.user"]
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.8
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert (
         vertexai_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"]
@@ -111,13 +122,18 @@ def test_vertexai_stream(exporter):
     ]
 
     vertexai_span = spans[0]
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "text-bison"
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "text-bison"
+    )
     assert (
         "Give me ten interview questions for the role of program manager."
         in vertexai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.user"]
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.8
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert vertexai_span.attributes[
         f"{SpanAttributes.LLM_COMPLETIONS}.0.content"
@@ -129,7 +145,9 @@ def test_vertexai_stream_async(exporter):
     async def async_streaming_prediction() -> list:
         """Streaming Text Example with a Large Language Model"""
 
-        text_generation_model = TextGenerationModel.from_pretrained("text-bison")
+        text_generation_model = TextGenerationModel.from_pretrained(
+            "text-bison"
+        )
         parameters = {
             "max_output_tokens": 256,
             "top_p": 0.8,
@@ -151,13 +169,18 @@ async def async_streaming_prediction() -> list:
     ]
 
     vertexai_span = spans[0]
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "text-bison"
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "text-bison"
+    )
     assert (
         "Give me ten interview questions for the role of program manager."
         in vertexai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.user"]
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.8
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert vertexai_span.attributes[
         f"{SpanAttributes.LLM_COMPLETIONS}.0.content"
@@ -197,14 +220,17 @@ def test_vertexai_chat(exporter):
 
     vertexai_span = spans[0]
     assert (
-        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "chat-bison@001"
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "chat-bison@001"
     )
     assert (
         "How many planets are there in the solar system?"
         in vertexai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.user"]
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.95
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert (
         vertexai_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"]
@@ -247,11 +273,16 @@ def test_vertexai_chat_stream(exporter):
 
     vertexai_span = spans[0]
     assert (
-        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "chat-bison@001"
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MODEL]
+        == "chat-bison@001"
     )
     assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.95
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] == 0.8
-    assert vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] == 0.8
+    )
+    assert (
+        vertexai_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 256
+    )
     assert vertexai_span.attributes[SpanAttributes.LLM_TOP_K] == 40
     assert vertexai_span.attributes[
         f"{SpanAttributes.LLM_COMPLETIONS}.0.content"

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/tests/disabled_test_gemini.py

@@ -1,8 +1,9 @@
 import pytest
 import vertexai
-from opentelemetry.semconv_ai import SpanAttributes
 from vertexai.preview.generative_models import GenerativeModel, Part
 
+from opentelemetry.semconv_ai import SpanAttributes
+
 vertexai.init()