From 5846251a5f7453b406ce2184939c56d4f8ce8cb9 Mon Sep 17 00:00:00 2001
From: Devis Lucato <devis@microsoft.com>
Date: Thu, 30 Nov 2023 23:31:55 -0800
Subject: [PATCH] Upgrade KM dependency and use new interfaces

---
 Directory.Packages.props                     |  2 +-
 PostgresMemoryStorage/DependencyInjection.cs |  2 +-
 PostgresMemoryStorage/PostgresMemory.cs      | 17 ++++++++++++++---
 TestApplication/TestApplication.csproj       |  2 +-
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/Directory.Packages.props b/Directory.Packages.props
index 0039c7a..a203447 100644
--- a/Directory.Packages.props
+++ b/Directory.Packages.props
@@ -3,7 +3,7 @@
     <ManagePackageVersionsCentrally>true</ManagePackageVersionsCentrally>
   </PropertyGroup>
   <ItemGroup>
-    <PackageVersion Include="Microsoft.KernelMemory.Abstractions" Version="0.12.231122.1-preview " />
+    <PackageVersion Include="Microsoft.KernelMemory.Abstractions" Version="0.15.231130.2-preview" />
     <PackageVersion Include="Microsoft.Extensions.DependencyInjection" Version="8.0.0" />
   </ItemGroup>
   <!-- Sources -->
diff --git a/PostgresMemoryStorage/DependencyInjection.cs b/PostgresMemoryStorage/DependencyInjection.cs
index 13320c5..e2213de 100644
--- a/PostgresMemoryStorage/DependencyInjection.cs
+++ b/PostgresMemoryStorage/DependencyInjection.cs
@@ -47,7 +47,7 @@ public static IServiceCollection AddPostgresAsVectorDb(this IServiceCollection s
     {
         return services
             .AddSingleton<PostgresConfig>(config)
-            .AddSingleton<IVectorDb, PostgresMemory>();
+            .AddSingleton<IMemoryDb, PostgresMemory>();
     }
 
     /// <summary>
diff --git a/PostgresMemoryStorage/PostgresMemory.cs b/PostgresMemoryStorage/PostgresMemory.cs
index ce4ee2d..0e380f3 100644
--- a/PostgresMemoryStorage/PostgresMemory.cs
+++ b/PostgresMemoryStorage/PostgresMemory.cs
@@ -7,26 +7,37 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.KernelMemory.Diagnostics;
 using Microsoft.KernelMemory.MemoryStorage;
+using Microsoft.SemanticKernel.AI.Embeddings;
 
 namespace Microsoft.KernelMemory.Postgres;
 
 /// <summary>
 /// Postgres connector for Kernel Memory.
 /// </summary>
-public class PostgresMemory : IVectorDb
+public class PostgresMemory : IMemoryDb
 {
     private readonly ILogger<PostgresMemory> _log;
+    private readonly ITextEmbeddingGeneration _embeddingGenerator;
 
     /// <summary>
     /// Create a new instance of Postgres KM connector
     /// </summary>
     /// <param name="config">Postgres configuration</param>
+    /// <param name="embeddingGenerator">Text embedding generator</param>
     /// <param name="log">Application logger</param>
     public PostgresMemory(
         PostgresConfig config,
+        ITextEmbeddingGeneration embeddingGenerator,
         ILogger<PostgresMemory>? log = null)
     {
         this._log = log ?? DefaultLogger<PostgresMemory>.Instance;
+
+        this._embeddingGenerator = embeddingGenerator;
+
+        if (this._embeddingGenerator == null)
+        {
+            throw new PostgresException("Embedding generator not configured");
+        }
     }
 
     /// <inheritdoc />
@@ -65,12 +76,12 @@ public class PostgresMemory : IVectorDb
     /// <inheritdoc />
     public IAsyncEnumerable<(MemoryRecord, double)> GetSimilarListAsync(
         string index,
-        Embedding embedding,
+        string text,
         ICollection<MemoryFilter>? filters = null,
         double minRelevance = 0,
         int limit = 1,
         bool withEmbeddings = false,
-        CancellationToken cancellationToken = default)
+        CancellationToken cancellationToken = new CancellationToken())
     {
         if (filters != null)
         {
diff --git a/TestApplication/TestApplication.csproj b/TestApplication/TestApplication.csproj
index a888574..0b155d6 100644
--- a/TestApplication/TestApplication.csproj
+++ b/TestApplication/TestApplication.csproj
@@ -15,7 +15,7 @@
     </ItemGroup>
 
     <ItemGroup>
-        <PackageReference Include="Microsoft.KernelMemory.Core" Version="0.12.231122.1-preview "/>
+        <PackageReference Include="Microsoft.KernelMemory.Core" Version="0.15.231130.2-preview" />
     </ItemGroup>
 
     <ItemGroup>