mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-17 11:57:59 +00:00
added retrieval with vector and bm25 search
This commit is contained in:
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()
|
||||
|
||||
+96
@@ -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);
|
||||
|
||||
Reference in new issue
Block a user