fix(instrumentations): add cache token span attributes for OpenAI, Bedrock, Together AI, and Vertex AI#1026
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
✅ Files skipped from review due to trivial changes (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughCache token span attribute guards are tightened from loose ChangesCache Token Guard and Calculation Fixes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~5 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/instrumentation-together/src/instrumentation.ts (1)
525-533: ⚡ Quick winConsider extending the Completion type instead of using a weak type cast.
The current implementation uses
as unknown as Record<string, unknown>to accesscached_tokens, which bypasses TypeScript's type safety. Following the pattern established in the OpenAI instrumentation (lines 117-119 inpackages/instrumentation-openai/src/instrumentation.ts), consider defining a local type extension at the top of the file:interface CompletionUsageWithCache extends Completion.CompletionUsage { cached_tokens?: number; }Then access it with better type safety:
- const cachedTokens = ( - result.usage as unknown as Record<string, unknown> - ).cached_tokens; - if (cachedTokens) { + const cachedTokens = (result.usage as CompletionUsageWithCache) + ?.cached_tokens; + if (typeof cachedTokens === 'number') { span.setAttribute( ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, - cachedTokens as number, + cachedTokens, ); }This approach provides better type safety and is more maintainable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/instrumentation-together/src/instrumentation.ts` around lines 525 - 533, The code uses a weak type cast (as unknown as Record<string, unknown>) to access cached_tokens from the result.usage object, which bypasses TypeScript's type safety. Define a local interface extension at the top of the instrumentation.ts file that extends Completion.CompletionUsage with an optional cached_tokens property (following the pattern in the OpenAI instrumentation), then replace the weak type cast with this properly typed interface when accessing cached_tokens before the span.setAttribute call for ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS. This improves type safety and maintainability without changing the functional behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/instrumentation-bedrock/src/instrumentation.ts`:
- Around line 706-717: Cache token attributes are currently being guarded with
truthy checks, which silently drops valid zero values from provider responses.
In packages/instrumentation-bedrock/src/instrumentation.ts at lines 706-717,
replace the truthy checks for usage["cache_read_input_tokens"] and
usage["cache_creation_input_tokens"] with null/undefined checks using != null or
!== undefined operators to preserve zero values in
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS and
ATTR_GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS attributes. In
packages/instrumentation-vertexai/src/vertexai-instrumentation.ts at lines
256-260, similarly replace the truthy check for cachedContentTokenCount with a
!== undefined or != null check when setting
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS to ensure zero values are faithfully
emitted.
---
Nitpick comments:
In `@packages/instrumentation-together/src/instrumentation.ts`:
- Around line 525-533: The code uses a weak type cast (as unknown as
Record<string, unknown>) to access cached_tokens from the result.usage object,
which bypasses TypeScript's type safety. Define a local interface extension at
the top of the instrumentation.ts file that extends Completion.CompletionUsage
with an optional cached_tokens property (following the pattern in the OpenAI
instrumentation), then replace the weak type cast with this properly typed
interface when accessing cached_tokens before the span.setAttribute call for
ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS. This improves type safety and
maintainability without changing the functional behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 10f66fbf-e369-4454-b281-f604f43ee0e1
⛔ Files ignored due to path filters (1)
pnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (10)
packages/instrumentation-bedrock/src/instrumentation.tspackages/instrumentation-bedrock/tests/anthropic.test.tspackages/instrumentation-openai/src/instrumentation.tspackages/instrumentation-openai/test/instrumentation.test.tspackages/instrumentation-together/package.jsonpackages/instrumentation-together/src/instrumentation.tspackages/instrumentation-together/test/instrumentation.test.tspackages/instrumentation-vertexai/package.jsonpackages/instrumentation-vertexai/src/vertexai-instrumentation.tspackages/instrumentation-vertexai/tests/gemini.test.ts
45a1278 to
ce12fa3
Compare
Summary by CodeRabbit
Bug Fixes
Tests