feat: include tool callback func

Piotr Stachaczynski · Piotr Stachaczynski · commit 4adaf0786948 · 2025-11-04T13:21:47.000+01:00
diff --git 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<Chat>(), _agentContext.GetAgentId(), It.IsAny<Knowledge>(), It.IsAny<bool>(), null))
+            .Setup(s => s.Process(It.IsAny<Chat>(), _agentContext.GetAgentId(), It.IsAny<Knowledge>(), It.IsAny<bool>(), null, 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
@@ -98,7 +98,7 @@ public async Task ProcessAsync_WithStringMessage_ShouldReturnChatResult()
             .ReturnsAsync(chat);
 
         _mockAgentService
-            .Setup(s => s.Process(It.IsAny<Chat>(), firstAgent.Id, It.IsAny<Knowledge>(), It.IsAny<bool>(), null))
+            .Setup(s => s.Process(It.IsAny<Chat>(), firstAgent.Id, It.IsAny<Knowledge>(), It.IsAny<bool>(), null, null))
             .ReturnsAsync(new Chat { 
                 Model = "test-model", 
                 Name = "test",
diff --git a/src/MaIN.Core/.nuspec b/src/MaIN.Core/.nuspec
@@ -2,7 +2,7 @@
 <package>
   <metadata>
     <id>MaIN.NET</id>
-    <version>0.7.7</version>
+    <version>0.7.8</version>
     <authors>Wisedev</authors>
     <owners>Wisedev</owners>
     <icon>favicon.png</icon>
diff --git a/src/MaIN.Core/Hub/Contexts/AgentContext.cs b/src/MaIN.Core/Hub/Contexts/AgentContext.cs
@@ -227,7 +227,11 @@ public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)
         };
     }
     
-    public async Task<ChatResult> ProcessAsync(string message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
+    public async Task<ChatResult> ProcessAsync(
+        string message,
+        bool translate = false,
+        Func<LLMTokenValue, Task>? tokenCallback = null,
+        Func<ToolInvocation, Task>? toolCallback = null)
     {
         if (_knowledge == null)
         {
@@ -241,7 +245,7 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals
             Type = MessageType.LocalLLM,
             Time = DateTime.Now
         });
-        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
+        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);
         var messageResult = result.Messages.LastOrDefault()!;
         return new ChatResult()
         {
@@ -252,15 +256,18 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals
         };
     }
     
-    public async Task<ChatResult> ProcessAsync(Message message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
+    public async Task<ChatResult> ProcessAsync(Message message, 
+        bool translate = false,
+        Func<LLMTokenValue, Task>? tokenCallback = null,
+        Func<ToolInvocation, Task>? toolCallback = null)
     {
         if (_knowledge == null)
         {
             LoadExistingKnowledgeIfExists();
         }
         var chat = await _agentService.GetChatByAgent(_agent.Id);
         chat.Messages.Add(message);
-        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
+        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;
         var messageResult = result.Messages.LastOrDefault()!;
         return new ChatResult()
         {
@@ -271,7 +278,11 @@ public async Task<ChatResult> ProcessAsync(Message message, bool translate = fal
         };
     }
     
-    public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
+    public async Task<ChatResult> ProcessAsync(
+        IEnumerable<Message> messages,
+        bool translate = false,
+        Func<LLMTokenValue, Task>? tokenCallback = null,
+        Func<ToolInvocation, Task>? toolCallback = null)
     {
         if (_knowledge == null)
         {
@@ -285,7 +296,7 @@ public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool t
             chat.Messages.Add(systemMsg);
         }
         chat.Messages.AddRange(messages);
-        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
+        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;
         var messageResult = result.Messages.LastOrDefault()!;
         return new ChatResult()
         {
diff --git a/src/MaIN.Domain/Entities/Tools/ToolCall.cs b/src/MaIN.Domain/Entities/Tools/ToolCall.cs
diff --git a/src/MaIN.Domain/Entities/Tools/ToolInvocation.cs b/src/MaIN.Domain/Entities/Tools/ToolInvocation.cs
@@ -0,0 +1,8 @@
+﻿namespace MaIN.Domain.Entities.Tools;
+
+public class ToolInvocation
+{
+    public string ToolName { get; set; }
+    public string Arguments { get; set; } = null!;
+    public bool Done { get; set; } = false;
+}
diff --git a/src/MaIN.Services/Services/Abstract/IAgentService.cs b/src/MaIN.Services/Services/Abstract/IAgentService.cs
@@ -1,15 +1,15 @@
 using MaIN.Domain.Entities;
 using MaIN.Domain.Entities.Agents;
 using MaIN.Domain.Entities.Agents.Knowledge;
+using MaIN.Domain.Entities.Tools;
 using MaIN.Domain.Models;
-using MaIN.Services.Services.LLMService;
 
 namespace MaIN.Services.Services.Abstract;
 
 public interface IAgentService
 {
     Task<Chat> Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false,
-        Func<LLMTokenValue, Task>? callback = null);
+        Func<LLMTokenValue, Task>? callbackToken = null, Func<ToolInvocation, Task>? callbackTool = null);
     Task<Agent> CreateAgent(Agent agent, bool flow = false, bool interactiveResponse = false,
         InferenceParams? inferenceParams = null, MemoryParams? memoryParams = null, bool disableCache = false);
     Task<Chat> GetChatByAgent(string agentId);
diff --git 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.Entities.Tools;
 using MaIN.Domain.Models;
 using MaIN.Infrastructure.Models;
 using Microsoft.Extensions.Logging;
@@ -13,6 +14,7 @@ Task<Chat> ProcessSteps(AgentContextDocument context,
         Knowledge? knowledge,
         Chat chat,
         Func<LLMTokenValue, Task>? callback,
+        Func<ToolInvocation, Task>? callbackTool,
         Func<string, string, string?, string, string, Task> notifyProgress,
         Func<Chat, Task> updateChat,
         ILogger logger);
diff --git a/src/MaIN.Services/Services/AgentService.cs b/src/MaIN.Services/Services/AgentService.cs
@@ -3,13 +3,13 @@
 using MaIN.Domain.Entities;
 using MaIN.Domain.Entities.Agents;
 using MaIN.Domain.Entities.Agents.Knowledge;
+using MaIN.Domain.Entities.Tools;
 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;
@@ -35,7 +35,8 @@ public async Task<Chat> Process(
         string agentId,
         Knowledge? knowledge,
         bool translatePrompt = false,
-        Func<LLMTokenValue, Task>? callback = null)
+        Func<LLMTokenValue, Task>? callbackToken = null,
+        Func<ToolInvocation, Task>? callbackTool = null)
     {
         var agent = await agentRepository.GetAgentById(agentId);
         if (agent == null) 
@@ -53,7 +54,8 @@ await notificationService.DispatchNotification(
                 agent,
                 knowledge,
                 chat,
-                callback,
+                callbackToken,
+                callbackTool,
                 async (status, id, progress, behaviour, details) =>
                 {
                     await notificationService.DispatchNotification(
diff --git a/src/MaIN.Services/Services/LLMService/AnthropicService.cs b/src/MaIN.Services/Services/LLMService/AnthropicService.cs
@@ -10,6 +10,7 @@
 using System.Text;
 using LLama.Common;
 using MaIN.Domain.Configuration;
+using MaIN.Domain.Entities.Tools;
 using MaIN.Services.Services.LLMService.Utils;
 using MaIN.Services.Services.LLMService;
 
@@ -237,7 +238,19 @@ await notificationService.DispatchNotification(
                 try
                 {
                     var inputJson = JsonSerializer.Serialize(toolUse.Input);
+                    options.ToolCallback?.Invoke(new ToolInvocation()
+                    {
+                        ToolName = toolUse.Name,
+                        Arguments = toolUse.Input.ToString() ?? string.Empty,
+                        Done = false
+                    });
                     var toolResult = await executor(inputJson);
+                    options.ToolCallback?.Invoke(new ToolInvocation()
+                    {
+                        ToolName = toolUse.Name,
+                        Arguments = toolUse.Input.ToString() ?? string.Empty,
+                        Done = true
+                    });
 
                     toolResults.Add(new
                     {
diff --git a/src/MaIN.Services/Services/LLMService/ChatRequestOptions.cs b/src/MaIN.Services/Services/LLMService/ChatRequestOptions.cs
@@ -11,4 +11,5 @@ public class ChatRequestOptions
     public Func<LLMTokenValue, Task>? TokenCallback { get; set; }
     public List<ToolDefinition>? Tools { get; set; }
     public string? ToolChoice { get; set; }
+    public Func<ToolInvocation, Task>? ToolCallback { get; set; }
 }
diff --git a/src/MaIN.Services/Services/LLMService/OpenAiCompatibleService.cs b/src/MaIN.Services/Services/LLMService/OpenAiCompatibleService.cs
@@ -222,8 +222,19 @@ await _notificationService.DispatchNotification(
 
                 try
                 {
+                    options.ToolCallback?.Invoke(new ToolInvocation()
+                    {
+                        ToolName = toolCall.Function.Name,
+                        Arguments = toolCall.Function.Arguments,
+                        Done = false
+                    });
                     var toolResult = await executor(toolCall.Function.Arguments);
-
+                    options.ToolCallback?.Invoke(new ToolInvocation()
+                    {
+                        ToolName = toolCall.Function.Name,
+                        Arguments = toolCall.Function.Arguments,
+                        Done = true
+                    });
                     var toolMessage = new ChatMessage(ServiceConstants.Roles.Tool, toolResult)
                     {
                         ToolCallId = toolCall.Id,
@@ -752,18 +763,15 @@ internal static async Task<object[]> BuildMessagesArray(List<ChatMessage> conver
                 ["content"] = content ?? string.Empty
             };
 
-            // Add tool calls if present (for assistant messages)
             if (msg.ToolCalls != null && msg.ToolCalls.Any())
             {
                 messageObj["tool_calls"] = msg.ToolCalls;
             }
 
-            // Add tool call id if present (for tool messages)
             if (!string.IsNullOrEmpty(msg.ToolCallId))
             {
                 messageObj["tool_call_id"] = msg.ToolCallId;
                 
-                // Tool messages also need the function name
                 if (!string.IsNullOrEmpty(msg.Name))
                 {
                     messageObj["name"] = msg.Name;
diff --git 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.Entities.Tools;
 using MaIN.Domain.Models;
 using MaIN.Services.Constants;
 using MaIN.Services.Services.Models.Commands.Base;
@@ -16,4 +17,5 @@ public class AnswerCommand : BaseCommand, ICommand<Message?>
     public Knowledge? Knowledge { get; init; }
     public string CommandName => "ANSWER";
     public Func<LLMTokenValue, Task>? Callback { get; set; }
+    public Func<ToolInvocation, Task>? ToolCallback { get; set; }
 }
diff --git 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.Entities.Tools;
 using MaIN.Domain.Models;
 using MaIN.Infrastructure.Models;
 
@@ -18,4 +19,5 @@ public class StepContext
     public required string StepName { get; init; }
     public Knowledge? Knowledge { get; set; }
     public Func<LLMTokenValue, Task>? Callback { get; set; }
+    public Func<ToolInvocation, Task>? ToolCallback { get; set; }
 }
diff --git 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.Entities.Tools;
 using MaIN.Domain.Models;
 using MaIN.Infrastructure.Models;
 using MaIN.Services.Services.Abstract;
@@ -25,12 +26,12 @@ public StepProcessor(IEnumerable<IStepHandler> stepHandlers)
         }
     }
 
-    public async Task<Chat> ProcessSteps(
-        AgentContextDocument context,
+    public async Task<Chat> ProcessSteps(AgentContextDocument context,
         AgentDocument agent,
         Knowledge? knowledge,
         Chat chat,
-        Func<LLMTokenValue, Task>? callback,
+        Func<LLMTokenValue, Task>? callbackToken,
+        Func<ToolInvocation, Task>? callbackTool,
         Func<string, string, string?, string, string, Task> notifyProgress,
         Func<Chat, Task> updateChat,
         ILogger logger)
@@ -58,7 +59,8 @@ public async Task<Chat> ProcessSteps(
                 McpConfig = context.McpConfig,
                 NotifyProgress = notifyProgress,
                 UpdateChat = updateChat,
-                Callback = lastStep ? callback : null,
+                Callback = lastStep ? callbackToken : null,
+                ToolCallback = callbackTool,
                 StepName = stepName
             };
 
diff --git a/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs b/src/MaIN.Services/Services/Steps/AnswerStepHandler.cs
@@ -30,7 +30,8 @@ public async Task<StepResult> Handle(StepContext context)
                         : KnowledgeUsage.None,
             Knowledge = context.Knowledge,
             AgentId = context.Agent.Id,
-            Callback = context.Callback
+            Callback = context.Callback,
+            ToolCallback = context.ToolCallback
         };
         
         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
@@ -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, TokenCallback = command.Callback });
+                new ChatRequestOptions { InteractiveUpdates = command.Chat.Interactive, TokenCallback = command.Callback, ToolCallback = command.ToolCallback });
 
         return result!.Message;
     }

Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,11 @@ public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)`
`227`	`227`	`};`
`228`	`228`	`}`
`229`	`229`
`230`		`- public async Task<ChatResult> ProcessAsync(string message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)`
	`230`	`+ public async Task<ChatResult> ProcessAsync(`
	`231`	`+ string message,`
	`232`	`+ bool translate = false,`
	`233`	`+ Func<LLMTokenValue, Task>? tokenCallback = null,`
	`234`	`+ Func<ToolInvocation, Task>? toolCallback = null)`
`231`	`235`	`{`
`232`	`236`	`if (_knowledge == null)`
`233`	`237`	`{`
`@@ -241,7 +245,7 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals`
`241`	`245`	`Type = MessageType.LocalLLM,`
`242`	`246`	`Time = DateTime.Now`
`243`	`247`	`});`
`244`		`- var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);`
	`248`	`+ var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);`
`245`	`249`	`var messageResult = result.Messages.LastOrDefault()!;`
`246`	`250`	`return new ChatResult()`
`247`	`251`	`{`
`@@ -252,15 +256,18 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals`
`252`	`256`	`};`
`253`	`257`	`}`
`254`	`258`
`255`		`- public async Task<ChatResult> ProcessAsync(Message message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)`
	`259`	`+ public async Task<ChatResult> ProcessAsync(Message message,`
	`260`	`+ bool translate = false,`
	`261`	`+ Func<LLMTokenValue, Task>? tokenCallback = null,`
	`262`	`+ Func<ToolInvocation, Task>? toolCallback = null)`
`256`	`263`	`{`
`257`	`264`	`if (_knowledge == null)`
`258`	`265`	`{`
`259`	`266`	`LoadExistingKnowledgeIfExists();`
`260`	`267`	`}`
`261`	`268`	`var chat = await _agentService.GetChatByAgent(_agent.Id);`
`262`	`269`	`chat.Messages.Add(message);`
`263`		`- var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);`
	`270`	`+ var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;`
`264`	`271`	`var messageResult = result.Messages.LastOrDefault()!;`
`265`	`272`	`return new ChatResult()`
`266`	`273`	`{`
`@@ -271,7 +278,11 @@ public async Task<ChatResult> ProcessAsync(Message message, bool translate = fal`
`271`	`278`	`};`
`272`	`279`	`}`
`273`	`280`
`274`		`- public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool translate = false, Func<LLMTokenValue, Task>? callback = null)`
	`281`	`+ public async Task<ChatResult> ProcessAsync(`
	`282`	`+ IEnumerable<Message> messages,`
	`283`	`+ bool translate = false,`
	`284`	`+ Func<LLMTokenValue, Task>? tokenCallback = null,`
	`285`	`+ Func<ToolInvocation, Task>? toolCallback = null)`
`275`	`286`	`{`
`276`	`287`	`if (_knowledge == null)`
`277`	`288`	`{`
`@@ -285,7 +296,7 @@ public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool t`
`285`	`296`	`chat.Messages.Add(systemMsg);`
`286`	`297`	`}`
`287`	`298`	`chat.Messages.AddRange(messages);`
`288`		`- var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);`
	`299`	`+ var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;`
`289`	`300`	`var messageResult = result.Messages.LastOrDefault()!;`
`290`	`301`	`return new ChatResult()`
`291`	`302`	`{`