Skip to content

Commit b6068e2

Browse files
michaelsafyanmichaelsafyan
committed
Add more tests.
1 parent 431fdf5 commit b6068e2

File tree

5 files changed

+297
-31
lines changed

5 files changed

+297
-31
lines changed

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from typing import Union, Optional, Iterator, AsyncIterator, Awaitable
22

3+
import logging
34
import time
45
import os
56
import functools
@@ -20,6 +21,9 @@
2021
from opentelemetry.semconv.attributes import error_attributes
2122

2223

24+
_logger = logging.getLogger(__name__)
25+
26+
2327
class _MethodsSnapshot:
2428

2529
def __init__(self):
@@ -189,6 +193,7 @@ def finalize_processing(self):
189193
span.set_attribute(gen_ai_attributes.GEN_AI_RESPONSE_FINISH_REASONS, sorted(list(self._finish_reasons_set)))
190194
self._record_token_usage_metric()
191195
self._record_duration_metric()
196+
self._otel_wrapper.done()
192197

193198
def _maybe_update_token_counts(self, response: GenerateContentResponse):
194199
input_tokens = _get_response_property(response, 'usage_metadata.prompt_token_count')
@@ -199,22 +204,51 @@ def _maybe_update_token_counts(self, response: GenerateContentResponse):
199204
self._output_tokens += output_tokens
200205

201206
def _maybe_update_error_type(self, response: GenerateContentResponse):
202-
pass
207+
if response.candidates:
208+
return
209+
if ((not response.prompt_feedback) or
210+
(not response.prompt_feedback.block_reason) or
211+
(block_reason == genai_types.BlockedReason.BLOCKED_REASON_UNSPECIFIED)):
212+
self._error_type = 'NO_CANDIDATES'
213+
return
214+
block_reason = response.prompt_feedback.block_reason
215+
self._error_type = 'BLOCKED_{}'.format(block_reason.name)
203216

204217
def _maybe_log_system_instruction(self, config: Optional[GenerateContentConfigOrDict]=None):
205218
if not self._content_recording_enabled:
206219
return
207-
pass
220+
system_instruction = _get_config_property(config, 'system_instruction')
221+
if not system_instruction:
222+
return
223+
self._otel_wrapper.log_system_prompt(
224+
attributes={
225+
gen_ai_attributes.GEN_AI_SYSTEM: self._genai_system,
226+
},
227+
body={
228+
'content': system_instruction,
229+
})
208230

209231
def _maybe_log_user_prompt(self, contents: Union[ContentListUnion, ContentListUnionDict]):
210232
if not self._content_recording_enabled:
211233
return
212-
pass
234+
self._otel_wrapper.log_user_prompt(
235+
attributes={
236+
gen_ai_attributes.GEN_AI_SYSTEM: self._genai_system,
237+
},
238+
body={
239+
'content': contents,
240+
})
213241

214242
def _maybe_log_response(self, response: GenerateContentResponse):
215243
if not self._content_recording_enabled:
216244
return
217-
pass
245+
self._otel_wrapper.log_response_content(
246+
attributes={
247+
gen_ai_attributes.GEN_AI_SYSTEM: self._genai_system,
248+
},
249+
body={
250+
'content': response.model_dump(),
251+
})
218252

219253
def _record_token_usage_metric(self):
220254
self._otel_wrapper.token_usage_metric.record(

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/otel_wrapper.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
1+
import logging
2+
3+
import json
14
import google.genai
25

36
from .version import __version__ as _LIBRARY_VERSION
7+
from opentelemetry._events import Event
48
from opentelemetry.semconv.schemas import Schemas
59
from opentelemetry.semconv._incubating.metrics import gen_ai_metrics
610

711

12+
_logger = logging.getLogger(__name__)
13+
14+
815
_LIBRARY_NAME = 'opentelemetry-instrumentation-google-genai'
916
_SCHEMA_URL = Schemas.V1_30_0.value
1017
_SCOPE_ATTRIBUTES = {
@@ -41,6 +48,9 @@ def from_providers(
4148
def start_as_current_span(self, *args, **kwargs):
4249
return self._tracer.start_as_current_span(*args, **kwargs)
4350

51+
def done(self):
52+
pass
53+
4454
@property
4555
def tracer(self):
4656
return self._tracer
@@ -60,3 +70,23 @@ def operation_duration_metric(self):
6070
@property
6171
def token_usage_metric(self):
6272
return self._token_usage_metric
73+
74+
def log_system_prompt(self, attributes, body):
75+
_logger.debug('Recording system prompt.')
76+
event_name = 'gen_ai.system.message'
77+
self._log_event(event_name, attributes, body)
78+
79+
def log_user_prompt(self, attributes, body):
80+
_logger.debug('Recording user prompt.')
81+
event_name = 'gen_ai.user.message'
82+
self._log_event(event_name, attributes, body)
83+
84+
def log_response_content(self, attributes, body):
85+
_logger.debug('Recording response.')
86+
event_name = 'gen_ai.assistant.message'
87+
self._log_event(event_name, attributes, body)
88+
89+
def _log_event(self, event_name, attributes, body):
90+
event = Event(event_name, body=body, attributes=attributes)
91+
self.event_logger.emit(event)
92+

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/base.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,18 @@ def setUp(self):
2121
self._otel.install()
2222
self._requests = RequestsMocker()
2323
self._requests.install()
24-
self._instrumentation_context = InstrumentationContext()
25-
self._instrumentation_context.install()
24+
self._instrumentation_context = None
2625
self._api_key = 'test-api-key'
2726
self._project = 'test-project'
2827
self._location = 'test-location'
2928
self._client = None
3029
self._uses_vertex = False
3130
self._credentials = _FakeCredentials()
3231

32+
def _lazy_init(self):
33+
self._instrumentation_context = InstrumentationContext()
34+
self._instrumentation_context.install()
35+
3336
@property
3437
def client(self):
3538
if self._client is None:
@@ -48,6 +51,7 @@ def set_use_vertex(self, use_vertex):
4851
self._uses_vertex = use_vertex
4952

5053
def _create_client(self):
54+
self._lazy_init()
5155
if self._uses_vertex:
5256
os.environ['GOOGLE_API_KEY'] = self._api_key
5357
return google.genai.Client(
@@ -58,6 +62,7 @@ def _create_client(self):
5862
return google.genai.Client(api_key=self._api_key)
5963

6064
def tearDown(self):
61-
self._instrumentation_context.uninstall()
65+
if self._instrumentation_context is not None:
66+
self._instrumentation_context.uninstall()
6267
self._requests.uninstall()
6368
self._otel.uninstall()

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/otel_mocker.py‎

Lines changed: 113 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import opentelemetry.trace
44
import opentelemetry._logs._internal
5+
import opentelemetry._events
56
import opentelemetry.metrics._internal
67
from opentelemetry.util._once import Once
78

@@ -13,25 +14,31 @@
1314
get_logger_provider,
1415
set_logger_provider
1516
)
17+
from opentelemetry._events import (
18+
get_event_logger_provider,
19+
set_event_logger_provider
20+
)
1621
from opentelemetry.metrics import (
1722
get_meter_provider,
1823
set_meter_provider
1924
)
2025

26+
from opentelemetry.sdk._events import EventLoggerProvider
2127
from opentelemetry.sdk._logs import LoggerProvider
2228
from opentelemetry.sdk._logs.export import SimpleLogRecordProcessor
23-
from opentelemetry.sdk._logs._internal import SynchronousMultiLogRecordProcessor
24-
from opentelemetry.sdk._logs._internal.export.in_memory_log_exporter import InMemoryLogExporter
29+
from opentelemetry.sdk._logs.export import InMemoryLogExporter
2530
from opentelemetry.sdk.metrics._internal.export import InMemoryMetricReader
2631
from opentelemetry.sdk.metrics import MeterProvider
2732
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
2833
from opentelemetry.sdk.trace import TracerProvider
2934
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
3035

3136

37+
3238
def _bypass_otel_once():
3339
opentelemetry.trace._TRACER_PROVIDER_SET_ONCE = Once()
3440
opentelemetry._logs._internal._LOGGER_PROVIDER_SET_ONCE = Once()
41+
opentelemetry._events._EVENT_LOGGER_PROVIDER_SET_ONCE = Once()
3542
opentelemetry.metrics._internal._METER_PROVIDER_SET_ONCE = Once()
3643

3744

@@ -41,15 +48,70 @@ class OTelProviderSnapshot:
4148
def __init__(self):
4249
self._tracer_provider = get_tracer_provider()
4350
self._logger_provider = get_logger_provider()
51+
self._event_logger_provider = get_event_logger_provider()
4452
self._meter_provider = get_meter_provider()
4553

4654
def restore(self):
4755
_bypass_otel_once()
4856
set_tracer_provider(self._tracer_provider)
4957
set_logger_provider(self._logger_provider)
58+
set_event_logger_provider(self._event_logger_provider)
5059
set_meter_provider(self._meter_provider)
5160

5261

62+
class _LogWrapper:
63+
64+
def __init__(self, log_data):
65+
self._log_data = log_data
66+
67+
@property
68+
def scope(self):
69+
return self._log_data.instrumentation_scope
70+
71+
@property
72+
def resource(self):
73+
return self._log_data.log_record.resource
74+
75+
@property
76+
def attributes(self):
77+
return self._log_data.log_record.attributes
78+
79+
@property
80+
def body(self):
81+
return self._log_data.log_record.body
82+
83+
def __str__(self):
84+
return self._log_data.log_record.to_json()
85+
86+
87+
class _MetricDataPointWrapper:
88+
89+
def __init__(self, resource, scope, metric):
90+
self._resource = resource
91+
self._scope = scope
92+
self._metric = metric
93+
94+
@property
95+
def resource(self):
96+
return self._resource
97+
98+
@property
99+
def scope(self):
100+
return self._scope
101+
102+
@property
103+
def metric(self):
104+
return self._metric
105+
106+
@property
107+
def name(self):
108+
return self._metric.name
109+
110+
@property
111+
def data(self):
112+
return self._metric.data
113+
114+
53115
class OTelMocker:
54116

55117
def __init__(self):
@@ -58,6 +120,8 @@ def __init__(self):
58120
self._traces = InMemorySpanExporter()
59121
self._metrics = InMemoryMetricReader()
60122
self._spans = []
123+
self._finished_logs = []
124+
self._metrics_data = []
61125

62126
def install(self):
63127
self._snapshot = OTelProviderSnapshot()
@@ -70,15 +134,26 @@ def uninstall(self):
70134
self._snapshot.restore()
71135

72136
def get_finished_logs(self):
73-
return self._logs.get_finished_logs()
74-
137+
for log_data in self._logs.get_finished_logs():
138+
self._finished_logs.append(_LogWrapper(log_data))
139+
return self._finished_logs
140+
75141
def get_finished_spans(self):
76142
for span in self._traces.get_finished_spans():
77143
self._spans.append(span)
78144
return self._spans
79145

80146
def get_metrics_data(self):
81-
return self._metrics.get_metrics_data()
147+
data = self._metrics.get_metrics_data()
148+
if data is not None:
149+
for resource_metric in data.resource_metrics:
150+
resource = resource_metric.resource
151+
for scope_metrics in resource_metric.scope_metrics:
152+
scope = scope_metrics.scope
153+
for metric in scope_metrics.metrics:
154+
wrapper = _MetricDataPointWrapper(resource, scope, metric)
155+
self._metrics_data.append(wrapper)
156+
return self._metrics_data
82157

83158
def get_span_named(self, name):
84159
for span in self.get_finished_spans():
@@ -91,11 +166,41 @@ def assert_has_span_named(self, name):
91166
finished_spans = self.get_finished_spans()
92167
assert span is not None, 'Could not find span named "{}"; finished spans: {}'.format(name, finished_spans)
93168

169+
def get_event_named(self, event_name):
170+
for event in self.get_finished_logs():
171+
event_name_attr = event.attributes.get('event.name')
172+
if event_name_attr is None:
173+
continue
174+
if event_name_attr == event_name:
175+
return event
176+
return None
177+
178+
def assert_has_event_named(self, name):
179+
event = self.get_event_named(name)
180+
finished_logs = self.get_finished_logs()
181+
assert event is not None, 'Could not find event named "{}"; finished logs: {}'.format(name, finished_logs)
182+
183+
def assert_does_not_have_event_named(self, name):
184+
event = self.get_event_named(name)
185+
assert event is None, 'Unexpected event: {}'.format(event)
186+
187+
def get_metrics_data_named(self, name):
188+
results = []
189+
for entry in self.get_metrics_data():
190+
if entry.name == name:
191+
results.append(entry)
192+
return results
193+
194+
def assert_has_metrics_data_named(self, name):
195+
data = self.get_metrics_data_named(name)
196+
assert len(data) > 0
197+
94198
def _install_logs(self):
95-
processor = SynchronousMultiLogRecordProcessor()
96-
processor.add_log_record_processor(SimpleLogRecordProcessor(self._logs))
97-
provider = LoggerProvider(processor)
199+
provider = LoggerProvider()
200+
provider.add_log_record_processor(SimpleLogRecordProcessor(self._logs))
98201
set_logger_provider(provider)
202+
event_provider = EventLoggerProvider(logger_provider=provider)
203+
set_event_logger_provider(event_provider)
99204

100205
def _install_metrics(self):
101206
provider = MeterProvider(metric_readers=[self._metrics])

0 commit comments

Comments
 (0)