Moved the first round rule of the token count into the tool calling loop

This commit is contained in:
Thorsten Sommer 2026-09-24 12:39:40 +02:00
parent 0c585144e2
commit a4ac6bc5a0
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
5 changed files with 87 additions and 31 deletions

View File

@ -17,8 +17,8 @@ public readonly record struct ChatCompletionStreamPart(string TextDelta, IList<I
/// Whether this part has anything to show at all.
/// </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.
/// The usage is not part of that: it is nothing to show, and whether it reaches the answer at
/// all is the tool calling loop's decision, which knows which round this is.
/// </remarks>
public bool HasContent => this.TextDelta.Length > 0 || this.Sources.Count > 0;
}

View File

@ -56,27 +56,17 @@ public sealed class ChatCompletionToolCallingAdapter<TRequest>(
ParallelToolCalls = requestDtoBase.Tools is null || !mayAskForSequentialToolCalls ? 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
// behind it, fragment by fragment.
// behind it, fragment by fragment. The usage goes out with every round: which of them
// describes the conversation is the loop's decision, which knows which round this is.
//
var accumulator = new ChatCompletionToolCallAccumulator(readSources);
await foreach (var serverSentEvent in streamRequestAsync(requestDto, token))
{
var part = accumulator.Process(serverSentEvent);
var usage = passesOnUsage ? part.Usage : TokenUsage.UNKNOWN;
if (part.HasContent || usage.IsKnown)
yield return ToolCallingStreamEvent.TextDelta(new ContentStreamChunk(part.TextDelta, part.Sources, Usage: usage));
if (part.HasContent || part.Usage.IsKnown)
yield return ToolCallingStreamEvent.TextDelta(new ContentStreamChunk(part.TextDelta, part.Sources, Usage: part.Usage));
}
var message = accumulator.Build();

View File

@ -39,6 +39,7 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
var toolResultCharacterCount = 0L;
var toolSources = new List<Source>();
var hasStreamedTextBefore = false;
var isFirstRound = true;
while (true)
{
@ -51,7 +52,21 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
ToolCallingRound? round = null;
var roundStreamedText = 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.
//
// Decided here rather than in the adapters, because every wire format reports its usage
// per request, and which request this is only the loop knows for all of them alike.
//
var passesOnUsage = isFirstRound;
isFirstRound = false;
//
// The model's words go out while the round is still running. That includes what it
// writes before a tool call -- "let me look that up" -- which used to be dropped on
@ -68,7 +83,17 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
if (streamEvent.Delta is null)
continue;
if (!string.IsNullOrWhiteSpace(streamEvent.Delta.Content))
var delta = streamEvent.Delta;
if (!passesOnUsage && delta.Usage.IsKnown)
{
// A delta which carried nothing but the usage has nothing left to hand over:
if (delta.Content.Length is 0 && delta.Sources.Count is 0)
continue;
delta = delta with { Usage = TokenUsage.UNKNOWN };
}
if (!string.IsNullOrWhiteSpace(delta.Content))
{
//
// The separator goes out once the new round actually has something to say:
@ -76,12 +101,12 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
//
if (!roundStreamedText && hasStreamedTextBefore)
yield return new ContentStreamChunk(ROUND_TEXT_SEPARATOR, []);
roundStreamedText = true;
hasStreamedTextBefore = true;
}
yield return streamEvent.Delta;
yield return delta;
}
//

View File

@ -13,11 +13,9 @@ namespace AIStudio.Tests.Provider.ToolCalling;
/// </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.
/// it cost. The adapter passes on each report as it arrives, the line without choices included.
/// Which of them describes the conversation is not its decision: that is the tool calling loop's,
/// which knows which round this is, and is checked in ToolCallingLoopTests.
///
/// What a round asks for is one tool call at a time, wherever the provider lets it ask: a provider
/// which rejects the question fails the whole request, so it is not asked at all.
@ -30,7 +28,7 @@ public sealed class ChatCompletionToolCallingAdapterTests
private const string SECOND_ROUND_USAGE = """{"choices":[],"usage":{"prompt_tokens":9800,"completion_tokens":150}}""";
[Test]
public async Task OnlyTheFirstRoundPassesOnWhatItsRequestCost()
public async Task EveryRoundPassesOnWhatItsRequestCost()
{
var adapter = Adapter(
[
@ -57,8 +55,8 @@ public sealed class ChatCompletionToolCallingAdapterTests
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.");
Assert.That(firstRound, Is.EqualTo(new[] { 1200 }), "The first round reports the conversation up to the question.");
Assert.That(secondRound, Is.EqualTo(new[] { 9800 }), "The second round reports its own request, tool result included, and leaves it to the loop to drop.");
});
}

View File

@ -215,7 +215,48 @@ public sealed class ToolCallingLoopTests
Assert.That(chunks.Select(x => x.Content), Has.None.Contains(NO_ANSWER), "A stop is not a failure to answer, so it is not reported as one.");
});
}
[Test]
public async Task OnlyTheFirstRoundsUsageReachesTheAnswer()
{
//
// Every round 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.
//
var adapter = new ScriptedAdapter(
[Usage(1200), Completed(string.Empty, [Call("call-1")])],
[Text(ANSWER), Usage(9800), Completed(ANSWER)]);
var usages = (await Collect(adapter)).Where(chunk => chunk.Usage.IsKnown).Select(chunk => chunk.Usage.PromptTokens);
Assert.That(usages, Is.EqualTo(new[] { 1200 }), "The second round's prompt holds the tool result as well, which the next question is not sent with.");
}
[Test]
public async Task ALaterRoundsUsageLeavesItsTextAndNothingElse()
{
//
// Some providers send the usage next to the last piece of text rather than on a line of
// its own. Dropping the usage must not drop that text, and a delta which carried nothing
// but the usage must not turn into an empty chunk of its own.
//
var withUsage = await Collect(new ScriptedAdapter(
[Completed(string.Empty, [Call("call-1")])],
[Usage(9800), Usage(9800, ANSWER), Completed(ANSWER)]));
var withoutUsage = await Collect(new ScriptedAdapter(
[Completed(string.Empty, [Call("call-1")])],
[Text(ANSWER), Completed(ANSWER)]));
Assert.Multiple(() =>
{
Assert.That(withUsage.Select(x => x.Content), Is.EqualTo(withoutUsage.Select(x => x.Content)), "The same chunks arrive as if the round had reported nothing.");
Assert.That(withUsage.Select(x => x.Usage.IsKnown), Has.None.True, "And none of them carries the usage on.");
});
}
/// <summary>
/// As many rounds calling one tool each as it takes to use up the tool budget.
/// </summary>
@ -225,7 +266,9 @@ public sealed class ToolCallingLoopTests
.ToList();
private static ToolCallingStreamEvent Text(string text) => ToolCallingStreamEvent.TextDelta(text);
private static ToolCallingStreamEvent Usage(int promptTokens, string text = "") => ToolCallingStreamEvent.TextDelta(new ContentStreamChunk(text, [], TokenUsage.Of(promptTokens)));
private static ToolCallingStreamEvent Completed(string text, IReadOnlyList<ToolCallingRequestedCall>? calls = null, IReadOnlyList<ISource>? sources = null)
=> ToolCallingStreamEvent.RoundCompleted(new ToolCallingRound(text, calls ?? [], sources ?? []));