From 395bbdbaf8e3940f2885dc59e02792abf0d71f4c Mon Sep 17 00:00:00 2001 From: NikitolProject Date: Thu, 26 Mar 2026 13:16:29 +0300 Subject: [PATCH] add gen_ai.usage span attributes from conversation response --- ai_core_api_client/api/ai_api.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/ai_core_api_client/api/ai_api.py b/ai_core_api_client/api/ai_api.py index 308ea18..b97d001 100644 --- a/ai_core_api_client/api/ai_api.py +++ b/ai_core_api_client/api/ai_api.py @@ -26,6 +26,27 @@ from ai_core_api_client.api_response import ApiResponse from ai_core_api_client.rest import RESTResponseType +def _record_usage_to_span(response) -> None: + """Record token usage from ai-core-api response as OTel span attributes.""" + try: + usage = getattr(response, 'usage', None) + if usage is None: + return + from opentelemetry import trace + span = trace.get_current_span() + if not span.is_recording(): + return + model = getattr(usage, 'model', None) + if model: + span.set_attribute("gen_ai.usage.model", model) + for attr in ('prompt_tokens', 'completion_tokens', 'total_tokens', 'cached_tokens'): + val = getattr(usage, attr, None) + if val is not None: + span.set_attribute(f"gen_ai.usage.{attr}", val) + except Exception: + pass + + class AiApi: """NOTE: This class is auto generated by OpenAPI Generator Ref: https://openapi-generator.tech @@ -103,11 +124,15 @@ class AiApi: _request_timeout=_request_timeout ) response_data.read() - return self.api_client.response_deserialize( + result = self.api_client.response_deserialize( response_data=response_data, response_types_map=_response_types_map, ).data + _record_usage_to_span(result) + + return result + @validate_call def conversation_post_with_http_info(