Skip to content

Commit

Permalink
Re-implementing the Postgres memory store (#1735)
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.
-->
Reimplement Postgres in-memory storage based on the discussion in #1338.


### Description
<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
- Re-implementing the Postgres memory store, mapping between SK
collection to Postgres table.
- PostgresMemoryStore no longer implements `IDisposable` pattern.
- No longer execute the enable pgvector extension statement, add it to
README.md to be executed by the user. It only needs to be executed once
in the database, and the extension may be enabled differently for
different hosting methods.

### 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 <[email protected]>
  • Loading branch information
JadynWong and dmytrostruk committed Jun 29, 2023
1 parent 537bc49 commit 07aa6a7
Show file tree
Hide file tree
Showing 12 changed files with 1,052 additions and 571 deletions.
13 changes: 10 additions & 3 deletions dotnet/samples/KernelSyntaxExamples/Example39_Postgres.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,22 +5,29 @@
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.Connectors.Memory.Postgres;
using Microsoft.SemanticKernel.Memory;
using Npgsql;
using Pgvector.Npgsql;
using RepoUtils;

// ReSharper disable once InconsistentNaming
public static class Example39_Postgres
{
private const string MemoryCollectionName = "postgres-test";
private const string MemoryCollectionName = "postgres_test";

public static async Task RunAsync()
{
string connectionString = Env.Var("POSTGRES_CONNECTIONSTRING");
using PostgresMemoryStore memoryStore = await PostgresMemoryStore.ConnectAsync(connectionString, vectorSize: 1536);
NpgsqlDataSourceBuilder dataSourceBuilder = new(Env.Var("POSTGRES_CONNECTIONSTRING"));
dataSourceBuilder.UseVector();
using NpgsqlDataSource dataSource = dataSourceBuilder.Build();

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

IKernel kernel = Kernel.Builder
.WithLogger(ConsoleLogger.Log)
.WithOpenAITextCompletionService("text-davinci-003", Env.Var("OPENAI_API_KEY"))
.WithOpenAITextEmbeddingGenerationService("text-embedding-ada-002", Env.Var("OPENAI_API_KEY"))
.WithMemoryStorage(memoryStore)
//.WithPostgresMemoryStore(dataSource, vectorSize: 1536, schema: "public") // This method offers an alternative approach to registering Postgres memory store.
.Build();

Console.WriteLine("== Printing Collections in DB ==");
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Pgvector;

namespace Microsoft.SemanticKernel.Connectors.Memory.Postgres;

/// <summary>
/// Interface for client managing postgres database operations.
/// </summary>
public interface IPostgresDbClient
{
/// <summary>
/// Check if a collection exists.
/// </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></returns>
Task<bool> DoesCollectionExistsAsync(string collectionName, CancellationToken cancellationToken = default);

/// <summary>
/// Create a collection.
/// </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></returns>
Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default);

/// <summary>
/// Get all collections.
/// </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);

/// <summary>
/// Delete a collection.
/// </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></returns>
Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default);

/// <summary>
/// Upsert entry into a collection.
/// </summary>
/// <param name="collectionName">The name assigned to a collection 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="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);

/// <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="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);

/// <summary>
/// Read a entry by its key.
/// </summary>
/// <param name="collectionName">The name assigned to a collection 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="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);

/// <summary>
/// Delete a entry by its key.
/// </summary>
/// <param name="collectionName">The name assigned to a collection 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);
}
Loading

0 comments on commit 07aa6a7

Please sign in to comment.