Skip to content

Commit

Permalink
feat: Add usage for Bedrock Anthropic Claude models (#157)
Browse files Browse the repository at this point in the history
* Add usage tracking to AnthropicClaude3ChatModel

Introduce a `usage` variable to track usage information throughout the `AnthropicClaude3ChatModel` class. Add a `GetUsage` method to extract usage data from JSON nodes. Initialize `usage` to `null` and update it with data from JSON responses, ensuring it is non-null before use with a fallback to `Usage.Empty`. Update `usage` incrementally in streaming responses and include it in `ChatResponse`, passing it to `AddUsage` and `provider.AddUsage` methods. Remove `CONTENT_BLOCK_STOP` type handling from streaming response processing.

* Improve null safety in usage variable assignment

Modified the assignment of the `usage` variable to use `GetUsage(response?["usage"])` instead of `GetUsage(response)`. This change ensures that the `usage` variable is assigned from the "usage" field of the `response` object, if it exists, preventing potential null reference exceptions.
  • Loading branch information
GregoireDaure authored Jan 29, 2025
1 parent 27a9346 commit f69fc1a
Showing 1 changed file with 37 additions and 8 deletions.
45 changes: 37 additions & 8 deletions src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
modelSettings: Settings,
providerSettings: provider.ChatSettings);

Usage? usage = null;

var bodyJson = CreateBodyJson(prompt, usedSettings, request.Image);

if (usedSettings.UseStreaming == true)
Expand All @@ -51,7 +53,15 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
var streamEvent = (PayloadPart)payloadPart;
var chunk = await JsonSerializer.DeserializeAsync<JsonObject>(streamEvent.Bytes, cancellationToken: cancellationToken)
.ConfigureAwait(false);

usage ??= GetUsage(chunk?["message"]?["usage"]);
var type = chunk?["type"]!.GetValue<string>().ToUpperInvariant();

if (type == "MESSAGE_DELTA")
{
usage += GetUsage(chunk?["usage"]);
}

if (type == "CONTENT_BLOCK_DELTA")
{
var delta = chunk?["delta"]?["text"]!.GetValue<string>();
Expand All @@ -62,10 +72,6 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
});
stringBuilder.Append(delta);
}
if (type == "CONTENT_BLOCK_STOP")
{
break;
}
}

OnDeltaReceived(new ChatResponseDelta
Expand All @@ -82,24 +88,27 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
else
{
var response = await provider.Api.InvokeModelAsync(Id, bodyJson, cancellationToken).ConfigureAwait(false);
usage = GetUsage(response?["usage"]);

var generatedText = response?["content"]?[0]?["text"]?.GetValue<string>() ?? "";

messages.Add(generatedText.AsAiMessage());
}

var usage = Usage.Empty with
usage ??= Usage.Empty;
usage = usage.Value with
{
Time = watch.Elapsed,
Messages = messages.Count,
};
AddUsage(usage);
provider.AddUsage(usage);
AddUsage(usage.Value);
provider.AddUsage(usage.Value);

var chatResponse = new ChatResponse
{
Messages = messages,
UsedSettings = usedSettings,
Usage = usage,
Usage = usage.Value,
};
OnResponseReceived(chatResponse);

Expand Down Expand Up @@ -162,4 +171,24 @@ private static JsonObject CreateBodyJson(

return bodyJson;
}

/// <summary>
/// Extracts usage information from the provided JSON node.
/// </summary>
/// <param name="usageNode">The JSON node containing usage information.</param>
/// <returns>A <see cref="Usage"/> object with the extracted usage data.</returns>
private Usage GetUsage(JsonNode? usageNode)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)

Check warning on line 180 in src/Amazon.Bedrock/src/Chat/AnthropicClaude3ChatModel.cs

View workflow job for this annotation

GitHub Actions / Build, test and publish / Build, test and publish

Member 'GetUsage' does not access instance data and can be marked as static (https://learn.microsoft.com/dotnet/fundamentals/code-analysis/quality-rules/ca1822)
{
var inputTokens = usageNode?["input_tokens"]?.GetValue<int>() ?? 0;
var outputTokens = usageNode?["output_tokens"]?.GetValue<int>() ?? 0;
var priceInUsd = 0.0;

return Usage.Empty with
{
InputTokens = inputTokens,
OutputTokens = outputTokens,
Messages = 0,
PriceInUsd = priceInUsd,
};
}
}

0 comments on commit f69fc1a

Please sign in to comment.