|
26 | 26 | GenerationSpanData, |
27 | 27 | GuardrailSpanData, |
28 | 28 | HandoffSpanData, |
| 29 | + ResponseSpanData, |
29 | 30 | ) |
30 | 31 |
|
31 | 32 | from opentelemetry.instrumentation.openai_agents._processor import ( |
@@ -380,3 +381,91 @@ def test_processor_does_not_throw(self, setup): |
380 | 381 | processor.on_span_end(None) |
381 | 382 | processor.on_trace_start(None) |
382 | 383 | processor.on_trace_end(None) |
| 384 | + |
| 385 | + |
| 386 | +class _FakeUsage: |
| 387 | + def __init__(self, input_tokens=100, output_tokens=50): |
| 388 | + self.input_tokens = input_tokens |
| 389 | + self.output_tokens = output_tokens |
| 390 | + |
| 391 | + |
| 392 | +class _FakeResponse: |
| 393 | + def __init__(self, model="gpt-4o", id="resp-001", usage=None): |
| 394 | + self.model = model |
| 395 | + self.id = id |
| 396 | + self.usage = usage or _FakeUsage() |
| 397 | + |
| 398 | + |
| 399 | +class TestResponseSpan: |
| 400 | + def test_response_creates_chat_span(self, setup): |
| 401 | + processor, exporter = setup |
| 402 | + trace = FakeTrace() |
| 403 | + processor.on_trace_start(trace) |
| 404 | + |
| 405 | + data = ResponseSpanData(response=_FakeResponse()) |
| 406 | + span = FakeSpan( |
| 407 | + span_data=data, |
| 408 | + span_id="resp_span_001", |
| 409 | + trace_id=trace.trace_id, |
| 410 | + ) |
| 411 | + processor.on_span_start(span) |
| 412 | + processor.on_span_end(span) |
| 413 | + processor.on_trace_end(trace) |
| 414 | + |
| 415 | + spans = exporter.get_finished_spans() |
| 416 | + chat_spans = [s for s in spans if "chat" in s.name] |
| 417 | + assert len(chat_spans) == 1 |
| 418 | + attrs = dict(chat_spans[0].attributes) |
| 419 | + assert attrs["gen_ai.operation.name"] == "chat" |
| 420 | + assert attrs["gen_ai.system"] == "openai" |
| 421 | + assert attrs["gen_ai.request.model"] == "gpt-4o" |
| 422 | + |
| 423 | + def test_response_extracts_usage(self, setup): |
| 424 | + processor, exporter = setup |
| 425 | + trace = FakeTrace() |
| 426 | + processor.on_trace_start(trace) |
| 427 | + |
| 428 | + resp = _FakeResponse( |
| 429 | + model="gpt-4o-mini", |
| 430 | + usage=_FakeUsage(input_tokens=200, output_tokens=80), |
| 431 | + ) |
| 432 | + data = ResponseSpanData(response=resp) |
| 433 | + span = FakeSpan( |
| 434 | + span_data=data, |
| 435 | + span_id="resp_span_002", |
| 436 | + trace_id=trace.trace_id, |
| 437 | + ) |
| 438 | + processor.on_span_start(span) |
| 439 | + processor.on_span_end(span) |
| 440 | + processor.on_trace_end(trace) |
| 441 | + |
| 442 | + spans = exporter.get_finished_spans() |
| 443 | + chat_spans = [s for s in spans if "chat" in s.name] |
| 444 | + assert len(chat_spans) == 1 |
| 445 | + attrs = dict(chat_spans[0].attributes) |
| 446 | + assert attrs["gen_ai.response.model"] == "gpt-4o-mini" |
| 447 | + assert attrs["gen_ai.response.id"] == "resp-001" |
| 448 | + assert attrs["gen_ai.usage.input_tokens"] == 200 |
| 449 | + assert attrs["gen_ai.usage.output_tokens"] == 80 |
| 450 | + |
| 451 | + def test_response_captures_input_content(self, setup): |
| 452 | + processor, exporter = setup |
| 453 | + trace = FakeTrace() |
| 454 | + processor.on_trace_start(trace) |
| 455 | + |
| 456 | + data = ResponseSpanData(response=_FakeResponse()) |
| 457 | + data.input = [{"role": "user", "content": "Hello"}] |
| 458 | + span = FakeSpan( |
| 459 | + span_data=data, |
| 460 | + span_id="resp_span_003", |
| 461 | + trace_id=trace.trace_id, |
| 462 | + ) |
| 463 | + processor.on_span_start(span) |
| 464 | + processor.on_span_end(span) |
| 465 | + processor.on_trace_end(trace) |
| 466 | + |
| 467 | + spans = exporter.get_finished_spans() |
| 468 | + chat_spans = [s for s in spans if "chat" in s.name] |
| 469 | + assert len(chat_spans) == 1 |
| 470 | + attrs = dict(chat_spans[0].attributes) |
| 471 | + assert "gen_ai.input.messages" in attrs |
0 commit comments