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(