frag-openai-usage.xml

<fragment>
	<!-- Usage logs for non-streaming requests only -->
	<choose>
		<when condition="@(context.Variables.GetValueOrDefault<string>("isStream","false").Equals("false", StringComparison.OrdinalIgnoreCase))">
			<!-- Set the response body as a variable to be used in various policies -->
			<set-variable name="responseBody" value="@(context.Response.Body.As<JObject>())" />
			<!-- Log OpenAI usage to EventHub -->
			<choose>
				<when condition="@(context.Response.StatusCode == 200)">
					<log-to-eventhub logger-id="usage-eventhub-logger">@{
                    //Avoid reading response body as it can only be ready once, instead, before calling this fragement, a variable call responseBody will be set in the outbound policy
                    //var responseBody = context.Response.Body?.As<JObject>(true); //Avoid this one
                    var responseBody = (JObject)context.Variables["responseBody"]; //It is set in the outbound policy before calling the fragment
                    return new JObject(
                        new JProperty("id", responseBody?["id"]?.ToString() ?? Guid.NewGuid().ToString()),
                        new JProperty("timestamp", DateTime.UtcNow.ToString()),
                        new JProperty("appId", context.Request.Headers.GetValueOrDefault("Authorization",string.Empty).Split(' ').LastOrDefault()?.AsJwt()?.Claims.GetValueOrDefault("appid", "NA")),
                        new JProperty("subscriptionId", context.Subscription?.Id?.ToString() ?? "Portal-Admin"),
                        new JProperty("productName", context.Product?.Name?.ToString() ?? "Portal-Admin"),
                        new JProperty("targetService", responseBody?["object"]?.ToString() ?? "NA"),
                        new JProperty("model", responseBody?["model"]?.ToString() ?? "NA"),
                        new JProperty("gatewayName", context.Deployment?.ServiceName ?? "NA"),
                        new JProperty("gatewayRegion", context.Deployment?.Region ?? "NA"),
                        new JProperty("aiGatewayId", context.Deployment?.Gateway?.Id ?? "NA"),
                        new JProperty("RequestIp", context.Request?.IpAddress ?? "NA"),
                        new JProperty("operationName", context.Operation?.Name ?? "NA"),
                        new JProperty("sessionId", (string)context.Variables.GetValueOrDefault<string>("sessionId", "NA")),
                        new JProperty("endUserId", (string)context.Variables.GetValueOrDefault<string>("endUserId", "NA")),
                        new JProperty("backendId", (string)context.Variables.GetValueOrDefault<string>("backendId", "NA")),
                        new JProperty("routeLocation", (string)context.Variables.GetValueOrDefault<string>("routeLocation", "NA")),
                        new JProperty("routeName", (string)context.Variables.GetValueOrDefault<string>("routeName", "NA")),
                        new JProperty("deploymentName", (string)context.Variables.GetValueOrDefault<string>("deploymentName", "NA")),
                        new JProperty("promptTokens", responseBody?["usage"]?["prompt_tokens"]?.Value<int>() ?? 1),
                        new JProperty("responseTokens", responseBody?["usage"]?["completion_tokens"]?.Value<int>() ?? 0),
                        new JProperty("totalTokens", responseBody?["usage"]?["total_tokens"]?.Value<int>() ?? 1)
                ).ToString();
                }</log-to-eventhub>
				</when>
			</choose>
			<set-body>@{
                    return ((JObject)context.Variables["responseBody"]).ToString();
            }</set-body>
		</when>
	</choose>
</fragment>