Skip to content

Commit 2dab021

Browse files
Prachig-MicrosoftCopilot
Prachig-Microsoft
and
Copilot
committed
fix: Handle UsageDetails object and add missing dependencies
- Fix extract_token_usage() to handle UsageDetails object (not just dict) from agent framework SDK, resolving 0-token reports for RAI, Summarize, and GapAnalysis agents - Add azure-monitor-events-extension dependency to ContentProcessor and ContentProcessorWorkflow pyproject.toml - Add unit test for UsageDetails object handling - Enable Application Insights monitoring in main.parameters.json Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent febc3da commit 2dab021

8 files changed

Lines changed: 2014 additions & 1931 deletions

File tree

‎infra/main.parameters.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
"existingFoundryProjectResourceId": {
3030
"value": "${AZURE_EXISTING_AIPROJECT_RESOURCE_ID}"
3131
},
32+
"enableMonitoring": {
33+
"value": true
34+
},
3235
"containerRegistryEndpoint": {
3336
"value": "${AZURE_ENV_CONTAINER_REGISTRY_ENDPOINT}"
3437
},

‎src/ContentProcessor/pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ dependencies = [
99
"azure-ai-inference==1.0.0b9",
1010
"azure-appconfiguration==1.8.0",
1111
"azure-identity==1.26.0b1",
12+
"azure-monitor-events-extension>=0.1.0",
1213
"azure-monitor-opentelemetry==1.8.7",
1314
"azure-storage-blob==12.29.0b1",
1415
"azure-storage-queue==12.16.0b1",

‎src/ContentProcessor/src/libs/token_usage_utils.py‎

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -56,21 +56,35 @@ def extract_token_usage(response: Any) -> dict[str, int]:
5656

5757
# Path 1: usage_details attribute (set by agent framework SDK)
5858
usage_details = getattr(response, "usage_details", None)
59-
if isinstance(usage_details, dict):
60-
input_tokens = _to_int(
61-
usage_details.get("input_token_count")
62-
or usage_details.get("prompt_tokens")
63-
or usage_details.get("input_tokens")
64-
)
65-
output_tokens = _to_int(
66-
usage_details.get("output_token_count")
67-
or usage_details.get("completion_tokens")
68-
or usage_details.get("output_tokens")
69-
)
70-
total_tokens = _to_int(
71-
usage_details.get("total_token_count")
72-
or usage_details.get("total_tokens")
73-
) or (input_tokens + output_tokens)
59+
if usage_details is not None:
60+
if isinstance(usage_details, dict):
61+
input_tokens = _to_int(
62+
usage_details.get("input_token_count")
63+
or usage_details.get("prompt_tokens")
64+
or usage_details.get("input_tokens")
65+
)
66+
output_tokens = _to_int(
67+
usage_details.get("output_token_count")
68+
or usage_details.get("completion_tokens")
69+
or usage_details.get("output_tokens")
70+
)
71+
total_tokens = _to_int(
72+
usage_details.get("total_token_count")
73+
or usage_details.get("total_tokens")
74+
) or (input_tokens + output_tokens)
75+
else:
76+
# UsageDetails object with attributes
77+
input_tokens = _to_int(
78+
getattr(usage_details, "input_token_count", 0)
79+
or getattr(usage_details, "prompt_tokens", 0)
80+
)
81+
output_tokens = _to_int(
82+
getattr(usage_details, "output_token_count", 0)
83+
or getattr(usage_details, "completion_tokens", 0)
84+
)
85+
total_tokens = _to_int(
86+
getattr(usage_details, "total_token_count", 0)
87+
) or (input_tokens + output_tokens)
7488

7589
# Path 2: raw_representation.usage (raw Azure OpenAI response)
7690
if total_tokens == 0:

‎src/ContentProcessor/tests/unit/libs/test_token_usage_utils.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,21 @@ def test_raw_representation_dict_usage(self):
111111
"total_tokens": 75,
112112
}
113113

114+
def test_usage_details_object_with_attributes(self):
115+
"""Handle UsageDetails object (not dict) from agent framework."""
116+
response = MagicMock()
117+
usage_obj = MagicMock()
118+
usage_obj.input_token_count = 400
119+
usage_obj.output_token_count = 150
120+
usage_obj.total_token_count = 550
121+
response.usage_details = usage_obj
122+
result = extract_token_usage(response)
123+
assert result == {
124+
"input_tokens": 400,
125+
"output_tokens": 150,
126+
"total_tokens": 550,
127+
}
128+
114129
def test_no_usage_returns_zeros(self):
115130
response = MagicMock()
116131
response.usage_details = None

‎src/ContentProcessor/uv.lock‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/ContentProcessorWorkflow/pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ dependencies = [
1414
"azure-appconfiguration==1.8.0",
1515
"azure-core==1.38.0",
1616
"azure-identity==1.26.0b1",
17+
"azure-monitor-events-extension>=0.1.0",
1718
"azure-monitor-opentelemetry==1.8.7",
1819
"azure-storage-blob==12.29.0b1",
1920
"azure-storage-file-datalake==12.23.0",

‎src/ContentProcessorWorkflow/src/libs/token_usage_utils.py‎

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -56,21 +56,35 @@ def extract_token_usage(response: Any) -> dict[str, int]:
5656

5757
# Path 1: usage_details attribute (set by agent framework SDK)
5858
usage_details = getattr(response, "usage_details", None)
59-
if isinstance(usage_details, dict):
60-
input_tokens = _to_int(
61-
usage_details.get("input_token_count")
62-
or usage_details.get("prompt_tokens")
63-
or usage_details.get("input_tokens")
64-
)
65-
output_tokens = _to_int(
66-
usage_details.get("output_token_count")
67-
or usage_details.get("completion_tokens")
68-
or usage_details.get("output_tokens")
69-
)
70-
total_tokens = _to_int(
71-
usage_details.get("total_token_count")
72-
or usage_details.get("total_tokens")
73-
) or (input_tokens + output_tokens)
59+
if usage_details is not None:
60+
if isinstance(usage_details, dict):
61+
input_tokens = _to_int(
62+
usage_details.get("input_token_count")
63+
or usage_details.get("prompt_tokens")
64+
or usage_details.get("input_tokens")
65+
)
66+
output_tokens = _to_int(
67+
usage_details.get("output_token_count")
68+
or usage_details.get("completion_tokens")
69+
or usage_details.get("output_tokens")
70+
)
71+
total_tokens = _to_int(
72+
usage_details.get("total_token_count")
73+
or usage_details.get("total_tokens")
74+
) or (input_tokens + output_tokens)
75+
else:
76+
# UsageDetails object with attributes
77+
input_tokens = _to_int(
78+
getattr(usage_details, "input_token_count", 0)
79+
or getattr(usage_details, "prompt_tokens", 0)
80+
)
81+
output_tokens = _to_int(
82+
getattr(usage_details, "output_token_count", 0)
83+
or getattr(usage_details, "completion_tokens", 0)
84+
)
85+
total_tokens = _to_int(
86+
getattr(usage_details, "total_token_count", 0)
87+
) or (input_tokens + output_tokens)
7488

7589
# Path 2: raw_representation.usage (raw Azure OpenAI response)
7690
if total_tokens == 0:

0 commit comments

Comments
 (0)