add agent max message count

Jicheng Lu · Jicheng Lu · commit 42ad511c7e3c · 2024-12-13T14:34:17.000-06:00
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/AgentField.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/AgentField.cs
@@ -17,7 +17,8 @@ public enum AgentField
     Response,
     Sample,
     LlmConfig,
-    Utility
+    Utility,
+    MaxMessageCount
 }
 
 public enum AgentTaskField
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Models/Agent.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Models/Agent.cs
@@ -104,6 +104,12 @@ public class Agent
     /// </summary>
     public string? InheritAgentId { get; set; }
 
+    /// <summary>
+    /// Maximum message count when load conversation
+    /// </summary>
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public int? MaxMessageCount { get; set; }
+
     public List<RoutingRule> RoutingRules { get; set; } = new();
 
     /// <summary>
@@ -133,6 +139,8 @@ public static Agent Clone(Agent agent)
             Knowledges = agent.Knowledges,
             IsPublic = agent.IsPublic,
             Disabled = agent.Disabled,
+            MergeUtility = agent.MergeUtility,
+            MaxMessageCount = agent.MaxMessageCount,
             Profiles = agent.Profiles,
             RoutingRules = agent.RoutingRules,
             LlmConfig = agent.LlmConfig,
diff --git a/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs b/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs
@@ -53,7 +53,7 @@ public interface IBotSharpRepository : IHaveServiceProvider
 
     #region Agent
     void UpdateAgent(Agent agent, AgentField field);
-    Agent? GetAgent(string agentId);
+    Agent? GetAgent(string agentId, bool basicsOnly = false);
     List<Agent> GetAgents(AgentFilter filter);
     List<UserAgent> GetUserAgents(string userId);
     void BulkInsertAgents(List<Agent> agents);
diff --git a/src/Infrastructure/BotSharp.Core/Agents/Hooks/BasicAgentHook.cs b/src/Infrastructure/BotSharp.Core/Agents/Hooks/BasicAgentHook.cs
@@ -47,7 +47,7 @@ public override void OnAgentUtilityLoaded(Agent agent)
             var entryAgentId = routing.EntryAgentId;
             if (!string.IsNullOrEmpty(entryAgentId))
             {
-                var entryAgent = db.GetAgent(entryAgentId);
+                var entryAgent = db.GetAgent(entryAgentId, basicsOnly: true);
                 var (fns, tps) = GetUniqueContent(entryAgent?.Utilities);
                 functionNames = functionNames.Concat(fns).Distinct().ToList();
                 templateNames = templateNames.Concat(tps).Distinct().ToList();
diff --git a/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.UpdateAgent.cs b/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.UpdateAgent.cs
@@ -28,6 +28,7 @@ public async Task UpdateAgent(Agent agent, AgentField updateField)
         record.IsPublic = agent.IsPublic;
         record.Disabled = agent.Disabled;
         record.MergeUtility = agent.MergeUtility;
+        record.MaxMessageCount = agent.MaxMessageCount;
         record.Type = agent.Type;
         record.Profiles = agent.Profiles ?? [];
         record.RoutingRules = agent.RoutingRules ?? [];
diff --git a/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.cs b/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.cs
@@ -153,7 +153,10 @@ public List<RoleDialogModel> GetDialogHistory(int lastCount = 100, bool fromBrea
             }
         }
 
-        return dialogs.TakeLast(lastCount).ToList();
+        var agentMsgCount = GetAgentMessageCount();
+        var count = agentMsgCount.HasValue && agentMsgCount.Value > 0 ? agentMsgCount.Value : lastCount;
+
+        return dialogs.TakeLast(count).ToList();
     }
 
     public void SetConversationId(string conversationId, List<MessageState> states, bool isReadOnly = false)
@@ -192,4 +195,16 @@ public bool IsConversationMode()
     {
         return !string.IsNullOrWhiteSpace(_conversationId);
     }
+
+
+    private int? GetAgentMessageCount()
+    {
+        var db = _services.GetRequiredService<IBotSharpRepository>();
+        var routingCtx = _services.GetRequiredService<IRoutingContext>();
+
+        if (string.IsNullOrEmpty(routingCtx.EntryAgentId)) return null;
+
+        var agent = db.GetAgent(routingCtx.EntryAgentId, basicsOnly: true);
+        return agent?.MaxMessageCount;
+    }
 }
diff --git a/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs b/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs
@@ -16,7 +16,7 @@ public class BotSharpDbContext : Database, IBotSharpRepository
     #endregion
 
     #region Agent
-    public Agent GetAgent(string agentId)
+    public Agent GetAgent(string agentId, bool basicsOnly = false)
         => throw new NotImplementedException();
 
     public List<Agent> GetAgents(AgentFilter filter)
diff --git a/src/Infrastructure/BotSharp.Core/Repository/FileRepository/FileRepository.Agent.cs b/src/Infrastructure/BotSharp.Core/Repository/FileRepository/FileRepository.Agent.cs
@@ -57,6 +57,9 @@ public void UpdateAgent(Agent agent, AgentField field)
                 case AgentField.Utility:
                     UpdateAgentUtilities(agent.Id, agent.MergeUtility, agent.Utilities);
                     break;
+                case AgentField.MaxMessageCount:
+                    UpdateAgentMaxMessageCount(agent.Id, agent.MaxMessageCount);
+                    break;
                 case AgentField.All:
                     UpdateAgentAllFields(agent);
                     break;
@@ -283,6 +286,17 @@ private void UpdateAgentLlmConfig(string agentId, AgentLlmConfig? config)
             File.WriteAllText(agentFile, json);
         }
 
+        private void UpdateAgentMaxMessageCount(string agentId, int? maxMessageCount)
+        {
+            var (agent, agentFile) = GetAgentFromFile(agentId);
+            if (agent == null) return;
+
+            agent.MaxMessageCount = maxMessageCount;
+            agent.UpdatedDateTime = DateTime.UtcNow;
+            var json = JsonSerializer.Serialize(agent, _options);
+            File.WriteAllText(agentFile, json);
+        }
+
         private void UpdateAgentAllFields(Agent inputAgent)
         {
             var (agent, agentFile) = GetAgentFromFile(inputAgent.Id);
@@ -298,6 +312,7 @@ private void UpdateAgentAllFields(Agent inputAgent)
             agent.Utilities = inputAgent.Utilities;
             agent.RoutingRules = inputAgent.RoutingRules;
             agent.LlmConfig = inputAgent.LlmConfig;
+            agent.MaxMessageCount = inputAgent.MaxMessageCount;
             agent.UpdatedDateTime = DateTime.UtcNow;
             var json = JsonSerializer.Serialize(agent, _options);
             File.WriteAllText(agentFile, json);
@@ -329,7 +344,7 @@ public List<string> GetAgentResponses(string agentId, string prefix, string inte
             return responses;
         }
 
-        public Agent? GetAgent(string agentId)
+        public Agent? GetAgent(string agentId, bool basicsOnly = false)
         {
             var agentDir = Path.Combine(_dbSettings.FileRepository, _agentSettings.DataDir);
             var dir = Directory.GetDirectories(agentDir).FirstOrDefault(x => x.Split(Path.DirectorySeparatorChar).Last() == agentId);
@@ -342,6 +357,8 @@ public List<string> GetAgentResponses(string agentId, string prefix, string inte
                 var record = JsonSerializer.Deserialize<Agent>(json, _options);
                 if (record == null) return null;
 
+                if (basicsOnly) return record;
+
                 var (defaultInstruction, channelInstructions) = FetchInstructions(dir);
                 var functions = FetchFunctions(dir);
                 var samples = FetchSamples(dir);
diff --git a/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentCreationModel.cs b/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentCreationModel.cs
@@ -51,6 +51,8 @@ public class AgentCreationModel
 
     public bool MergeUtility { get; set; }
 
+    public int? MaxMessageCount { get; set; }
+
     public List<AgentUtility> Utilities { get; set; } = new();
     public List<RoutingRuleUpdateModel> RoutingRules { get; set; } = new();
     public AgentLlmConfig? LlmConfig { get; set; }
@@ -72,6 +74,7 @@ public Agent ToAgent()
             Type = Type,
             Disabled = Disabled,
             MergeUtility = MergeUtility,
+            MaxMessageCount = MaxMessageCount,
             Profiles = Profiles,
             RoutingRules = RoutingRules?.Select(x => RoutingRuleUpdateModel.ToDomainElement(x))?.ToList() ?? new List<RoutingRule>(),
             LlmConfig = LlmConfig
diff --git a/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentUpdateModel.cs b/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentUpdateModel.cs
@@ -57,6 +57,9 @@ public class AgentUpdateModel
 
     public bool Disabled { get; set; }
 
+    [JsonPropertyName("max_message_count")]
+    public int? MaxMessageCount { get; set; }
+
     /// <summary>
     /// Profile by channel
     /// </summary>
@@ -77,6 +80,7 @@ public Agent ToAgent()
             IsPublic = IsPublic,
             Disabled = Disabled,
             MergeUtility = MergeUtility,
+            MaxMessageCount = MaxMessageCount,
             Type = Type,
             Profiles = Profiles ?? new List<string>(),
             RoutingRules = RoutingRules?.Select(x => RoutingRuleUpdateModel.ToDomainElement(x))?.ToList() ?? new List<RoutingRule>(),
diff --git a/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentViewModel.cs b/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentViewModel.cs
@@ -46,6 +46,10 @@ public class AgentViewModel
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
     public AgentLlmConfig? LlmConfig { get; set; }
 
+    [JsonPropertyName("max_message_count")]
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public int? MaxMessageCount { get; set; }
+
     public PluginDef Plugin { get; set; }
 
     public IEnumerable<string>? Actions { get; set; }
@@ -75,6 +79,7 @@ public static AgentViewModel FromAgent(Agent agent)
             Disabled = agent.Disabled,
             MergeUtility = agent.MergeUtility,
             IconUrl = agent.IconUrl,
+            MaxMessageCount = agent.MaxMessageCount,
             Profiles = agent.Profiles ?? new List<string>(),
             RoutingRules = agent.RoutingRules,
             LlmConfig = agent.LlmConfig,
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Collections/AgentDocument.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Collections/AgentDocument.cs
@@ -11,6 +11,7 @@ public class AgentDocument : MongoBase
     public bool IsPublic { get; set; }
     public bool Disabled { get; set; }
     public bool MergeUtility { get; set; }
+    public int? MaxMessageCount { get; set; }
     public List<ChannelInstructionMongoElement> ChannelInstructions { get; set; }
     public List<AgentTemplateMongoElement> Templates { get; set; }
     public List<FunctionDefMongoElement> Functions { get; set; }
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.Agent.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.Agent.cs
@@ -9,7 +9,7 @@ public partial class MongoRepository
 {
     public void UpdateAgent(Agent agent, AgentField field)
     {
-        if (agent == null || string.IsNullOrEmpty(agent.Id)) return;
+        if (agent == null || string.IsNullOrWhiteSpace(agent.Id)) return;
 
         switch (field)
         {
@@ -58,6 +58,9 @@ public void UpdateAgent(Agent agent, AgentField field)
             case AgentField.Utility:
                 UpdateAgentUtilities(agent.Id, agent.MergeUtility, agent.Utilities);
                 break;
+            case AgentField.MaxMessageCount:
+                UpdateAgentMaxMessageCount(agent.Id, agent.MaxMessageCount);
+                break;
             case AgentField.All:
                 UpdateAgentAllFields(agent);
                 break;
@@ -158,10 +161,8 @@ private void UpdateAgentRoutingRules(string agentId, List<RoutingRule> rules)
 
     private void UpdateAgentInstructions(string agentId, string instruction, List<ChannelInstruction>? channelInstructions)
     {
-        if (string.IsNullOrWhiteSpace(agentId)) return;
-
         var instructionElements = channelInstructions?.Select(x => ChannelInstructionMongoElement.ToMongoElement(x))?
-                                                      .ToList() ?? new List<ChannelInstructionMongoElement>();
+                                                      .ToList() ?? [];
 
         var filter = Builders<AgentDocument>.Filter.Eq(x => x.Id, agentId);
         var update = Builders<AgentDocument>.Update
@@ -200,7 +201,7 @@ private void UpdateAgentTemplates(string agentId, List<AgentTemplate> templates)
 
     private void UpdateAgentResponses(string agentId, List<AgentResponse> responses)
     {
-        if (responses == null) return;
+        if (responses == null || string.IsNullOrWhiteSpace(agentId)) return;
 
         var responsesToUpdate = responses.Select(r => AgentResponseMongoElement.ToMongoElement(r)).ToList();
         var filter = Builders<AgentDocument>.Filter.Eq(x => x.Id, agentId);
@@ -249,6 +250,16 @@ private void UpdateAgentLlmConfig(string agentId, AgentLlmConfig? config)
         _dc.Agents.UpdateOne(filter, update);
     }
 
+    private void UpdateAgentMaxMessageCount(string agentId, int? maxMessageCount)
+    {
+        var filter = Builders<AgentDocument>.Filter.Eq(x => x.Id, agentId);
+        var update = Builders<AgentDocument>.Update
+            .Set(x => x.MaxMessageCount, maxMessageCount)
+            .Set(x => x.UpdatedTime, DateTime.UtcNow);
+
+        _dc.Agents.UpdateOne(filter, update);
+    }
+
     private void UpdateAgentAllFields(Agent agent)
     {
         var filter = Builders<AgentDocument>.Filter.Eq(x => x.Id, agent.Id);
@@ -258,6 +269,7 @@ private void UpdateAgentAllFields(Agent agent)
             .Set(x => x.Disabled, agent.Disabled)
             .Set(x => x.MergeUtility, agent.MergeUtility)
             .Set(x => x.Type, agent.Type)
+            .Set(x => x.MaxMessageCount, agent.MaxMessageCount)
             .Set(x => x.Profiles, agent.Profiles)
             .Set(x => x.RoutingRules, agent.RoutingRules.Select(r => RoutingRuleMongoElement.ToMongoElement(r)).ToList())
             .Set(x => x.Instruction, agent.Instruction)
@@ -277,7 +289,7 @@ private void UpdateAgentAllFields(Agent agent)
     #endregion
 
 
-    public Agent? GetAgent(string agentId)
+    public Agent? GetAgent(string agentId, bool basicsOnly = false)
     {
         var agent = _dc.Agents.AsQueryable().FirstOrDefault(x => x.Id == agentId);
         if (agent == null) return null;
@@ -420,6 +432,7 @@ public void BulkInsertAgents(List<Agent> agents)
             InheritAgentId = x.InheritAgentId,
             Disabled = x.Disabled,
             MergeUtility = x.MergeUtility,
+            MaxMessageCount = x.MaxMessageCount,
             Profiles = x.Profiles,
             RoutingRules = x.RoutingRules?.Select(r => RoutingRuleMongoElement.ToMongoElement(r))?.ToList() ?? [],
             LlmConfig = AgentLlmConfigMongoElement.ToMongoElement(x.LlmConfig),
@@ -513,6 +526,7 @@ private Agent TransformAgentDocument(AgentDocument? agentDoc)
             Type = agentDoc.Type,
             InheritAgentId = agentDoc.InheritAgentId,
             Profiles = agentDoc.Profiles,
+            MaxMessageCount = agentDoc.MaxMessageCount
         };
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,8 @@ public enum AgentField`
`17`	`17`	`Response,`
`18`	`18`	`Sample,`
`19`	`19`	`LlmConfig,`
`20`		`- Utility`
	`20`	`+ Utility,`
	`21`	`+ MaxMessageCount`
`21`	`22`	`}`
`22`	`23`
`23`	`24`	`public enum AgentTaskField`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public override void OnAgentUtilityLoaded(Agent agent)`
`47`	`47`	`var entryAgentId = routing.EntryAgentId;`
`48`	`48`	`if (!string.IsNullOrEmpty(entryAgentId))`
`49`	`49`	`{`
`50`		`- var entryAgent = db.GetAgent(entryAgentId);`
	`50`	`+ var entryAgent = db.GetAgent(entryAgentId, basicsOnly: true);`
`51`	`51`	`var (fns, tps) = GetUniqueContent(entryAgent?.Utilities);`
`52`	`52`	`functionNames = functionNames.Concat(fns).Distinct().ToList();`
`53`	`53`	`templateNames = templateNames.Concat(tps).Distinct().ToList();`
Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,10 @@ public List<RoleDialogModel> GetDialogHistory(int lastCount = 100, bool fromBrea`
`153`	`153`	`}`
`154`	`154`	`}`
`155`	`155`
`156`		`- return dialogs.TakeLast(lastCount).ToList();`
	`156`	`+ var agentMsgCount = GetAgentMessageCount();`
	`157`	`+ var count = agentMsgCount.HasValue && agentMsgCount.Value > 0 ? agentMsgCount.Value : lastCount;`
	`158`	`+`
	`159`	`+ return dialogs.TakeLast(count).ToList();`
`157`	`160`	`}`
`158`	`161`
`159`	`162`	`public void SetConversationId(string conversationId, List<MessageState> states, bool isReadOnly = false)`
`@@ -192,4 +195,16 @@ public bool IsConversationMode()`
`192`	`195`	`{`
`193`	`196`	`return !string.IsNullOrWhiteSpace(_conversationId);`
`194`	`197`	`}`
	`198`	`+`
	`199`	`+`
	`200`	`+ private int? GetAgentMessageCount()`
	`201`	`+ {`
	`202`	`+ var db = _services.GetRequiredService<IBotSharpRepository>();`
	`203`	`+ var routingCtx = _services.GetRequiredService<IRoutingContext>();`
	`204`	`+`
	`205`	`+ if (string.IsNullOrEmpty(routingCtx.EntryAgentId)) return null;`
	`206`	`+`
	`207`	`+ var agent = db.GetAgent(routingCtx.EntryAgentId, basicsOnly: true);`
	`208`	`+ return agent?.MaxMessageCount;`
	`209`	`+ }`
`195`	`210`	`}`