Skip to content

Commit 8cf3189

Browse files
zarinn3palzarinn3pal
committed
feat(py): Added metrics for observability and implemented AIM telemetry for firebase
1 parent f639667 commit 8cf3189

File tree

9 files changed

+576
-32
lines changed

9 files changed

+576
-32
lines changed

‎py/plugins/firebase/src/genkit/plugins/firebase/__init__.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
"""Firebase Plugin for Genkit."""
1919

20+
from genkit.plugins.google_cloud.telemetry.tracing import add_gcp_telemetry
21+
2022

2123
def package_name() -> str:
2224
"""Get the package name for the Firebase plugin.
@@ -27,4 +29,13 @@ def package_name() -> str:
2729
return 'genkit.plugins.firebase'
2830

2931

30-
__all__ = ['package_name']
32+
def add_firebase_telemetry() -> None:
33+
"""Add Firebase telemetry export to Google Cloud Observability.
34+
35+
Exports traces to Cloud Trace and metrics to Cloud Monitoring.
36+
In development (GENKIT_ENV=dev), telemetry is disabled by default.
37+
"""
38+
add_gcp_telemetry(force_export=False)
39+
40+
41+
__all__ = ['package_name', 'add_firebase_telemetry']

‎py/plugins/firebase/src/genkit/plugins/firebase/tests/test_telemetry.py‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# SPDX-License-Identifier: Apache-2.0
16+
17+
"""Tests for Firebase telemetry functionality."""
18+
19+
import json
20+
from unittest.mock import MagicMock, Mock, patch
21+
22+
from opentelemetry.sdk.trace import ReadableSpan
23+
from opentelemetry.trace import Status, StatusCode
24+
25+
from genkit.plugins.firebase import add_firebase_telemetry
26+
27+
28+
def test_add_firebase_telemetry():
29+
"""Test that add_firebase_telemetry calls add_gcp_telemetry with force_export=False."""
30+
with patch('genkit.plugins.firebase.add_gcp_telemetry') as mock_gcp_telemetry:
31+
add_firebase_telemetry()
32+
mock_gcp_telemetry.assert_called_once_with(force_export=False)
33+
34+
35+
class TestMetrics:
36+
"""Tests for metrics recording functionality."""
37+
38+
def test_record_generate_metrics_with_valid_span(self):
39+
"""Test metrics recording for a valid model span."""
40+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
41+
42+
# Create mock span with model execution data
43+
span = Mock(spec=ReadableSpan)
44+
span.attributes = {
45+
'genkit:type': 'action',
46+
'genkit:metadata:subtype': 'model',
47+
'genkit:name': 'gemini-2.0-flash',
48+
'genkit:path': '/{myFlow,t:flow}',
49+
'genkit:output': json.dumps({
50+
'usage': {
51+
'inputTokens': 100,
52+
'outputTokens': 50,
53+
'inputCharacters': 500,
54+
'outputCharacters': 250,
55+
}
56+
}),
57+
}
58+
span.status = Status(StatusCode.OK)
59+
span.start_time = 1000000000
60+
span.end_time = 1100000000 # 100ms later
61+
62+
# Should not raise exception
63+
record_generate_metrics(span)
64+
65+
def test_record_generate_metrics_with_error(self):
66+
"""Test metrics recording for failed model execution."""
67+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
68+
69+
span = Mock(spec=ReadableSpan)
70+
span.attributes = {
71+
'genkit:type': 'action',
72+
'genkit:metadata:subtype': 'model',
73+
'genkit:name': 'gemini-2.0-flash',
74+
'genkit:path': '/{myFlow,t:flow}',
75+
}
76+
span.status = Status(StatusCode.ERROR)
77+
span.start_time = 1000000000
78+
span.end_time = 1100000000
79+
80+
# Should not raise exception
81+
record_generate_metrics(span)
82+
83+
def test_record_generate_metrics_skips_non_model_spans(self):
84+
"""Test that non-model spans are skipped."""
85+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
86+
87+
span = Mock(spec=ReadableSpan)
88+
span.attributes = {
89+
'genkit:type': 'action',
90+
'genkit:metadata:subtype': 'flow', # Not a model
91+
}
92+
93+
# Should not raise exception
94+
record_generate_metrics(span)
95+
96+
def test_record_generate_metrics_with_no_attributes(self):
97+
"""Test that spans without attributes are handled gracefully."""
98+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
99+
100+
span = Mock(spec=ReadableSpan)
101+
span.attributes = None
102+
103+
# Should not raise exception
104+
record_generate_metrics(span)
105+
106+
def test_record_generate_metrics_with_all_usage_types(self):
107+
"""Test metrics recording with all usage types."""
108+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
109+
110+
span = Mock(spec=ReadableSpan)
111+
span.attributes = {
112+
'genkit:type': 'action',
113+
'genkit:metadata:subtype': 'model',
114+
'genkit:name': 'gemini-2.0-flash',
115+
'genkit:path': '/{testFlow,t:flow}',
116+
'genkit:output': json.dumps({
117+
'usage': {
118+
'inputTokens': 100,
119+
'outputTokens': 50,
120+
'inputCharacters': 500,
121+
'outputCharacters': 250,
122+
'inputImages': 2,
123+
'outputImages': 1,
124+
'inputVideos': 1,
125+
'outputVideos': 0,
126+
'inputAudio': 1,
127+
'outputAudio': 1,
128+
}
129+
}),
130+
}
131+
span.status = Status(StatusCode.OK)
132+
span.start_time = 1000000000
133+
span.end_time = 1200000000
134+
135+
# Should not raise exception
136+
record_generate_metrics(span)
137+
138+
def test_record_generate_metrics_with_invalid_usage(self):
139+
"""Test metrics recording handles invalid usage data gracefully."""
140+
from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics
141+
142+
span = Mock(spec=ReadableSpan)
143+
span.attributes = {
144+
'genkit:type': 'action',
145+
'genkit:metadata:subtype': 'model',
146+
'genkit:name': 'test-model',
147+
'genkit:path': '/{flow,t:flow}',
148+
'genkit:output': json.dumps({
149+
'usage': {
150+
'inputTokens': 'invalid', # Invalid type
151+
'outputTokens': None,
152+
}
153+
}),
154+
}
155+
span.status = Status(StatusCode.OK)
156+
span.start_time = 1000000000
157+
span.end_time = 1100000000
158+
159+
# Should not raise exception
160+
record_generate_metrics(span)
161+
162+
def test_extract_feature_name(self):
163+
"""Test feature name extraction from paths."""
164+
from genkit.plugins.google_cloud.telemetry.metrics import _extract_feature_name
165+
166+
assert _extract_feature_name('/{myFlow,t:flow}') == 'myFlow'
167+
assert _extract_feature_name('/{outer,t:flow}/{inner,t:flow}') == 'outer'
168+
assert _extract_feature_name('') == '<unknown>'
169+
assert _extract_feature_name('/invalid') == '<unknown>'
170+
assert _extract_feature_name('/{test123,t:flow}') == 'test123'
171+
172+
173+
class TestGCPTelemetry:
174+
"""Tests for GCP telemetry configuration."""
175+
176+
def test_add_gcp_telemetry_in_dev_environment(self):
177+
"""Test that telemetry is not added in dev environment by default."""
178+
from genkit.plugins.google_cloud.telemetry.tracing import add_gcp_telemetry
179+
180+
with patch(
181+
'genkit.plugins.google_cloud.telemetry.tracing.is_dev_environment',
182+
return_value=True,
183+
):
184+
with patch('genkit.plugins.google_cloud.telemetry.tracing.add_custom_exporter') as mock_add:
185+
add_gcp_telemetry(force_export=False)
186+
mock_add.assert_not_called()
187+
188+
def test_add_gcp_telemetry_force_export_in_dev(self):
189+
"""Test that force_export=True works in dev environment."""
190+
from genkit.plugins.google_cloud.telemetry.tracing import add_gcp_telemetry
191+
192+
with patch(
193+
'genkit.plugins.google_cloud.telemetry.tracing.is_dev_environment',
194+
return_value=True,
195+
):
196+
with patch('genkit.plugins.google_cloud.telemetry.tracing.add_custom_exporter') as mock_add:
197+
with patch('genkit.plugins.google_cloud.telemetry.tracing.metrics'):
198+
with patch('genkit.plugins.google_cloud.telemetry.tracing.CloudMonitoringMetricsExporter'):
199+
with patch('genkit.plugins.google_cloud.telemetry.tracing.GoogleCloudResourceDetector'):
200+
add_gcp_telemetry(force_export=True)
201+
mock_add.assert_called_once()
202+
203+
def test_add_gcp_telemetry_in_prod_environment(self):
204+
"""Test that telemetry is added in production environment."""
205+
from genkit.plugins.google_cloud.telemetry.tracing import add_gcp_telemetry
206+
207+
with patch(
208+
'genkit.plugins.google_cloud.telemetry.tracing.is_dev_environment',
209+
return_value=False,
210+
):
211+
with patch('genkit.plugins.google_cloud.telemetry.tracing.add_custom_exporter') as mock_add:
212+
with patch('genkit.plugins.google_cloud.telemetry.tracing.metrics'):
213+
with patch('genkit.plugins.google_cloud.telemetry.tracing.CloudMonitoringMetricsExporter'):
214+
with patch('genkit.plugins.google_cloud.telemetry.tracing.GoogleCloudResourceDetector'):
215+
add_gcp_telemetry()
216+
mock_add.assert_called_once()
217+
218+
def test_genkitgcp_exporter_export_success(self):
219+
"""Test GenkitGCPExporter exports spans successfully."""
220+
from opentelemetry.sdk.trace.export import SpanExportResult
221+
222+
from genkit.plugins.google_cloud.telemetry.tracing import GenkitGCPExporter
223+
224+
with patch('genkit.plugins.google_cloud.telemetry.tracing.record_generate_metrics'):
225+
exporter = GenkitGCPExporter(project_id='test-project')
226+
227+
# Mock the client
228+
exporter.client = MagicMock()
229+
exporter._translate_to_cloud_trace = MagicMock(return_value=[])
230+
231+
span = Mock(spec=ReadableSpan)
232+
span.attributes = {}
233+
234+
result = exporter.export([span])
235+
assert result == SpanExportResult.SUCCESS
236+
237+
def test_genkitgcp_exporter_export_failure(self):
238+
"""Test GenkitGCPExporter handles export failures."""
239+
from opentelemetry.sdk.trace.export import SpanExportResult
240+
241+
from genkit.plugins.google_cloud.telemetry.tracing import GenkitGCPExporter
242+
243+
with patch('genkit.plugins.google_cloud.telemetry.tracing.record_generate_metrics'):
244+
exporter = GenkitGCPExporter(project_id='test-project')
245+
246+
# Mock the client to raise exception
247+
exporter.client = MagicMock()
248+
exporter.client.batch_write_spans.side_effect = Exception('Export failed')
249+
exporter._translate_to_cloud_trace = MagicMock(return_value=[])
250+
251+
span = Mock(spec=ReadableSpan)
252+
span.attributes = {}
253+
254+
result = exporter.export([span])
255+
assert result == SpanExportResult.FAILURE

‎py/plugins/google-cloud/pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ classifiers = [
1919
dependencies = [
2020
"genkit",
2121
"opentelemetry-exporter-gcp-trace>=1.9.0",
22+
"opentelemetry-exporter-gcp-monitoring>=1.9.0",
2223
"strenum>=0.4.15; python_version < '3.11'",
2324
]
2425
description = "Genkit Google Cloud Plugin"

‎py/plugins/google-cloud/src/genkit/plugins/google_cloud/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
"""Google Cloud Plugin for Genkit."""
1919

20+
from .telemetry import add_gcp_telemetry
21+
2022

2123
def package_name() -> str:
2224
"""Get the package name for the Google Cloud plugin.
@@ -27,4 +29,4 @@ def package_name() -> str:
2729
return 'genkit.plugins.google_cloud'
2830

2931

30-
__all__ = ['package_name']
32+
__all__ = ['package_name', 'add_gcp_telemetry']

‎py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/__init__.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# SPDX-License-Identifier: Apache-2.0
16+
17+
"""Telemetry exports for Google Cloud plugin."""
18+
19+
from .tracing import add_gcp_telemetry
20+
21+
__all__ = ['add_gcp_telemetry']

0 commit comments

Comments
 (0)