<policies>
  <inbound>
    <base />
    <set-variable name="attemptCount" value="@(0)" />
    <rewrite-uri template="/openai/deployments/YOUR_DEPLOYMENT_NAME/chat/completions" copy-unmatched-params="false" />
    <set-query-parameter name="api-version" exists-action="override"><value>2024-10-21</value></set-query-parameter>
    <set-header name="api-key" exists-action="delete" />
    <set-header name="Ocp-Apim-Subscription-Key" exists-action="delete" />
    <set-header name="Cookie" exists-action="delete" />
    <set-header name="x-ms-client-request-id" exists-action="override"><value>@(context.RequestId.ToString())</value></set-header>
    <authentication-managed-identity resource="https://cognitiveservices.azure.com" ignore-error="false" />
  </inbound>
  <backend>
    <retry condition="@(true &amp;&amp; context.Response != null &amp;&amp; context.Response.StatusCode == 429)" count="1" interval="1" first-fast-retry="false">
      <set-variable name="attemptCount" value="@(context.Variables.GetValueOrDefault&lt;int&gt;(&quot;attemptCount&quot;, 0) + 1)" />
      <set-backend-service backend-id="YOUR_POOL_ID" />
      <forward-request timeout="60" http-version="1" follow-redirects="false" buffer-request-body="true" buffer-response="false" fail-on-error-status-code="false" />
      <trace source="YOUR_TRACE_SOURCE" severity="information">
        <message>backend-attempt</message>
        <metadata name="requestId" value="@(context.RequestId.ToString())" />
        <metadata name="operation" value="@(context.Operation.Id)" />
        <metadata name="attempt" value="@(context.Variables.GetValueOrDefault&lt;int&gt;(&quot;attemptCount&quot;, 0).ToString())" />
        <metadata name="configuredBackend" value="YOUR_POOL_ID" />
        <metadata name="backendId" value="@(context.Backend.Id)" />
        <metadata name="backendHost" value="@(context.Request.Url.Host)" />
        <metadata name="status" value="@(context.Response.StatusCode.ToString())" />
        <metadata name="elapsedMs" value="@(context.Elapsed.TotalMilliseconds.ToString())" />
        <metadata name="retryAfter" value="@{ var header = context.Response.Headers.GetValueOrDefault(&quot;Retry-After&quot;, &quot;&quot;); return string.IsNullOrWhiteSpace(header) ? &quot;not-present&quot; : header; }" />
        <metadata name="retryAfterMs" value="@{ var header = context.Response.Headers.GetValueOrDefault(&quot;retry-after-ms&quot;, &quot;&quot;); return string.IsNullOrWhiteSpace(header) ? &quot;not-present&quot; : header; }" />
      </trace>
      <choose id="legacy-token-usage">
        <when condition="@(context.Response.StatusCode == 200 &amp;&amp; context.Response.Headers.GetValueOrDefault(&quot;Content-Type&quot;, &quot;&quot;).StartsWith(&quot;application/json&quot;, StringComparison.OrdinalIgnoreCase))">
          <set-variable name="llmUsage" value="@{
            var body = context.Response.Body.As&lt;JObject&gt;(preserveContent: true);
            var usage = body[&quot;usage&quot;] as JObject;
            var input = usage == null ? null : (usage[&quot;prompt_tokens&quot;] ?? usage[&quot;input_tokens&quot;]);
            var output = usage == null ? null : (usage[&quot;completion_tokens&quot;] ?? usage[&quot;output_tokens&quot;]);
            var total = usage == null ? null : usage[&quot;total_tokens&quot;];
            var reported = input != null &amp;&amp; input.Type == JTokenType.Integer &amp;&amp;
              output != null &amp;&amp; output.Type == JTokenType.Integer &amp;&amp;
              total != null &amp;&amp; total.Type == JTokenType.Integer;
            return new JObject {
              [&quot;model&quot;] = (string)body[&quot;model&quot;] ?? &quot;not-reported&quot;,
              [&quot;usageStatus&quot;] = reported ? &quot;reported&quot; : &quot;not-reported&quot;,
              [&quot;inputTokens&quot;] = reported ? input : null,
              [&quot;outputTokens&quot;] = reported ? output : null,
              [&quot;totalTokens&quot;] = reported ? total : null
            }.ToString();
          }" />
          <trace source="YOUR_TRACE_SOURCE" severity="information">
            <message>llm-usage</message>
            <metadata name="requestId" value="@(context.RequestId.ToString())" />
            <metadata name="subscriptionId" value="@(context.Subscription == null ? &quot;no-subscription&quot; : context.Subscription.Id)" />
            <metadata name="backendHost" value="@(context.Request.Url.Host)" />
            <metadata name="attempt" value="@(context.Variables.GetValueOrDefault&lt;int&gt;(&quot;attemptCount&quot;, 0).ToString())" />
            <metadata name="usage" value="@((string)context.Variables[&quot;llmUsage&quot;])" />
          </trace>
        </when>
      </choose>
    </retry>
  </backend>
  <outbound>
    <base />
    <set-header name="X-PoC-Request-Id" exists-action="override"><value>@(context.RequestId.ToString())</value></set-header>
    <set-header name="X-PoC-Attempts" exists-action="override"><value>@(context.Variables.GetValueOrDefault&lt;int&gt;("attemptCount", 0).ToString())</value></set-header>
    <set-header name="X-PoC-Backend-Host" exists-action="override"><value>@(context.Request.Url.Host)</value></set-header>
    <set-header name="X-PoC-Backend-Id" exists-action="override"><value>@(context.Backend.Id)</value></set-header>
    <set-header name="X-PoC-Error-Source" exists-action="override"><value>@(context.Response.StatusCode == 429 ? "backend" : "none")</value></set-header>
  </outbound>
  <on-error>
    <base />
    <set-header name="X-PoC-Request-Id" exists-action="override"><value>@(context.RequestId.ToString())</value></set-header>
    <set-header name="X-PoC-Attempts" exists-action="override"><value>@(context.Variables.GetValueOrDefault&lt;int&gt;("attemptCount", 0).ToString())</value></set-header>
    <set-header name="X-PoC-Error-Source" exists-action="override"><value>gateway</value></set-header>
  </on-error>
</policies>
