Skip to content

Commit 9610550

Browse files
Merge pull request #101 from wisedev-code/feat/accept_any_stream_in_files
feat: callback on agents
2 parents 4e7d9dc + 00201c0 commit 9610550

File tree

12 files changed

+37
-13
lines changed

12 files changed

+37
-13
lines changed

src/MaIN.Core.UnitTests/AgentContextTests.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ public async Task ProcessAsync_WithStringMessage_ShouldReturnChatResult()
173173
.ReturnsAsync(chat);
174174

175175
_mockAgentService
176-
.Setup(s => s.Process(It.IsAny<Chat>(), _agentContext.GetAgentId(), It.IsAny<Knowledge>(), It.IsAny<bool>()))
176+
.Setup(s => s.Process(It.IsAny<Chat>(), _agentContext.GetAgentId(), It.IsAny<Knowledge>(), It.IsAny<bool>(), null))
177177
.ReturnsAsync(new Chat {
178178
Model = "test-model",
179179
Name = "test",

src/MaIN.Core.UnitTests/FlowContextTests.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ public async Task ProcessAsync_WithStringMessage_ShouldReturnChatResult()
9898
.ReturnsAsync(chat);
9999

100100
_mockAgentService
101-
.Setup(s => s.Process(It.IsAny<Chat>(), firstAgent.Id, It.IsAny<Knowledge>(), It.IsAny<bool>()))
101+
.Setup(s => s.Process(It.IsAny<Chat>(), firstAgent.Id, It.IsAny<Knowledge>(), It.IsAny<bool>(), null))
102102
.ReturnsAsync(new Chat {
103103
Model = "test-model",
104104
Name = "test",

src/MaIN.Core/.nuspec

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
<package>
33
<metadata>
44
<id>MaIN.NET</id>
5-
<version>0.7.2</version>
5+
<version>0.7.3</version>
66
<authors>Wisedev</authors>
77
<owners>Wisedev</owners>
88
<icon>favicon.png</icon>

src/MaIN.Core/Hub/Contexts/AgentContext.cs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)
226226
};
227227
}
228228

229-
public async Task<ChatResult> ProcessAsync(string message, bool translate = false)
229+
public async Task<ChatResult> ProcessAsync(string message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
230230
{
231231
if (_knowledge == null)
232232
{
@@ -240,7 +240,7 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals
240240
Type = MessageType.LocalLLM,
241241
Time = DateTime.Now
242242
});
243-
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate);
243+
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
244244
var messageResult = result.Messages.LastOrDefault()!;
245245
return new ChatResult()
246246
{
@@ -251,15 +251,15 @@ public async Task<ChatResult> ProcessAsync(string message, bool translate = fals
251251
};
252252
}
253253

254-
public async Task<ChatResult> ProcessAsync(Message message, bool translate = false)
254+
public async Task<ChatResult> ProcessAsync(Message message, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
255255
{
256256
if (_knowledge == null)
257257
{
258258
LoadExistingKnowledgeIfExists();
259259
}
260260
var chat = await _agentService.GetChatByAgent(_agent.Id);
261261
chat.Messages.Add(message);
262-
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate);
262+
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
263263
var messageResult = result.Messages.LastOrDefault()!;
264264
return new ChatResult()
265265
{
@@ -270,7 +270,7 @@ public async Task<ChatResult> ProcessAsync(Message message, bool translate = fal
270270
};
271271
}
272272

273-
public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool translate = false)
273+
public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool translate = false, Func<LLMTokenValue, Task>? callback = null)
274274
{
275275
if (_knowledge == null)
276276
{
@@ -279,7 +279,7 @@ public async Task<ChatResult> ProcessAsync(IEnumerable<Message> messages, bool t
279279
var chat = await _agentService.GetChatByAgent(_agent.Id);
280280
chat.Messages.Clear();
281281
chat.Messages.AddRange(messages);
282-
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate);
282+
var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, callback);
283283
var messageResult = result.Messages.LastOrDefault()!;
284284
return new ChatResult()
285285
{

src/MaIN.Services/Services/Abstract/IAgentService.cs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
using MaIN.Domain.Entities;
22
using MaIN.Domain.Entities.Agents;
33
using MaIN.Domain.Entities.Agents.Knowledge;
4+
using MaIN.Domain.Models;
5+
using MaIN.Services.Services.LLMService;
46

57
namespace MaIN.Services.Services.Abstract;
68

79
public interface IAgentService
810
{
9-
Task<Chat> Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false);
11+
Task<Chat> Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false,
12+
Func<LLMTokenValue, Task>? callback = null);
1013
Task<Agent> CreateAgent(Agent agent, bool flow = false, bool interactiveResponse = false,
1114
InferenceParams? inferenceParams = null, MemoryParams? memoryParams = null, bool disableCache = false);
1215
Task<Chat> GetChatByAgent(string agentId);

src/MaIN.Services/Services/Abstract/IStepProcessor.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using MaIN.Domain.Entities;
22
using MaIN.Domain.Entities.Agents.Knowledge;
3+
using MaIN.Domain.Models;
34
using MaIN.Infrastructure.Models;
45
using Microsoft.Extensions.Logging;
56

@@ -11,6 +12,7 @@ Task<Chat> ProcessSteps(AgentContextDocument context,
1112
AgentDocument agent,
1213
Knowledge? knowledge,
1314
Chat chat,
15+
Func<LLMTokenValue, Task>? callback,
1416
Func<string, string, string?, string, string, Task> notifyProgress,
1517
Func<Chat, Task> updateChat,
1618
ILogger logger);

src/MaIN.Services/Services/AgentService.cs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33
using MaIN.Domain.Entities;
44
using MaIN.Domain.Entities.Agents;
55
using MaIN.Domain.Entities.Agents.Knowledge;
6+
using MaIN.Domain.Models;
67
using MaIN.Infrastructure.Repositories.Abstract;
78
using MaIN.Services.Constants;
89
using MaIN.Services.Mappers;
910
using MaIN.Services.Services.Abstract;
1011
using MaIN.Services.Services.ImageGenServices;
12+
using MaIN.Services.Services.LLMService;
1113
using MaIN.Services.Services.LLMService.Factory;
1214
using MaIN.Services.Services.Models.Commands;
1315
using MaIN.Services.Services.Steps.Commands;
@@ -28,7 +30,12 @@ public class AgentService(
2830
MaINSettings maInSettings)
2931
: IAgentService
3032
{
31-
public async Task<Chat> Process(Chat chat, string agentId, Knowledge? knowledge, bool translatePrompt = false)
33+
public async Task<Chat> Process(
34+
Chat chat,
35+
string agentId,
36+
Knowledge? knowledge,
37+
bool translatePrompt = false,
38+
Func<LLMTokenValue, Task>? callback = null)
3239
{
3340
var agent = await agentRepository.GetAgentById(agentId);
3441
if (agent == null)
@@ -46,6 +53,7 @@ await notificationService.DispatchNotification(
4653
agent,
4754
knowledge,
4855
chat,
56+
callback,
4957
async (status, id, progress, behaviour, details) =>
5058
{
5159
await notificationService.DispatchNotification(

src/MaIN.Services/Services/Models/Commands/AnswerCommand.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using MaIN.Domain.Entities;
22
using MaIN.Domain.Entities.Agents.Knowledge;
3+
using MaIN.Domain.Models;
34
using MaIN.Services.Constants;
45
using MaIN.Services.Services.Models.Commands.Base;
56
using MaIN.Services.Services.Steps.Commands;
@@ -14,4 +15,5 @@ public class AnswerCommand : BaseCommand, ICommand<Message?>
1415
public KnowledgeUsage KnowledgeUsage { get; init; }
1516
public Knowledge? Knowledge { get; init; }
1617
public string CommandName => "ANSWER";
18+
public Func<LLMTokenValue, Task>? Callback { get; set; }
1719
}

src/MaIN.Services/Services/Models/StepContext.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using MaIN.Domain.Entities;
22
using MaIN.Domain.Entities.Agents.Knowledge;
3+
using MaIN.Domain.Models;
34
using MaIN.Infrastructure.Models;
45

56
namespace MaIN.Services.Services.Models;
@@ -16,4 +17,5 @@ public class StepContext
1617
public required Func<Chat, Task> UpdateChat { get; init; }
1718
public required string StepName { get; init; }
1819
public Knowledge? Knowledge { get; set; }
20+
public Func<LLMTokenValue, Task>? Callback { get; set; }
1921
}

src/MaIN.Services/Services/StepProcessor.cs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using MaIN.Domain.Entities;
22
using MaIN.Domain.Entities.Agents.Knowledge;
3+
using MaIN.Domain.Models;
34
using MaIN.Infrastructure.Models;
45
using MaIN.Services.Services.Abstract;
56
using MaIN.Services.Services.Models;
@@ -29,14 +30,18 @@ public async Task<Chat> ProcessSteps(
2930
AgentDocument agent,
3031
Knowledge? knowledge,
3132
Chat chat,
33+
Func<LLMTokenValue, Task>? callback,
3234
Func<string, string, string?, string, string, Task> notifyProgress,
3335
Func<Chat, Task> updateChat,
3436
ILogger logger)
3537
{
3638
Message redirectMessage = chat.Messages.Last();
39+
var stepCount = 0;
3740
var tagsToReplaceWithFilter = new List<string>();
3841
foreach (var step in context.Steps!)
3942
{
43+
stepCount++;
44+
var lastStep = stepCount.Equals(context.Steps.Count);
4045
logger.LogInformation("Processing step: {Step} on agent {agent}", step, agent.Name);
4146

4247
var (stepName, arguments) = ParseStep(step);
@@ -53,6 +58,7 @@ public async Task<Chat> ProcessSteps(
5358
McpConfig = context.McpConfig,
5459
NotifyProgress = notifyProgress,
5560
UpdateChat = updateChat,
61+
Callback = lastStep ? callback : null,
5662
StepName = stepName
5763
};
5864

0 commit comments

Comments
 (0)