Skip to content

Commit 4c5226c

Browse files
srebreksrebrek
authored
Chat with files and tools conflict (#114) (#144)
* feat: No support for chat with file and tools at the same time (#114) If tools and files are configured at the same time, file-RAG is provided as a tool that model can call. Previously files had higher priority and tools were silently ignored. Tools only and files only pipelines remains as before. fixes #114 * refactor: remove dead code * fix: correct assertion * refactor: make IIngestedMemory internal * refactor: Use CreateIngestedMemoryAsync as it duplicates code in the AskMemory * fix: remove cancellation tokens from cleanup calls * fix: minor fixes - VertexService has proper ovverride pethod instead of new - ProcessWithToolsAsync in the OpenAiCompatibleService returns full response intead of truncated * fix: improve FileSearchTool error handling * fix: Cloud models not supporting semantic search correctly fallbacks to the prompt injection
1 parent 25bb0b6 commit 4c5226c

19 files changed

Lines changed: 917 additions & 85 deletions

‎MaIN.Core.IntegrationTests/Fakes/FakeHttpMessageHandler.cs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@ public sealed class FakeHttpMessageHandler : HttpMessageHandler
1212
public HttpStatusCode ResponseStatusCode { get; set; } = HttpStatusCode.OK;
1313
public string ResponseBody { get; set; } = string.Empty;
1414

15+
private readonly Queue<string> _responseQueue = new();
16+
17+
public void EnqueueResponse(string body) => _responseQueue.Enqueue(body);
18+
1519
protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken ct)
1620
{
1721
LastRequest = request;
@@ -25,9 +29,11 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
2529
catch { }
2630
}
2731

32+
var body = _responseQueue.Count > 0 ? _responseQueue.Dequeue() : ResponseBody;
33+
2834
return new HttpResponseMessage(ResponseStatusCode)
2935
{
30-
Content = new StringContent(ResponseBody, Encoding.UTF8, "application/json")
36+
Content = new StringContent(body, Encoding.UTF8, "application/json")
3137
};
3238
}
3339
}

‎MaIN.Core.IntegrationTests/Fakes/FakeKernelMemory.cs‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
using System.Runtime.CompilerServices;
2+
using Microsoft.KernelMemory;
3+
using Microsoft.KernelMemory.Context;
4+
5+
namespace MaIN.Core.IntegrationTests.Fakes;
6+
7+
internal sealed class FakeKernelMemory : IKernelMemory
8+
{
9+
public int SearchCallCount { get; private set; }
10+
public List<string> SearchQueries { get; } = [];
11+
12+
public Task<SearchResult> SearchAsync(
13+
string query,
14+
string? index = null,
15+
MemoryFilter? filter = null,
16+
ICollection<MemoryFilter>? filters = null,
17+
double minRelevance = 0,
18+
int limit = -1,
19+
IContext? context = null,
20+
CancellationToken ct = default)
21+
{
22+
SearchCallCount++;
23+
SearchQueries.Add(query);
24+
return Task.FromResult(new SearchResult
25+
{
26+
Results =
27+
[
28+
new Citation
29+
{
30+
SourceName = "test-doc",
31+
Partitions = [new Citation.Partition
32+
{ Text = "Copernicus proposed the heliocentric model of the solar system." }
33+
]
34+
}
35+
]
36+
});
37+
}
38+
39+
public Task DeleteIndexAsync(
40+
string? index = null,
41+
CancellationToken ct = default)
42+
=> Task.CompletedTask;
43+
44+
public async IAsyncEnumerable<MemoryAnswer> AskStreamingAsync(
45+
string question,
46+
string? index = null,
47+
MemoryFilter? filter = null,
48+
ICollection<MemoryFilter>? filters = null,
49+
double minRelevance = 0,
50+
SearchOptions? options = null,
51+
IContext? context = null,
52+
[EnumeratorCancellation] CancellationToken ct = default)
53+
{
54+
yield return new MemoryAnswer { Question = question, Result = "Fake answer.", NoResult = false };
55+
await Task.CompletedTask;
56+
}
57+
58+
public Task<string> ImportDocumentAsync(
59+
Stream content,
60+
string? fileName = null,
61+
string? documentId = null,
62+
TagCollection? tags = null,
63+
string? index = null,
64+
IEnumerable<string>? steps = null,
65+
IContext? context = null,
66+
CancellationToken ct = default)
67+
=> throw new NotImplementedException();
68+
69+
public Task<string> ImportDocumentAsync(
70+
string filePath,
71+
string? documentId = null,
72+
TagCollection? tags = null,
73+
string? index = null,
74+
IEnumerable<string>? steps = null,
75+
IContext? context = null,
76+
CancellationToken ct = default)
77+
=> throw new NotImplementedException();
78+
79+
public Task<string> ImportDocumentAsync(
80+
DocumentUploadRequest uploadRequest,
81+
IContext? context = null,
82+
CancellationToken ct = default)
83+
=> throw new NotImplementedException();
84+
85+
public Task<string> ImportDocumentAsync(
86+
Document document,
87+
string? index = null,
88+
IEnumerable<string>? steps = null,
89+
IContext? context = null,
90+
CancellationToken ct = default)
91+
=> throw new NotImplementedException();
92+
93+
public Task<string> ImportTextAsync(
94+
string text,
95+
string? documentId = null,
96+
TagCollection? tags = null,
97+
string? index = null,
98+
IEnumerable<string>? steps = null,
99+
IContext? context = null,
100+
CancellationToken ct = default)
101+
=> throw new NotImplementedException();
102+
103+
public Task<string> ImportWebPageAsync(
104+
string url,
105+
string? documentId = null,
106+
TagCollection? tags = null,
107+
string? index = null,
108+
IEnumerable<string>? steps = null,
109+
IContext? context = null,
110+
CancellationToken ct = default)
111+
=> throw new NotImplementedException();
112+
113+
public Task<StreamableFileContent> ExportFileAsync(
114+
string documentId,
115+
string fileName,
116+
string? index = null,
117+
CancellationToken ct = default)
118+
=> throw new NotImplementedException();
119+
120+
public Task<bool> IsDocumentReadyAsync(
121+
string documentId,
122+
string? index = null,
123+
CancellationToken ct = default)
124+
=> throw new NotImplementedException();
125+
126+
public Task<DataPipelineStatus?> GetDocumentStatusAsync(
127+
string documentId,
128+
string? index = null,
129+
CancellationToken ct = default)
130+
=> throw new NotImplementedException();
131+
132+
public Task DeleteDocumentAsync(
133+
string documentId,
134+
string? index = null,
135+
CancellationToken ct = default)
136+
=> throw new NotImplementedException();
137+
138+
public Task<IEnumerable<IndexDetails>> ListIndexesAsync(
139+
CancellationToken ct = default)
140+
=> throw new NotImplementedException();
141+
142+
}

‎MaIN.Core.IntegrationTests/Fakes/FakeMemoryFactory.cs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
using System.Diagnostics.CodeAnalysis;
2+
using LLama;
3+
using MaIN.Domain.Entities;
4+
using MaIN.Services.Services.LLMService.Memory;
5+
using MaIN.Services.Services.LLMService.Memory.Embeddings;
6+
using Microsoft.KernelMemory;
7+
8+
namespace MaIN.Core.IntegrationTests.Fakes;
9+
10+
internal sealed class FakeMemoryFactory : IMemoryFactory
11+
{
12+
public FakeKernelMemory KernelMemory { get; } = new();
13+
14+
[Experimental("KMEXP00")]
15+
public (IKernelMemory km, LLamaSharpTextEmbeddingMaINClone generator, LlamaSharpTextGen textGenerator)
16+
CreateMemoryWithModel(string modelsPath, LLamaWeights llmModel, string modelName, MemoryParams memoryParams)
17+
=> throw new NotImplementedException("Local model memory not used in this test.");
18+
19+
public IKernelMemory CreateMemoryWithOpenAi(string openAiKey, MemoryParams memoryParams)
20+
=> KernelMemory;
21+
22+
public IKernelMemory CreateMemoryWithGemini(string geminiKey, MemoryParams memoryParams)
23+
=> KernelMemory;
24+
25+
public IKernelMemory CreateMemoryWithVertex(Func<ValueTask<string>> bearerTokenProvider, string location,
26+
string projectId, MemoryParams memoryParams)
27+
=> KernelMemory;
28+
}

‎MaIN.Core.IntegrationTests/Fakes/FakeMemoryService.cs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
using MaIN.Services.Services.LLMService;
2+
using MaIN.Services.Services.LLMService.Memory;
3+
using Microsoft.KernelMemory;
4+
using Microsoft.KernelMemory.AI;
5+
6+
namespace MaIN.Core.IntegrationTests.Fakes;
7+
8+
internal sealed class FakeMemoryService : IMemoryService
9+
{
10+
public Task ImportDataToMemory((IKernelMemory km, ITextEmbeddingGenerator? generator) memory,
11+
ChatMemoryOptions options, CancellationToken cancellationToken)
12+
=> Task.CompletedTask;
13+
14+
public string CleanResponseText(string text) => text;
15+
}

‎MaIN.Core.IntegrationTests/FilesAndToolsIntegrationTests.cs‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
using MaIN.Core.Hub.Utils;
2+
using MaIN.Core.IntegrationTests.Fakes;
3+
using MaIN.Domain.Configuration;
4+
using MaIN.Domain.Models.Abstract;
5+
using MaIN.Services.Services.LLMService.Memory;
6+
using Microsoft.Extensions.DependencyInjection;
7+
using FileInfo = MaIN.Domain.Entities.FileInfo;
8+
9+
namespace MaIN.Core.IntegrationTests;
10+
11+
[Collection("IntegrationTests")]
12+
public class FilesAndToolsIntegrationTests : LLMServiceTestBase
13+
{
14+
private const string TestModelId = "files-tools-integration-model";
15+
16+
private readonly FakeMemoryFactory _fakeMemoryFactory = new();
17+
private readonly FakeMemoryService _fakeMemoryService = new();
18+
19+
public FilesAndToolsIntegrationTests()
20+
{
21+
ModelRegistry.RegisterOrReplace(new GenericCloudModel(TestModelId, BackendType.OpenAi));
22+
}
23+
24+
protected override void ConfigureServices(IServiceCollection services)
25+
{
26+
base.ConfigureServices(services);
27+
services.AddSingleton<IMemoryFactory>(_fakeMemoryFactory);
28+
services.AddSingleton<IMemoryService>(_fakeMemoryService);
29+
}
30+
31+
[Fact]
32+
public async Task FilesAndTools_BothToolsAreExecuted_AndNoExceptionIsThrown()
33+
{
34+
// Arrange
35+
HttpHandler.EnqueueResponse(ToolCallResponse("search_documents", """{"query":"Copernicus astronomy"}"""));
36+
HttpHandler.EnqueueResponse(ToolCallResponse("get_current_time", "{}"));
37+
HttpHandler.EnqueueResponse(FinalResponse(
38+
"Copernicus proposed the heliocentric model. The time is 2026-06-29."));
39+
40+
var userToolCalled = false;
41+
var tools = new ToolsConfigurationBuilder()
42+
.AddTool("get_current_time", "Returns the current date", () =>
43+
{
44+
userToolCalled = true;
45+
return "2026-06-29";
46+
})
47+
.Build();
48+
49+
var files = new List<FileInfo>
50+
{
51+
new() { Name = "copernicus.txt", Extension = ".txt", Content = "Copernicus was a Polish astronomer." }
52+
};
53+
54+
// Act
55+
var result = await AIHub.Chat()
56+
.WithModel(TestModelId)
57+
.WithMessage("What did Copernicus contribute to astronomy? Also, what is the current time?")
58+
.WithFiles(files)
59+
.WithTools(tools)
60+
.CompleteAsync();
61+
62+
// Assert
63+
Assert.NotNull(result);
64+
Assert.True(_fakeMemoryFactory.KernelMemory.SearchCallCount > 0);
65+
Assert.True(userToolCalled);
66+
Assert.False(string.IsNullOrEmpty(result.Message.Content));
67+
}
68+
69+
[Fact]
70+
public async Task FilesOnly_DoesNotInvokeMemoryFactory()
71+
{
72+
// Arrange
73+
HttpHandler.ResponseBody = FinalResponse("Summary of Copernicus.");
74+
75+
var files = new List<FileInfo>
76+
{
77+
new() { Name = "copernicus.txt", Extension = ".txt", Content = "Copernicus was a Polish astronomer." }
78+
};
79+
80+
var searchCountBefore = _fakeMemoryFactory.KernelMemory.SearchCallCount;
81+
82+
// Act
83+
var result = await AIHub.Chat()
84+
.WithModel(TestModelId)
85+
.WithMessage("Summarise the document.")
86+
.WithFiles(files)
87+
.CompleteAsync();
88+
89+
// Assert
90+
Assert.NotNull(result);
91+
Assert.Equal(searchCountBefore, _fakeMemoryFactory.KernelMemory.SearchCallCount);
92+
}
93+
94+
private static string ToolCallResponse(string toolName, string arguments, string callId = "call_001") =>
95+
$$"""
96+
{
97+
"choices": [{
98+
"message": {
99+
"role": "assistant",
100+
"content": null,
101+
"tool_calls": [{
102+
"id": "{{callId}}",
103+
"type": "function",
104+
"function": {
105+
"name": "{{toolName}}",
106+
"arguments": "{{EscapeJson(arguments)}}"
107+
}
108+
}]
109+
}
110+
}]
111+
}
112+
""";
113+
114+
private static string FinalResponse(string content) =>
115+
$$"""
116+
{
117+
"choices": [{
118+
"message": {
119+
"role": "assistant",
120+
"content": "{{content}}"
121+
}
122+
}]
123+
}
124+
""";
125+
126+
private static string EscapeJson(string s) => s.Replace("\"", "\\\"");
127+
}

0 commit comments

Comments
 (0)