diff --git a/tests/agents/core/test_llm_processor.py b/tests/agents/core/test_llm_processor.py index 5779bea65..1b3762616 100644 --- a/tests/agents/core/test_llm_processor.py +++ b/tests/agents/core/test_llm_processor.py @@ -10,6 +10,7 @@ import asyncio from typing import List from unittest.mock import Mock +from unittest.mock import patch import pytest @@ -23,6 +24,7 @@ class _StubAgent(BaseAgent): + async def _run_async_impl(self, ctx): yield @@ -53,21 +55,17 @@ def register_test_model(): @pytest.fixture def model(): m = MockLLMModel(model_name="test-llmproc-model") - m._responses = [ - LlmResponse( - content=Content(parts=[Part(text="hello")]), - partial=False, - ) - ] + m._responses = [LlmResponse( + content=Content(parts=[Part(text="hello")]), + partial=False, + )] return m @pytest.fixture def invocation_context(): service = InMemorySessionService() - session = asyncio.run( - service.create_session(app_name="test", user_id="u1", session_id="s1") - ) + session = asyncio.run(service.create_session(app_name="test", user_id="u1", session_id="s1")) agent = _StubAgent(name="test_agent") ctx = InvocationContext( session_service=service, @@ -86,6 +84,7 @@ def invocation_context(): class TestCreateEventFromResponse: + def test_maps_response_fields(self, model, invocation_context): proc = LlmProcessor(model) response = LlmResponse( @@ -121,6 +120,7 @@ def test_preserves_error_fields(self, model, invocation_context): class TestCreateErrorEvent: + def test_creates_error_event(self, model, invocation_context): proc = LlmProcessor(model) event = proc._create_error_event(invocation_context, "err_code", "err_msg") @@ -136,6 +136,7 @@ def test_creates_error_event(self, model, invocation_context): class TestProcessPlanningResponse: + def test_no_planner_returns_event_unchanged(self, model, invocation_context): proc = LlmProcessor(model) event = Event( @@ -159,6 +160,7 @@ def test_event_without_content_skips_planning(self, model, invocation_context): class TestCallLlmAsync: + def test_yields_events_for_responses(self, model, invocation_context): proc = LlmProcessor(model) request = LlmRequest() @@ -210,3 +212,33 @@ async def run(): assert len(content_events) == 2 assert content_events[0].partial is True assert content_events[1].partial is False + + def test_error_response_is_traced_before_consumer_stops(self, invocation_context): + m = MockLLMModel(model_name="test-llmproc-model") + m._responses = [ + LlmResponse( + error_code="STREAMING_ERROR", + error_message="rate limit exceeded", + partial=False, + ) + ] + proc = LlmProcessor(m) + request = LlmRequest() + + async def run(): + stream = proc.call_llm_async(request, invocation_context, stream=True) + event = await anext(stream) + # The downstream LlmAgent returns immediately for an error event, + # so tracing and span-context cleanup must be complete at this point. + mock_trace.assert_called_once() + span_context.__exit__.assert_called_once() + await stream.aclose() + return event + + with patch("trpc_agent_sdk.agents.core._llm_processor.trace_call_llm") as mock_trace, \ + patch("trpc_agent_sdk.agents.core._llm_processor.tracer") as mock_tracer: + span_context = mock_tracer.start_as_current_span.return_value + event = asyncio.run(run()) + + assert event.error_code == "STREAMING_ERROR" + assert mock_trace.call_args.args[3].error_message == "rate limit exceeded" diff --git a/tests/models/test_retry.py b/tests/models/test_retry.py index fbd4e284c..a70a37984 100644 --- a/tests/models/test_retry.py +++ b/tests/models/test_retry.py @@ -10,8 +10,11 @@ from collections.abc import AsyncGenerator from typing import Optional from unittest.mock import AsyncMock +from unittest.mock import MagicMock from unittest.mock import patch +from opentelemetry import trace + from trpc_agent_sdk.configs import ExponentialBackoffConfig from trpc_agent_sdk.configs import ModelRetryConfig from trpc_agent_sdk.models._llm_response import LlmResponse @@ -33,7 +36,6 @@ def __init__(self, status_code: int | str, headers: Optional[dict] = None): self.response = type("Resp", (), {"headers": headers})() - class _HeadersError(Exception): def __init__(self, headers: dict): @@ -73,13 +75,11 @@ async def _collect( *, get_retry_info=None, ) -> list[LlmResponse]: - return [ - response async for response in retry_model_call( - call_model, - config, - get_retry_info=get_retry_info, - ) - ] + return [response async for response in retry_model_call( + call_model, + config, + get_retry_info=get_retry_info, + )] class TestRetryHelpers: @@ -192,12 +192,46 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]: yield _content_response("ok") with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep: - responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) + responses = await _collect(call_model, + self._retry_cfg(), + get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) assert attempts == 2 assert sleep.await_count == 1 assert responses[-1].content.parts[0].text == "ok" assert all(response.error_code is None for response in responses) + async def test_retry_records_failed_attempt_span(self): + attempts = 0 + error = _StatusError(429) + + async def call_model() -> AsyncGenerator[LlmResponse, None]: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise error + yield _content_response("ok") + + span = MagicMock() + with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()), \ + patch("trpc_agent_sdk.models._retry._retry_tracer") as retry_tracer: + retry_tracer.start_as_current_span.return_value.__enter__.return_value = span + responses = await _collect( + call_model, + self._retry_cfg(), + get_retry_info=lambda _: ModelRetryInfo(should_retry=True), + ) + + assert responses[-1].content.parts[0].text == "ok" + retry_tracer.start_as_current_span.assert_called_once_with("model_retry") + span.record_exception.assert_called_once_with(error) + span.set_status.assert_called_once_with(trace.StatusCode.ERROR, "status 429") + span.set_attribute.assert_any_call("gen_ai.operation.name", "model_retry") + span.set_attribute.assert_any_call("error.type", "_StatusError") + span.set_attribute.assert_any_call("gen_ai.retry.number", 1) + span.set_attribute.assert_any_call("gen_ai.retry.max_retries", 2) + span.set_attribute.assert_any_call("gen_ai.retry.delay_seconds", 0.0) + span.set_attribute.assert_any_call("http.response.status_code", 429) + async def test_exhausts_budget_then_yields_error(self): attempts = 0 @@ -208,7 +242,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]: yield with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep: - responses = await _collect(call_model, self._retry_cfg(num_retries=2), get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) + responses = await _collect(call_model, + self._retry_cfg(num_retries=2), + get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) assert attempts == 3 assert sleep.await_count == 2 assert responses[-1].error_code == "API_ERROR" @@ -224,7 +260,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]: yield with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep: - responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=False)) + responses = await _collect(call_model, + self._retry_cfg(), + get_retry_info=lambda _: ModelRetryInfo(should_retry=False)) assert attempts == 1 assert sleep.await_count == 0 assert responses[-1].error_code == "API_ERROR" @@ -258,7 +296,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]: raise _StatusError(429) with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep: - responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) + responses = await _collect(call_model, + self._retry_cfg(), + get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) assert attempts == 1 assert sleep.await_count == 0 assert responses[0].content.parts[0].text == "partial" @@ -279,6 +319,8 @@ async def second_attempt() -> AsyncGenerator[LlmResponse, None]: attempts = iter([first_attempt, second_attempt]) with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()): - responses = await _collect(lambda: next(attempts)(), self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) + responses = await _collect(lambda: next(attempts)(), + self._retry_cfg(), + get_retry_info=lambda _: ModelRetryInfo(should_retry=True)) assert closed_attempts == ["first"] assert responses[-1].content.parts[0].text == "ok" diff --git a/tests/telemetry/test_custom_trace.py b/tests/telemetry/test_custom_trace.py index 02f27ab14..f5e0c961b 100644 --- a/tests/telemetry/test_custom_trace.py +++ b/tests/telemetry/test_custom_trace.py @@ -30,14 +30,17 @@ _SyntheticTool, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -def _make_invocation_context(agent_name="test_agent", session_id="sess-1", - user_id="user-1", user_content=None, - invocation_id="inv-1", instruction=None): + +def _make_invocation_context(agent_name="test_agent", + session_id="sess-1", + user_id="user-1", + user_content=None, + invocation_id="inv-1", + instruction=None): ctx = MagicMock() ctx.agent = MagicMock() ctx.agent.name = agent_name @@ -72,13 +75,18 @@ def _make_event( text="", event_id="evt-1", content=None, + error_code=None, error_message=None, + custom_metadata=None, ): event = MagicMock() event.partial = partial event.id = event_id event.content = content + event.error_code = error_code event.error_message = error_message + event.custom_metadata = custom_metadata + event.is_error = MagicMock(return_value=error_code is not None) event.get_function_calls = MagicMock(return_value=function_calls or []) event.get_function_responses = MagicMock(return_value=function_responses or []) event.get_text = MagicMock(return_value=text) @@ -89,7 +97,9 @@ def _make_event( # Tests: _SyntheticTool # --------------------------------------------------------------------------- + class TestSyntheticTool: + def test_init_with_name_and_description(self): tool = _SyntheticTool(name="my_tool", description="My tool desc") assert tool.name == "my_tool" @@ -114,7 +124,9 @@ async def test_run_async_impl_raises(self): # Tests: CustomTraceReporter.__init__ # --------------------------------------------------------------------------- + class TestCustomTraceReporterInit: + def test_default_init(self): reporter = CustomTraceReporter(agent_name="agent_a") assert reporter.agent_name == "agent_a" @@ -140,7 +152,9 @@ def test_custom_params(self): # Tests: _create_synthetic_llm_request # --------------------------------------------------------------------------- + class TestCreateSyntheticLlmRequest: + @patch("trpc_agent_sdk.telemetry._custom_trace.LlmRequest") @patch("trpc_agent_sdk.telemetry._custom_trace.GenerateContentConfig") def test_with_user_content(self, MockConfig, MockLlmRequest): @@ -181,17 +195,26 @@ def test_without_user_content(self, MockConfig, MockLlmRequest): # Tests: _create_synthetic_llm_response # --------------------------------------------------------------------------- + class TestCreateSyntheticLlmResponse: + @patch("trpc_agent_sdk.telemetry._custom_trace.LlmResponse") def test_with_event(self, MockLlmResponse): reporter = CustomTraceReporter(agent_name="a") - event = _make_event(content="content_obj", error_message="err") + event = _make_event( + content="content_obj", + error_code="REMOTE_ERROR", + error_message="err", + custom_metadata={"error_type": "RemoteError"}, + ) reporter._create_synthetic_llm_response(event) MockLlmResponse.assert_called_once_with( content="content_obj", + error_code="REMOTE_ERROR", error_message="err", + custom_metadata={"error_type": "RemoteError"}, ) @patch("trpc_agent_sdk.telemetry._custom_trace.LlmResponse") @@ -203,7 +226,9 @@ def test_with_none_event(self, MockLlmResponse): MockLlmResponse.assert_called_once_with( content=None, + error_code=None, error_message=None, + custom_metadata=None, ) @@ -211,7 +236,9 @@ def test_with_none_event(self, MockLlmResponse): # Tests: _trace_function_call # --------------------------------------------------------------------------- + class TestTraceFunctionCall: + def test_single_function_call(self): reporter = CustomTraceReporter(agent_name="a") fc = _make_function_call(name="tool_1", fc_id="fc-1", args={"k": "v"}) @@ -248,7 +275,9 @@ def test_args_none_becomes_empty_dict(self): # Tests: _trace_function_response # --------------------------------------------------------------------------- + class TestTraceFunctionResponse: + @patch("trpc_agent_sdk.telemetry._custom_trace.trace_tool_call") @patch("trpc_agent_sdk.telemetry._custom_trace.tracer") def test_matched_response(self, mock_tracer, mock_trace_tool_call): @@ -262,7 +291,9 @@ def test_matched_response(self, mock_tracer, mock_trace_tool_call): ) reporter.pending_function_calls["fc-1"] = { "name": "tool_x", - "args": {"input": "val"}, + "args": { + "input": "val" + }, "id": "fc-1", } @@ -298,10 +329,14 @@ def test_multiple_responses(self, mock_tracer, mock_trace_tool_call): reporter = CustomTraceReporter(agent_name="a") reporter.pending_function_calls["fc-1"] = { - "name": "t1", "args": {}, "id": "fc-1", + "name": "t1", + "args": {}, + "id": "fc-1", } reporter.pending_function_calls["fc-2"] = { - "name": "t2", "args": {}, "id": "fc-2", + "name": "t2", + "args": {}, + "id": "fc-2", } fr1 = _make_function_response(resp_id="fc-1") @@ -318,7 +353,9 @@ def test_multiple_responses(self, mock_tracer, mock_trace_tool_call): # Tests: _trace_llm_response # --------------------------------------------------------------------------- + class TestTraceLlmResponse: + @patch("trpc_agent_sdk.telemetry._custom_trace.trace_call_llm") @patch("trpc_agent_sdk.telemetry._custom_trace.tracer") def test_traces_llm_call(self, mock_tracer, mock_trace_call_llm): @@ -396,7 +433,9 @@ def test_no_instruction(self, mock_tracer, mock_trace_call_llm): # Tests: _should_trace_text # --------------------------------------------------------------------------- + class TestShouldTraceText: + def test_empty_text_returns_false(self): reporter = CustomTraceReporter(agent_name="a") assert reporter._should_trace_text("") is False @@ -431,7 +470,9 @@ def test_filter_with_none_text(self): # Tests: trace_event # --------------------------------------------------------------------------- + class TestTraceEvent: + def test_skip_partial_event(self): reporter = CustomTraceReporter(agent_name="a") ctx = _make_invocation_context() @@ -502,6 +543,20 @@ def test_empty_text_event_skips_llm_trace(self): m_llm.assert_not_called() + def test_error_event_without_text_traces_llm(self): + reporter = CustomTraceReporter(agent_name="a") + ctx = _make_invocation_context() + event = _make_event( + text="", + error_code="REMOTE_ERROR", + error_message="remote model failed", + ) + + with patch.object(reporter, "_trace_llm_response") as m_llm: + reporter.trace_event(ctx, event) + + m_llm.assert_called_once_with(ctx, event) + def test_text_filtered_out_skips_llm_trace(self): reporter = CustomTraceReporter( agent_name="a", @@ -533,7 +588,9 @@ def test_function_call_takes_priority_over_text(self): # Tests: reset # --------------------------------------------------------------------------- + class TestReset: + def test_reset_clears_pending(self): reporter = CustomTraceReporter(agent_name="a") reporter.pending_function_calls["fc-1"] = {"name": "t", "args": {}, "id": "fc-1"} @@ -554,13 +611,13 @@ def test_reset_idempotent(self): # Tests: Integration-like end-to-end flow # --------------------------------------------------------------------------- + class TestEndToEndFlow: + @patch("trpc_agent_sdk.telemetry._custom_trace.trace_call_llm") @patch("trpc_agent_sdk.telemetry._custom_trace.trace_tool_call") @patch("trpc_agent_sdk.telemetry._custom_trace.tracer") - def test_full_flow_fc_then_fr_then_text( - self, mock_tracer, mock_trace_tool_call, mock_trace_call_llm - ): + def test_full_flow_fc_then_fr_then_text(self, mock_tracer, mock_trace_tool_call, mock_trace_call_llm): mock_tracer.start_as_current_span = MagicMock() mock_tracer.start_as_current_span.return_value.__enter__ = MagicMock() mock_tracer.start_as_current_span.return_value.__exit__ = MagicMock(return_value=False) @@ -593,9 +650,7 @@ def test_full_flow_fc_then_fr_then_text( @patch("trpc_agent_sdk.telemetry._custom_trace.trace_call_llm") @patch("trpc_agent_sdk.telemetry._custom_trace.trace_tool_call") @patch("trpc_agent_sdk.telemetry._custom_trace.tracer") - def test_reset_between_invocations( - self, mock_tracer, mock_trace_tool_call, mock_trace_call_llm - ): + def test_reset_between_invocations(self, mock_tracer, mock_trace_tool_call, mock_trace_call_llm): mock_tracer.start_as_current_span = MagicMock() mock_tracer.start_as_current_span.return_value.__enter__ = MagicMock() mock_tracer.start_as_current_span.return_value.__exit__ = MagicMock(return_value=False) diff --git a/tests/telemetry/test_trace.py b/tests/telemetry/test_trace.py index cce8dad0d..51bbd4dfe 100644 --- a/tests/telemetry/test_trace.py +++ b/tests/telemetry/test_trace.py @@ -23,6 +23,7 @@ from unittest.mock import MagicMock, patch import pytest +from opentelemetry import trace from trpc_agent_sdk.telemetry._trace import ( _build_llm_request_for_trace, @@ -883,10 +884,12 @@ def _make_llm_request(self, model="test-model", contents=None): req.config.model_dump = MagicMock(return_value={"temperature": 0.7}) return req - def _make_llm_response(self, content=None, usage=None, error_message=None): + def _make_llm_response(self, content=None, usage=None, error_code=None, error_message=None, custom_metadata=None): resp = MagicMock() resp.content = content + resp.error_code = error_code resp.error_message = error_message + resp.custom_metadata = custom_metadata resp.model_dump_json = MagicMock(return_value='{"content": "response"}') resp.usage_metadata = usage return resp @@ -909,6 +912,35 @@ def test_basic_llm_trace(self, mock_get_span): span.set_attribute.assert_any_call("trpc.python.agent.session_id", "sess-1") span.set_attribute.assert_any_call("trpc.python.agent.event_id", "e-1") + @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") + def test_error_response_marks_span_and_records_exception(self, mock_get_span): + span = _mock_span() + mock_get_span.return_value = span + ctx = _make_invocation_context() + + req = self._make_llm_request() + resp = self._make_llm_response( + error_code="STREAMING_ERROR", + error_message="rate limit exceeded", + custom_metadata={"error_type": "RateLimitError"}, + ) + + trace_call_llm(ctx, event_id="e-1", llm_request=req, llm_response=resp) + + span.set_status.assert_called_once_with(trace.StatusCode.ERROR, "rate limit exceeded") + span.set_attribute.assert_any_call("error.type", "RateLimitError") + span.set_attribute.assert_any_call( + "trpc.python.agent.llm.error_code", + "STREAMING_ERROR", + ) + span.add_event.assert_called_once_with( + "exception", + { + "exception.type": "RateLimitError", + "exception.message": "rate limit exceeded", + }, + ) + @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") def test_with_usage_metadata(self, mock_get_span): span = _mock_span() diff --git a/trpc_agent_sdk/agents/core/_llm_processor.py b/trpc_agent_sdk/agents/core/_llm_processor.py index 6d3ca3730..b0ee51c8e 100644 --- a/trpc_agent_sdk/agents/core/_llm_processor.py +++ b/trpc_agent_sdk/agents/core/_llm_processor.py @@ -82,11 +82,14 @@ async def call_llm_async(self, return # Step 2: Call the model and process responses with telemetry tracing. + terminal_event: Optional[Event] = None with tracer.start_as_current_span('call_llm'): event_id = Event.new_id() final_llm_response = None aggregated_raw_function_calls: list[dict] = [] aggregated_event_function_calls: list[dict] = [] + instruction = getattr(context.agent, 'instruction', None) + instruction_metadata = getattr(instruction, 'metadata', None) def _append_function_calls(target: list[dict], calls: list) -> None: for call in calls or []: @@ -125,15 +128,40 @@ def _append_function_calls(target: list[dict], calls: list) -> None: # Process response with planner if available event = self._process_planning_response(event, context) - # Track the latest non-partial response for tracing - # In streaming mode, only the final (non-partial) response - # contains complete data suitable for telemetry reporting. if not llm_response.partial: final_llm_response = llm_response + # Trace before yielding because consumers stop + # immediately after receiving an error event. + trace_call_llm( + context, + event_id, + request, + llm_response, + instruction_metadata=instruction_metadata, + stream_function_calls_raw=aggregated_raw_function_calls, + stream_function_calls_post_planner=aggregated_event_function_calls, + ) + terminal_event = event + # Finish the model stream and exit the span context + # before exposing the terminal event downstream. + continue yield event except Exception as ex: metrics_error_type = type(ex).__name__ + trace_call_llm( + context, + event_id, + request, + LlmResponse( + error_code="LLM_CALL_ERROR", + error_message=str(ex), + custom_metadata={"error_type": type(ex).__name__}, + ), + instruction_metadata=instruction_metadata, + stream_function_calls_raw=aggregated_raw_function_calls, + stream_function_calls_post_planner=aggregated_event_function_calls, + ) raise finally: duration_s = time.monotonic() - t_start @@ -148,18 +176,8 @@ def _append_function_calls(target: list[dict], calls: list) -> None: error_type=metrics_error_type, ) - # Trace the LLM call once after the stream completes, - # using the final complete response to avoid attribute - # overwrites from multiple partial trace_call_llm calls. - if final_llm_response is not None: - instruction_metadata = getattr(context.agent.instruction, 'metadata', None) - trace_call_llm(context, - event_id, - request, - final_llm_response, - instruction_metadata=instruction_metadata, - stream_function_calls_raw=aggregated_raw_function_calls, - stream_function_calls_post_planner=aggregated_event_function_calls) + if terminal_event is not None: + yield terminal_event except Exception as ex: # pylint: disable=broad-except logger.error("LLM call failed for agent %s: %s", author, ex) diff --git a/trpc_agent_sdk/models/_retry.py b/trpc_agent_sdk/models/_retry.py index a62cdca1f..48ba3f496 100644 --- a/trpc_agent_sdk/models/_retry.py +++ b/trpc_agent_sdk/models/_retry.py @@ -17,6 +17,8 @@ from typing import Any from typing import Optional +from opentelemetry import trace + from trpc_agent_sdk.configs import ExponentialBackoffConfig from trpc_agent_sdk.configs import ModelRetryConfig from trpc_agent_sdk.log import logger @@ -24,6 +26,7 @@ from ._llm_response import LlmResponse _MAX_RETRY_AFTER_SECONDS = 60.0 +_retry_tracer = trace.get_tracer("trpc.python.agent") @dataclass(frozen=True) @@ -152,6 +155,27 @@ def _build_error_response(ex: Exception, error_code: str) -> LlmResponse: ) +def _trace_retry_failure( + ex: Exception, + *, + retry_number: int, + max_retries: int, + delay_seconds: float, +) -> None: + """Record a retryable model failure as a short child span.""" + with _retry_tracer.start_as_current_span("model_retry") as span: + span.record_exception(ex) + span.set_status(trace.StatusCode.ERROR, str(ex)) + span.set_attribute("gen_ai.operation.name", "model_retry") + span.set_attribute("error.type", type(ex).__name__) + span.set_attribute("gen_ai.retry.number", retry_number) + span.set_attribute("gen_ai.retry.max_retries", max_retries) + span.set_attribute("gen_ai.retry.delay_seconds", delay_seconds) + status_code = _extract_status_code(ex) + if status_code is not None: + span.set_attribute("http.response.status_code", status_code) + + async def retry_model_call( call_model: Callable[[], AsyncGenerator[LlmResponse, None]], retry_config: Optional[ModelRetryConfig], @@ -195,6 +219,12 @@ async def retry_model_call( return delay = _compute_exponential_backoff(retry_config.backoff, attempt, retry_info.retry_after) + _trace_retry_failure( + ex, + retry_number=attempt + 1, + max_retries=retry_config.num_retries, + delay_seconds=delay, + ) logger.warning( "Model call failed (exception=%s); retrying in %.2fs (attempt %d/%d).", type(ex).__name__, diff --git a/trpc_agent_sdk/telemetry/_custom_trace.py b/trpc_agent_sdk/telemetry/_custom_trace.py index e28286952..04653c3ee 100644 --- a/trpc_agent_sdk/telemetry/_custom_trace.py +++ b/trpc_agent_sdk/telemetry/_custom_trace.py @@ -146,7 +146,9 @@ def _create_synthetic_llm_response(self, event: Event) -> LlmResponse: """ return LlmResponse( content=event.content if event else None, + error_code=event.error_code if event else None, error_message=event.error_message if event else None, + custom_metadata=event.custom_metadata if event else None, ) def _trace_function_call(self, event: Event) -> None: @@ -242,6 +244,13 @@ def trace_event( if event.partial: return + # Error events commonly have no text content. Trace them before the + # content-based routing below so failed custom-model invocations remain + # visible and are marked as errors by trace_call_llm. + if event.is_error(): + self._trace_llm_response(ctx, event) + return + # Check for function_call (tool invocation request) if event.get_function_calls(): self._trace_function_call(event) diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py index 9c7259066..6c18a59ea 100644 --- a/trpc_agent_sdk/telemetry/_trace.py +++ b/trpc_agent_sdk/telemetry/_trace.py @@ -457,6 +457,26 @@ def trace_call_llm( llm_response_json, ) + error_code = getattr(llm_response, "error_code", None) + if error_code: + error_message = getattr(llm_response, "error_message", None) + custom_metadata = getattr(llm_response, "custom_metadata", None) + error_type = custom_metadata.get("error_type") if isinstance(custom_metadata, dict) else None + error_type = str(error_type or error_code) + status_description = str(error_message or error_code) + + span.set_status(trace.StatusCode.ERROR, status_description) + span.set_attribute("error.type", error_type) + span.set_attribute(f"{_trpc_agent_span_name}.llm.error_code", str(error_code)) + if error_message: + span.set_attribute(f"{_trpc_agent_span_name}.llm.error_message", str(error_message)) + + exception_attributes = { + "exception.type": error_type, + "exception.message": status_description, + } + span.add_event("exception", exception_attributes) + if stream_function_calls_raw: span.set_attribute( f"{_trpc_agent_span_name}.stream_function_calls.raw", @@ -517,8 +537,10 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: """ # Some fields in LlmRequest are function pointers and can not be serialized. result = { - "model": llm_request.model, - "config": llm_request.config.model_dump(exclude_none=True, exclude="response_schema"), + "model": + llm_request.model, + "config": (llm_request.config.model_dump(exclude_none=True, exclude="response_schema") + if llm_request.config is not None else {}), "contents": [], } # We do not want to send bytes data to the trace.