Skip to content

Commit 858019c

Browse files
M-HietalaM-Hietala
authored
agent tracing runid to feature branch (#40888)
1 parent 4d185aa commit 858019c

3 files changed

Lines changed: 30 additions & 2 deletions

File tree

‎sdk/ai/azure-ai-agents/azure/ai/agents/telemetry/_ai_agents_instrumentor.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast, TYPE_CHECKING
1515
from urllib.parse import urlparse
1616

17-
from azure.ai.agents.models import AgentRunStream, AsyncAgentRunStream, _models
17+
from azure.ai.agents.models import AgentRunStream, AsyncAgentRunStream, RunStepMessageCreationDetails, _models
1818
from azure.ai.agents.models._enums import (
1919
AgentsResponseFormatMode,
2020
MessageRole,
@@ -596,8 +596,9 @@ def _add_tool_assistant_message_event(self, span, step: RunStep) -> None:
596596
def set_end_run(self, span: "AbstractSpan", run: Optional[ThreadRun]) -> None:
597597
if run and span and span.span_instance.is_recording:
598598
span.add_attribute(GEN_AI_THREAD_RUN_STATUS, self._status_to_string(run.status))
599+
span.add_attribute(GEN_AI_THREAD_RUN_ID, self._status_to_string(run.id))
599600
span.add_attribute(GEN_AI_RESPONSE_MODEL, run.model)
600-
if run and run.usage:
601+
if run.usage:
601602
span.add_attribute(GEN_AI_USAGE_INPUT_TOKENS, run.usage.prompt_tokens)
602603
span.add_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, run.usage.completion_tokens)
603604

@@ -1121,6 +1122,7 @@ async def trace_create_run_async(self, operation_name, function, *args, **kwargs
11211122
result = await function(*args, **kwargs)
11221123
if span.span_instance.is_recording:
11231124
span.add_attribute(GEN_AI_THREAD_RUN_STATUS, self._status_to_string(result.status))
1125+
span.add_attribute(GEN_AI_THREAD_RUN_ID, self._status_to_string(result.id))
11241126
span.add_attribute(GEN_AI_RESPONSE_MODEL, result.model)
11251127
if result.usage:
11261128
span.add_attribute(GEN_AI_USAGE_INPUT_TOKENS, result.usage.prompt_tokens)
@@ -2040,6 +2042,13 @@ def on_run_step(self, step: "RunStep") -> None: # type: ignore[func-returns-val
20402042
)
20412043
elif step.type == "message_creation" and step.status == RunStepStatus.COMPLETED:
20422044
self.instrumentor.add_thread_message_event(self.span, cast(ThreadMessage, self.last_message), step.usage)
2045+
if (
2046+
self.span
2047+
and self.span.span_instance.is_recording
2048+
and isinstance(step.step_details, RunStepMessageCreationDetails)
2049+
):
2050+
self.span.add_attribute(GEN_AI_MESSAGE_ID, step.step_details.message_creation.message_id)
2051+
20432052
self.last_message = None
20442053

20452054
return retval # type: ignore
@@ -2152,6 +2161,13 @@ async def on_run_step(self, step: "RunStep") -> None: # type: ignore[func-retur
21522161
)
21532162
elif step.type == "message_creation" and step.status == RunStepStatus.COMPLETED:
21542163
self.instrumentor.add_thread_message_event(self.span, cast(ThreadMessage, self.last_message), step.usage)
2164+
if (
2165+
self.span
2166+
and self.span.span_instance.is_recording
2167+
and isinstance(step.step_details, RunStepMessageCreationDetails)
2168+
):
2169+
self.span.add_attribute(GEN_AI_MESSAGE_ID, step.step_details.message_creation.message_id)
2170+
21552171
self.last_message = None
21562172

21572173
return retval # type: ignore

‎sdk/ai/azure-ai-agents/tests/test_ai_agents_instrumentor.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ def test_agent_chat_with_tracing_content_recording_enabled(self, **kwargs):
269269
("gen_ai.operation.name", "start_thread_run"),
270270
("server.address", ""),
271271
("gen_ai.thread.id", ""),
272+
("gen_ai.thread.run.id", ""),
272273
("gen_ai.agent.id", ""),
273274
("gen_ai.thread.run.status", "queued"),
274275
("gen_ai.response.model", "gpt-4o"),
@@ -414,6 +415,7 @@ def test_agent_chat_with_tracing_content_recording_disabled(self, **kwargs):
414415
("gen_ai.operation.name", "start_thread_run"),
415416
("server.address", ""),
416417
("gen_ai.thread.id", ""),
418+
("gen_ai.thread.run.id", ""),
417419
("gen_ai.agent.id", ""),
418420
("gen_ai.thread.run.status", "queued"),
419421
("gen_ai.response.model", "gpt-4o"),
@@ -594,6 +596,8 @@ def fetch_weather(location: str) -> str:
594596
("server.address", ""),
595597
("gen_ai.thread.id", ""),
596598
("gen_ai.agent.id", ""),
599+
("gen_ai.thread.run.id", ""),
600+
("gen_ai.message.id", ""),
597601
("gen_ai.thread.run.status", "completed"),
598602
("gen_ai.response.model", "gpt-4o"),
599603
("gen_ai.usage.input_tokens", "+"),
@@ -927,6 +931,8 @@ def fetch_weather(location: str) -> str:
927931
("server.address", ""),
928932
("gen_ai.thread.id", ""),
929933
("gen_ai.agent.id", ""),
934+
("gen_ai.thread.run.id", ""),
935+
("gen_ai.message.id", ""),
930936
("gen_ai.thread.run.status", "completed"),
931937
("gen_ai.response.model", "gpt-4o"),
932938
("gen_ai.usage.input_tokens", "+"),

‎sdk/ai/azure-ai-agents/tests/test_ai_agents_instrumentor_async.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,7 @@ async def test_agent_chat_with_tracing_content_recording_enabled(self, **kwargs)
215215
("gen_ai.operation.name", "start_thread_run"),
216216
("server.address", ""),
217217
("gen_ai.thread.id", ""),
218+
("gen_ai.thread.run.id", ""),
218219
("gen_ai.agent.id", ""),
219220
("gen_ai.thread.run.status", "queued"),
220221
("gen_ai.response.model", "gpt-4o"),
@@ -358,6 +359,7 @@ async def test_agent_chat_with_tracing_content_recording_disabled(self, **kwargs
358359
("gen_ai.operation.name", "start_thread_run"),
359360
("server.address", ""),
360361
("gen_ai.thread.id", ""),
362+
("gen_ai.thread.run.id", ""),
361363
("gen_ai.agent.id", ""),
362364
("gen_ai.thread.run.status", "queued"),
363365
("gen_ai.response.model", "gpt-4o"),
@@ -559,6 +561,8 @@ def fetch_weather(location: str) -> str:
559561
("server.address", ""),
560562
("gen_ai.thread.id", ""),
561563
("gen_ai.agent.id", ""),
564+
("gen_ai.thread.run.id", ""),
565+
("gen_ai.message.id", ""),
562566
("gen_ai.thread.run.status", "completed"),
563567
("gen_ai.response.model", "gpt-4o"),
564568
("gen_ai.usage.input_tokens", "+"),
@@ -782,6 +786,8 @@ def fetch_weather(location: str) -> str:
782786
("server.address", ""),
783787
("gen_ai.thread.id", ""),
784788
("gen_ai.agent.id", ""),
789+
("gen_ai.thread.run.id", ""),
790+
("gen_ai.message.id", ""),
785791
("gen_ai.thread.run.status", "completed"),
786792
("gen_ai.response.model", "gpt-4o"),
787793
("gen_ai.usage.input_tokens", "+"),

0 commit comments

Comments
 (0)