Added the exact token count to chats which offer tools

This commit is contained in:
Thorsten Sommer 2026-09-23 18:58:23 +02:00
parent 3277f20f60
commit ffc6536ce3
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
8 changed files with 198 additions and 17 deletions

View File

@ -401,8 +401,7 @@ public sealed record ChatThread
/// asked, and no earlier answer ever stands in for it. Whatever came after an older report -- a /// asked, and no earlier answer ever stands in for it. Whatever came after an older report -- a
/// message whose request was turned down, an answer which is still being written, an answer /// message whose request was turned down, an answer which is still being written, an answer
/// without a report of its own -- is missing from that report's number, and the estimate is /// without a report of its own -- is missing from that report's number, and the estimate is
/// closer to the truth than a figure which leaves it out. Answers written while tools were /// closer to the truth than a figure which leaves it out.<br/><br/>
/// offered carry no report yet, so for them the estimate always takes over.<br/><br/>
/// ///
/// A report also stops counting when the thread holds a different number of blocks than the /// A report also stops counting when the thread holds a different number of blocks than the
/// request did, which means an earlier message was deleted, and when the next request goes to /// request did, which means an earlier message was deleted, and when the next request goes to

View File

@ -5,7 +5,8 @@ namespace AIStudio.Provider.OpenAI;
/// </summary> /// </summary>
/// <param name="TextDelta">The text this line carried, empty when it carried none.</param> /// <param name="TextDelta">The text this line carried, empty when it carried none.</param>
/// <param name="Sources">The sources this line announced, empty when it announced none.</param> /// <param name="Sources">The sources this line announced, empty when it announced none.</param>
public readonly record struct ChatCompletionStreamPart(string TextDelta, IList<ISource> Sources) /// <param name="Usage">What the provider said the request cost, unknown on every line but the one which carries it.</param>
public readonly record struct ChatCompletionStreamPart(string TextDelta, IList<ISource> Sources, TokenUsage Usage = default)
{ {
/// <summary> /// <summary>
/// The part of a line which says nothing to the user, such as a fragment of a tool call. /// The part of a line which says nothing to the user, such as a fragment of a tool call.
@ -15,5 +16,9 @@ public readonly record struct ChatCompletionStreamPart(string TextDelta, IList<I
/// <summary> /// <summary>
/// Whether this part has anything to show at all. /// Whether this part has anything to show at all.
/// </summary> /// </summary>
/// <remarks>
/// The usage is not part of that: it is nothing to show, and whether it is passed on at all is
/// the adapter's decision, which knows which round this is.
/// </remarks>
public bool HasContent => this.TextDelta.Length > 0 || this.Sources.Count > 0; public bool HasContent => this.TextDelta.Length > 0 || this.Sources.Count > 0;
} }

View File

@ -61,9 +61,15 @@ public sealed class ChatCompletionToolCallAccumulator(Func<ServerSentEvent, ILis
// //
var sources = readSources?.Invoke(serverSentEvent) ?? []; var sources = readSources?.Invoke(serverSentEvent) ?? [];
//
// The usage arrives on a line without choices at most providers, and next to the last
// piece of text at some. Read before the choices are looked at, it is not lost in either.
//
var usage = line?.Usage?.ToTokenUsage() ?? TokenUsage.UNKNOWN;
var delta = line?.Choices?.FirstOrDefault()?.Delta; var delta = line?.Choices?.FirstOrDefault()?.Delta;
if (delta is null) if (delta is null)
return WithSources(string.Empty, sources); return WithSources(string.Empty, sources, usage);
this.hasReadAnything = true; this.hasReadAnything = true;
@ -80,10 +86,10 @@ public sealed class ChatCompletionToolCallAccumulator(Func<ServerSentEvent, ILis
var textDelta = delta.Content; var textDelta = delta.Content;
if (textDelta.Length is 0) if (textDelta.Length is 0)
return WithSources(string.Empty, sources); return WithSources(string.Empty, sources, usage);
this.text.Append(textDelta); this.text.Append(textDelta);
return new ChatCompletionStreamPart(textDelta, sources); return new ChatCompletionStreamPart(textDelta, sources, usage);
} }
/// <summary> /// <summary>
@ -183,10 +189,10 @@ public sealed class ChatCompletionToolCallAccumulator(Func<ServerSentEvent, ILis
private static string? Coalesce(string? value) => string.IsNullOrWhiteSpace(value) ? null : value; private static string? Coalesce(string? value) => string.IsNullOrWhiteSpace(value) ? null : value;
/// <summary> /// <summary>
/// A part for a line which brought sources but no text, or nothing at all. /// A part for a line which brought sources or a usage but no text, or nothing at all.
/// </summary> /// </summary>
private static ChatCompletionStreamPart WithSources(string text, IList<ISource> sources) private static ChatCompletionStreamPart WithSources(string text, IList<ISource> sources, TokenUsage usage)
=> sources.Count is 0 ? ChatCompletionStreamPart.Nothing : new ChatCompletionStreamPart(text, sources); => sources.Count is 0 && !usage.IsKnown ? ChatCompletionStreamPart.Nothing : new ChatCompletionStreamPart(text, sources, usage);
/// <summary> /// <summary>
/// One tool call while its fragments are still arriving. /// One tool call while its fragments are still arriving.

View File

@ -54,6 +54,16 @@ public sealed class ChatCompletionToolCallingAdapter<TRequest>(
ParallelToolCalls = requestDtoBase.Tools is null ? null : false, ParallelToolCalls = requestDtoBase.Tools is null ? null : false,
}; };
//
// Only the first round passes on what its request cost. Its prompt is the conversation up
// to the question, which is exactly what the next question will be sent after. Every later
// round carries the tool calls and their results on top, and none of that is sent again
// once the answer stands -- a report of such a round would count a chat far larger than
// the one the next request carries. What the answer adds, all rounds of text together, is
// counted from its text afterwards, cf. ReportedHistory.
//
var passesOnUsage = this.internalMessages.Count is 0;
// //
// The text goes out while it is being written; the tool calls are put back together // The text goes out while it is being written; the tool calls are put back together
// behind it, fragment by fragment. // behind it, fragment by fragment.
@ -62,8 +72,9 @@ public sealed class ChatCompletionToolCallingAdapter<TRequest>(
await foreach (var serverSentEvent in streamRequestAsync(requestDto, token)) await foreach (var serverSentEvent in streamRequestAsync(requestDto, token))
{ {
var part = accumulator.Process(serverSentEvent); var part = accumulator.Process(serverSentEvent);
if (part.HasContent) var usage = passesOnUsage ? part.Usage : TokenUsage.UNKNOWN;
yield return ToolCallingStreamEvent.TextDelta(new ContentStreamChunk(part.TextDelta, part.Sources)); if (part.HasContent || usage.IsKnown)
yield return ToolCallingStreamEvent.TextDelta(new ContentStreamChunk(part.TextDelta, part.Sources, Usage: usage));
} }
var message = accumulator.Build(); var message = accumulator.Build();

View File

@ -11,4 +11,15 @@ namespace AIStudio.Provider.OpenAI;
/// </remarks> /// </remarks>
/// <param name="Id">The ID of the answer.</param> /// <param name="Id">The ID of the answer.</param>
/// <param name="Choices">The choices this line adds to.</param> /// <param name="Choices">The choices this line adds to.</param>
public sealed record ChatCompletionToolStreamLine(string? Id, IList<ChatCompletionToolStreamChoice?>? Choices); public sealed record ChatCompletionToolStreamLine(string? Id, IList<ChatCompletionToolStreamChoice?>? Choices)
{
/// <summary>
/// What the provider says the request cost, on the one line which carries it.
/// </summary>
/// <remarks>
/// The same block the plain text path reads, and asked for the same way: every streamed
/// ChatCompletionAPIRequest asks for it, the requests of the tool rounds included. Not a
/// positional parameter, because nobody but the serializer ever builds this line.
/// </remarks>
public ChatCompletionUsage? Usage { get; init; }
}

View File

@ -88,8 +88,8 @@ public sealed class ChatThreadReportedHistoryTests
public void AnAnswerWithoutAReportDoesNotBorrowTheReportBeforeIt() public void AnAnswerWithoutAReportDoesNotBorrowTheReportBeforeIt()
{ {
// //
// What an answer written while tools were offered looks like today: finished, but without // What an answer looks like whose provider reports nothing, or whose API is not read for a
// a report of its own. // report yet: finished, but without a report of its own.
// //
var withoutReport = Block(ChatRole.AI, new ContentText { Text = "Second answer" }, 4); var withoutReport = Block(ChatRole.AI, new ContentText { Text = "Second answer" }, 4);
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), withoutReport); var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), withoutReport);

View File

@ -184,6 +184,37 @@ public sealed class ChatCompletionToolCallAccumulatorTests
Assert.That(part.Sources.Select(x => x.URL), Is.EqualTo(new[] { "https://example.org/" }), "Whatever the provider announced on that line reaches the user with it."); Assert.That(part.Sources.Select(x => x.URL), Is.EqualTo(new[] { "https://example.org/" }), "Whatever the provider announced on that line reaches the user with it.");
} }
[Test]
public void TheLineWithoutChoicesStatesWhatTheRequestCost()
{
var accumulator = new ChatCompletionToolCallAccumulator();
var part = accumulator.Process(Event("""{"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":345,"total_tokens":1545}}"""));
Assert.Multiple(() =>
{
Assert.That(part.Usage.IsKnown, Is.True, "The last line of the stream has no choices, and it must not be dropped for that.");
Assert.That(part.Usage.PromptTokens, Is.EqualTo(1200));
Assert.That(part.HasContent, Is.False, "It has nothing to show, though.");
});
}
[Test]
public void AUsageNextToTheLastTextIsReadAsWell()
{
//
// Some providers put the usage on the line which carries the last piece of the answer
// rather than on a line of its own.
//
var accumulator = new ChatCompletionToolCallAccumulator();
var part = accumulator.Process(Event("""{"choices":[{"index":0,"delta":{"content":"Bye"}}],"usage":{"prompt_tokens":1200,"completion_tokens":2}}"""));
Assert.Multiple(() =>
{
Assert.That(part.TextDelta, Is.EqualTo("Bye"));
Assert.That(part.Usage.PromptTokens, Is.EqualTo(1200));
});
}
private static ChatCompletionResponseMessage? Read(params string[] data) private static ChatCompletionResponseMessage? Read(params string[] data)
{ {
var accumulator = new ChatCompletionToolCallAccumulator(); var accumulator = new ChatCompletionToolCallAccumulator();

View File

@ -0,0 +1,118 @@
using System.Runtime.CompilerServices;
using AIStudio.Provider;
using AIStudio.Provider.OpenAI;
using Microsoft.Extensions.Logging.Abstractions;
namespace AIStudio.Tests.Provider.ToolCalling;
/// <summary>
/// Checks which round of a tool calling conversation passes on what its request cost.
/// </summary>
/// <remarks>
/// Every round of a tool conversation is a request of its own, and every one of them reports what
/// it cost. Only the first one describes what the next question will be sent after: every later
/// round carries the tool calls and their results on top, none of which is sent again once the
/// answer stands. Passing on the last report instead would put the chat at the size of everything
/// the tools returned, which is the one number a person watching their context window must not
/// see as exact.
/// </remarks>
[TestFixture]
public sealed class ChatCompletionToolCallingAdapterTests
{
private const string FIRST_ROUND_USAGE = """{"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":20}}""";
private const string SECOND_ROUND_USAGE = """{"choices":[],"usage":{"prompt_tokens":9800,"completion_tokens":150}}""";
[Test]
public async Task OnlyTheFirstRoundPassesOnWhatItsRequestCost()
{
var adapter = Adapter(
[
"""{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"web_search","arguments":"{\"query\":\"weather\"}"}}]}}]}""",
FIRST_ROUND_USAGE,
"[DONE]",
],
[
"""{"choices":[{"index":0,"delta":{"content":"It is sunny."}}]}""",
SECOND_ROUND_USAGE,
"[DONE]",
]);
var firstRound = await Usages(adapter);
//
// What the loop does between two rounds: the model's turn and the tool's result become part
// of the next request.
//
adapter.RecordAssistantTurn();
adapter.RecordToolResult("call_1", "Sunny, 24 degrees.");
var secondRound = await Usages(adapter);
Assert.Multiple(() =>
{
Assert.That(firstRound, Is.EqualTo(new[] { 1200 }), "The first round's prompt is the conversation up to the question.");
Assert.That(secondRound, Is.Empty, "The second round's prompt holds the tool result as well, which the next question is not sent with.");
});
}
[Test]
public async Task ARoundWithoutToolCallsPassesItOnAsWell()
{
//
// Offering tools does not mean the model uses them. Then the first round is the only one,
// and its report is as good as the one of a request which offered none.
//
var adapter = Adapter(
[
"""{"choices":[{"index":0,"delta":{"content":"Hello."}}]}""",
FIRST_ROUND_USAGE,
"[DONE]",
]);
Assert.That(await Usages(adapter), Is.EqualTo(new[] { 1200 }));
}
/// <summary>
/// Runs the next round and returns the prompt of every usage it passed on.
/// </summary>
private static async Task<List<int>> Usages(ChatCompletionToolCallingAdapter<ChatCompletionAPIRequest> adapter)
{
var usages = new List<int>();
await foreach (var streamEvent in adapter.ExecuteRoundAsync(null, true))
if (streamEvent.Delta is { Usage.IsKnown: true } delta)
usages.Add(delta.Usage.PromptTokens);
return usages;
}
/// <summary>
/// Builds an adapter whose requests are answered by the given rounds, one after another.
/// </summary>
private static ChatCompletionToolCallingAdapter<ChatCompletionAPIRequest> Adapter(params string[][] rounds)
{
var nextRound = 0;
return new(
(_, _, _) => Task.FromResult(new ChatCompletionAPIRequest("model-a", [], true)),
new TextMessage("You are a helpful assistant.", "system"),
new Dictionary<string, object>(),
[],
[],
(_, token) => Lines(rounds[nextRound++], token),
_ => [],
NullLogger.Instance);
}
private static async IAsyncEnumerable<ServerSentEvent> Lines(string[] data, [EnumeratorCancellation] CancellationToken token = default)
{
foreach (var line in data)
{
token.ThrowIfCancellationRequested();
yield return new ServerSentEvent($"data: {line}", line);
}
await Task.CompletedTask;
}
}