Adding contentType to document (#473)

## Motivation and Context (Why the change? What's the scenario?) Address #471: [Bug] "contentType" is not populated by WriteFileAsync in MongoDbAtlasStorage ## High level description (Approach, Design) Added a `contentType` to the document. --------- Co-authored-by: Pradeep Rajendran <[email protected]> Co-authored-by: Devis Lucato <[email protected]>
microsoft · May 11, 2024 · 6b0d347 · 6b0d347
1 parent 0239be6
commit 6b0d347
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/extensions/MongoDbAtlas/MongoDbAtlas/MongoDbAtlasStorage.cs b/extensions/MongoDbAtlas/MongoDbAtlas/MongoDbAtlasStorage.cs
@@ -6,6 +6,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.KernelMemory.ContentStorage;
+using Microsoft.KernelMemory.Pipeline;
 using MongoDB.Bson;
 using MongoDB.Driver;
 using MongoDB.Driver.GridFS;
@@ -15,8 +16,13 @@ namespace Microsoft.KernelMemory.MongoDbAtlas;
 [Experimental("KMEXP03")]
 public sealed class MongoDbAtlasStorage : MongoDbAtlasBaseStorage, IContentStorage
 {
-    public MongoDbAtlasStorage(MongoDbAtlasConfig config) : base(config)
+    private readonly IMimeTypeDetection _mimeTypeDetection;
+
+    public MongoDbAtlasStorage(
+        MongoDbAtlasConfig config,
+        IMimeTypeDetection? mimeTypeDetection = null) : base(config)
     {
+        this._mimeTypeDetection = mimeTypeDetection ?? new MimeTypesDetection();
     }
 
     public Task CreateIndexDirectoryAsync(string index, CancellationToken cancellationToken = default)
@@ -73,12 +79,13 @@ public async Task WriteFileAsync(
         if (extension == ".txt")
         {
             using var reader = new StreamReader(streamContent);
-            var doc = new BsonDocument()
+            var doc = new BsonDocument
             {
                 { "_id", id },
                 { "documentId", documentId },
                 { "fileName", fileName },
-                { "content", new BsonString(await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false)) }
+                { "content", new BsonString(await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false)) },
+                { "contentType", MimeTypes.PlainText }
             };
             await this.SaveDocumentAsync(index, id, doc, cancellationToken).ConfigureAwait(false);
         }
@@ -94,6 +101,7 @@ public async Task WriteFileAsync(
             doc["documentId"] = documentId;
             doc["fileName"] = fileName;
             doc["content"] = content;
+            doc["contentType"] = MimeTypes.PlainText;
             await this.SaveDocumentAsync(index, id, doc, cancellationToken).ConfigureAwait(false);
         }
         else
@@ -105,14 +113,15 @@ public async Task WriteFileAsync(
                 {
                     { "index", index },
                     { "documentId", documentId },
-                    { "fileName", fileName }
+                    { "fileName", fileName },
+                    { "contentType", this._mimeTypeDetection.GetFileType(fileName) }
                 }
             };
 
             // Since the pattern of usage is that you can upload a file for a document id and then update, we need to delete
             // any existing file with the same id check if the file exists and delete it
             IAsyncCursor<GridFSFileInfo<string>> existingFile = await GetFromBucketByIdAsync(id, bucket, cancellationToken).ConfigureAwait(false);
-            if (existingFile.Any(cancellationToken))
+            if (await existingFile.AnyAsync(cancellationToken).ConfigureAwait(false))
             {
                 await bucket.DeleteAsync(id, cancellationToken).ConfigureAwait(false);
             }