ElasticsearchMemory.cs

// Copyright (c) Free Mind Labs, Inc. All rights reserved.

usingSystem;
usingSystem.Collections.Generic;
usingSystem.Linq;
usingSystem.Threading;
usingSystem.Threading.Tasks;
usingElastic.Clients.Elasticsearch;
usingElastic.Clients.Elasticsearch.Mapping;
usingMicrosoft.Extensions.Logging;
usingMicrosoft.KernelMemory;
usingMicrosoft.KernelMemory.Diagnostics;
usingMicrosoft.KernelMemory.MemoryStorage;

namespaceFreeMindLabs.KernelMemory.Elasticsearch;

/// <summary>
/// Elasticsearch connector for Kernel Memory.
/// </summary>
publicclassElasticsearchMemory:IMemoryDb
{
privatereadonlyElasticsearchConfig_config;
privatereadonlyILogger<ElasticsearchMemory>_log;
privatereadonlyElasticsearchClient_client;

/// <summary>
/// Create a new instance of Elasticsearch KM connector
/// </summary>
/// <param name="config">Elasticsearch configuration</param>
/// <param name="log">Application logger</param>
publicElasticsearchMemory(
ElasticsearchConfigconfig,
ILogger<ElasticsearchMemory>?log=null)
{
this._config=config??thrownewArgumentNullException(nameof(config));
this._client=newElasticsearchClient(this._config.ToElasticsearchClientSettings());
this._log=log??DefaultLogger<ElasticsearchMemory>.Instance;
}

/// <inheritdoc />
publicasyncTaskCreateIndexAsync(
stringindex,
intvectorSize,
CancellationTokencancellationToken=default)
{
index=index??thrownewArgumentNullException(nameof(index));
varexistsResponse=awaitthis._client.Indices.ExistsAsync(index,cancellationToken).ConfigureAwait(false);
if(existsResponse.Exists)
{
//this._log.LogInformation($"Index {index} already exists");
return;
}

varcreateIdxResponse=awaitthis._client.Indices.CreateAsync(index,cancellationToken).ConfigureAwait(false);

constintDimensions=1536;

varnp=newNestedProperty()
{
Properties=newProperties()
{
{ElasticsearchTag.NameField,newKeywordProperty()},
{ElasticsearchTag.ValueField,newTextProperty()}
}
};

varmapResponse=awaitthis._client.Indices.PutMappingAsync(index, x =>x
.Properties<ElasticsearchMemoryRecord>(p =>
{
p.Keyword(x =>x.Id);
p.Nested(ElasticsearchMemoryRecord.TagsField,np);
p.Text(x =>x.Payload);
p.DenseVector(x =>x.Vector, d =>d.Index(true).Dims(Dimensions).Similarity("cosine"));
}),
cancellationToken).ConfigureAwait(false);
}

/// <inheritdoc />
publicasyncTask<IEnumerable<string>>GetIndexesAsync(
CancellationTokencancellationToken=default)
{
Indicesindices="";
varresp=awaitthis._client.Indices.GetAsync(indices,cancellationToken).ConfigureAwait(false);

varnames=resp.Indices.Select(x =>x.Key.ToString());
returnnames;
}

/// <inheritdoc />
publicasyncTaskDeleteIndexAsync(
stringindex,
CancellationTokencancellationToken=default)
{
vardelResponse=awaitthis._client.Indices.DeleteAsync(index,cancellationToken).ConfigureAwait(false);
}


/// <inheritdoc />
publicasyncTaskDeleteAsync(
stringindex,
MemoryRecordrecord,
CancellationTokencancellationToken=default)
{
record=record??thrownewArgumentNullException(nameof(record));
vardelResponse=awaitthis._client.DeleteAsync<ElasticsearchMemoryRecord>(
index,
record.Id,
(delReq)=>
{
},
cancellationToken)
.ConfigureAwait(false);

}

/// <inheritdoc />
publicasyncTask<string>UpsertAsync(
stringindex,
MemoryRecordrecord,
CancellationTokencancellationToken=default)
{
varresponse=awaitthis._client.UpdateAsync<ElasticsearchMemoryRecord,ElasticsearchMemoryRecord>(
index,
record.Id,
(updateReq)=>
{
varmemRec=ElasticsearchMemoryRecord.FromMemoryRecord(record);
updateReq.Doc(memRec);
updateReq.DocAsUpsert(true);
},
cancellationToken)
.ConfigureAwait(false);

returnresponse.Id;
}

/// <inheritdoc />
publicasyncIAsyncEnumerable<(MemoryRecord,double)>GetSimilarListAsync(
stringindex,
stringtext,
ICollection<MemoryFilter>?filters=null,
doubleminRelevance=0,intlimit=1,boolwithEmbeddings=false,CancellationTokencancellationToken=default)
{
if(filters!=null)
{
foreach(MemoryFilterfilterinfilters)
{
if(filterisElasticsearchMemoryFilterextendedFilter)
{
// use ElasticsearchMemoryFilter filtering logic
}

// use MemoryFilter filtering logic
}
}

varresp=awaitthis._client.SearchAsync<ElasticsearchMemoryRecord>(s =>
s.Index(index)
.Query(q =>q.MatchAll()),
cancellationToken)
.ConfigureAwait(false);

foreach(varhitinresp.Hits)
{
yieldreturn(hit.Source!.ToMemoryRecord(),hit.Score??0);
}
}

/// <inheritdoc />
publicIAsyncEnumerable<MemoryRecord>GetListAsync(
stringindex,
ICollection<MemoryFilter>?filters=null,
intlimit=1,
boolwithEmbeddings=false,
CancellationTokencancellationToken=default)
{
if(filters!=null)
{
foreach(MemoryFilterfilterinfilters)
{
if(filterisElasticsearchMemoryFilterextendedFilter)
{
// use ElasticsearchMemoryFilter filtering logic
}

// use MemoryFilter filtering logic
}
}

thrownewNotImplementedException();
}
}