Skip to content

Commit f42c239

Browse files
alkampfergitalkampfergit
committed
added disk based prompt store
also added prompt store to query rewriter
1 parent efae48b commit f42c239

File tree

5 files changed

+359
-58
lines changed

5 files changed

+359
-58
lines changed

‎src/KernelMemory.Extensions.FunctionalTests/Helper/LocalFolderPromptStoreTests.cs

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
using System;
2+
using System.IO;
3+
using System.Threading.Tasks;
4+
using Microsoft.Extensions.Logging;
5+
using Moq;
6+
using Xunit;
7+
8+
namespace KernelMemory.Extensions.FunctionalTests.Helper;
9+
10+
public class LocalFolderPromptStoreTests : IDisposable
11+
{
12+
private readonly string _testDirectory;
13+
private readonly Mock<ILogger<LocalFolderPromptStore>> _loggerMock;
14+
private readonly LocalFolderPromptStore _store;
15+
16+
public LocalFolderPromptStoreTests()
17+
{
18+
_testDirectory = Path.Combine(Path.GetTempPath(), $"promptstore_tests_{Guid.NewGuid()}");
19+
_loggerMock = new Mock<ILogger<LocalFolderPromptStore>>();
20+
_store = new LocalFolderPromptStore(_testDirectory, _loggerMock.Object);
21+
}
22+
23+
public void Dispose()
24+
{
25+
if (Directory.Exists(_testDirectory))
26+
{
27+
Directory.Delete(_testDirectory, true);
28+
}
29+
}
30+
31+
[Fact]
32+
public async Task GetPromptAsync_NonExistentKey_ReturnsNull()
33+
{
34+
// Act
35+
var result = await _store.GetPromptAsync("nonexistent");
36+
37+
// Assert
38+
Assert.Null(result);
39+
}
40+
41+
[Fact]
42+
public async Task SetAndGetPromptAsync_ValidKey_ReturnsStoredPrompt()
43+
{
44+
// Arrange
45+
const string key = "test-key";
46+
const string expectedPrompt = "This is a test prompt";
47+
48+
// Act
49+
await _store.SetPromptAsync(key, expectedPrompt);
50+
var result = await _store.GetPromptAsync(key);
51+
52+
// Assert
53+
Assert.Equal(expectedPrompt, result);
54+
}
55+
56+
[Fact]
57+
public async Task GetPromptAndSetDefaultAsync_NonExistentKey_SetsAndReturnsDefault()
58+
{
59+
// Arrange
60+
const string key = "default-key";
61+
const string defaultPrompt = "Default prompt value";
62+
63+
// Act
64+
var result = await _store.GetPromptAndSetDefaultAsync(key, defaultPrompt);
65+
var storedPrompt = await _store.GetPromptAsync(key);
66+
67+
// Assert
68+
Assert.Equal(defaultPrompt, result);
69+
Assert.Equal(defaultPrompt, storedPrompt);
70+
}
71+
72+
[Fact]
73+
public async Task GetPromptAndSetDefaultAsync_ExistingKey_ReturnsExistingPrompt()
74+
{
75+
// Arrange
76+
const string key = "existing-key";
77+
const string existingPrompt = "Existing prompt";
78+
const string defaultPrompt = "Default prompt";
79+
await _store.SetPromptAsync(key, existingPrompt);
80+
81+
// Act
82+
var result = await _store.GetPromptAndSetDefaultAsync(key, defaultPrompt);
83+
84+
// Assert
85+
Assert.Equal(existingPrompt, result);
86+
}
87+
88+
[Fact]
89+
public async Task SetPromptAsync_KeyWithSpecialCharacters_HandlesCorrectly()
90+
{
91+
// Arrange
92+
const string key = "special/\\*:?\"<>|characters";
93+
const string expectedPrompt = "Prompt with special characters";
94+
95+
// Act
96+
await _store.SetPromptAsync(key, expectedPrompt);
97+
var result = await _store.GetPromptAsync(key);
98+
99+
// Assert
100+
Assert.Equal(expectedPrompt, result);
101+
}
102+
103+
[Fact]
104+
public async Task GetPromptAndSetDefaultAsync_MissingPlaceholder_LogsError()
105+
{
106+
// Arrange
107+
const string key = "test-placeholder";
108+
const string existingPrompt = "A prompt without placeholder";
109+
const string defaultPrompt = "Default prompt with {{$placeholder}}";
110+
await _store.SetPromptAsync(key, existingPrompt);
111+
112+
// Act
113+
var result = await _store.GetPromptAndSetDefaultAsync(key, defaultPrompt);
114+
115+
// Assert
116+
_loggerMock.Verify(
117+
x => x.Log(
118+
LogLevel.Error,
119+
It.IsAny<EventId>(),
120+
It.Is<It.IsAnyType>((v, t) => v.ToString().Contains("{{$placeholder}}")),
121+
It.IsAny<Exception>(),
122+
It.IsAny<Func<It.IsAnyType, Exception, string>>()
123+
),
124+
Times.Once);
125+
}
126+
127+
[Fact]
128+
public async Task GetPromptAndSetDefaultAsync_MultipleMissingPlaceholders_LogsMultipleErrors()
129+
{
130+
// Arrange
131+
const string key = "test-multiple-placeholders";
132+
const string existingPrompt = "A prompt without any placeholders";
133+
const string defaultPrompt = "Default with {{$first}} and {{$second}}";
134+
await _store.SetPromptAsync(key, existingPrompt);
135+
136+
// Act
137+
var result = await _store.GetPromptAndSetDefaultAsync(key, defaultPrompt);
138+
139+
// Assert
140+
_loggerMock.Verify(
141+
x => x.Log(
142+
LogLevel.Error,
143+
It.IsAny<EventId>(),
144+
It.Is<It.IsAnyType>((v, t) => true),
145+
It.IsAny<Exception>(),
146+
It.IsAny<Func<It.IsAnyType, Exception, string>>()
147+
),
148+
Times.Exactly(2));
149+
}
150+
151+
[Fact]
152+
public async Task GetPromptAndSetDefaultAsync_ValidPlaceholders_NoErrors()
153+
{
154+
// Arrange
155+
const string key = "test-valid-placeholders";
156+
const string existingPrompt = "A prompt with {{$placeholder}} correctly set";
157+
const string defaultPrompt = "Default with {{$placeholder}}";
158+
await _store.SetPromptAsync(key, existingPrompt);
159+
160+
// Act
161+
var result = await _store.GetPromptAndSetDefaultAsync(key, defaultPrompt);
162+
163+
// Assert
164+
_loggerMock.Verify(
165+
x => x.Log(
166+
LogLevel.Error,
167+
It.IsAny<EventId>(),
168+
It.Is<It.IsAnyType>((v, t) => true),
169+
It.IsAny<Exception>(),
170+
It.IsAny<Func<It.IsAnyType, Exception, string>>()
171+
),
172+
Times.Never);
173+
}
174+
}

‎src/KernelMemory.Extensions/BasePromptStore.cs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
using System.Linq;
2+
using System.Threading.Tasks;
3+
using Microsoft.Extensions.Logging;
4+
using System.Text.RegularExpressions;
5+
6+
namespace KernelMemory.Extensions;
7+
8+
public abstract class BasePromptStore : IPromptStore
9+
{
10+
private static readonly Regex PlaceholderRegex = new(@"\{\{\$\w+\}\}", RegexOptions.Compiled);
11+
12+
protected readonly ILogger _log;
13+
14+
protected BasePromptStore(ILogger log)
15+
{
16+
_log = log;
17+
}
18+
19+
protected void ValidatePlaceholders(string defaultPrompt, string loadedPrompt)
20+
{
21+
var placeholders = PlaceholderRegex.Matches(defaultPrompt)
22+
.Cast<Match>()
23+
.Select(m => m.Value)
24+
.ToList();
25+
26+
foreach (var placeholder in placeholders)
27+
{
28+
if (!loadedPrompt.Contains(placeholder))
29+
{
30+
_log.LogError("The prompt does not contain {Placeholder} placeholder, the prompt will not work correctly", placeholder);
31+
}
32+
}
33+
}
34+
35+
public abstract Task<string?> GetPromptAsync(string key);
36+
public abstract Task SetPromptAsync(string key, string prompt);
37+
38+
public virtual async Task<string> GetPromptAndSetDefaultAsync(string key, string defaultPrompt)
39+
{
40+
var prompt = await GetPromptAsync(key);
41+
if (prompt == null)
42+
{
43+
await SetPromptAsync(key, defaultPrompt);
44+
return defaultPrompt;
45+
}
46+
47+
ValidatePlaceholders(defaultPrompt, prompt);
48+
return prompt;
49+
}
50+
}

‎src/KernelMemory.Extensions/IPromptStore.cs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@ public interface IPromptStore
2323
/// various components will use some default prompts.</returns>
2424
Task<string?> GetPromptAsync(string key);
2525

26+
/// <summary>
27+
/// Get the prompt for the given key and set the default prompt if the prompt is not present.
28+
/// </summary>
29+
/// <param name="key"></param>
30+
/// <param name="defaultPrompt"></param>
31+
/// <returns></returns>
32+
Task<string> GetPromptAndSetDefaultAsync(string key, string defaultPrompt);
33+
2634
/// <summary>
2735
/// Allow setting prompt value.
2836
/// </summary>
@@ -45,6 +53,17 @@ public class NullPromptStore : IPromptStore
4553
return Task.FromResult<string?>(null);
4654
}
4755

56+
/// <summary>
57+
/// Return the default prompt always
58+
/// </summary>
59+
/// <param name="key"></param>
60+
/// <param name="defaultPrompt"></param>
61+
/// <returns></returns>
62+
public Task<string> GetPromptAndSetDefaultAsync(string key, string defaultPrompt)
63+
{
64+
return Task.FromResult<string>(defaultPrompt);
65+
}
66+
4867
/// <summary>
4968
/// Allow setting prompt value.
5069
/// </summary>

‎src/KernelMemory.Extensions/LocalFolderPromptStore.cs

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
using System;
2+
using System.IO;
3+
using System.Threading.Tasks;
4+
using Microsoft.Extensions.Logging;
5+
using System.Linq;
6+
using Microsoft.KernelMemory.Diagnostics;
7+
8+
namespace KernelMemory.Extensions;
9+
10+
public class LocalFolderPromptStore : BasePromptStore
11+
{
12+
private readonly string _promptDirectory;
13+
14+
public LocalFolderPromptStore(string promptDirectory, ILogger<LocalFolderPromptStore>? log = null)
15+
: base(log ?? DefaultLogger<LocalFolderPromptStore>.Instance)
16+
{
17+
_promptDirectory = promptDirectory;
18+
Directory.CreateDirectory(promptDirectory);
19+
}
20+
21+
private string GetPromptFilePath(string key)
22+
{
23+
var sanitizedKey = string.Join("_", key.Split(Path.GetInvalidFileNameChars()));
24+
return Path.Combine(_promptDirectory, $"{sanitizedKey}.prompt");
25+
}
26+
27+
public override async Task<string?> GetPromptAsync(string key)
28+
{
29+
var filePath = GetPromptFilePath(key);
30+
if (!File.Exists(filePath))
31+
{
32+
return null;
33+
}
34+
35+
return await File.ReadAllTextAsync(filePath);
36+
}
37+
38+
public override async Task SetPromptAsync(string key, string prompt)
39+
{
40+
var filePath = GetPromptFilePath(key);
41+
await File.WriteAllTextAsync(filePath, prompt);
42+
}
43+
}

0 commit comments

Comments
 (0)