Skip to content

Commit

Permalink
VertexAI emit user events
Browse files Browse the repository at this point in the history
  • Loading branch information
aabmass committed Jan 23, 2025
1 parent ec3c51d commit 208af92
Show file tree
Hide file tree
Showing 7 changed files with 294 additions and 8 deletions.
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))
- VertexAI emit user events
([#3203](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3203))
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
# 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.

"""
Factories for event types described in
https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-events.md#system-event.
Hopefully this code can be autogenerated by Weaver once Gen AI semantic conventions are
schematized in YAML and the Weaver tool supports it.
"""

from opentelemetry._events import Event
from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
from opentelemetry.util.types import AnyValue


def user_event(
*,
role: str = "user",
content: AnyValue = None,
) -> Event:
"""Creates a User event
https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#user-event
"""
body: dict[str, AnyValue] = {
"role": role,
}
if content is not None:
body["content"] = content
return Event(
name="gen_ai.user.message",
attributes={
gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
},
body=body,
)
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
GenerateContentParams,
get_genai_request_attributes,
get_span_name,
request_to_events,
)
from opentelemetry.trace import SpanKind, Tracer

Expand Down Expand Up @@ -107,13 +108,12 @@ def traced_method(
name=span_name,
kind=SpanKind.CLIENT,
attributes=span_attributes,
) as _span:
# TODO: emit request events
# if span.is_recording():
# for message in kwargs.get("messages", []):
# event_logger.emit(
# message_to_event(message, capture_content)
# )
) as span:
if span.is_recording():
for event in request_to_events(
params=params, capture_content=capture_content
):
event_logger.emit(event)

# TODO: set error.type attribute
# https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,18 @@
from os import environ
from typing import (
TYPE_CHECKING,
Iterable,
Mapping,
Sequence,
cast,
)

from opentelemetry._events import Event
from opentelemetry.instrumentation.vertexai.events import user_event
from opentelemetry.semconv._incubating.attributes import (
gen_ai_attributes as GenAIAttributes,
)
from opentelemetry.util.types import AttributeValue
from opentelemetry.util.types import AnyValue, AttributeValue

if TYPE_CHECKING:
from google.cloud.aiplatform_v1.types import content, tool
Expand Down Expand Up @@ -137,3 +141,21 @@ def get_span_name(span_attributes: Mapping[str, AttributeValue]) -> str:
if not model:
return f"{name}"
return f"{name} {model}"


def request_to_events(
*, params: GenerateContentParams, capture_content: bool
) -> Iterable[Event]:
for content in params.contents or []:
if content.role == "model":
# TODO: handle assistant message
pass
# Assume user event but role should be "user"
else:
request_content = None
if capture_content:
request_content = [
cast(dict[str, AnyValue], type(part).to_dict(part)) # type: ignore[reportUnknownMemberType]
for part in content.parts
]
yield user_event(role=content.role, content=request_content)
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
interactions:
- request:
body: |-
{
"contents": [
{
"role": "invalid_role",
"parts": [
{
"text": "Say this is a test"
}
]
}
]
}
headers:
Accept:
- '*/*'
Accept-Encoding:
- gzip, deflate
Connection:
- keep-alive
Content-Length:
- '149'
Content-Type:
- application/json
User-Agent:
- python-requests/2.32.3
method: POST
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-1.5-flash-002:generateContent?%24alt=json%3Benum-encoding%3Dint
response:
body:
string: |-
{
"error": {
"code": 400,
"message": "Please use a valid role: user, model.",
"status": "INVALID_ARGUMENT",
"details": []
}
}
headers:
Content-Type:
- application/json; charset=UTF-8
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
content-length:
- '416'
status:
code: 400
message: Bad Request
version: 1
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
interactions:
- request:
body: |-
{
"contents": [
{
"role": "user",
"parts": [
{
"text": "Say this is a test"
}
]
}
]
}
headers:
Accept:
- '*/*'
Accept-Encoding:
- gzip, deflate
Connection:
- keep-alive
Content-Length:
- '141'
Content-Type:
- application/json
User-Agent:
- python-requests/2.32.3
method: POST
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-1.5-flash-002:generateContent?%24alt=json%3Benum-encoding%3Dint
response:
body:
string: |-
{
"candidates": [
{
"content": {
"role": "model",
"parts": [
{
"text": "Okay, I understand. I'm ready for your test. Please proceed.\n"
}
]
},
"finishReason": 1,
"avgLogprobs": -0.005519990466142956
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 19,
"totalTokenCount": 24
},
"modelVersion": "gemini-1.5-flash-002"
}
headers:
Content-Type:
- application/json; charset=UTF-8
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
content-length:
- '453'
status:
code: 200
message: OK
version: 1
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
)

from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
from opentelemetry.sdk._logs._internal.export.in_memory_log_exporter import (
InMemoryLogExporter,
)
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
Expand All @@ -18,6 +21,7 @@
@pytest.mark.vcr
def test_generate_content(
span_exporter: InMemorySpanExporter,
log_exporter: InMemoryLogExporter,
instrument_with_content: VertexAIInstrumentor,
):
model = GenerativeModel("gemini-1.5-flash-002")
Expand All @@ -27,6 +31,48 @@ def test_generate_content(
]
)

# Emits span
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
assert spans[0].name == "chat gemini-1.5-flash-002"
assert dict(spans[0].attributes) == {
"gen_ai.operation.name": "chat",
"gen_ai.request.model": "gemini-1.5-flash-002",
"gen_ai.system": "vertex_ai",
}

# Emits content event
logs = log_exporter.get_finished_logs()
assert len(logs) == 1
log_record = logs[0].log_record
span_context = spans[0].get_span_context()
assert log_record.trace_id == span_context.trace_id
assert log_record.span_id == span_context.span_id
assert log_record.trace_flags == span_context.trace_flags
assert log_record.attributes == {
"gen_ai.system": "vertex_ai",
"event.name": "gen_ai.user.message",
}
assert log_record.body == {
"content": [{"text": "Say this is a test"}],
"role": "user",
}


@pytest.mark.vcr
def test_generate_content_without_events(
span_exporter: InMemorySpanExporter,
log_exporter: InMemoryLogExporter,
instrument_no_content: VertexAIInstrumentor,
):
model = GenerativeModel("gemini-1.5-flash-002")
model.generate_content(
[
Content(role="user", parts=[Part.from_text("Say this is a test")]),
]
)

# Emits span
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
assert spans[0].name == "chat gemini-1.5-flash-002"
Expand All @@ -36,6 +82,16 @@ def test_generate_content(
"gen_ai.system": "vertex_ai",
}

# Emits event without body.content
logs = log_exporter.get_finished_logs()
assert len(logs) == 1
log_record = logs[0].log_record
assert log_record.attributes == {
"gen_ai.system": "vertex_ai",
"event.name": "gen_ai.user.message",
}
assert log_record.body == {"role": "user"}


@pytest.mark.vcr
def test_generate_content_empty_model(
Expand Down Expand Up @@ -98,6 +154,7 @@ def test_generate_content_missing_model(
@pytest.mark.vcr
def test_generate_content_invalid_temperature(
span_exporter: InMemorySpanExporter,
log_exporter: InMemoryLogExporter,
instrument_with_content: VertexAIInstrumentor,
):
model = GenerativeModel("gemini-1.5-flash-002")
Expand Down Expand Up @@ -126,6 +183,38 @@ def test_generate_content_invalid_temperature(
assert_span_error(spans[0])


@pytest.mark.vcr
def test_generate_content_invalid_role(
log_exporter: InMemoryLogExporter,
instrument_with_content: VertexAIInstrumentor,
):
model = GenerativeModel("gemini-1.5-flash-002")
try:
# Fails because role must be "user" or "model"
model.generate_content(
[
Content(
role="invalid_role",
parts=[Part.from_text("Say this is a test")],
)
]
)
except BadRequest:
pass

# Emits the faulty content which caused the request to fail
logs = log_exporter.get_finished_logs()
assert len(logs) == 1
assert logs[0].log_record.attributes == {
"gen_ai.system": "vertex_ai",
"event.name": "gen_ai.user.message",
}
assert logs[0].log_record.body == {
"content": [{"text": "Say this is a test"}],
"role": "invalid_role",
}


@pytest.mark.vcr()
def test_generate_content_extra_params(span_exporter, instrument_no_content):
generation_config = GenerationConfig(
Expand Down

0 comments on commit 208af92

Please sign in to comment.