|
14 | 14 | ) |
15 | 15 |
|
16 | 16 | from ..common.base import TestCase as CommonTestCaseBase |
17 | | -from .base import _HAS_INTERACTIONS, AsyncInteractionsResource, InteractionsResource |
| 17 | +from .base import ( |
| 18 | + _HAS_INTERACTIONS, |
| 19 | + AsyncInteractionsResource, |
| 20 | + InteractionsResource, |
| 21 | +) |
18 | 22 | from .util import ( |
19 | 23 | FakeAsyncStream, |
20 | 24 | FakeStream, |
| 25 | + create_mock_content_event, |
21 | 26 | create_mock_error_event, |
| 27 | + create_mock_fetched_interaction, |
| 28 | + create_mock_sse_completed_event, |
22 | 29 | create_mock_step_event, |
23 | 30 | parse_sse_events, |
24 | | - create_mock_sse_completed_event, |
25 | | - create_mock_content_event, |
26 | | - create_mock_fetched_interaction, |
27 | 31 | ) |
28 | 32 |
|
29 | 33 | _SPAN_NAME = "fetch_response" |
@@ -239,9 +243,7 @@ def _result(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: |
239 | 243 | "index": 0, |
240 | 244 | "step": { |
241 | 245 | "type": "model_output", |
242 | | - "content": [ |
243 | | - {"type": "text", "text": "Hello"} |
244 | | - ], |
| 246 | + "content": [{"type": "text", "text": "Hello"}], |
245 | 247 | }, |
246 | 248 | }, |
247 | 249 | { |
@@ -357,7 +359,9 @@ def _result(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: |
357 | 359 | interaction.steps = None |
358 | 360 | interaction.output_text = None |
359 | 361 | return [ |
360 | | - create_mock_step_event("step.start", 0, step_type="thought"), |
| 362 | + create_mock_step_event( |
| 363 | + "step.start", 0, step_type="thought" |
| 364 | + ), |
361 | 365 | create_mock_step_event("step.delta", 0, text=None), |
362 | 366 | create_mock_step_event( |
363 | 367 | "step.start", 1, step_type="model_output" |
@@ -408,7 +412,9 @@ def run_streaming_fetch_with_caller_error( |
408 | 412 | ) -> None: |
409 | 413 | raise NotImplementedError() |
410 | 414 |
|
411 | | - def run_streaming_fetch_closed_early(self, *args: Any, **kwargs: Any) -> Any: |
| 415 | + def run_streaming_fetch_closed_early( |
| 416 | + self, *args: Any, **kwargs: Any |
| 417 | + ) -> Any: |
412 | 418 | raise NotImplementedError() |
413 | 419 |
|
414 | 420 | def drain_stream(self, *args: Any, **kwargs: Any) -> Any: |
@@ -465,7 +471,9 @@ def test_response_model_and_status_recorded(self) -> None: |
465 | 471 | self.assertEqual( |
466 | 472 | span.attributes["gen_ai.response.model"], "gemini-2.5-pro" |
467 | 473 | ) |
468 | | - self.assertEqual(span.attributes["gen_ai.response.status"], "completed") |
| 474 | + self.assertEqual( |
| 475 | + span.attributes["gen_ai.response.status"], "completed" |
| 476 | + ) |
469 | 477 | self.assertEqual( |
470 | 478 | span.attributes["gen_ai.response.finish_reasons"], ("stop",) |
471 | 479 | ) |
@@ -645,8 +653,12 @@ def test_streaming_fetch_generates_span(self) -> None: |
645 | 653 | self.assertEqual(events[1].interaction.id, "stream-id-1") |
646 | 654 | span = self.otel.get_span_named(_SPAN_NAME) |
647 | 655 | self.assertEqual(span.attributes["gen_ai.response.id"], "stream-id-1") |
648 | | - self.assertIs(span.attributes["gen_ai.request.stream"], True) |
649 | | - self.assertEqual(span.attributes["gen_ai.response.status"], "completed") |
| 656 | + self.assertEqual( |
| 657 | + span.attributes["gen_ai.response.status"], "completed" |
| 658 | + ) |
| 659 | + # fetch_response has no gen_ai.request.stream attribute; only a resumed |
| 660 | + # fetch is distinguished, by its cursor. |
| 661 | + self.assertNotIn("gen_ai.request.stream", span.attributes) |
650 | 662 | self.assertNotIn("gen_ai.request.stream_cursor", span.attributes) |
651 | 663 |
|
652 | 664 | def test_resumed_streaming_fetch_records_cursor(self) -> None: |
@@ -693,9 +705,7 @@ def test_mid_stream_error_records_error(self) -> None: |
693 | 705 | def test_caller_side_error_records_error_and_closes_stream(self) -> None: |
694 | 706 | self.configure_valid_fetch() |
695 | 707 | with self.assertRaises(RuntimeError): |
696 | | - self.run_streaming_fetch_with_caller_error( |
697 | | - "abc-123", stream=True |
698 | | - ) |
| 708 | + self.run_streaming_fetch_with_caller_error("abc-123", stream=True) |
699 | 709 | span = self.otel.get_span_named(_SPAN_NAME) |
700 | 710 | self.assertEqual(span.attributes["error.type"], "RuntimeError") |
701 | 711 | self.assertTrue(self._last_stream.closed) |
@@ -728,9 +738,7 @@ def test_in_band_stream_error_fails_the_span(self) -> None: |
728 | 738 | events = self.run_streaming_fetch("abc-123", stream=True) |
729 | 739 | self.assertEqual(len(events), 2) |
730 | 740 | span = self.otel.get_span_named(_SPAN_NAME) |
731 | | - self.assertEqual( |
732 | | - span.attributes["error.type"], "service_unavailable" |
733 | | - ) |
| 741 | + self.assertEqual(span.attributes["error.type"], "service_unavailable") |
734 | 742 | # No completion event arrived, so nothing is known about the response. |
735 | 743 | self.assertNotIn("gen_ai.response.status", span.attributes) |
736 | 744 |
|
@@ -952,8 +960,6 @@ def test_early_close_finalizes_span(self) -> None: |
952 | 960 | self.configure_valid_fetch(interaction_id="closed-early") |
953 | 961 | self.run_streaming_fetch_closed_early("closed-early", stream=True) |
954 | 962 | span = self.otel.get_span_named(_SPAN_NAME) |
955 | | - self.assertEqual( |
956 | | - span.attributes["gen_ai.response.id"], "closed-early" |
957 | | - ) |
| 963 | + self.assertEqual(span.attributes["gen_ai.response.id"], "closed-early") |
958 | 964 | self.assertNotIn("error.type", span.attributes) |
959 | 965 | self.assertTrue(self._last_stream.closed) |
0 commit comments