Merge pull request #68 from tryAGI/Add-metadata-searching-for-InMemory

Add inmemory searching for metadata and also add tests
tryAGI · Oct 21, 2024 · dad3e84 · dad3e84
2 parents 6bf112e + 6499dda
commit dad3e84
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 3 deletions.
diff --git a/src/InMemory/src/InMemoryVectorCollection.cs b/src/InMemory/src/InMemoryVectorCollection.cs
@@ -92,8 +92,29 @@ public Task<bool> IsEmptyAsync(CancellationToken cancellationToken = default)
         return Task.FromResult(_vectors.GetValueOrDefault(id));
     }
 
-    Task<List<Vector>> IVectorCollection.SearchByMetadata(Dictionary<string, object> filters, CancellationToken cancellationToken)
+    public Task<List<Vector>> SearchByMetadata(Dictionary<string, object> filters, CancellationToken cancellationToken = default)
     {
-        throw new NotImplementedException();
+        filters = filters ?? throw new ArgumentNullException(nameof(filters));
+
+        var filteredVectors = _vectors.Values.Where(vector =>
+        {
+            // Check if all filters match
+            foreach (var filter in filters)
+            {
+                object? metadataValue = null;
+                if (vector.Metadata != null && !vector.Metadata.TryGetValue(filter.Key, out metadataValue) || metadataValue == null)
+                {
+                    return false;
+                }
+                else if (!metadataValue.Equals(filter.Value)) // Convert metadata value to string and compare
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }).ToList();
+
+        return Task.FromResult(filteredVectors);
     }
 }
diff --git a/src/IntegrationTests/DatabaseTests.cs b/src/IntegrationTests/DatabaseTests.cs
@@ -1,5 +1,5 @@
-using LangChain.Extensions;
 using LangChain.DocumentLoaders;
+using LangChain.Extensions;
 
 namespace LangChain.Databases.IntegrationTests;
 
@@ -296,4 +296,47 @@ public async Task SimilaritySearchWithScores_Ok(SupportedDatabase database)
             first.Distance.Should().BeGreaterOrEqualTo(1f);
         }
     }
+
+    [TestCase(SupportedDatabase.InMemory)]
+    //[TestCase(SupportedDatabase.Chroma)]
+    //[TestCase(SupportedDatabase.OpenSearch)]
+    //[TestCase(SupportedDatabase.Postgres)]
+    [TestCase(SupportedDatabase.SqLite)]
+    //[TestCase(SupportedDatabase.DuckDb)]
+    //[TestCase(SupportedDatabase.Weaviate)]
+    //[TestCase(SupportedDatabase.Elasticsearch)]
+    //[TestCase(SupportedDatabase.Milvus)]
+    public async Task MetadataSearch_Ok(SupportedDatabase database)
+    {
+        await using var environment = await StartEnvironmentForAsync(database);
+        var vectorCollection = await environment.VectorDatabase.GetOrCreateCollectionAsync(environment.CollectionName, dimensions: environment.Dimensions);
+
+        var texts = new[] { "apple", "orange" };
+
+        var metadatas = new Dictionary<string, object>[2];
+        metadatas[0] = new Dictionary<string, object>
+        {
+            ["color"] = "red"
+        };
+
+        metadatas[1] = new Dictionary<string, object>
+        {
+            ["color"] = "orange"
+        };
+
+        var totalItems = await vectorCollection.AddTextsAsync(environment.EmbeddingModel, texts, metadatas);
+
+        // Define the filters to get the orange entry
+        var filters = new Dictionary<string, object>
+        {
+            { "color", "orange" }
+        };
+
+
+        var items = await vectorCollection.SearchByMetadata(filters);
+
+        totalItems.Should().HaveCount(2);
+        items.Should().HaveCount(1);
+        var vector = items.SingleOrDefault()?.Metadata?["color"].Should().Be("orange");
+    }
 }