From 00201c09518bcc23503a0d806439b37deb4b1c1f Mon Sep 17 00:00:00 2001 From: Piotr Stachaczynski Date: Mon, 27 Oct 2025 14:01:52 +0100 Subject: [PATCH] feat: callback on agents --- src/MaIN.Core.UnitTests/AgentContextTests.cs | 2 +- src/MaIN.Core.UnitTests/FlowContextTests.cs | 2 +- src/MaIN.Core/.nuspec | 2 +- src/MaIN.Core/Hub/Contexts/AgentContext.cs | 12 ++++++------ src/MaIN.Services/Services/Abstract/IAgentService.cs | 5 ++++- .../Services/Abstract/IStepProcessor.cs | 2 ++ src/MaIN.Services/Services/AgentService.cs | 10 +++++++++- .../Services/Models/Commands/AnswerCommand.cs | 2 ++ src/MaIN.Services/Services/Models/StepContext.cs | 2 ++ src/MaIN.Services/Services/StepProcessor.cs | 6 ++++++ .../Services/Steps/AnswerStepHandler.cs | 3 ++- .../Services/Steps/Commands/AnswerCommandHandler.cs | 2 +- 12 files changed, 37 insertions(+), 13 deletions(-) diff --git a/src/MaIN.Core.UnitTests/AgentContextTests.cs b/src/MaIN.Core.UnitTests/AgentContextTests.cs index 7d117a10..c6dbfe1e 100644 --- a/src/MaIN.Core.UnitTests/AgentContextTests.cs +++ b/src/MaIN.Core.UnitTests/AgentContextTests.cs @@ -173,7 +173,7 @@ public async Task ProcessAsync_WithStringMessage_ShouldReturnChatResult() .ReturnsAsync(chat); _mockAgentService - .Setup(s => s.Process(It.IsAny(), _agentContext.GetAgentId(), It.IsAny(), It.IsAny())) + .Setup(s => s.Process(It.IsAny(), _agentContext.GetAgentId(), It.IsAny(), It.IsAny(), null)) .ReturnsAsync(new Chat { Model = "test-model", Name = "test", diff --git a/src/MaIN.Core.UnitTests/FlowContextTests.cs b/src/MaIN.Core.UnitTests/FlowContextTests.cs index 54d7a87a..f88bd9bf 100644 --- a/src/MaIN.Core.UnitTests/FlowContextTests.cs +++ b/src/MaIN.Core.UnitTests/FlowContextTests.cs @@ -98,7 +98,7 @@ public async Task ProcessAsync_WithStringMessage_ShouldReturnChatResult() .ReturnsAsync(chat); _mockAgentService - .Setup(s => s.Process(It.IsAny(), firstAgent.Id, It.IsAny(), It.IsAny())) + .Setup(s => s.Process(It.IsAny(), firstAgent.Id, It.IsAny(), It.IsAny(), null)) .ReturnsAsync(new Chat { Model = "test-model", Name = "test", diff --git a/src/MaIN.Core/.nuspec b/src/MaIN.Core/.nuspec index 046eb3b9..01b2cd00 100644 --- a/src/MaIN.Core/.nuspec +++ b/src/MaIN.Core/.nuspec @@ -2,7 +2,7 @@ MaIN.NET - 0.7.2 + 0.7.3 Wisedev Wisedev favicon.png diff --git a/src/MaIN.Core/Hub/Contexts/AgentContext.cs b/src/MaIN.Core/Hub/Contexts/AgentContext.cs index 3d53354d..860eb754 100644 --- a/src/MaIN.Core/Hub/Contexts/AgentContext.cs +++ b/src/MaIN.Core/Hub/Contexts/AgentContext.cs @@ -226,7 +226,7 @@ public async Task ProcessAsync(Chat chat, bool translate = false) }; } - public async Task ProcessAsync(string message, bool translate = false) + public async Task ProcessAsync(string message, bool translate = false, Func? callback = null) { if (_knowledge == null) { @@ -240,7 +240,7 @@ public async Task ProcessAsync(string message, bool translate = fals Type = MessageType.LocalLLM, Time = DateTime.Now }); - var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate); + var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback); var messageResult = result.Messages.LastOrDefault()!; return new ChatResult() { @@ -251,7 +251,7 @@ public async Task ProcessAsync(string message, bool translate = fals }; } - public async Task ProcessAsync(Message message, bool translate = false) + public async Task ProcessAsync(Message message, bool translate = false, Func? callback = null) { if (_knowledge == null) { @@ -259,7 +259,7 @@ public async Task ProcessAsync(Message message, bool translate = fal } var chat = await _agentService.GetChatByAgent(_agent.Id); chat.Messages.Add(message); - var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate); + var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback); var messageResult = result.Messages.LastOrDefault()!; return new ChatResult() { @@ -270,7 +270,7 @@ public async Task ProcessAsync(Message message, bool translate = fal }; } - public async Task ProcessAsync(IEnumerable messages, bool translate = false) + public async Task ProcessAsync(IEnumerable messages, bool translate = false, Func? callback = null) { if (_knowledge == null) { @@ -279,7 +279,7 @@ public async Task ProcessAsync(IEnumerable messages, bool t var chat = await _agentService.GetChatByAgent(_agent.Id); chat.Messages.Clear(); chat.Messages.AddRange(messages); - var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate); + var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback); var messageResult = result.Messages.LastOrDefault()!; return new ChatResult() { diff --git a/src/MaIN.Services/Services/Abstract/IAgentService.cs b/src/MaIN.Services/Services/Abstract/IAgentService.cs index d62951c6..15938d97 100644 --- a/src/MaIN.Services/Services/Abstract/IAgentService.cs +++ b/src/MaIN.Services/Services/Abstract/IAgentService.cs @@ -1,12 +1,15 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; +using MaIN.Services.Services.LLMService; namespace MaIN.Services.Services.Abstract; public interface IAgentService { - Task Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false); + Task Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false, + Func? callback = null); Task CreateAgent(Agent agent, bool flow = false, bool interactiveResponse = false, InferenceParams? inferenceParams = null, MemoryParams? memoryParams = null, bool disableCache = false); Task GetChatByAgent(string agentId); diff --git a/src/MaIN.Services/Services/Abstract/IStepProcessor.cs b/src/MaIN.Services/Services/Abstract/IStepProcessor.cs index 832c0955..53d2c69a 100644 --- a/src/MaIN.Services/Services/Abstract/IStepProcessor.cs +++ b/src/MaIN.Services/Services/Abstract/IStepProcessor.cs @@ -1,5 +1,6 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; using MaIN.Infrastructure.Models; using Microsoft.Extensions.Logging; @@ -11,6 +12,7 @@ Task ProcessSteps(AgentContextDocument context, AgentDocument agent, Knowledge? knowledge, Chat chat, + Func? callback, Func notifyProgress, Func updateChat, ILogger logger); diff --git a/src/MaIN.Services/Services/AgentService.cs b/src/MaIN.Services/Services/AgentService.cs index 4270f9eb..c3436cf6 100644 --- a/src/MaIN.Services/Services/AgentService.cs +++ b/src/MaIN.Services/Services/AgentService.cs @@ -3,11 +3,13 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; using MaIN.Infrastructure.Repositories.Abstract; using MaIN.Services.Constants; using MaIN.Services.Mappers; using MaIN.Services.Services.Abstract; using MaIN.Services.Services.ImageGenServices; +using MaIN.Services.Services.LLMService; using MaIN.Services.Services.LLMService.Factory; using MaIN.Services.Services.Models.Commands; using MaIN.Services.Services.Steps.Commands; @@ -28,7 +30,12 @@ public class AgentService( MaINSettings maInSettings) : IAgentService { - public async Task Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false) + public async Task Process( + Chat chat, + string agentId, + Knowledge? knowledge, + bool translatePrompt = false, + Func? callback = null) { var agent = await agentRepository.GetAgentById(agentId); if (agent == null) @@ -46,6 +53,7 @@ await notificationService.DispatchNotification( agent, knowledge, chat, + callback, async (status, id, progress, behaviour, details) => { await notificationService.DispatchNotification( diff --git a/src/MaIN.Services/Services/Models/Commands/AnswerCommand.cs b/src/MaIN.Services/Services/Models/Commands/AnswerCommand.cs index 6e56303b..8ec09cb4 100644 --- a/src/MaIN.Services/Services/Models/Commands/AnswerCommand.cs +++ b/src/MaIN.Services/Services/Models/Commands/AnswerCommand.cs @@ -1,5 +1,6 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; using MaIN.Services.Constants; using MaIN.Services.Services.Models.Commands.Base; using MaIN.Services.Services.Steps.Commands; @@ -14,4 +15,5 @@ public class AnswerCommand : BaseCommand, ICommand public KnowledgeUsage KnowledgeUsage { get; init; } public Knowledge? Knowledge { get; init; } public string CommandName => "ANSWER"; + public Func? Callback { get; set; } } \ No newline at end of file diff --git a/src/MaIN.Services/Services/Models/StepContext.cs b/src/MaIN.Services/Services/Models/StepContext.cs index 1e254f7a..c1e11caa 100644 --- a/src/MaIN.Services/Services/Models/StepContext.cs +++ b/src/MaIN.Services/Services/Models/StepContext.cs @@ -1,5 +1,6 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; using MaIN.Infrastructure.Models; namespace MaIN.Services.Services.Models; @@ -16,4 +17,5 @@ public class StepContext public required Func UpdateChat { get; init; } public required string StepName { get; init; } public Knowledge? Knowledge { get; set; } + public Func? Callback { get; set; } } \ No newline at end of file diff --git a/src/MaIN.Services/Services/StepProcessor.cs b/src/MaIN.Services/Services/StepProcessor.cs index 43094ca3..ce0da3d9 100644 --- a/src/MaIN.Services/Services/StepProcessor.cs +++ b/src/MaIN.Services/Services/StepProcessor.cs @@ -1,5 +1,6 @@ using MaIN.Domain.Entities; using MaIN.Domain.Entities.Agents.Knowledge; +using MaIN.Domain.Models; using MaIN.Infrastructure.Models; using MaIN.Services.Services.Abstract; using MaIN.Services.Services.Models; @@ -29,14 +30,18 @@ public async Task ProcessSteps( AgentDocument agent, Knowledge? knowledge, Chat chat, + Func? callback, Func notifyProgress, Func updateChat, ILogger logger) { Message redirectMessage = chat.Messages.Last(); + var stepCount = 0; var tagsToReplaceWithFilter = new List(); foreach (var step in context.Steps!) { + stepCount++; + var lastStep = stepCount.Equals(context.Steps.Count); logger.LogInformation("Processing step: {Step} on agent {agent}", step, agent.Name); var (stepName, arguments) = ParseStep(step); @@ -53,6 +58,7 @@ public async Task ProcessSteps( McpConfig = context.McpConfig, NotifyProgress = notifyProgress, UpdateChat = updateChat, + Callback = lastStep ? callback : null, StepName = stepName }; diff --git a/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs b/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs index 998d4bcb..3c0d56fe 100644 --- a/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs +++ b/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs @@ -29,7 +29,8 @@ public async Task Handle(StepContext context) : useMemory ? KnowledgeUsage.UseMemory : KnowledgeUsage.None, Knowledge = context.Knowledge, - AgentId = context.Agent.Id + AgentId = context.Agent.Id, + Callback = context.Callback }; var answerResponse = await commandDispatcher.DispatchAsync(answerCommand); diff --git a/src/MaIN.Services/Services/Steps/Commands/AnswerCommandHandler.cs b/src/MaIN.Services/Services/Steps/Commands/AnswerCommandHandler.cs index 0f2cac06..c187bdf6 100644 --- a/src/MaIN.Services/Services/Steps/Commands/AnswerCommandHandler.cs +++ b/src/MaIN.Services/Services/Steps/Commands/AnswerCommandHandler.cs @@ -53,7 +53,7 @@ public class AnswerCommandHandler( result = command.Chat.Visual ? await imageGenService!.Send(command.Chat) : await llmService.Send(command.Chat, - new ChatRequestOptions { InteractiveUpdates = command.Chat.Interactive }); + new ChatRequestOptions { InteractiveUpdates = command.Chat.Interactive, TokenCallback = command.Callback }); return result!.Message; }