Skip to content

Conversation

@steven10a
Copy link
Collaborator

Implements an update for issue 41

  • For Guardrails that use LLM's, return token usage data
  • Per guardrail token usage is returned via its info_dict and can be accessed with
const response = await client.guardrails.responses.create(...);
for (const gr of response.guardrail_results.allResults) {
  const usage = gr.info.token_usage;
  if (usage) {
    console.log(`${gr.info.guardrail_name}: ${usage.total_tokens} tokens`);
  }
}
  • Additionally, the total tokens used in a response can be returned with the totalGuardrailTokenUsage helper
    totalGuardrailTokenUsage(response). This works for all clients (GuardrailAgent, GuardrailAsyncOpenAI, ...) and works with streaming and non-streaming
  • Updated the documentation to reflect this new functionality
  • Updated example scripts to show implementation
  • Added tests

Note: Known limitation, we do not return the token count for GuardrailAgent. OpenAI's Agent SDK JS version is slightly different from the Python implementation. The RunResult does not contain details about the guardrail results so we can not pass token count back. This is noted in our Quickstart.

Copilot AI review requested due to automatic review settings December 2, 2025 17:16
Copilot finished reviewing on behalf of steven10a December 2, 2025 17:21
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds token usage tracking for LLM-based guardrails, enabling developers to monitor and analyze the token costs of running guardrails. The implementation provides both per-guardrail token counts and aggregated totals across all guardrails in a response.

Key Changes:

  • Introduced TokenUsage and TokenUsageSummary types with helper functions to extract and aggregate token usage
  • Updated all LLM-based guardrails (Jailbreak, Prompt Injection Detection, Hallucination Detection, and LLM-base) to return token usage data
  • Added totalGuardrailTokenUsage() helper function for easy access to aggregated token counts from any result object
  • Updated examples and documentation to demonstrate token usage tracking

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
src/types.ts Core token usage types (TokenUsage, TokenUsageSummary) and utility functions for extraction, aggregation, and conversion
src/index.ts Exported new token usage types and helper function
src/checks/llm-base.ts Modified runLLM to return token usage alongside analysis results; updated createErrorResult to include token usage
src/checks/prompt_injection_detection.ts Integrated token usage tracking with proper fallback handling for skipped checks and LLM failures
src/checks/jailbreak.ts Added token usage tracking to jailbreak guardrail results
src/checks/hallucination-detection.ts Added token usage extraction from Responses API and included in results
src/base-client.ts Added totalTokenUsage property to GuardrailResults interface and implementation
src/tests/unit/types.test.ts Added tests for token usage extraction, aggregation, and helper functions
src/tests/unit/prompt_injection_detection.test.ts Updated test to verify token_usage field presence
src/tests/unit/llm-base.test.ts Updated tests to verify token usage in success and error scenarios
src/tests/unit/checks/jailbreak.test.ts Updated mocks and assertions to handle token usage in return tuples
src/tests/unit/base-client.test.ts Added tests for token usage aggregation in GuardrailResultsImpl
src/tests/unit/agents.test.ts Added test for token usage propagation in agent guardrails; indentation issue on line 454
examples/basic/streaming.ts Added totalGuardrailTokenUsage() call to display token counts after streaming
examples/basic/multiturn_with_prompt_injection_detection.ts Added token usage display for initial and final guardrail runs
examples/basic/local_model.ts Added token usage display after processing input
examples/basic/hello_world.ts Added token usage display in guardrail results section
docs/quickstart.md Added comprehensive Token Usage Tracking section with examples for all client types
docs/agents_sdk_integration.md Documented JavaScript Agents SDK limitation and workaround for token usage tracking

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

@steven10a steven10a requested a review from Copilot December 2, 2025 17:31
@steven10a
Copy link
Collaborator Author

@codex review

Copilot finished reviewing on behalf of steven10a December 2, 2025 17:33
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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


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

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines +304 to +308
if (completion !== null) {
totalCompletion += completion;
}
if (total !== null) {
totalTokens += total;

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid zero totals when only prompt/completion counts exist

Aggregating token usage currently increments totalTokens only when a total_tokens field is present, so guardrails that report just prompt_tokens/completion_tokens (e.g., providers that emit input/output counts without a precomputed total) produce totalTokenUsage.total_tokens === 0 even though tokens were consumed. This under-reports usage for those models; the total should be derived from prompt+completion or left null rather than defaulting to zero.

Useful? React with 👍 / 👎.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WAI. The client interfaces our package supports all return prompt/completion/total tokens so we will not be receiving partial token dicts like this

@steven10a steven10a requested a review from Copilot December 2, 2025 17:44
Copilot finished reviewing on behalf of steven10a December 2, 2025 17:48
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.


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

@gabor-openai gabor-openai merged commit 0cdaaaf into main Dec 2, 2025
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants