Count the tool schemas the request carries

This commit is contained in:
Thorsten Sommer 2026-09-14 19:23:32 +02:00
parent a3d6d6abc3
commit 21a6244720
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
3 changed files with 76 additions and 29 deletions

View File

@ -1,3 +1,7 @@
using System.Text.Json;
using AIStudio.Tools.ToolCallingSystem;
namespace AIStudio.Chat;
/// <summary>
@ -6,9 +10,10 @@ namespace AIStudio.Chat;
/// <remarks>
/// Collected here rather than while counting, so that what counts towards a token budget is one
/// question with one answer which a test can ask. It follows what the message builder actually
/// sends: the system prompt, the text of every block, and the attachments hanging off those
/// blocks -- plus whatever is standing in the composer but has not been sent yet, because that is
/// the part a person is deciding about while they look at the number.
/// sends: the system prompt, the schema of every tool the model may call, the text of every block,
/// and the attachments hanging off those blocks -- plus whatever is standing in the composer but
/// has not been sent yet, because that is the part a person is deciding about while they look at
/// the number.
///
/// And, while a request is running, what its tools have returned so far. That is the one part
/// which is not about the next request but about the one in flight: it is what the model is
@ -70,8 +75,12 @@ public sealed record ConversationParts
/// <param name="draft">What stands in the composer.</param>
/// <param name="draftAttachments">What is attached to the composer.</param>
/// <param name="imagesAreSent">Whether the model takes images at all. When it does not, none are sent.</param>
/// <param name="toolDefinitions">
/// The tools the model may call, filtered for the provider the same way they are before
/// sending, or null when there are none.
/// </param>
/// <returns>The parts of the conversation.</returns>
public static ConversationParts Of(ChatThread? thread, string systemPrompt, string draft, IEnumerable<FileAttachment>? draftAttachments, bool imagesAreSent)
public static ConversationParts Of(ChatThread? thread, string systemPrompt, string draft, IEnumerable<FileAttachment>? draftAttachments, bool imagesAreSent, IEnumerable<ToolDefinition>? toolDefinitions)
{
var texts = new List<string>();
var growing = new List<string>();
@ -81,6 +90,15 @@ public sealed record ConversationParts
if (!string.IsNullOrWhiteSpace(systemPrompt))
texts.Add(systemPrompt);
//
// The tools ride along beside the messages, one schema each, in every single request of a
// conversation. Counted with the lasting texts rather than with the growing ones: a schema
// is the same string all session long, so measuring it once and remembering it is exactly
// what the cache is for.
//
foreach (var definition in toolDefinitions ?? [])
texts.Add(Describe(definition));
if (thread is not null)
{
//
@ -128,6 +146,27 @@ public sealed record ConversationParts
};
}
/// <summary>
/// What one tool costs the request it is offered in.
/// </summary>
/// <remarks>
/// Its name, what it tells the model it does, and the arguments it takes -- that is what the
/// provider adapters put into the tool list of the request body. The wire shape differs
/// between the APIs: they name the fields differently, and a strict schema is rewritten for
/// the OpenAI ones. None of that changes the length by an amount which matters next to a
/// conversation, and the number is reported as an estimate anyway.
/// </remarks>
/// <param name="definition">The tool as it was declared.</param>
/// <returns>The text to count for it.</returns>
private static string Describe(ToolDefinition definition)
{
var parameters = definition.Function.Parameters.ValueKind is JsonValueKind.Undefined
? string.Empty
: definition.Function.Parameters.GetRawText();
return $"{definition.Function.Name}{definition.Function.DescriptionForLLM}{parameters}";
}
/// <summary>
/// Puts attachments into the two groups they are counted in.
/// </summary>

View File

@ -1475,8 +1475,9 @@ public partial class ChatComponent : MSGComponentBase
// of it would tell a person their window is empty while their first message is not.
//
var thread = this.ChatThread ?? this.NewChatThread(string.Empty);
var toolDefinitions = this.GetRunnableToolDefinitions();
provider = this.Provider;
parts = ConversationParts.Of(thread, this.BuildSystemPromptFor(thread), this.UserInput, this.ComposerState.FileAttachments, provider.SupportsImageInput());
parts = ConversationParts.Of(thread, this.BuildSystemPromptFor(thread, toolDefinitions), this.UserInput, this.ComposerState.FileAttachments, provider.SupportsImageInput(), toolDefinitions);
});
var counted = await this.ConversationTokenCounter.CountAsync(provider, parts, token);
@ -1501,22 +1502,29 @@ public partial class ChatComponent : MSGComponentBase
/// source is appended to it, the selected profile adds a paragraph, and the policy of the
/// selected tools adds another. Switching a profile while writing therefore moves the number,
/// which is the whole reason this is asked rather than read off the thread.
///
/// The tools are filtered for the provider the same way they are before sending, so that a tool
/// the provider is not trusted enough to receive does not count either.
/// </remarks>
/// <param name="thread">The thread to build the prompt for.</param>
/// <param name="toolDefinitions">The tools whose policy the prompt states.</param>
/// <returns>The system prompt as it would be sent.</returns>
private string BuildSystemPromptFor(ChatThread thread)
{
var toolDefinitions = this.ToolRegistry.FilterToolIdsForProvider(this.Provider, this.selectedToolIds)
.Select(this.ToolRegistry.GetDefinition)
.Where(definition => definition is not null)
.Select(definition => definition!)
.ToList();
private string BuildSystemPromptFor(ChatThread thread, IReadOnlyList<ToolDefinition> toolDefinitions) => thread.BuildSystemPrompt(this.SettingsManager, toolDefinitions).Text;
return thread.BuildSystemPrompt(this.SettingsManager, toolDefinitions).Text;
}
/// <summary>
/// The tools the next request would offer the model.
/// </summary>
/// <remarks>
/// Filtered for the provider the same way they are before sending, so that a tool the provider
/// is not trusted enough to receive does not count either.
///
/// Asked for once and used twice: their policy goes into the system prompt, and their schemas
/// travel next to it in the request body. Both cost tokens, and both change the moment somebody
/// switches a tool on.
/// </remarks>
/// <returns>The definitions of the selected tools.</returns>
private IReadOnlyList<ToolDefinition> GetRunnableToolDefinitions() => this.ToolRegistry.FilterToolIdsForProvider(this.Provider, this.selectedToolIds)
.Select(this.ToolRegistry.GetDefinition)
.Where(definition => definition is not null)
.Select(definition => definition!)
.ToList();
/// <summary>
/// The thread a new chat starts with, as the selections made so far decide it.

View File

@ -44,7 +44,7 @@ public sealed class ConversationPartsTests
],
};
var parts = ConversationParts.Of(thread, "You are helpful.", "And of Italy?", null, imagesAreSent: true);
var parts = ConversationParts.Of(thread, "You are helpful.", "And of Italy?", null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -65,7 +65,7 @@ public sealed class ConversationPartsTests
((ContentText)streaming.Content!).IsStreaming = true;
var thread = new ChatThread { Blocks = [Block("A question."), streaming] };
var parts = ConversationParts.Of(thread, string.Empty, "a draft", null, imagesAreSent: true);
var parts = ConversationParts.Of(thread, string.Empty, "a draft", null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -80,7 +80,7 @@ public sealed class ConversationPartsTests
var finished = Block("The whole answer.");
((ContentText)finished.Content!).IsStreaming = false;
var parts = ConversationParts.Of(new() { Blocks = [finished] }, string.Empty, string.Empty, null, imagesAreSent: true);
var parts = ConversationParts.Of(new() { Blocks = [finished] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -100,7 +100,7 @@ public sealed class ConversationPartsTests
//
var thread = new ChatThread { SystemPrompt = "What the person typed." };
var parts = ConversationParts.Of(thread, "What the request carries.", string.Empty, null, imagesAreSent: true);
var parts = ConversationParts.Of(thread, "What the request carries.", string.Empty, null, imagesAreSent: true, toolDefinitions: null);
Assert.That(parts.Texts, Is.EqualTo(new[] { "What the request carries." }));
}
@ -115,7 +115,7 @@ public sealed class ConversationPartsTests
var hidden = Block("An instruction the user does not see.");
var thread = new ChatThread { Blocks = [new() { ContentType = hidden.ContentType, Role = hidden.Role, Content = hidden.Content, HideFromUser = true }] };
var parts = ConversationParts.Of(thread, string.Empty, string.Empty, null, imagesAreSent: true);
var parts = ConversationParts.Of(thread, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);
Assert.That(parts.Texts, Is.EqualTo(new[] { "An instruction the user does not see." }));
}
@ -123,7 +123,7 @@ public sealed class ConversationPartsTests
[Test]
public void WithoutAConversationOnlyTheDraftCounts()
{
var parts = ConversationParts.Of(null, string.Empty, "Hello", null, imagesAreSent: true);
var parts = ConversationParts.Of(null, string.Empty, "Hello", null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -136,7 +136,7 @@ public sealed class ConversationPartsTests
[TestCase(" ")]
public void NothingWrittenIsNothingToCount(string draft)
{
var parts = ConversationParts.Of(null, string.Empty, draft, null, imagesAreSent: true);
var parts = ConversationParts.Of(null, string.Empty, draft, null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -157,7 +157,7 @@ public sealed class ConversationPartsTests
var empty = Block(string.Empty);
((ContentText)empty.Content!).FileAttachments.Add(FileAttachment.FromPath(document));
var parts = ConversationParts.Of(new() { Blocks = [empty] }, string.Empty, string.Empty, null, imagesAreSent: true);
var parts = ConversationParts.Of(new() { Blocks = [empty] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -179,7 +179,7 @@ public sealed class ConversationPartsTests
var block = Block("Please read this.");
((ContentText)block.Content!).FileAttachments.Add(FileAttachment.FromPath(older));
var parts = ConversationParts.Of(new() { Blocks = [block] }, string.Empty, "And this one.", [FileAttachment.FromPath(draft)], imagesAreSent: true);
var parts = ConversationParts.Of(new() { Blocks = [block] }, string.Empty, "And this one.", [FileAttachment.FromPath(draft)], imagesAreSent: true, toolDefinitions: null);
Assert.That(parts.Documents.Select(document => document.FileName), Is.EqualTo(new[] { "older.txt", "draft.txt" }));
}
@ -192,7 +192,7 @@ public sealed class ConversationPartsTests
//
var attachment = FileAttachment.FromPath(Path.Combine(this.directory, "never-existed.txt"));
var parts = ConversationParts.Of(null, string.Empty, "Here", [attachment], imagesAreSent: true);
var parts = ConversationParts.Of(null, string.Empty, "Here", [attachment], imagesAreSent: true, toolDefinitions: null);
Assert.That(parts.Documents, Is.Empty);
}
@ -203,7 +203,7 @@ public sealed class ConversationPartsTests
var document = this.WriteFile("notes.txt", "content");
var image = this.WriteFile("photo.png", "not really a png");
var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(document), FileAttachment.FromPath(image)], imagesAreSent: true);
var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(document), FileAttachment.FromPath(image)], imagesAreSent: true, toolDefinitions: null);
Assert.Multiple(() =>
{
@ -221,7 +221,7 @@ public sealed class ConversationPartsTests
//
var image = this.WriteFile("photo.png", "not really a png");
var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(image)], imagesAreSent: false);
var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(image)], imagesAreSent: false, toolDefinitions: null);
Assert.That(parts.Images, Is.Zero);
}