small fixes, get CI passing

aabmass · Jan 16, 2025 · 5dfc282 · 5dfc282
1 parent 816212b
commit 5dfc282
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 15 deletions.
diff --git a/...pentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py b/...pentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py
@@ -52,7 +52,9 @@ def traced_method(
         args: Any,
         kwargs: Any,
     ):
-        # Use exact parameter signature to handle named vs positional args robustly
+        # Use parameter signature from
+        # https://github.com/googleapis/python-aiplatform/blob/v1.76.0/vertexai/generative_models/_generative_models.py#L595
+        # to handle named vs positional args robustly
         def extract_params(
             contents: ContentsType,
             *,
@@ -62,6 +64,7 @@ def extract_params(
             tool_config: Optional[ToolConfig] = None,
             labels: Optional[dict[str, str]] = None,
             stream: bool = False,
+            **_kwargs: Any,
         ) -> GenerateContentParams:
             return GenerateContentParams(
                 contents=contents,

diff --git a/...pentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py b/...pentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py
@@ -70,7 +70,6 @@ class GenerationConfigDict(TypedDict, total=False):
 
 
 def get_genai_request_attributes(
-    # TODO: use types
     instance: _GenerativeModel,
     params: GenerateContentParams,
     operation_name: GenAIAttributes.GenAiOperationNameValues = GenAIAttributes.GenAiOperationNameValues.CHAT,
@@ -130,10 +129,6 @@ def _get_model_name(instance: _GenerativeModel) -> str:
     return model_name
 
 
-# TODO: Everything below here should be replaced with
-# opentelemetry.instrumentation.genai_utils instead once it is released.
-# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3191
-
 OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = (
     "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
 )

diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py
@@ -81,19 +81,17 @@ def fixture_meter_provider(metric_reader):
 
 @pytest.fixture(autouse=True)
 def vertexai_init(vcr: VCR) -> None:
-    # Unfortunately I couldn't find a nice way to globally reset the global_config for each
-    # test because different vertex submodules reference the global instance directly
-    # https://github.com/googleapis/python-aiplatform/blob/v1.74.0/google/cloud/aiplatform/initializer.py#L687
-    # so this config will leak if we don't call init() for each test.
-
     # When not recording (in CI), don't do any auth. That prevents trying to read application
     # default credentials from the filesystem or metadata server and oauth token exchange. This
     # is not the interesting part of our instrumentation to test.
-    vertex_init_kwargs = {"api_transport": "rest"}
+    credentials = None
+    project = None
     if vcr.record_mode == RecordMode.NONE:
-        vertex_init_kwargs["credentials"] = AnonymousCredentials()
-        vertex_init_kwargs["project"] = FAKE_PROJECT
-    vertexai.init(**vertex_init_kwargs)
+        credentials = AnonymousCredentials()
+        project = FAKE_PROJECT
+    vertexai.init(
+        api_transport="rest", credentials=credentials, project=project
+    )
 
 
 @pytest.fixture