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 e164c54
Show file tree
Hide file tree
Showing 6 changed files with 315 additions and 8 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
# 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,
)


def assistant_event(
*,
role: str = "model",
content: AnyValue = None,
) -> Event:
"""Creates an Assistant event
https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#assistant-event
"""
# TODO: add tool_calls once instrumentation supports it

return Event(
name="gen_ai.assistant.message",
attributes={
gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
},
body={
"role": role,
"content": content,
},
)
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 @@ -16,17 +16,22 @@

import re
from dataclasses import dataclass
from enum import Enum
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 +142,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
Loading

0 comments on commit e164c54

Please sign in to comment.