Function test_trace_call_llm

tests/unittests/telemetry/test_spans.py:155–221

Test trace_call_llm sets all telemetry attributes correctly with normal content.

(monkeypatch, mock_span_fixture)

@pytest.mark.asyncio
async def test_trace_call_llm(monkeypatch, mock_span_fixture):
  """Test trace_call_llm sets all telemetry attributes correctly with normal content."""
  monkeypatch.setattr(
      'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
  )

  agent = LlmAgent(name='test_agent')
  invocation_context = await _create_invocation_context(agent)
  llm_request = LlmRequest(
      model='gemini-pro',
      contents=[
          types.Content(
              role='user',
              parts=[types.Part(text='Hello, how are you?')],
          ),
      ],
      config=types.GenerateContentConfig(
          top_p=0.95,
          max_output_tokens=1024,
          thinking_config=types.ThinkingConfig(thinking_budget=10),
      ),
  )
  llm_response = LlmResponse(
      turn_complete=True,
      finish_reason=types.FinishReason.STOP,
      usage_metadata=types.GenerateContentResponseUsageMetadata(
          total_token_count=100,
          prompt_token_count=50,
          candidates_token_count=50,
          thoughts_token_count=10,
      ),
  )
  # We dynamically assign system_instruction_tokens rather than passing it
  # to the GenerateContentResponseUsageMetadata constructor to ensure backward
  # compatibility with older versions of the google-genai SDK that do not have
  # this property defined in their Pydantic models.
  try:
    llm_response.usage_metadata.system_instruction_tokens = 5
  except Exception:
    pass

  trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)

  expected_calls = [
      mock.call('gen_ai.system', 'gcp.vertex.agent'),
      mock.call('gen_ai.request.top_p', 0.95),
      mock.call('gen_ai.request.max_tokens', 1024),
      mock.call('gcp.vertex.agent.llm_response', mock.ANY),
      mock.call('gen_ai.usage.experimental.reasoning_tokens_limit', 10),
      mock.call('gen_ai.response.finish_reasons', ['stop']),
  ]

  expected_usage_attrs = {
      'gen_ai.usage.input_tokens': 50,
      'gen_ai.usage.output_tokens': 60,
      'gen_ai.usage.reasoning.output_tokens': 10,
  }
  if hasattr(llm_response.usage_metadata, 'system_instruction_tokens'):