Skip to content

Commit 5ca98aa

Browse files
sdcbsdcb
committed
Implement anthropic count tokens api
1 parent ff4d1c2 commit 5ca98aa

4 files changed

Lines changed: 128 additions & 3 deletions

File tree

‎src/BE/Controllers/Api/AnthropicCompatible/AnthropicCountTokenController.cs‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
using Chats.BE.Controllers.Api.AnthropicCompatible.Dtos;
2+
using Chats.BE.DB;
3+
using Chats.BE.DB.Enums;
4+
using Chats.BE.Services;
5+
using Chats.BE.Services.Models;
6+
using Chats.BE.Services.Models.ChatServices;
7+
using Chats.BE.Services.OpenAIApiKeySession;
8+
using Microsoft.AspNetCore.Authorization;
9+
using Microsoft.AspNetCore.Mvc;
10+
using System.Text.Json.Nodes;
11+
12+
namespace Chats.BE.Controllers.Api.AnthropicCompatible;
13+
14+
[Authorize(AuthenticationSchemes = "OpenAIApiKey")]
15+
public class AnthropicCountTokenController(
16+
CurrentApiKey currentApiKey,
17+
ChatFactory cf,
18+
UserModelManager userModelManager,
19+
ILogger<AnthropicCountTokenController> logger) : ControllerBase
20+
{
21+
private static readonly DBApiType[] AllowedApiTypes = [DBApiType.OpenAIChatCompletion, DBApiType.OpenAIResponse, DBApiType.AnthropicMessages];
22+
23+
[HttpPost("v1/messages/count_tokens")]
24+
public async Task<ActionResult> CountTokens([FromBody] JsonObject json, CancellationToken cancellationToken)
25+
{
26+
AnthropicCountTokenRequestWrapper request = new(json);
27+
28+
if (!request.SeemsValid())
29+
{
30+
return ErrorMessage(AnthropicErrorTypes.InvalidRequestError, "Invalid request: model and messages are required.");
31+
}
32+
33+
if (string.IsNullOrWhiteSpace(request.Model))
34+
{
35+
return ErrorMessage(AnthropicErrorTypes.InvalidRequestError, "model is required.");
36+
}
37+
38+
UserModel? userModel = await userModelManager.GetUserModel(currentApiKey.ApiKey, request.Model, cancellationToken);
39+
if (userModel == null)
40+
{
41+
return ErrorMessage(AnthropicErrorTypes.NotFoundError, $"The model `{request.Model}` does not exist or you do not have access to it.");
42+
}
43+
44+
if (!AllowedApiTypes.Contains(userModel.Model.ApiType))
45+
{
46+
return ErrorMessage(AnthropicErrorTypes.InvalidRequestError, $"The model `{request.Model}` does not support messages API.");
47+
}
48+
49+
try
50+
{
51+
Model cm = userModel.Model;
52+
using ChatService s = cf.CreateChatService(cm);
53+
ChatRequest chatRequest = request.ToChatRequest(currentApiKey.User.Id.ToString(), cm);
54+
int inputTokens = await s.CountTokenAsync(chatRequest, cancellationToken);
55+
56+
return Ok(new AnthropicCountTokenResponse { InputTokens = inputTokens });
57+
}
58+
catch (Exception e)
59+
{
60+
logger.LogError(e, "Error counting tokens");
61+
return ErrorMessage(AnthropicErrorTypes.ApiError, "Internal server error");
62+
}
63+
}
64+
65+
private BadRequestObjectResult ErrorMessage(string errorType, string message)
66+
{
67+
return BadRequest(new AnthropicErrorResponse
68+
{
69+
Error = new AnthropicErrorDetail
70+
{
71+
Type = errorType,
72+
Message = message
73+
}
74+
});
75+
}
76+
}
77+
78+
public class AnthropicCountTokenRequestWrapper(JsonObject json)
79+
{
80+
public string? Model => (string?)json["model"];
81+
82+
public bool SeemsValid()
83+
{
84+
return Model != null && json["messages"] != null;
85+
}
86+
87+
public ChatRequest ToChatRequest(string userId, Model model)
88+
{
89+
// For count tokens, we reuse the AnthropicRequestWrapper logic
90+
// but we need to ensure max_tokens has a default value
91+
json["max_tokens"] ??= model.MaxResponseTokens;
92+
json["stream"] ??= false;
93+
94+
AnthropicRequestWrapper wrapper = new(json);
95+
return wrapper.ToChatRequest(userId, model);
96+
}
97+
}

‎src/BE/Controllers/Api/AnthropicCompatible/Dtos/AnthropicResponse.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,9 @@ public record AnthropicUsage
103103
[JsonPropertyName("cache_read_input_tokens"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
104104
public int? CacheReadInputTokens { get; init; }
105105
}
106+
107+
public record AnthropicCountTokenResponse
108+
{
109+
[JsonPropertyName("input_tokens")]
110+
public required int InputTokens { get; init; }
111+
}

‎src/BE/Services/Models/ChatServices/Anthropic/AnthropicChatService.cs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,26 @@ public override async Task<string[]> ListModels(ModelKey modelKey, CancellationT
179179
return [.. result.Data.Select(x => x.ID)];
180180
}
181181

182+
public override async Task<int> CountTokenAsync(ChatRequest request, CancellationToken cancellationToken)
183+
{
184+
AnthropicClient anthropicClient = CreateAnthropicClient(request.ChatConfig.Model.ModelKey);
185+
MessageCreateParams messageParams = ConvertOptions(request);
186+
187+
MessageCountTokensParams countParams = new()
188+
{
189+
Messages = messageParams.Messages,
190+
Model = messageParams.Model,
191+
System = request.ChatConfig.SystemPrompt is { } systemPrompt ? systemPrompt : null!,
192+
Thinking = messageParams.Thinking,
193+
Tools = messageParams.Tools != null
194+
? [.. messageParams.Tools.Select(x => new MessageCountTokensTool(x.Json))]
195+
: null,
196+
};
197+
198+
MessageTokensCount result = await anthropicClient.Messages.CountTokens(countParams, cancellationToken);
199+
return (int)result.InputTokens;
200+
}
201+
182202
static MessageCreateParams ConvertOptions(ChatRequest request)
183203
{
184204
// Anthropic has a very strict policy on thinking blocks - they need pass back thinking AND signature together

‎src/BE/Services/Models/ChatServices/ChatService.cs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,9 @@
33
using Chats.BE.DB.Enums;
44
using Chats.BE.Services.FileServices;
55
using Chats.BE.Services.Models.ChatServices;
6-
using Chats.BE.Services.Models.ChatServices.OpenAI;
76
using Chats.BE.Services.Models.Dtos;
87
using Microsoft.ML.Tokenizers;
98
using OpenAI;
10-
using OpenAI.Models;
11-
using System.ClientModel;
129
using System.Runtime.CompilerServices;
1310
using Tokenizer = Microsoft.ML.Tokenizers.Tokenizer;
1411

@@ -24,6 +21,11 @@ public abstract partial class ChatService : IDisposable
2421

2522
public virtual Task<string[]> ListModels(ModelKey modelKey, CancellationToken cancellationToken) => Task.FromResult(Array.Empty<string>());
2623

24+
public virtual Task<int> CountTokenAsync(ChatRequest request, CancellationToken cancellationToken)
25+
{
26+
return Task.FromResult(request.EstimatePromptTokens(Tokenizer));
27+
}
28+
2729
public virtual async Task<ChatSegment> Chat(ChatRequest request, CancellationToken cancellationToken)
2830
{
2931
List<ChatSegmentItem> segments = [];

0 commit comments

Comments
 (0)