feat: Add token usage to OpenAIPrompt and OpenAIEmbedding

[ranadeepsingh]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Extending usage of OpenAIPrompt and OpenAIEmbedding which are the core usage functions powering PySpark AI functions to have a new boolean returnUsage paramater.
This parameter is False by default returning just string or list response.
When it is true, it returns a strict with the following members:

1. "response": string for OpenAIPrompt() or DenseVector for OpenAIEmbedding()
2. "usage": usage tokens as stuct with the following fields:

• input_tokens: Long
• output_tokens: Long
• total_tokens: Long
• input_token_details: Map - This field is dynamic and depends on the model. It can contain details like "cached_tokens"
• output_token_details: Map - This field is dynamic and depends on the model. It can contain details like "reasoning_tokens"

Screenshot of output types and usage:
Works across all three different openai APIs

• OpenAIPrompt

• OpenAIEmbedding

How is this patch tested?

• I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.

Does this PR change any dependencies?

• No. You can skip this section.
• Yes. Make sure the dependencies are resolved correctly, and list changes here.

Does this PR add a new feature? If so, have you added samples on website?

• No. You can skip this section.
• Yes. Make sure you have added samples following below steps.

1. Find the corresponding markdown file for your new feature in website/docs/documentation folder.
   Make sure you choose the correct class estimators/transformers and namespace.
2. Follow the pattern in markdown file and add another section for your new API, including pyspark, scala (and .NET potentially) samples.
3. Make sure the DocTable points to correct API link.
4. Navigate to website folder, and run yarn run start to make sure the website renders correctly.
5. Don't forget to add <!--pytest-codeblocks:cont--> before each python code blocks to enable auto-tests for python samples.
6. Make sure the WebsiteSamplesTests job pass in the pipeline.

[ranadeepsingh]

/azp run

[github-actions]

Hey @ranadeepsingh 👋!
Thank you so much for contributing to our repository 🙌.
Someone from SynapseML Team will be reviewing this pull request soon.

We use semantic commit messages to streamline the release process.
Before your pull request can be merged, you should make sure your first commit and PR title start with a semantic prefix.
This helps us to create release messages and credit you for your hard work!

Examples of commit messages with semantic prefixes:

• fix: Fix LightGBM crashes with empty partitions
• feat: Make HTTP on Spark back-offs configurable
• docs: Update Spark Serving usage
• build: Add codecov support
• perf: improve LightGBM memory usage
• refactor: make python code generation rely on classes
• style: Remove nulls from CNTKModel
• test: Add test coverage for CNTKModel

To test your commit locally, please follow our guild on building from source.
Check out the developer guide for additional guidance on testing your change.

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[codecov-commenter]

Codecov Report

❌ Patch coverage is 94.91525% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.43%. Comparing base (df189c8) to head (f0a384b).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
...zure/synapse/ml/services/openai/OpenAIPrompt.scala	83.33%	5 Missing ⚠️
...azure/synapse/ml/services/openai/ReturnUsage.scala	98.38%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2444      +/-   ##
==========================================
- Coverage   84.48%   82.43%   -2.05%     
==========================================
  Files         334      335       +1     
  Lines       17591    17690      +99     
  Branches     1601     1619      +18     
==========================================
- Hits        14862    14583     -279     
- Misses       2729     3107     +378     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Copilot]

Pull Request Overview

This PR adds token usage tracking functionality to OpenAIPrompt and OpenAIEmbedding transformers. When enabled, these components return structured output containing both the response and usage statistics (input/output tokens, total tokens, and detailed token breakdowns).

• Introduces a returnUsage boolean parameter (default: false) to control whether usage statistics are included
• Adds a new HasReturnUsage trait and ReturnUsage.scala to standardize usage tracking across different OpenAI APIs
• Updates response schemas to include optional usage fields for Chat Completions, Responses, and Embeddings APIs

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
ReturnUsage.scala	New trait and utilities for standardizing usage statistics across different OpenAI API types
OpenAIPrompt.scala	Integrates HasReturnUsage trait; modifies transform and transformSchema to conditionally return usage data
OpenAIEmbedding.scala	Integrates HasReturnUsage trait; overrides transform and transformSchema to support usage tracking
OpenAISchemas.scala	Updates response case classes to include optional usage fields and token details
OpenAIPromptSuite.scala	Adds tests for default, false, and true returnUsage parameter behaviors
OpenAIEmbeddingsSuite.scala	Adds tests for default, false, and true returnUsage parameter behaviors

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ranadeepsingh]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).