added retrieval with vector and bm25 search

This commit is contained in:
PaulKoudelka committed 2026-08-03 20:56:19 +02:00
1 parent c78a8f05ef
commit 010f35b08a
21 files changed
+775 -49

No files matched your search

@@ -26,6 +26,8 @@ public abstract class EmbeddingStateClient(string name, string path) : DatabaseC
public abstract Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token);
public abstract Task<IReadOnlyList<EmbeddingStateSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token);
public abstract Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token);
}
@@ -51,3 +53,26 @@ public sealed record EmbeddingStateChunk(
int ChunkIndex,
string ChunkText,
DateTime EmbeddedAtUtc);
public sealed record EmbeddingStateSearchResult(
string ChunkId,
string ParentFileId,
string DataSourceId,
string DataSourceName,
string DataSourceType,
string AbsolutePath,
string FileName,
string RelativePath,
string FileType,
int? PageNumber,
int ChunkIndex,
string ChunkText,
double Score,
string Fingerprint,
long FileSize,
DateTime CreationUtc,
DateTime LastWriteUtc,
DateTime EmbeddedAtUtc,
int ChunkCount,
string ComplianceLevel,
int ComplianceLevelRank);
@@ -46,6 +46,9 @@ public sealed class NoEmbeddingStateClient(string name, string? unavailableReaso
public override Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token) => Task.CompletedTask;
public override Task<IReadOnlyList<EmbeddingStateSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token) =>
Task.FromResult<IReadOnlyList<EmbeddingStateSearchResult>>([]);
public override Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token) => Task.CompletedTask;
public override void Dispose()
@@ -1,4 +1,5 @@
using System.Globalization;
using System.Text.RegularExpressions;
using AIStudio.Settings;
using AIStudio.Tools.PluginSystem;
@@ -16,6 +17,9 @@ public sealed class SqliteEmbeddingStateClientImplementation(
{
private const string DATABASE_NAME = "Local RAG Index";
private const string DATABASE_FILENAME = "rag-index.sqlite3";
private const int MAX_FTS_QUERY_TERMS = 32;
private static readonly Regex FTS_TOKEN_REGEX = new(@"[\p{L}\p{Nd}_]+", RegexOptions.Compiled | RegexOptions.CultureInvariant);
private readonly string databasePath = databasePath;
private readonly string connectionString = new SqliteConnectionStringBuilder
@@ -316,6 +320,84 @@ public sealed class SqliteEmbeddingStateClientImplementation(
await transaction.CommitAsync(token);
}
public override async Task<IReadOnlyList<EmbeddingStateSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token)
{
if (maxMatches <= 0)
return [];
var ftsQuery = BuildFtsQuery(query);
if (string.IsNullOrWhiteSpace(ftsQuery))
return [];
var results = new List<EmbeddingStateSearchResult>(maxMatches);
await using var connection = await this.OpenConnectionAsync(token);
await using var command = CreateCommand(connection, """
SELECT
c.chunk_id,
c.parent_file_id,
ds.data_source_id,
ds.data_source_name,
ds.data_source_type,
f.absolute_path,
f.file_name,
f.relative_path,
f.file_type,
c.page_number,
c.chunk_index,
c.chunk_text,
bm25(embedding_chunks_fts) AS score,
f.fingerprint,
f.file_size,
f.creation_utc,
f.last_write_utc,
c.embedded_at_utc,
f.chunk_count,
f.compliance_level,
f.compliance_level_rank
FROM embedding_chunks_fts
JOIN embedding_chunks c ON c.id = embedding_chunks_fts.rowid
JOIN embedded_files f ON f.parent_file_id = c.parent_file_id
JOIN data_sources ds ON ds.data_source_id = f.data_source_id
WHERE ds.data_source_id = $dataSourceId
AND embedding_chunks_fts MATCH $query
ORDER BY score
LIMIT $maxMatches
""");
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
command.Parameters.AddWithValue("$query", ftsQuery);
command.Parameters.AddWithValue("$maxMatches", maxMatches);
await using var reader = await command.ExecuteReaderAsync(token);
while (await reader.ReadAsync(token))
{
results.Add(new EmbeddingStateSearchResult(
reader.GetString(0),
reader.GetString(1),
reader.GetString(2),
reader.GetString(3),
reader.GetString(4),
reader.GetString(5),
reader.GetString(6),
reader.GetString(7),
reader.GetString(8),
reader.IsDBNull(9) ? null : reader.GetInt32(9),
reader.GetInt32(10),
reader.GetString(11),
reader.GetDouble(12),
reader.GetString(13),
reader.GetInt64(14),
ParseUtc(reader.GetString(15)),
ParseUtc(reader.GetString(16)),
ParseUtc(reader.GetString(17)),
reader.GetInt32(18),
reader.GetString(19),
reader.GetInt32(20)));
}
return results;
}
public override async Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
@@ -525,6 +607,20 @@ public sealed class SqliteEmbeddingStateClientImplementation(
: DateTime.UnixEpoch;
}
private static string BuildFtsQuery(string query)
{
var terms = FTS_TOKEN_REGEX
.Matches(query)
.Select(match => match.Value)
.Where(term => !string.IsNullOrWhiteSpace(term))
.Distinct(StringComparer.OrdinalIgnoreCase)
.Take(MAX_FTS_QUERY_TERMS)
.Select(term => $"\"{term.Replace("\"", "\"\"", StringComparison.Ordinal)}\"")
.ToList();
return terms.Count == 0 ? string.Empty : string.Join(" OR ", terms);
}
private static NoEmbeddingStateClient CreateNoEmbeddingStateClient(string name, string? unavailableReason, DatabaseClientStatus status, ILogger<DatabaseClient> databaseClientLogger)
{
var client = new NoEmbeddingStateClient(name, unavailableReason, status);
@@ -28,6 +28,9 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason,
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token) =>
Task.FromException<IReadOnlyList<VectorSearchResult>>(this.CreateUnavailableException());
public override Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
@@ -15,6 +15,7 @@ public sealed class QdrantEdgeClientImplementation(
private const string INFO_PATH = "/system/qdrant-edge/info";
private const string ENSURE_PATH = "/system/qdrant-edge/ensure";
private const string INSERT_PATH = "/system/qdrant-edge/insert";
private const string SEARCH_PATH = "/system/qdrant-edge/search";
private const string DELETE_FILE_PATH = "/system/qdrant-edge/delete-file";
private const string DELETE_STORE_PATH = "/system/qdrant-edge/delete-store";
@@ -86,6 +87,18 @@ public sealed class QdrantEdgeClientImplementation(
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token);
public override async Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token)
{
if (maxMatches <= 0)
return [];
return await rustService.ExecuteDatabaseQuery<SearchEmbeddingRequest, List<VectorSearchResult>>(
DATABASE_NAME,
SEARCH_PATH,
new SearchEmbeddingRequest(storeName, vector, maxMatches),
token) ?? [];
}
public override Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, DELETE_FILE_PATH, new DeleteEmbeddingByFileRequest(storeName, filePath), token);
@@ -108,6 +121,8 @@ public sealed class QdrantEdgeClientImplementation(
private sealed record InsertEmbeddingRequest(string StoreName, IReadOnlyList<VectorStoragePoint> Points);
private sealed record SearchEmbeddingRequest(string StoreName, IReadOnlyList<float> Vector, int MaxMatches);
private sealed record DeleteEmbeddingByFileRequest(string StoreName, string FilePath);
private sealed record DeleteVectorStoreRequest(string StoreName);
@@ -0,0 +1,24 @@
namespace AIStudio.Tools.Databases.VectorStore;
public sealed record VectorSearchResult(
string PointId,
double Score,
string DataSourceId,
string DataSourceName,
string DataSourceType,
string ChunkId,
string ParentFileId,
string FilePath,
string AbsolutePath,
string FileName,
string RelativePath,
string FileType,
int? PageNumber,
int ChunkIndex,
string Text,
string Fingerprint,
string CreationUtc,
string LastWriteUtc,
string EmbeddedAtUtc,
string ComplianceLevel,
int ComplianceLevelRank);
@@ -6,6 +6,8 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClien
public abstract Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token);
public abstract Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token);
public abstract Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token);
public abstract Task DeleteVectorStore(string storeName, CancellationToken token);