6
6
using System . Threading ;
7
7
using System . Threading . Tasks ;
8
8
using Microsoft . KernelMemory . ContentStorage ;
9
+ using Microsoft . KernelMemory . Pipeline ;
9
10
using MongoDB . Bson ;
10
11
using MongoDB . Driver ;
11
12
using MongoDB . Driver . GridFS ;
@@ -15,8 +16,13 @@ namespace Microsoft.KernelMemory.MongoDbAtlas;
15
16
[ Experimental ( "KMEXP03" ) ]
16
17
public sealed class MongoDbAtlasStorage : MongoDbAtlasBaseStorage , IContentStorage
17
18
{
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 )
19
24
{
25
+ this . _mimeTypeDetection = mimeTypeDetection ?? new MimeTypesDetection ( ) ;
20
26
}
21
27
22
28
public Task CreateIndexDirectoryAsync ( string index , CancellationToken cancellationToken = default )
@@ -73,12 +79,13 @@ public async Task WriteFileAsync(
73
79
if ( extension == ".txt" )
74
80
{
75
81
using var reader = new StreamReader ( streamContent ) ;
76
- var doc = new BsonDocument ( )
82
+ var doc = new BsonDocument
77
83
{
78
84
{ "_id" , id } ,
79
85
{ "documentId" , documentId } ,
80
86
{ "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 }
82
89
} ;
83
90
await this . SaveDocumentAsync ( index , id , doc , cancellationToken ) . ConfigureAwait ( false ) ;
84
91
}
@@ -94,6 +101,7 @@ public async Task WriteFileAsync(
94
101
doc [ "documentId" ] = documentId ;
95
102
doc [ "fileName" ] = fileName ;
96
103
doc [ "content" ] = content ;
104
+ doc [ "contentType" ] = MimeTypes . PlainText ;
97
105
await this . SaveDocumentAsync ( index , id , doc , cancellationToken ) . ConfigureAwait ( false ) ;
98
106
}
99
107
else
@@ -105,14 +113,15 @@ public async Task WriteFileAsync(
105
113
{
106
114
{ "index" , index } ,
107
115
{ "documentId" , documentId } ,
108
- { "fileName" , fileName }
116
+ { "fileName" , fileName } ,
117
+ { "contentType" , this . _mimeTypeDetection . GetFileType ( fileName ) }
109
118
}
110
119
} ;
111
120
112
121
// Since the pattern of usage is that you can upload a file for a document id and then update, we need to delete
113
122
// any existing file with the same id check if the file exists and delete it
114
123
IAsyncCursor < GridFSFileInfo < string > > existingFile = await GetFromBucketByIdAsync ( id , bucket , cancellationToken ) . ConfigureAwait ( false ) ;
115
- if ( existingFile . Any ( cancellationToken ) )
124
+ if ( await existingFile . AnyAsync ( cancellationToken ) . ConfigureAwait ( false ) )
116
125
{
117
126
await bucket . DeleteAsync ( id , cancellationToken ) . ConfigureAwait ( false ) ;
118
127
}
0 commit comments