mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-09-27 00:53:37 +00:00
Update the token count while tools are running
This commit is contained in:
parent
b8b8e3f66d
commit
a3d6d6abc3
@ -28,6 +28,15 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
|
||||
|
||||
public DateTimeOffset LastCheckpoint { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// When the chat was last told that something happened which was not a streamed chunk.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Kept on the job rather than in the loop which streams, because the tool calling reports
|
||||
/// from outside that loop: it runs inside the provider call the loop is waiting on.
|
||||
/// </remarks>
|
||||
public DateTimeOffset LastActivityNotification { get; set; }
|
||||
|
||||
public bool IsCompletionStarted { get; set; }
|
||||
|
||||
public readonly Lock SyncRoot = new();
|
||||
@ -79,6 +88,44 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
|
||||
return this.jobs.TryGetValue(jobId, out var job) ? job.ChatGenerationRequest?.ChatThread : null;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Says that the answer of a chat has moved without a chunk having arrived.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A model which calls tools asks several times before it says anything, and while it does,
|
||||
/// this service sits in the provider call and hands nothing to the screen. But the request is
|
||||
/// growing the whole time -- every tool result travels with the next round -- and the chat is
|
||||
/// what recounts the tokens when it renders. Without this, the only thing which would ever ask
|
||||
/// again is the ten-second heartbeat of the token tracker.
|
||||
///
|
||||
/// Throttled like the streamed chunks, and by the same setting: a round which calls five tools
|
||||
/// in a row must not turn into five renders of the whole chat when somebody asked us to go easy
|
||||
/// on their battery.
|
||||
///
|
||||
/// A chat without a running job is not an error. The same tool calling loop runs for the
|
||||
/// assistants, which have no job behind them and no token count to update.
|
||||
/// </remarks>
|
||||
/// <param name="chatId">The chat whose answer moved.</param>
|
||||
public async Task NotifyChatActivityAsync(Guid chatId)
|
||||
{
|
||||
if (!this.activeChatJobsByChatId.TryGetValue(chatId, out var jobId))
|
||||
return;
|
||||
|
||||
if (!this.jobs.TryGetValue(jobId, out var job))
|
||||
return;
|
||||
|
||||
lock (job.SyncRoot)
|
||||
{
|
||||
var now = DateTimeOffset.Now;
|
||||
if (settingsManager.ConfigurationData.App.IsSavingEnergy && now - job.LastActivityNotification < STREAMING_EVENT_MIN_TIME)
|
||||
return;
|
||||
|
||||
job.LastActivityNotification = now;
|
||||
}
|
||||
|
||||
await this.NotifyChangedAsync(job);
|
||||
}
|
||||
|
||||
public async Task<AIJobSnapshot?> TryStartChatGenerationAsync(ChatGenerationRequest request)
|
||||
{
|
||||
if (this.activeChatJobsByChatId.TryGetValue(request.ChatThread.ChatId, out var existingJobId))
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
using AIStudio.Chat;
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Tools.AIJobs;
|
||||
|
||||
namespace AIStudio.Tools.ToolCallingSystem.Harness;
|
||||
|
||||
@ -55,7 +56,7 @@ public sealed class ToolCallingLoopContext
|
||||
return;
|
||||
|
||||
this.CurrentAssistantContent.ToolInvocations.Add(trace);
|
||||
await this.CurrentAssistantContent.StreamingEvent();
|
||||
await this.AnnounceAsync(this.CurrentAssistantContent);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@ -73,7 +74,7 @@ public sealed class ToolCallingLoopContext
|
||||
return;
|
||||
|
||||
this.CurrentAssistantContent.PendingToolConversation = [..adapter.RecordedRequestTexts];
|
||||
await this.CurrentAssistantContent.StreamingEvent();
|
||||
await this.AnnounceAsync(this.CurrentAssistantContent);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@ -90,7 +91,7 @@ public sealed class ToolCallingLoopContext
|
||||
ToolNames = [.. toolNames],
|
||||
};
|
||||
|
||||
await this.CurrentAssistantContent.StreamingEvent();
|
||||
await this.AnnounceAsync(this.CurrentAssistantContent);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@ -106,6 +107,32 @@ public sealed class ToolCallingLoopContext
|
||||
return;
|
||||
|
||||
this.CurrentAssistantContent.ToolRuntimeStatus = new();
|
||||
await this.CurrentAssistantContent.StreamingEvent();
|
||||
await this.AnnounceAsync(this.CurrentAssistantContent);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Says that something about the running answer has changed.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Two receivers, because the screen is built from two of them. The content's own event
|
||||
/// renders the message block, which is what shows a running tool and the calls it has made.
|
||||
/// The job service renders the chat around it, and that is what recounts the tokens -- which
|
||||
/// nothing else would ask for during a tool run: the chat hears about progress one streamed
|
||||
/// chunk at a time, and a tool run produces none until it is over.<br/><br/>
|
||||
/// One method rather than two calls at each of the four places above, because the second of
|
||||
/// them is the one which is easy to forget.
|
||||
/// </remarks>
|
||||
/// <param name="content">The assistant message which changed.</param>
|
||||
private async Task AnnounceAsync(ContentText content)
|
||||
{
|
||||
await content.StreamingEvent();
|
||||
|
||||
//
|
||||
// Asked for here rather than taken as a dependency: the same loop runs for the assistants,
|
||||
// where there is no job to tell and nothing which counts tokens.
|
||||
//
|
||||
var jobService = Program.SERVICE_PROVIDER.GetService<AIJobService>();
|
||||
if (jobService is not null)
|
||||
await jobService.NotifyChatActivityAsync(this.ChatThread.ChatId);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user