Skip to content

Commit

Permalink
Improve postgres memory (#1794)
Browse files Browse the repository at this point in the history
### Motivation and Context
<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->
Follow
#1735 (review)
advice to rename methods and separate responsibilities.

### Description
<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
- Separate the responsibilities of PostgresMemoryStore and
PostgresDbClient.
- Add the real batch methods to PostgresDbClient.
- Save timestamp with `TIMESTAMP WITH TIME ZONE` type
- Stop creating index

### Contribution Checklist
<!-- Before submitting this PR, please make sure: -->
- [x] The code builds clean without any errors or warnings
- [x] The PR follows SK Contribution Guidelines
(https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
- [x] The code follows the .NET coding conventions
(https://learn.microsoft.com/dotnet/csharp/fundamentals/coding-style/coding-conventions)
verified with `dotnet format`
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

---------

Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
  • Loading branch information
JadynWong and dmytrostruk committed Jul 3, 2023
1 parent 93a172b commit b139203
Show file tree
Hide file tree
Showing 9 changed files with 343 additions and 319 deletions.
2 changes: 1 addition & 1 deletion dotnet/samples/KernelSyntaxExamples/Example39_Postgres.cs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ public static async Task RunAsync()
dataSourceBuilder.UseVector();
using NpgsqlDataSource dataSource = dataSourceBuilder.Build();

PostgresMemoryStore memoryStore = new(dataSource, vectorSize: 1536, schema: "public", numberOfLists: 100);
PostgresMemoryStore memoryStore = new(dataSource, vectorSize: 1536, schema: "public");

IKernel kernel = Kernel.Builder
.WithLogger(ConsoleLogger.Log)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
using System.Net.Http;
using System.Text;
using System.Text.Json.Serialization;
using Microsoft.SemanticKernel.Connectors.Memory.Pinecone.Http;

namespace Microsoft.SemanticKernel.Connectors.Memory.Pinecone.Model;

Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
Expand All @@ -13,76 +14,95 @@ namespace Microsoft.SemanticKernel.Connectors.Memory.Postgres;
public interface IPostgresDbClient
{
/// <summary>
/// Check if a collection exists.
/// Check if a table exists.
/// </summary>
/// <param name="collectionName">The name assigned to a collection of entries.</param>
/// <param name="tableName">The name assigned to a 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 table.
/// </summary>
/// <param name="collectionName">The name assigned to a collection of entries.</param>
/// <param name="tableName">The name assigned to a 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.
/// Get all tables.
/// </summary>
/// <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);
/// <returns>A group of tables.</returns>
IAsyncEnumerable<string> GetTablesAsync(CancellationToken cancellationToken = default);

/// <summary>
/// Delete a collection.
/// Delete a table.
/// </summary>
/// <param name="collectionName">The name assigned to a collection of entries.</param>
/// <param name="tableName">The name assigned to a 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 table.
/// </summary>
/// <param name="collectionName">The name assigned to a collection of entries.</param>
/// <param name="tableName">The name assigned to a 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. Its 'DateTimeKind' must be <see cref="DateTimeKind.Utc"/></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, DateTime? 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 table of entries.</param>
/// <param name="embedding">The <see cref="Vector"/> to compare the table'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 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 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 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 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);
}
Loading

0 comments on commit b139203

Please sign in to comment.