ElasticsearchMemory.cs

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

usingSystem.Runtime.CompilerServices;
usingElastic.Clients.Elasticsearch;
usingElastic.Clients.Elasticsearch.Mapping;
usingElastic.Clients.Elasticsearch.QueryDsl;
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;
privatereadonlyIIndexNameHelper_indexNameHelper;
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="client">Elasticsearch client</param>
/// <param name="log">Application logger</param>
/// <param name="embeddingGenerator">Embedding generator</param>
/// <param name="indexNameHelper">Index name helper</param>
publicElasticsearchMemory(
ElasticsearchConfigconfig,
ElasticsearchClientclient,
ITextEmbeddingGeneratorembeddingGenerator,
IIndexNameHelperindexNameHelper,
ILogger<ElasticsearchMemory>?log=null)
{
this._embeddingGenerator=embeddingGenerator??thrownewArgumentNullException(nameof(embeddingGenerator));
this._indexNameHelper=indexNameHelper??thrownewArgumentNullException(nameof(indexNameHelper));
this._config=config??thrownewArgumentNullException(nameof(config));
this._client=client;// new ElasticsearchClient(this._config.ToElasticsearchClientSettings()); // TODO: inject
this._log=log??DefaultLogger<ElasticsearchMemory>.Instance;
}

/// <inheritdoc />
publicasyncTaskCreateIndexAsync(
stringindex,
intvectorSize,
CancellationTokencancellationToken=default)
{
index=this._indexNameHelper.Convert(index);

varexistsResponse=awaitthis._client.Indices.ExistsAsync(index,cancellationToken).ConfigureAwait(false);
if(existsResponse.Exists)
{
this._log.LogTrace("{MethodName}: Index {Index} already exists.",nameof(CreateIndexAsync),index);
return;
}

varcreateIdxResponse=awaitthis._client.Indices.CreateAsync(index,
 cfg =>
{
cfg.Settings(setts =>
{
setts.NumberOfShards(this._config.ShardCount);
setts.NumberOfReplicas(this._config.ReplicaCount);
});
},
cancellationToken).ConfigureAwait(false);

constintDimensions=1536;// TODO: make not hardcoded

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

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

this._config.ConfigureProperties?.Invoke(propDesc);
}),
cancellationToken).ConfigureAwait(false);

this._log.LogTrace("{MethodName}: Index {Index} creeated.",nameof(CreateIndexAsync),index);
}

/// <inheritdoc />
publicasyncTask<IEnumerable<string>>GetIndexesAsync(
CancellationTokencancellationToken=default)
{
varresp=awaitthis._client.Indices.GetAsync(this._config.IndexPrefix+"*",cancellationToken).ConfigureAwait(false);

varnames=resp.Indices
.Select(x =>x.Key.ToString().Replace(this._config.IndexPrefix,string.Empty,StringComparison.Ordinal))
.ToHashSet(StringComparer.OrdinalIgnoreCase);

this._log.LogTrace("{MethodName}: Returned {IndexCount} indices: {Indices}.",nameof(GetIndexesAsync),names.Count,string.Join(", ",names));

returnnames;
}

/// <inheritdoc />
publicasyncTaskDeleteIndexAsync(
stringindex,
CancellationTokencancellationToken=default)
{
index=this._indexNameHelper.Convert(index);

vardelResponse=awaitthis._client.Indices.DeleteAsync(
index,
cancellationToken).ConfigureAwait(false);

if(delResponse.IsSuccess())
{
this._log.LogTrace("{MethodName}: Index {Index} deleted.",nameof(DeleteIndexAsync),index);
}
else
{
this._log.LogWarning("{MethodName}: Index {Index} delete failed.",nameof(DeleteIndexAsync),index);
}
}

/// <inheritdoc />
publicasyncTaskDeleteAsync(
stringindex,
MemoryRecordrecord,
CancellationTokencancellationToken=default)
{
index=this._indexNameHelper.Convert(index);

record=record??thrownewArgumentNullException(nameof(record));

vardelResponse=awaitthis._client.DeleteAsync<ElasticsearchMemoryRecord>(
index,
record.Id,
(delReq)=>
{
delReq.Refresh(Refresh.WaitFor);
},
cancellationToken)
.ConfigureAwait(false);

if(delResponse.IsSuccess())
{
this._log.LogTrace("{MethodName}: Record {RecordId} deleted.",nameof(DeleteAsync),record.Id);
}
else
{
this._log.LogWarning("{MethodName}: Record {RecordId} delete failed.",nameof(DeleteAsync),record.Id);
}
}

/// <inheritdoc />
publicasyncTask<string>UpsertAsync(
stringindex,
MemoryRecordrecord,
CancellationTokencancellationToken=default)
{
index=this._indexNameHelper.Convert(index);

varmemRec=ElasticsearchMemoryRecord.FromMemoryRecord(record);

varresponse=awaitthis._client.UpdateAsync<ElasticsearchMemoryRecord,ElasticsearchMemoryRecord>(
index,
memRec.Id,
(updateReq)=>
{
updateReq.Refresh(Refresh.WaitFor);

varmemRec2=memRec;
updateReq.Doc(memRec2);
updateReq.DocAsUpsert(true);
},
cancellationToken)
.ConfigureAwait(false);

if(response.IsSuccess())
{
this._log.LogTrace("{MethodName}: Record {RecordId} upserted.",nameof(UpsertAsync),memRec.Id);
}
else
{
this._log.LogError("{MethodName}: Record {RecordId} upsert failed.",nameof(UpsertAsync),memRec.Id);
}

returnresponse.Id;
}

/// <inheritdoc />
publicasyncIAsyncEnumerable<(MemoryRecord,double)>GetSimilarListAsync(
stringindex,
stringtext,
ICollection<MemoryFilter>?filters=null,
doubleminRelevance=0,intlimit=1,boolwithEmbeddings=false,[EnumeratorCancellation]CancellationTokencancellationToken=default)
{
if(limit<0)
{
limit=10;
}

index=this._indexNameHelper.Convert(index);

this._log.LogTrace("{MethodName}: Searching for '{Text}' on index '{IndexName}' with filters {Filters}. {MinRelevance} {Limit} {WithEmbeddings}",
nameof(GetSimilarListAsync),text,index,filters.ToDebugString(),minRelevance,limit,withEmbeddings);

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)
.Filter(q =>this.ConvertTagFilters(q,filters))
.NumCandidates(limit+100)
.Field(x =>x.Vector)
.QueryVector(coll);
}),
cancellationToken)
.ConfigureAwait(false);

if((resp.HitsMetadataisnull)||(resp.HitsMetadata.Hitsisnull))
{
this._log.LogWarning("The search returned a null result. Should retry?");
yieldbreak;
}

foreach(varhitinresp.HitsMetadata.Hits)
{
if(hit?.Source==null)
{
continue;
}

this._log.LogTrace("{MethodName} Hit: {HitScore}, {HitId}",nameof(GetSimilarListAsync),hit.Score,hit.Id);
yieldreturn(hit.Source!.ToMemoryRecord(),hit.Score??0);
}
}

/// <inheritdoc />
publicasyncIAsyncEnumerable<MemoryRecord>GetListAsync(
stringindex,
ICollection<MemoryFilter>?filters=null,
intlimit=1,
boolwithEmbeddings=false,
[EnumeratorCancellation]
CancellationTokencancellationToken=default)
{
this._log.LogTrace("{MethodName}: querying index '{IndexName}' with filters {Filters}. {Limit} {WithEmbeddings}",
nameof(GetListAsync),index,filters.ToDebugString(),limit,withEmbeddings);

if(limit<0)
{
limit=10;
}

index=this._indexNameHelper.Convert(index);

varresp=awaitthis._client.SearchAsync<ElasticsearchMemoryRecord>(s =>
s.Index(index)
.Size(limit)
.Query(qd =>
{
this.ConvertTagFilters(qd,filters);
}),
cancellationToken)
.ConfigureAwait(false);

if((resp.HitsMetadataisnull)||(resp.HitsMetadata.Hitsisnull))
{
yieldbreak;
}

foreach(varhitinresp.Hits)
{
if(hit?.Source==null)
{
continue;
}

this._log.LogTrace("{MethodName} Hit: {HitScore}, {HitId}",nameof(GetListAsync),hit.Score,hit.Id);
yieldreturnhit.Source!.ToMemoryRecord();
}
}

//private string ConvertIndexName(string index) => ESIndexName.Convert(this._config.IndexPrefix + index);

privateQueryDescriptor<ElasticsearchMemoryRecord>ConvertTagFilters(
QueryDescriptor<ElasticsearchMemoryRecord>qd,
ICollection<MemoryFilter>?filters=null)
{
if((filters==null)||(filters.Count==0))
{
qd.MatchAll();
returnqd;
}

filters=filters.Where(f =>f.Keys.Count>0)
.ToList();// Remove empty filters

if(filters.Count==0)
{
qd.MatchAll();
returnqd;
}

foreach(MemoryFilterfilterinfilters)
{
List<Query>all=new();

// Each tag collection is an element of a List<string, List<string?>>>
foreach(vartagNameinfilter.Keys)
{
List<string?>tagValues=filter[tagName];
List<FieldValue>terms=tagValues.Select(x =>(FieldValue)(x??FieldValue.Null))
.ToList();
// ----------------
QuerynewTagQuery=newTermQuery(ElasticsearchMemoryRecord.Tags_Name){Value=tagName};
newTagQuery&=newTermsQuery()
{
Field=ElasticsearchMemoryRecord.Tags_Value,
Terms=newTermsQueryField(terms)
};
varnestedQd=newNestedQuery();
nestedQd.Path=ElasticsearchMemoryRecord.TagsField;
nestedQd.Query=newTagQuery;

all.Add(nestedQd);
qd.Bool(bq =>bq.Must(all.ToArray()));
}
}

// ---------------------

//qd.Nested(nqd =>
//{
// nqd.Path(ElasticsearchMemoryRecord.TagsField);

// nqd.Query(nq =>
// {
// // Each filter is a tag collection.
// foreach (MemoryFilter filter in filters)
// {
// List<Query> all = new();

// // Each tag collection is an element of a List<string, List<string?>>>
// foreach (var tagName in filter.Keys)
// {
// List<string?> tagValues = filter[tagName];
// List<FieldValue> terms = tagValues.Select(x => (FieldValue)(x ?? FieldValue.Null))
// .ToList();
// // ---------------- 

// Query newTagQuery = new TermQuery(ElasticsearchMemoryRecord.Tags_Name) { Value = tagName };
// newTagQuery &= new TermsQuery() {
// Field = ElasticsearchMemoryRecord.Tags_Value,
// Terms = new TermsQueryField(terms)
// };

// all.Add(newTagQuery);
// }

// nq.Bool(bq => bq.Must(all.ToArray()));
// }
// });
//});

returnqd;
}
}