Improve postgres memory

microsoft · Jun 30, 2023 · 7d73b22 · 7d73b22
1 parent 348c1c8
commit 7d73b22
Show file tree

Hide file tree

Showing 5 changed files with 320 additions and 266 deletions.
diff --git a/dotnet/src/Connectors/Connectors.Memory.Postgres/IPostgresDbClient.cs b/dotnet/src/Connectors/Connectors.Memory.Postgres/IPostgresDbClient.cs
@@ -13,76 +13,103 @@ namespace Microsoft.SemanticKernel.Connectors.Memory.Postgres;
 public interface IPostgresDbClient
 {
     /// <summary>
-    /// Check if a collection exists.
+    /// Check if a collection table exists.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task<bool> DoesCollectionExistsAsync(string collectionName, CancellationToken cancellationToken = default);
+    Task<bool> DoesTableExistsAsync(string tableName, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Create a collection.
+    /// Create a collection table.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default);
+    Task CreateTableAsync(string tableName, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Get all collections.
+    /// Create a index for collection table.
     /// </summary>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    IAsyncEnumerable<string> GetCollectionsAsync(CancellationToken cancellationToken = default);
+    Task CreateIndexAsync(string tableName, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Delete a collection.
+    /// Get all collection tables.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A group of collection tables.</returns>
+    IAsyncEnumerable<string> GetTablesAsync(CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Delete a collection table.
+    /// </summary>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default);
+    Task DeleteTableAsync(string tableName, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Upsert entry into a collection.
+    /// Upsert entry into a collection table.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="key">The key of the entry to upsert.</param>
     /// <param name="metadata">The metadata of the entry.</param>
     /// <param name="embedding">The embedding of the entry.</param>
-    /// <param name="timestamp">The timestamp of the entry</param>
+    /// <param name="timestamp">The timestamp of the entry.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task UpsertAsync(string collectionName, string key, string? metadata, Vector? embedding, long? timestamp, CancellationToken cancellationToken = default);
+    Task UpsertAsync(string tableName, string key, string? metadata, Vector? embedding, long? timestamp, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Gets the nearest matches to the <see cref="Vector"/>.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
-    /// <param name="embeddingFilter">The <see cref="Vector"/> to compare the collection's embeddings with.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
+    /// <param name="embedding">The <see cref="Vector"/> to compare the collection's embeddings with.</param>
     /// <param name="limit">The maximum number of similarity results to return.</param>
     /// <param name="minRelevanceScore">The minimum relevance threshold for returned results.</param>
     /// <param name="withEmbeddings">If true, the embeddings will be returned in the entries.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns></returns>
-    IAsyncEnumerable<(PostgresMemoryEntry, double)> GetNearestMatchesAsync(string collectionName, Vector embeddingFilter, int limit, double minRelevanceScore = 0, bool withEmbeddings = false, CancellationToken cancellationToken = default);
+    /// <returns>An asynchronous stream of <see cref="PostgresMemoryEntry"/> objects that the nearest matches to the <see cref="Vector"/>.</returns>
+    IAsyncEnumerable<(PostgresMemoryEntry, double)> GetNearestMatchesAsync(string tableName, Vector embedding, int limit, double minRelevanceScore = 0, bool withEmbeddings = false, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Read a entry by its key.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="key">The key of the entry to read.</param>
-    /// <param name="withEmbeddings">If true, the embeddings will be returned in the entries.</param>
+    /// <param name="withEmbeddings">If true, the embeddings will be returned in the entry.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task<PostgresMemoryEntry?> ReadAsync(string collectionName, string key, bool withEmbeddings = false, CancellationToken cancellationToken = default);
+    Task<PostgresMemoryEntry?> ReadAsync(string tableName, string key, bool withEmbeddings = false, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Read multiple entries by their keys.
+    /// </summary>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
+    /// <param name="keys">The keys of the entries to read.</param>
+    /// <param name="withEmbeddings">If true, the embeddings will be returned in the entries.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>An asynchronous stream of <see cref="PostgresMemoryEntry"/> objects that match the given keys.</returns>
+    IAsyncEnumerable<PostgresMemoryEntry> ReadBatchAsync(string tableName, IEnumerable<string> keys, bool withEmbeddings = false, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Delete a entry by its key.
     /// </summary>
-    /// <param name="collectionName">The name assigned to a collection of entries.</param>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
     /// <param name="key">The key of the entry to delete.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns></returns>
-    Task DeleteAsync(string collectionName, string key, CancellationToken cancellationToken = default);
+    Task DeleteAsync(string tableName, string key, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Delete multiple entries by their key.
+    /// </summary>
+    /// <param name="tableName">The name assigned to a collection table of entries.</param>
+    /// <param name="keys">The keys of the entries to delete.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns></returns>
+    Task DeleteBatchAsync(string tableName, IEnumerable<string> keys, CancellationToken cancellationToken = default);
 }