2323
2424
2525def _create_model_span (
26- model_name : str = " gemini-pro" ,
27- path : str = " /{myflow,t:flow}" ,
26+ model_name : str = ' gemini-pro' ,
27+ path : str = ' /{myflow,t:flow}' ,
2828 output : str = '{"usage": {"inputTokens": 100, "outputTokens": 50}}' ,
2929 is_ok : bool = True ,
3030 start_time : int = 1000000000 ,
@@ -45,30 +45,30 @@ def _create_model_span(
4545 """
4646 mock_span = MagicMock (spec = ReadableSpan )
4747 mock_span .attributes = {
48- " genkit:type" : " action" ,
49- " genkit:metadata:subtype" : " model" ,
50- " genkit:name" : model_name ,
51- " genkit:path" : path ,
52- " genkit:output" : output ,
48+ ' genkit:type' : ' action' ,
49+ ' genkit:metadata:subtype' : ' model' ,
50+ ' genkit:name' : model_name ,
51+ ' genkit:path' : path ,
52+ ' genkit:output' : output ,
5353 }
5454 mock_span .status .is_ok = is_ok
5555 mock_span .start_time = start_time
5656 mock_span .end_time = end_time
5757 return mock_span
5858
5959
60- @patch (" genkit.plugins.firebase.add_gcp_telemetry" )
60+ @patch (' genkit.plugins.firebase.add_gcp_telemetry' )
6161def test_firebase_telemetry_delegates_to_gcp (mock_add_gcp_telemetry ):
6262 """Test that Firebase telemetry delegates to GCP telemetry."""
6363 add_firebase_telemetry ()
6464 mock_add_gcp_telemetry .assert_called_once_with (force_export = False )
6565
6666
67- @patch (" genkit.plugins.google_cloud.telemetry.metrics._output_tokens" )
68- @patch (" genkit.plugins.google_cloud.telemetry.metrics._input_tokens" )
69- @patch (" genkit.plugins.google_cloud.telemetry.metrics._latency" )
70- @patch (" genkit.plugins.google_cloud.telemetry.metrics._failures" )
71- @patch (" genkit.plugins.google_cloud.telemetry.metrics._requests" )
67+ @patch (' genkit.plugins.google_cloud.telemetry.metrics._output_tokens' )
68+ @patch (' genkit.plugins.google_cloud.telemetry.metrics._input_tokens' )
69+ @patch (' genkit.plugins.google_cloud.telemetry.metrics._latency' )
70+ @patch (' genkit.plugins.google_cloud.telemetry.metrics._failures' )
71+ @patch (' genkit.plugins.google_cloud.telemetry.metrics._requests' )
7272def test_record_generate_metrics_with_model_action (
7373 mock_requests ,
7474 mock_failures ,
@@ -91,16 +91,16 @@ def test_record_generate_metrics_with_model_action(
9191
9292 # Create test span using helper
9393 mock_span = _create_model_span (
94- model_name = " gemini-pro" ,
95- path = " /{myflow,t:flow}" ,
94+ model_name = ' gemini-pro' ,
95+ path = ' /{myflow,t:flow}' ,
9696 output = '{"usage": {"inputTokens": 100, "outputTokens": 50}}' ,
9797 )
9898
9999 # Execute
100100 record_generate_metrics (mock_span )
101101
102102 # Verify dimensions
103- expected_dimensions = {" model" : " gemini-pro" , " source" : " myflow" , " error" : " none" }
103+ expected_dimensions = {' model' : ' gemini-pro' , ' source' : ' myflow' , ' error' : ' none' }
104104
105105 # Verify requests counter
106106 mock_request_counter .add .assert_called_once_with (1 , expected_dimensions )
0 commit comments