Skip to content

Commit

Permalink
fix Bedrock usage
Browse files Browse the repository at this point in the history
  • Loading branch information
Daniel OBrien committed Nov 1, 2024
1 parent 7c5c585 commit 9c6fdd3
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 18 deletions.
15 changes: 12 additions & 3 deletions packages/sdk/ai/src/LDAIConfigTracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { LDClient, LDContext } from '@launchdarkly/node-server-sdk';

import {
BedrockTokenUsage,
createBedrockTokenUsage,
FeedbackKind,
OpenAITokenUsage,
TokenUsage,
Expand Down Expand Up @@ -66,17 +67,25 @@ export class LDAIConfigTracker {
return result;
}

async trackBedrockConverse(res: any): Promise<any> {
async trackBedrockConverse(res: {
$metadata?: { httpStatusCode: number };
metrics?: { latencyMs: number };
usage?: {
inputTokens: number;
outputTokens: number;
totalTokens: number;
};
}): Promise<any> {
if (res.$metadata?.httpStatusCode === 200) {
this.trackGeneration(1);
} else if (res.$metadata?.httpStatusCode >= 400) {
} else if (res.$metadata?.httpStatusCode && res.$metadata.httpStatusCode >= 400) {
this.trackError(res.$metadata.httpStatusCode);
}
if (res.metrics) {
this.trackDuration(res.metrics.latencyMs);
}
if (res.usage) {
this.trackTokens(new BedrockTokenUsage(res.usage));
this.trackTokens(createBedrockTokenUsage(res.usage));
}
return res;
}
Expand Down
22 changes: 7 additions & 15 deletions packages/sdk/ai/src/api/metrics/BedrockTokenUsage.ts
Original file line number Diff line number Diff line change
@@ -1,19 +1,11 @@
export class BedrockTokenUsage {
export function createBedrockTokenUsage(data: {
totalTokens: number;
inputTokens: number;
outputTokens: number;

constructor(data: any) {
this.totalTokens = data.totalTokens || 0;
this.inputTokens = data.inputTokens || 0;
this.outputTokens = data.outputTokens || 0;
}

toMetrics() {
return {
total: this.totalTokens,
input: this.inputTokens,
output: this.outputTokens,
};
}
}) {
return {
total: data.totalTokens || 0,
input: data.inputTokens || 0,
output: data.outputTokens || 0,
};
}

0 comments on commit 9c6fdd3

Please sign in to comment.