refine state search

This commit is contained in:
Jicheng Lu 2024-07-09 02:34:49 -05:00
parent 8a6a52ce58
commit d4ff9e0308
7 changed files with 58 additions and 58 deletions

View file

@ -21,11 +21,16 @@ public class ConversationFilter
/// <summary>
/// Check whether each key in the list is in the conversation states and its value equals to target value if not empty
/// </summary>
public IEnumerable<KeyValue> States { get; set; } = new List<KeyValue>();
public List<KeyValue> States { get; set; } = new List<KeyValue>();
}
public class KeyValue
{
public string Key { get; set; }
public string? Value { get; set; }
public override string ToString()
{
return $"Key: {Key}, Value: {Value}";
}
}

View file

@ -1,5 +1,3 @@
using BotSharp.Plugin.MongoStorage.Models;
namespace BotSharp.Plugin.MongoStorage.Collections;
public class AgentDocument : MongoBase

View file

@ -1,5 +1,3 @@
using BotSharp.Plugin.MongoStorage.Models;
namespace BotSharp.Plugin.MongoStorage.Collections;
public class ConversationDialogDocument : MongoBase

View file

@ -1,5 +1,3 @@
using BotSharp.Plugin.MongoStorage.Models;
namespace BotSharp.Plugin.MongoStorage.Collections;
public class ConversationStateDocument : MongoBase

View file

@ -1,5 +1,3 @@
using BotSharp.Plugin.MongoStorage.Models;
namespace BotSharp.Plugin.MongoStorage.Collections;
public class LlmCompletionLogDocument : MongoBase

View file

@ -40,6 +40,19 @@ public class MongoDbContext
return collection;
}
private IMongoCollection<ConversationStateDocument> CreateConversationStateIndex()
{
var collection = Database.GetCollection<ConversationStateDocument>($"{_collectionPrefix}_ConversationStates");
var indexes = collection.Indexes.List().ToList();
var stateIndex = indexes.FirstOrDefault(x => x.GetElement("name").ToString().StartsWith("States.Key"));
if (stateIndex == null)
{
var indexDef = Builders<ConversationStateDocument>.IndexKeys.Ascending("States.Key");
collection.Indexes.CreateOne(new CreateIndexModel<ConversationStateDocument>(indexDef));
}
return collection;
}
private IMongoCollection<AgentTaskDocument> CreateAgentTaskIndex()
{
var collection = Database.GetCollection<AgentTaskDocument>($"{_collectionPrefix}_AgentTasks");
@ -93,7 +106,7 @@ public class MongoDbContext
=> Database.GetCollection<ConversationDialogDocument>($"{_collectionPrefix}_ConversationDialogs");
public IMongoCollection<ConversationStateDocument> ConversationStates
=> Database.GetCollection<ConversationStateDocument>($"{_collectionPrefix}_ConversationStates");
=> CreateConversationStateIndex();
public IMongoCollection<ExecutionLogDocument> ExectionLogs
=> Database.GetCollection<ExecutionLogDocument>($"{_collectionPrefix}_ExecutionLogs");

View file

@ -1,5 +1,9 @@
using Amazon.Util.Internal;
using BotSharp.Abstraction.Conversations.Models;
using BotSharp.Abstraction.Repositories.Filters;
using MongoDB.Bson.Serialization;
using MongoDB.Driver;
using System.Collections.Immutable;
namespace BotSharp.Plugin.MongoStorage.Repository;
@ -225,91 +229,77 @@ public partial class MongoRepository
public PagedItems<Conversation> GetConversations(ConversationFilter filter)
{
var conversations = new List<Conversation>();
var builder = Builders<ConversationDocument>.Filter;
var filters = new List<FilterDefinition<ConversationDocument>>() { builder.Empty };
var convBuilder = Builders<ConversationDocument>.Filter;
var convFilters = new List<FilterDefinition<ConversationDocument>>() { convBuilder.Empty };
// Filter conversations
if (!string.IsNullOrEmpty(filter?.Id))
{
filters.Add(builder.Eq(x => x.Id, filter.Id));
convFilters.Add(convBuilder.Eq(x => x.Id, filter.Id));
}
if (!string.IsNullOrEmpty(filter?.AgentId))
{
filters.Add(builder.Eq(x => x.AgentId, filter.AgentId));
convFilters.Add(convBuilder.Eq(x => x.AgentId, filter.AgentId));
}
if (!string.IsNullOrEmpty(filter?.Status))
{
filters.Add(builder.Eq(x => x.Status, filter.Status));
convFilters.Add(convBuilder.Eq(x => x.Status, filter.Status));
}
if (!string.IsNullOrEmpty(filter?.Channel))
{
filters.Add(builder.Eq(x => x.Channel, filter.Channel));
convFilters.Add(convBuilder.Eq(x => x.Channel, filter.Channel));
}
if (!string.IsNullOrEmpty(filter?.UserId))
{
filters.Add(builder.Eq(x => x.UserId, filter.UserId));
convFilters.Add(convBuilder.Eq(x => x.UserId, filter.UserId));
}
if (!string.IsNullOrEmpty(filter?.TaskId))
{
filters.Add(builder.Eq(x => x.TaskId, filter.TaskId));
convFilters.Add(convBuilder.Eq(x => x.TaskId, filter.TaskId));
}
if (filter?.StartTime != null)
{
filters.Add(builder.Gte(x => x.CreatedTime, filter.StartTime.Value));
convFilters.Add(convBuilder.Gte(x => x.CreatedTime, filter.StartTime.Value));
}
// Check states
if (filter != null && !filter.States.IsNullOrEmpty())
// Filter states
var stateFilters = new List<FilterDefinition<ConversationStateDocument>>();
if (filter != null && string.IsNullOrEmpty(filter.Id) && !filter.States.IsNullOrEmpty())
{
var targetConvIds = new List<string>();
foreach (var pair in filter.States)
{
if (pair == null || string.IsNullOrWhiteSpace(pair.Key)) continue;
var query = _dc.ConversationStates.AsQueryable();
var convIds = query.AsEnumerable().Where(x =>
var elementFilters = new List<FilterDefinition<StateMongoElement>> { Builders<StateMongoElement>.Filter.Eq(x => x.Key, pair.Key) };
if (!string.IsNullOrEmpty(pair.Value))
{
var foundState = x.States.FirstOrDefault(s => s.Key.IsEqualTo(pair.Key));
if (foundState == null) return false;
if (!string.IsNullOrWhiteSpace(pair.Value))
{
return pair.Value.IsEqualTo(foundState.Values.LastOrDefault()?.Data);
}
return true;
}).Select(x => x.ConversationId).ToList();
targetConvIds = targetConvIds.Concat(convIds).Distinct().ToList();
elementFilters.Add(Builders<StateMongoElement>.Filter.Eq("Values.Data", pair.Value));
}
stateFilters.Add(Builders<ConversationStateDocument>.Filter.ElemMatch(x => x.States, Builders<StateMongoElement>.Filter.And(elementFilters)));
}
filters.Add(builder.In(x => x.Id, targetConvIds));
var targetConvIds = _dc.ConversationStates.Find(Builders<ConversationStateDocument>.Filter.And(stateFilters)).ToEnumerable().Select(x => x.ConversationId).Distinct().ToList();
convFilters.Add(convBuilder.In(x => x.Id, targetConvIds));
}
var filterDef = builder.And(filters);
// Sort and paginate
var filterDef = convBuilder.And(convFilters);
var sortDef = Builders<ConversationDocument>.Sort.Descending(x => x.CreatedTime);
var pager = filter?.Pager ?? new Pagination();
var conversationDocs = _dc.Conversations.Find(filterDef).Sort(sortDef).Skip(pager.Offset).Limit(pager.Size).ToList();
var count = _dc.Conversations.CountDocuments(filterDef);
foreach (var conv in conversationDocs)
var conversations = conversationDocs.Select(x => new Conversation
{
var convId = conv.Id.ToString();
conversations.Add(new Conversation
{
Id = convId,
AgentId = conv.AgentId.ToString(),
UserId = conv.UserId.ToString(),
TaskId = conv.TaskId,
Title = conv.Title,
Channel = conv.Channel,
Status = conv.Status,
DialogCount = conv.DialogCount,
CreatedTime = conv.CreatedTime,
UpdatedTime = conv.UpdatedTime
});
}
Id = x.Id.ToString(),
AgentId = x.AgentId.ToString(),
UserId = x.UserId.ToString(),
TaskId = x.TaskId,
Title = x.Title,
Channel = x.Channel,
Status = x.Status,
DialogCount = x.DialogCount,
CreatedTime = x.CreatedTime,
UpdatedTime = x.UpdatedTime
}).ToList();
return new PagedItems<Conversation>
{