[Feature Request] Support for llm-emit-token-metric in AIProjectClient ChatCompletionsClient for Azure OpenAI Token Tracking

[o1100]

Feature Request: Enable llm-emit-token-metric or azure-openai-emit-token-metric in AIProjectClient's ChatCompletionsClient

---

Is your feature request related to a problem? Please describe.

I am looking to use the llm-emit-token-metric or azure-openai-emit-token-metric headers within Azure API Management (APIM) to capture token usage. My setup uses Azure Foundry alongside the new azure.ai.projects AIProjectClient and its corresponding ChatCompletionsClient.

While I understand how to enable token metric headers when performing regular POST requests to Azure OpenAI endpoints, I have been unable to find functionality to do this within the ChatCompletionsClient.

The goal is to configure an APIM instance for the endpoint https://<deployment>.openai.azure.com/ and log token metrics to Azure Log Analytics.

---

Describe the solution you'd like

I would like the ChatCompletionsClient (from azure.ai.projects) to include functionality for enabling and utilizing the llm-emit-token-metric or azure-openai-emit-token-metric headers when interacting with Azure OpenAI services.

This would allow token usage tracking to seamlessly integrate with APIM and Log Analytics while leveraging the new client library.

---

Describe alternatives you've considered

• I have been able to use regular POST requests to Azure OpenAI endpoints and successfully capture token metrics with the llm-emit-token-metric or azure-openai-emit-token-metric headers.
• However, I cannot find similar functionality within the ChatCompletionsClient of azure.ai.projects.

---

Additional context

• I have reviewed the Journey of the Geek blog post for guidance on this matter. While their solution effectively captures token metrics, it does not utilize the new AIProjectClient or ChatCompletionsClient from the azure.ai.projects library.

[github-actions]

Thank you for your feedback. Tagging and routing to the team member best able to assist.

[github-actions]

Thanks for the feedback! We are routing this to the appropriate team for follow-up. cc @achauhan-scc @kingernupur @luigiw @needuv @paulshealy1 @singankit.

[o1100]

Hi @achauhan-scc @kingernupur @luigiw @needuv @paulshealy1 @singankit,

I wanted to check in on this request. It has been two weeks since it was opened, and I haven't seen any updates. Could you provide any timeline or status on this?

Thanks for your time!

[O-EAI]

I second this, have been looking for a way to track token usage with APIM and the new chat completion clients.

[kingernupur]

Tagging @dargilco

[dargilco]

@O-EAI @o1100 the issue just now got assigned to the right team. I'll look at it today and respond here. If/when you open another GitHub issue, please mention the package name (azure-ai-projects in this case) for fast triage.

[lmolkova]

Is the ask specific to APIM in some way? Azure AI Services report token usage metrics on the server-side - https://learn.microsoft.com/en-us/azure/azure-monitor/reference/supported-metrics/microsoft-cognitiveservices-accounts-metrics by default. Is there a case when something else is necessary?

[dargilco]

@O-EAI @o1100 in addition to @lmolkova question above, if you say you got it "to work" with REST API calls (per spec https://learn.microsoft.com/en-us/azure/ai-foundry/model-inference/reference/reference-model-inference-api?tabs=rest) and you are asking how to do the equivalent with the SDK, if you send me the details of your REST call I can look at how to configure the Python ChatCompletionsClient to make similar REST API call. It may be possible to do it right now, as you can always add additional HTTP request headers and JSON elements in the request payload.

[dargilco]

For my reference, adding this public doc: "Overview of generative AI gateway capabilities in Azure API Management" https://learn.microsoft.com/en-us/azure/api-management/genai-gateway-capabilities

[o1100]

Note:
I think my original post could have been clearer. I want to use the ChatCompletionsClient but forward requests through an Azure API Management instance so that I can collect azure-openai-emit-token-metric. I can easily do this with direct POST requests, but project_client.inference.get_chat_completions_client() only lets me use predefined connections.

@lmolkova Azure AI Services can emit token metrics per deployed model/resource but in the instance where I have a shared endpoint for multiple customers/departments, I can't distinguish how many tokens each department uses. I can only see aggregate token metrics.

@dargilco I am able to do a REST API call to the APIM endpoint directly and communicate with the AI Services Endpoint (extracting token metrics and logging results using APIM policies), however the ChatCompletionsClient only lets you communicate with predetermined Azure AI connections.

What I am trying to emulate with ChatCompletionsClient:

headers = {
    "api-key": api_key,
    "Content-Type": "application/json",
    "billingID": "12345"
}

payload = {
    "messages": [
        {"role": "user", "content": "Hi"}
    ],
    "max_tokens": 50
}

response = requests.post(apim_endpoint, headers=headers, json=payload)
print(response)

Returns 200 OK, and I can see token metrics:

[dargilco]

@o1100 You can create an instance of the ChatCompletionsClient from azure-ai-inference, provide the inference URL endpoint, provide the api-key, and add any additional headers you want using the code below:

client = ChatCompletionsClient(
    endpoint="your-inference-endpoint ", 
    credentials=AzureKeyCredentials("your-api-key"),
    headers={"billingID": "12345"}  # A dict of additional HTTP request headers
)

You can turn on client console logging, to see the actual request payload by doing the following:

# Put this at the top of your Python souce file:
import sys
import logging
logger = logging.getLogger("azure")
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler(stream=sys.stdout))

# And add this as an additional input argument to the ChatComletionsClient constructor:
logging_enable=True

If you are using AIProjectClient, you can use the .connection operations to get the properties of the relevant connection, including the endpoint to be used for inference, and create an instance of the ChatCompletionsClient yourself.

If you want to use the .inference.get_chat_completions_client() to get your authenticated ChatCompletionsClient for the default AIServices connection, you can call:

project_client.inference.get_chat_completions_client(headers={"billingID": "12345"})

as the keywords that are entered in the call to get_chat_completions_client method get passed directly to the constructor of the ChatCompletionsClient.

Note that you also have the option to explicitly specify the connection name if the default does not work for you:

project_client.inference.get_chat_completions_client(connection_name="my-connection-name", headers={"billingID": "12345"})