Skip to content

Commit

Permalink
Add server attributes to Vertex AI spans
Browse files Browse the repository at this point in the history
  • Loading branch information
aabmass committed Jan 24, 2025
1 parent 2756c1e commit 152fb20
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,5 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#3192](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3192))
- Initial VertexAI instrumentation
([#3123](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3123))
- Add server attributes to Vertex AI spans
([#3208](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3208))
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from opentelemetry.instrumentation.vertexai.utils import (
GenerateContentParams,
get_genai_request_attributes,
get_server_attributes,
get_span_name,
)
from opentelemetry.trace import SpanKind, Tracer
Expand Down Expand Up @@ -100,7 +101,11 @@ def traced_method(
kwargs: Any,
):
params = _extract_params(*args, **kwargs)
span_attributes = get_genai_request_attributes(params)
api_endpoint: str = instance.api_endpoint # type: ignore[reportUnknownMemberType]
span_attributes = {
**get_genai_request_attributes(params),
**get_server_attributes(api_endpoint),
}

span_name = get_span_name(span_attributes)
with tracer.start_as_current_span(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,12 @@
Mapping,
Sequence,
)
from urllib.parse import urlparse

from opentelemetry.semconv._incubating.attributes import (
gen_ai_attributes as GenAIAttributes,
)
from opentelemetry.semconv.attributes import server_attributes
from opentelemetry.util.types import AttributeValue

if TYPE_CHECKING:
Expand Down Expand Up @@ -58,6 +60,24 @@ class GenerateContentParams:
) = None


def get_server_attributes(
endpoint: str,
) -> dict[str, AttributeValue]:
"""Get server.* attributes from the endpoint, which is a hostname with optional port e.g.
- ``us-central1-aiplatform.googleapis.com``
- ``us-central1-aiplatform.googleapis.com:5431``
"""
parsed = urlparse(f"scheme://{endpoint}")

if not parsed.hostname:
return {}

return {
server_attributes.SERVER_ADDRESS: parsed.hostname,
server_attributes.SERVER_PORT: parsed.port or 443,
}


def get_genai_request_attributes(
params: GenerateContentParams,
operation_name: GenAIAttributes.GenAiOperationNameValues = GenAIAttributes.GenAiOperationNameValues.CHAT,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ def test_generate_content(
"gen_ai.operation.name": "chat",
"gen_ai.request.model": "gemini-1.5-flash-002",
"gen_ai.system": "vertex_ai",
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}


Expand Down Expand Up @@ -62,6 +64,8 @@ def test_generate_content_empty_model(
"gen_ai.operation.name": "chat",
"gen_ai.request.model": "",
"gen_ai.system": "vertex_ai",
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}
assert_span_error(spans[0])

Expand Down Expand Up @@ -91,6 +95,8 @@ def test_generate_content_missing_model(
"gen_ai.operation.name": "chat",
"gen_ai.request.model": "gemini-does-not-exist",
"gen_ai.system": "vertex_ai",
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}
assert_span_error(spans[0])

Expand Down Expand Up @@ -122,6 +128,8 @@ def test_generate_content_invalid_temperature(
"gen_ai.request.model": "gemini-1.5-flash-002",
"gen_ai.request.temperature": 1000.0,
"gen_ai.system": "vertex_ai",
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}
assert_span_error(spans[0])

Expand Down Expand Up @@ -158,6 +166,8 @@ def test_generate_content_extra_params(span_exporter, instrument_no_content):
"gen_ai.request.temperature": 0.20000000298023224,
"gen_ai.request.top_p": 0.949999988079071,
"gen_ai.system": "vertex_ai",
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}


Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# Copyright The OpenTelemetry Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from opentelemetry.instrumentation.vertexai.utils import get_server_attributes


def test_get_server_attributes() -> None:
# without port
assert get_server_attributes("us-central1-aiplatform.googleapis.com") == {
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 443,
}

# with port
assert get_server_attributes(
"us-central1-aiplatform.googleapis.com:5432"
) == {
"server.address": "us-central1-aiplatform.googleapis.com",
"server.port": 5432,
}

0 comments on commit 152fb20

Please sign in to comment.