Skip to content

Commit 98ae8f6

Browse files
KSemenenkoKSemenenko
committed
storage
1 parent d93b494 commit 98ae8f6

10 files changed

Lines changed: 502 additions & 37 deletions

‎README.md‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ dotnet add package ManagedCode.MCPGateway
2525
- one registry for local tools, stdio MCP servers, HTTP MCP servers, or prebuilt `McpClient` instances
2626
- descriptor indexing that enriches search with tool name, description, required arguments, and input schema
2727
- vector search when an `IEmbeddingGenerator<string, Embedding<float>>` is registered
28+
- optional persisted tool embeddings through `IMcpGatewayToolEmbeddingStore`
2829
- lexical fallback when embeddings are unavailable
2930
- one invoke surface for both local `AIFunction` tools and MCP tools
3031
- optional meta-tools you can hand back to another model as normal `AITool` instances
@@ -182,6 +183,63 @@ services.AddManagedCodeMcpGateway(options =>
182183

183184
The keyed registration is the preferred one, so you can dedicate a specific embedder to the gateway without affecting other app services.
184185

186+
## Persistent Tool Embeddings
187+
188+
For process-local caching, the package already includes `McpGatewayInMemoryToolEmbeddingStore`:
189+
190+
```csharp
191+
var services = new ServiceCollection();
192+
193+
services.AddKeyedSingleton<IEmbeddingGenerator<string, Embedding<float>>, MyEmbeddingGenerator>(
194+
McpGatewayServiceKeys.EmbeddingGenerator);
195+
services.AddSingleton<IMcpGatewayToolEmbeddingStore, McpGatewayInMemoryToolEmbeddingStore>();
196+
197+
services.AddManagedCodeMcpGateway(options =>
198+
{
199+
options.AddTool(
200+
"local",
201+
AIFunctionFactory.Create(
202+
static (string query) => $"github:{query}",
203+
new AIFunctionFactoryOptions
204+
{
205+
Name = "github_search_repositories",
206+
Description = "Search GitHub repositories by user query."
207+
}));
208+
});
209+
```
210+
211+
If you want to keep descriptor embeddings in a database or another persistent store, register your own `IMcpGatewayToolEmbeddingStore` implementation instead:
212+
213+
```csharp
214+
var services = new ServiceCollection();
215+
216+
services.AddKeyedSingleton<IEmbeddingGenerator<string, Embedding<float>>, MyEmbeddingGenerator>(
217+
McpGatewayServiceKeys.EmbeddingGenerator);
218+
services.AddSingleton<IMcpGatewayToolEmbeddingStore, MyToolEmbeddingStore>();
219+
220+
services.AddManagedCodeMcpGateway(options =>
221+
{
222+
options.AddTool(
223+
"local",
224+
AIFunctionFactory.Create(
225+
static (string query) => $"github:{query}",
226+
new AIFunctionFactoryOptions
227+
{
228+
Name = "github_search_repositories",
229+
Description = "Search GitHub repositories by user query."
230+
}));
231+
});
232+
```
233+
234+
During `BuildIndexAsync()` the gateway:
235+
236+
- computes a descriptor-document hash per tool
237+
- asks `IMcpGatewayToolEmbeddingStore` for matching stored vectors
238+
- generates embeddings only for tools that are missing in the store
239+
- upserts the newly generated vectors back into the store
240+
241+
This avoids recalculating tool embeddings on every rebuild while still refreshing them automatically when the descriptor document changes. Query embeddings are still generated at search time from the registered `IEmbeddingGenerator<string, Embedding<float>>`.
242+
185243
## Supported Sources
186244

187245
- local `AITool` / `AIFunction`

‎src/ManagedCode.MCPGateway/Abstractions/IMcpGatewayToolEmbeddingStore.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
namespace ManagedCode.MCPGateway.Abstractions;
2+
3+
public interface IMcpGatewayToolEmbeddingStore
4+
{
5+
Task<IReadOnlyList<McpGatewayToolEmbedding>> GetAsync(
6+
IReadOnlyList<McpGatewayToolEmbeddingLookup> lookups,
7+
CancellationToken cancellationToken = default);
8+
9+
Task UpsertAsync(
10+
IReadOnlyList<McpGatewayToolEmbedding> embeddings,
11+
CancellationToken cancellationToken = default);
12+
}

‎src/ManagedCode.MCPGateway/McpGateway.cs‎

Lines changed: 161 additions & 36 deletions
Large diffs are not rendered by default.

‎src/ManagedCode.MCPGateway/McpGatewayInMemoryToolEmbeddingStore.cs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
using System.Collections.Concurrent;
2+
using ManagedCode.MCPGateway.Abstractions;
3+
4+
namespace ManagedCode.MCPGateway;
5+
6+
public sealed class McpGatewayInMemoryToolEmbeddingStore : IMcpGatewayToolEmbeddingStore
7+
{
8+
private readonly ConcurrentDictionary<McpGatewayToolEmbeddingLookup, McpGatewayToolEmbedding> _embeddings = new();
9+
10+
public Task<IReadOnlyList<McpGatewayToolEmbedding>> GetAsync(
11+
IReadOnlyList<McpGatewayToolEmbeddingLookup> lookups,
12+
CancellationToken cancellationToken = default)
13+
{
14+
cancellationToken.ThrowIfCancellationRequested();
15+
16+
var results = lookups
17+
.Where(_embeddings.ContainsKey)
18+
.Select(lookup => Clone(_embeddings[lookup]))
19+
.ToList();
20+
21+
return Task.FromResult<IReadOnlyList<McpGatewayToolEmbedding>>(results);
22+
}
23+
24+
public Task UpsertAsync(
25+
IReadOnlyList<McpGatewayToolEmbedding> embeddings,
26+
CancellationToken cancellationToken = default)
27+
{
28+
cancellationToken.ThrowIfCancellationRequested();
29+
30+
foreach (var embedding in embeddings)
31+
{
32+
var clone = Clone(embedding);
33+
_embeddings[new McpGatewayToolEmbeddingLookup(clone.ToolId, clone.DocumentHash)] = clone;
34+
}
35+
36+
return Task.CompletedTask;
37+
}
38+
39+
private static McpGatewayToolEmbedding Clone(McpGatewayToolEmbedding embedding)
40+
=> embedding with
41+
{
42+
Vector = [.. embedding.Vector]
43+
};
44+
}

‎src/ManagedCode.MCPGateway/Models/McpGatewayToolEmbedding.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
namespace ManagedCode.MCPGateway;
2+
3+
public sealed record McpGatewayToolEmbedding(
4+
string ToolId,
5+
string SourceId,
6+
string ToolName,
7+
string DocumentHash,
8+
float[] Vector);

‎src/ManagedCode.MCPGateway/Models/McpGatewayToolEmbeddingLookup.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
namespace ManagedCode.MCPGateway;
2+
3+
public sealed record McpGatewayToolEmbeddingLookup(
4+
string ToolId,
5+
string DocumentHash);

‎tests/ManagedCode.MCPGateway.Tests/Search/McpGatewayInMemoryToolEmbeddingStoreTests.cs‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
namespace ManagedCode.MCPGateway.Tests;
2+
3+
public sealed class McpGatewayInMemoryToolEmbeddingStoreTests
4+
{
5+
[TUnit.Core.Test]
6+
public async Task GetAsync_ReturnsEmbeddingsForMatchingLookups()
7+
{
8+
var store = new McpGatewayInMemoryToolEmbeddingStore();
9+
await store.UpsertAsync(
10+
[
11+
new McpGatewayToolEmbedding(
12+
"local:github_search_issues",
13+
"local",
14+
"github_search_issues",
15+
"hash-1",
16+
[1f, 2f, 3f]),
17+
new McpGatewayToolEmbedding(
18+
"local:weather_search_forecast",
19+
"local",
20+
"weather_search_forecast",
21+
"hash-2",
22+
[4f, 5f, 6f])
23+
]);
24+
25+
var result = await store.GetAsync(
26+
[
27+
new McpGatewayToolEmbeddingLookup("local:weather_search_forecast", "hash-2"),
28+
new McpGatewayToolEmbeddingLookup("local:missing", "hash-3"),
29+
new McpGatewayToolEmbeddingLookup("local:github_search_issues", "hash-1")
30+
]);
31+
32+
await Assert.That(result.Count).IsEqualTo(2);
33+
await Assert.That(result.Any(static item => item.ToolId == "local:github_search_issues")).IsTrue();
34+
await Assert.That(result.Any(static item => item.ToolId == "local:weather_search_forecast")).IsTrue();
35+
}
36+
37+
[TUnit.Core.Test]
38+
public async Task UpsertAsync_ClonesVectorsOnWriteAndRead()
39+
{
40+
var store = new McpGatewayInMemoryToolEmbeddingStore();
41+
var inputVector = new[] { 1f, 2f, 3f };
42+
43+
await store.UpsertAsync(
44+
[
45+
new McpGatewayToolEmbedding(
46+
"local:github_search_issues",
47+
"local",
48+
"github_search_issues",
49+
"hash-1",
50+
inputVector)
51+
]);
52+
53+
inputVector[0] = 99f;
54+
55+
var firstRead = await store.GetAsync(
56+
[
57+
new McpGatewayToolEmbeddingLookup("local:github_search_issues", "hash-1")
58+
]);
59+
firstRead[0].Vector[1] = 77f;
60+
61+
var secondRead = await store.GetAsync(
62+
[
63+
new McpGatewayToolEmbeddingLookup("local:github_search_issues", "hash-1")
64+
]);
65+
66+
await Assert.That(secondRead[0].Vector[0]).IsEqualTo(1f);
67+
await Assert.That(secondRead[0].Vector[1]).IsEqualTo(2f);
68+
await Assert.That(secondRead[0].Vector[2]).IsEqualTo(3f);
69+
}
70+
}

‎tests/ManagedCode.MCPGateway.Tests/Search/McpGatewaySearchTests.cs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,77 @@ public async Task SearchAsync_UsesContextOnlyInputWhenQueryIsMissing()
9494
await Assert.That(embeddingGenerator.Calls[1].Single()).IsEqualTo("context summary: weather forecast");
9595
}
9696

97+
[TUnit.Core.Test]
98+
public async Task BuildIndexAsync_ReusesStoredToolEmbeddingsOnNextBuild()
99+
{
100+
var embeddingStore = new McpGatewayInMemoryToolEmbeddingStore();
101+
var firstEmbeddingGenerator = new TestEmbeddingGenerator();
102+
103+
await using (var firstServiceProvider = GatewayTestServiceProviderFactory.Create(
104+
ConfigureSearchTools,
105+
firstEmbeddingGenerator,
106+
embeddingStore))
107+
{
108+
var gateway = firstServiceProvider.GetRequiredService<IMcpGateway>();
109+
var buildResult = await gateway.BuildIndexAsync();
110+
111+
await Assert.That(buildResult.VectorizedToolCount).IsEqualTo(2);
112+
await Assert.That(firstEmbeddingGenerator.Calls.Count).IsEqualTo(1);
113+
}
114+
115+
var secondEmbeddingGenerator = new TestEmbeddingGenerator(new TestEmbeddingGeneratorOptions
116+
{
117+
ThrowOnInput = static _ => true
118+
});
119+
120+
await using var secondServiceProvider = GatewayTestServiceProviderFactory.Create(
121+
ConfigureSearchTools,
122+
secondEmbeddingGenerator,
123+
embeddingStore);
124+
var secondGateway = secondServiceProvider.GetRequiredService<IMcpGateway>();
125+
126+
var secondBuildResult = await secondGateway.BuildIndexAsync();
127+
128+
await Assert.That(secondBuildResult.VectorizedToolCount).IsEqualTo(2);
129+
await Assert.That(secondBuildResult.Diagnostics.Any(static diagnostic => diagnostic.Code == "embedding_failed")).IsFalse();
130+
await Assert.That(secondEmbeddingGenerator.Calls.Count).IsEqualTo(0);
131+
}
132+
133+
[TUnit.Core.Test]
134+
public async Task BuildIndexAsync_GeneratesAndPersistsOnlyMissingStoredToolEmbeddings()
135+
{
136+
var embeddingStore = new TestToolEmbeddingStore();
137+
var initialEmbeddingGenerator = new TestEmbeddingGenerator();
138+
139+
await using (var initialServiceProvider = GatewayTestServiceProviderFactory.Create(
140+
ConfigureSearchTools,
141+
initialEmbeddingGenerator,
142+
embeddingStore))
143+
{
144+
var gateway = initialServiceProvider.GetRequiredService<IMcpGateway>();
145+
await gateway.BuildIndexAsync();
146+
}
147+
148+
embeddingStore.Remove("local:weather_search_forecast");
149+
150+
var incrementalEmbeddingGenerator = new TestEmbeddingGenerator();
151+
await using var incrementalServiceProvider = GatewayTestServiceProviderFactory.Create(
152+
ConfigureSearchTools,
153+
incrementalEmbeddingGenerator,
154+
embeddingStore);
155+
var incrementalGateway = incrementalServiceProvider.GetRequiredService<IMcpGateway>();
156+
157+
var buildResult = await incrementalGateway.BuildIndexAsync();
158+
159+
await Assert.That(buildResult.VectorizedToolCount).IsEqualTo(2);
160+
await Assert.That(incrementalEmbeddingGenerator.Calls.Count).IsEqualTo(1);
161+
await Assert.That(incrementalEmbeddingGenerator.Calls[0].Count).IsEqualTo(1);
162+
await Assert.That(incrementalEmbeddingGenerator.Calls[0].Single().Contains("weather_search_forecast", StringComparison.Ordinal)).IsTrue();
163+
await Assert.That(embeddingStore.UpsertCalls.Count).IsEqualTo(2);
164+
await Assert.That(embeddingStore.UpsertCalls[1].Count).IsEqualTo(1);
165+
await Assert.That(embeddingStore.UpsertCalls[1].Single().ToolId).IsEqualTo("local:weather_search_forecast");
166+
}
167+
97168
[TUnit.Core.Test]
98169
public async Task SearchAsync_PrefersKeyedEmbeddingGeneratorOverUnkeyedRegistration()
99170
{

‎tests/ManagedCode.MCPGateway.Tests/TestSupport/GatewayTestServiceProviderFactory.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
using Microsoft.Extensions.AI;
22
using Microsoft.Extensions.DependencyInjection;
33
using Microsoft.Extensions.Logging;
4+
using ManagedCode.MCPGateway.Abstractions;
45

56
namespace ManagedCode.MCPGateway.Tests;
67

78
internal static class GatewayTestServiceProviderFactory
89
{
910
public static ServiceProvider Create(
1011
Action<McpGatewayOptions> configure,
11-
IEmbeddingGenerator<string, Embedding<float>>? embeddingGenerator = null)
12+
IEmbeddingGenerator<string, Embedding<float>>? embeddingGenerator = null,
13+
IMcpGatewayToolEmbeddingStore? embeddingStore = null)
1214
{
1315
var services = new ServiceCollection();
1416
services.AddLogging(static logging => logging.SetMinimumLevel(LogLevel.Debug));
@@ -18,6 +20,11 @@ public static ServiceProvider Create(
1820
services.AddSingleton(embeddingGenerator);
1921
}
2022

23+
if (embeddingStore is not null)
24+
{
25+
services.AddSingleton(embeddingStore);
26+
}
27+
2128
services.AddManagedCodeMcpGateway(configure);
2229
return services.BuildServiceProvider();
2330
}

‎tests/ManagedCode.MCPGateway.Tests/TestSupport/TestToolEmbeddingStore.cs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
using ManagedCode.MCPGateway.Abstractions;
2+
3+
namespace ManagedCode.MCPGateway.Tests;
4+
5+
internal sealed class TestToolEmbeddingStore : IMcpGatewayToolEmbeddingStore
6+
{
7+
private readonly Dictionary<McpGatewayToolEmbeddingLookup, McpGatewayToolEmbedding> _embeddings = [];
8+
9+
public List<IReadOnlyList<McpGatewayToolEmbeddingLookup>> GetCalls { get; } = [];
10+
11+
public List<IReadOnlyList<McpGatewayToolEmbedding>> UpsertCalls { get; } = [];
12+
13+
public Task<IReadOnlyList<McpGatewayToolEmbedding>> GetAsync(
14+
IReadOnlyList<McpGatewayToolEmbeddingLookup> lookups,
15+
CancellationToken cancellationToken = default)
16+
{
17+
cancellationToken.ThrowIfCancellationRequested();
18+
19+
GetCalls.Add(lookups.ToList());
20+
21+
var matches = lookups
22+
.Where(_embeddings.ContainsKey)
23+
.Select(lookup => Clone(_embeddings[lookup]))
24+
.ToList();
25+
26+
return Task.FromResult<IReadOnlyList<McpGatewayToolEmbedding>>(matches);
27+
}
28+
29+
public Task UpsertAsync(
30+
IReadOnlyList<McpGatewayToolEmbedding> embeddings,
31+
CancellationToken cancellationToken = default)
32+
{
33+
cancellationToken.ThrowIfCancellationRequested();
34+
35+
var clonedBatch = embeddings
36+
.Select(Clone)
37+
.ToList();
38+
UpsertCalls.Add(clonedBatch);
39+
40+
foreach (var embedding in clonedBatch)
41+
{
42+
_embeddings[new McpGatewayToolEmbeddingLookup(embedding.ToolId, embedding.DocumentHash)] = embedding;
43+
}
44+
45+
return Task.CompletedTask;
46+
}
47+
48+
public void Remove(string toolId)
49+
{
50+
var keys = _embeddings.Keys
51+
.Where(lookup => string.Equals(lookup.ToolId, toolId, StringComparison.OrdinalIgnoreCase))
52+
.ToList();
53+
54+
foreach (var key in keys)
55+
{
56+
_embeddings.Remove(key);
57+
}
58+
}
59+
60+
private static McpGatewayToolEmbedding Clone(McpGatewayToolEmbedding embedding)
61+
=> embedding with
62+
{
63+
Vector = [.. embedding.Vector]
64+
};
65+
}

0 commit comments

Comments
 (0)