Add support for Count Message tokens endpoint

tghamm · Jan 28, 2025 · 2dca3f9 · 2dca3f9
1 parent 170bbbc
commit 2dca3f9
Show file tree

Hide file tree

Showing 7 changed files with 72 additions and 24 deletions.
diff --git a/Anthropic.SDK.Tests/Messages.cs b/Anthropic.SDK.Tests/Messages.cs
@@ -81,6 +81,21 @@ public async Task TestBasicClaude35HaikuMessage()
             Assert.IsNotNull(res.Message.ToString());
         }
 
+        [TestMethod]
+        public async Task TestBasicTokenCountMessage()
+        {
+            var client = new AnthropicClient();
+            var messages = new List<Message>();
+            messages.Add(new Message(RoleType.User, "Write me a sonnet about the Statue of Liberty"));
+            var parameters = new MessageCountTokenParameters
+            {
+                Messages = messages,
+                Model = AnthropicModels.Claude35Haiku
+            };
+            var res = await client.Messages.CountMessageTokensAsync(parameters);
+            Assert.IsTrue(res.InputTokens > 0);
+        }
+
         [TestMethod]
         public async Task TestStreamingClaude3HaikuMessage()
         {

diff --git a/Anthropic.SDK/EndpointBase.cs b/Anthropic.SDK/EndpointBase.cs
@@ -1,4 +1,4 @@
-using Anthropic.SDK.Extensions;
+using Anthropic.SDK.Extensions;
 using System;
 using System.Collections.Generic;
 using System.IO;
@@ -111,7 +111,7 @@ private string GetErrorMessage(string resultAsString, HttpResponseMessage respon
             return $"{resultAsString ?? "<no content>"}";
         }
 
-        protected async Task<MessageResponse> HttpRequestMessages(string url = null, HttpMethod verb = null,
+        protected async Task<TResponse> HttpRequestMessages<TResponse>(string url = null, HttpMethod verb = null,
             object postData = null, CancellationToken ctx = default)
         {
             var response = await HttpRequestRaw(url, verb, postData, ctx: ctx).ConfigureAwait(false);
@@ -124,10 +124,13 @@ protected async Task<MessageResponse> HttpRequestMessages(string url = null, Htt
             {
                 Converters = { ContentConverter.Instance }
             };
-            var res = await JsonSerializer.DeserializeAsync<MessageResponse>(
+            var res = await JsonSerializer.DeserializeAsync<TResponse>(
                 new MemoryStream(Encoding.UTF8.GetBytes(resultAsString)), options, cancellationToken: ctx);
 
-            res.RateLimits = GetRateLimits(response);
+            if (res is MessageResponse messageResponse)
+            {
+                messageResponse.RateLimits = GetRateLimits(response);
+            }
 
             return res;
         }

diff --git a/Anthropic.SDK/Extensions/MessageParameterConverter.cs b/Anthropic.SDK/Extensions/MessageParameterConverter.cs
@@ -1,26 +1,24 @@
 using System;
-using System.Collections.Generic;
 using System.Reflection;
-using System.Text;
 using System.Text.Json;
 using System.Text.Json.Serialization;
 using Anthropic.SDK.Messaging;
 
 namespace Anthropic.SDK.Extensions
 {
-    public class MessageParametersConverter : JsonConverter<MessageParameters>
+    public class MessageParametersConverter<T> : JsonConverter<T> where T: MessageCountTokenParameters
     {
-        public override MessageParameters Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+        public override T Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
         {
             // Implement the Read method if deserialization is needed
             throw new NotImplementedException();
         }
 
-        public override void Write(Utf8JsonWriter writer, MessageParameters value, JsonSerializerOptions options)
+        public override void Write(Utf8JsonWriter writer, T value, JsonSerializerOptions options)
         {
             writer.WriteStartObject();
 
-            var properties = typeof(MessageParameters).GetProperties(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance);
+            var properties = typeof(T).GetProperties(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance);
 
             foreach (var property in properties)
             {

diff --git a/Anthropic.SDK/Messaging/MessageCountTokenResponse.cs b/Anthropic.SDK/Messaging/MessageCountTokenResponse.cs
@@ -0,0 +1,9 @@
+using System.Text.Json.Serialization;
+
+namespace Anthropic.SDK.Messaging;
+
+public class MessageCountTokenResponse
+{
+    [JsonPropertyName("input_tokens")]
+    public int InputTokens { get; set; }
+}
diff --git a/Anthropic.SDK/Messaging/MessageParameters.cs b/Anthropic.SDK/Messaging/MessageParameters.cs
@@ -1,21 +1,30 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Linq;
-using System.Text;
 using System.Text.Json.Serialization;
 using Anthropic.SDK.Extensions;
 
 namespace Anthropic.SDK.Messaging
 {
-    [JsonConverter(typeof(MessageParametersConverter))]
-    public class MessageParameters
+    [JsonConverter(typeof(MessageParametersConverter<MessageCountTokenParameters>))]
+    public class MessageCountTokenParameters
     {
         [JsonPropertyName("model")]
         public string Model { get; set; }
         [JsonPropertyName("messages")]
         public List<Message> Messages { get; set; }
         [JsonPropertyName("system")]
         public List<SystemMessage> System { get; set; }
+        [JsonPropertyName("tools")]
+        private List<Common.Function> ToolsForClaude => Tools?.Select(p => p.Function).ToList();
+        [JsonIgnore]
+        public IList<Common.Tool> Tools { get; set; }
+        [JsonPropertyName("tool_choice")]
+        public ToolChoice ToolChoice { get; set; }
+    }
+
+    [JsonConverter(typeof(MessageParametersConverter<MessageParameters>))]
+    public class MessageParameters : MessageCountTokenParameters
+    {
         [JsonPropertyName("max_tokens")]
         public int MaxTokens { get; set; }
         [JsonPropertyName("metadata")]
@@ -30,18 +39,11 @@ public class MessageParameters
         public int? TopK { get; set; }
         [JsonPropertyName("top_p")]
         public decimal? TopP { get; set; }
-        [JsonPropertyName("tools")] 
-        private List<Common.Function> ToolsForClaude => Tools?.Select(p => p.Function).ToList();
-        [JsonIgnore]
-        public IList<Common.Tool> Tools { get; set; }
 
         /// <summary>
         /// Prompt Cache Type Definitions. Designed to be used as a bitwise assignment if you want to cache multiple types and are caching enough context.
         /// </summary>
         [JsonIgnore]
         public PromptCacheType PromptCaching { get; set; } = PromptCacheType.None;
-        [JsonPropertyName("tool_choice")]
-        public ToolChoice ToolChoice { get; set; }
     }
-
 }
diff --git a/Anthropic.SDK/Messaging/MessagesEndpoint.cs b/Anthropic.SDK/Messaging/MessagesEndpoint.cs
@@ -29,7 +29,7 @@ public async Task<MessageResponse> GetClaudeMessageAsync(MessageParameters param
             SetCacheControls(parameters);
 
             parameters.Stream = false;
-            var response = await HttpRequestMessages(Url, HttpMethod.Post, parameters, ctx).ConfigureAwait(false);
+            var response = await HttpRequestMessages<MessageResponse>(Url, HttpMethod.Post, parameters, ctx).ConfigureAwait(false);
 
             var toolCalls = new List<Function>();
             foreach (var message in response.Content)
@@ -147,7 +147,11 @@ public async IAsyncEnumerable<MessageResponse> StreamClaudeMessageAsync(MessageP
 
                 yield return result;
             }
-
+        }
+
+        public async Task<MessageCountTokenResponse> CountMessageTokensAsync(MessageCountTokenParameters parameters, CancellationToken ctx = default)
+        {
+            return await HttpRequestMessages<MessageCountTokenResponse>($"{Url}/count_tokens", HttpMethod.Post, parameters, ctx).ConfigureAwait(false);
         }
     }
 }
diff --git a/README.md b/README.md
@@ -13,6 +13,7 @@ Anthropic.SDK is an unofficial C# client designed for interacting with the Claud
 - [Examples](#examples)
   - [Non-Streaming Call](#non-streaming-call)
   - [Streaming Call](#streaming-call)
+  - [Token Count](#token-count)
   - [IChatClient](#ichatclient)
   - [Prompt Caching](#prompt-caching)
   - [PDF Support](#pdf-support)
@@ -177,6 +178,22 @@ Console.WriteLine($@"Used Tokens - Input:{outputs.First().StreamStartMessage.Usa
                             Output: {outputs.Last().Usage.OutputTokens}");
 ```
 
+### Token Count
+The `AnthropicClient` has support for the [message token count endpoint](https://docs.anthropic.com/en/docs/build-with-claude/token-counting). Below is an example of how to use it.
+
+```csharp
+ var client = new AnthropicClient();
+ var messages = new List<Message>();
+ messages.Add(new Message(RoleType.User, "Write me a sonnet about the Statue of Liberty"));
+ var parameters = new MessageCountTokenParameters
+ {
+     Messages = messages,
+     Model = AnthropicModels.Claude35Haiku
+ };
+ var response = await client.Messages.CountMessageTokensAsync(parameters);
+ Assert.IsTrue(res.InputTokens > 0);
+```
+
 ### IChatClient
 
 The `AnthropicClient` has support for the new `IChatClient` from Microsoft and offers a slightly different mechanism for using the `AnthropicClient`.  Below are a few examples.