ElasticsearchMemory.cs

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

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

namespaceFreeMindLabs.KernelMemory.Elasticsearch;

/// <summary>
/// Elasticsearch connector for Kernel Memory.
/// </summary>
publicclassElasticsearchMemory:IMemoryDb
{
privatereadonlyITextEmbeddingGenerator_embeddingGenerator;
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>
/// <param name="embeddingGenerator">Embedding generator</param>
publicElasticsearchMemory(
ElasticsearchConfigconfig,
ITextEmbeddingGeneratorembeddingGenerator,
ILogger<ElasticsearchMemory>?log=null)
{
this._embeddingGenerator=embeddingGenerator??thrownewArgumentNullException(nameof(embeddingGenerator));
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=Indexname.Convert(index);

varexistsResponse=awaitthis._client.Indices.ExistsAsync(index,cancellationToken).ConfigureAwait(false);
if(existsResponse.Exists)
{
this._log.LogTrace("Index {Index} already exists.",index);
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, pd =>pd.Index(false));
p.Text(x =>x.Content);
p.DenseVector(x =>x.Vector, d =>d.Index(true).Dims(Dimensions).Similarity("cosine"));
}),
cancellationToken).ConfigureAwait(false);

this._log.LogTrace("Index {Index} creeated.",index);
}

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

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

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

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

/// <inheritdoc />
publicasyncTask<string>UpsertAsync(
stringindex,
MemoryRecordrecord,
CancellationTokencancellationToken=default)
{
varmemRec=ElasticsearchMemoryRecord.FromMemoryRecord(record);

varresponse=awaitthis._client.UpdateAsync<ElasticsearchMemoryRecord,ElasticsearchMemoryRecord>(
Indexname.Convert(index),
memRec.Id,
(updateReq)=>
{
varmemRec2=memRec;
updateReq.Doc(memRec2);
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,[EnumeratorCancellation]CancellationTokencancellationToken=default)
{
index=Indexname.Convert(index);

if(filters!=null)
{
foreach(MemoryFilterfilterinfilters)
{
if(filterisElasticsearchMemoryFilterextendedFilter)
{
// use ElasticsearchMemoryFilter filtering logic
}

// use MemoryFilter filtering logic
}
}

Embeddingqembed=awaitthis._embeddingGenerator.GenerateEmbeddingAsync(text,cancellationToken).ConfigureAwait(false);
varcoll=qembed.Data.ToArray();

varresp=awaitthis._client.SearchAsync<ElasticsearchMemoryRecord>(s =>
s.Index(index)
.Knn(qd =>
{
qd.k(limit)
.NumCandidates(limit+100)
.Field(x =>x.Vector)
.QueryVector(coll);
}),
//.Query(q => q.MatchAll()),
cancellationToken)
.ConfigureAwait(false);

foreach(varhitinresp.Hits)
{
this._log.LogTrace("Hit: {HitId}, {HitScore}",hit.Id,hit.Score);
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();
}
}