Skip to content

Commit 6b0d347

Browse files
pradeepr-roboticistPradeep Rajendrandluc
authored
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]>
1 parent 0239be6 commit 6b0d347

File tree

1 file changed

+14
-5
lines changed

1 file changed

+14
-5
lines changed

extensions/MongoDbAtlas/MongoDbAtlas/MongoDbAtlasStorage.cs

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
using System.Threading;
77
using System.Threading.Tasks;
88
using Microsoft.KernelMemory.ContentStorage;
9+
using Microsoft.KernelMemory.Pipeline;
910
using MongoDB.Bson;
1011
using MongoDB.Driver;
1112
using MongoDB.Driver.GridFS;
@@ -15,8 +16,13 @@ namespace Microsoft.KernelMemory.MongoDbAtlas;
1516
[Experimental("KMEXP03")]
1617
public sealed class MongoDbAtlasStorage : MongoDbAtlasBaseStorage, IContentStorage
1718
{
18-
public MongoDbAtlasStorage(MongoDbAtlasConfig config) : base(config)
19+
private readonly IMimeTypeDetection _mimeTypeDetection;
20+
21+
public MongoDbAtlasStorage(
22+
MongoDbAtlasConfig config,
23+
IMimeTypeDetection? mimeTypeDetection = null) : base(config)
1924
{
25+
this._mimeTypeDetection = mimeTypeDetection ?? new MimeTypesDetection();
2026
}
2127

2228
public Task CreateIndexDirectoryAsync(string index, CancellationToken cancellationToken = default)
@@ -73,12 +79,13 @@ public async Task WriteFileAsync(
7379
if (extension == ".txt")
7480
{
7581
using var reader = new StreamReader(streamContent);
76-
var doc = new BsonDocument()
82+
var doc = new BsonDocument
7783
{
7884
{ "_id", id },
7985
{ "documentId", documentId },
8086
{ "fileName", fileName },
81-
{ "content", new BsonString(await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false)) }
87+
{ "content", new BsonString(await reader.ReadToEndAsync(cancellationToken).ConfigureAwait(false)) },
88+
{ "contentType", MimeTypes.PlainText }
8289
};
8390
await this.SaveDocumentAsync(index, id, doc, cancellationToken).ConfigureAwait(false);
8491
}
@@ -94,6 +101,7 @@ public async Task WriteFileAsync(
94101
doc["documentId"] = documentId;
95102
doc["fileName"] = fileName;
96103
doc["content"] = content;
104+
doc["contentType"] = MimeTypes.PlainText;
97105
await this.SaveDocumentAsync(index, id, doc, cancellationToken).ConfigureAwait(false);
98106
}
99107
else
@@ -105,14 +113,15 @@ public async Task WriteFileAsync(
105113
{
106114
{ "index", index },
107115
{ "documentId", documentId },
108-
{ "fileName", fileName }
116+
{ "fileName", fileName },
117+
{ "contentType", this._mimeTypeDetection.GetFileType(fileName) }
109118
}
110119
};
111120

112121
// Since the pattern of usage is that you can upload a file for a document id and then update, we need to delete
113122
// any existing file with the same id check if the file exists and delete it
114123
IAsyncCursor<GridFSFileInfo<string>> existingFile = await GetFromBucketByIdAsync(id, bucket, cancellationToken).ConfigureAwait(false);
115-
if (existingFile.Any(cancellationToken))
124+
if (await existingFile.AnyAsync(cancellationToken).ConfigureAwait(false))
116125
{
117126
await bucket.DeleteAsync(id, cancellationToken).ConfigureAwait(false);
118127
}

0 commit comments

Comments
 (0)