Merge pull request #806 from ChenGong-lessen/cgong/agent-knowledge_base

Oceania2018 · web-flow · commit bcb2b0bf5643 · 2024-12-27T16:59:34.000Z
add agent level knowledge base
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/AgentField.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/AgentField.cs
@@ -18,6 +18,7 @@ public enum AgentField
     Sample,
     LlmConfig,
     Utility,
+    KnowledgeBase,
     MaxMessageCount
 }
 
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Models/Agent.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Models/Agent.cs
@@ -99,6 +99,11 @@ public class Agent
     /// </summary>
     public List<AgentUtility> Utilities { get; set; } = new();
 
+    /// <summary>
+    /// Agent knowledge bases
+    /// </summary>
+    public List<AgentKnowledgeBase> KnowledgeBases { get; set; } = [];
+
     /// <summary>
     /// Inherit from agent
     /// </summary>
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Models/AgentKnowledgeBase.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Models/AgentKnowledgeBase.cs
@@ -0,0 +1,23 @@
+namespace BotSharp.Abstraction.Agents.Models;
+
+public class AgentKnowledgeBase
+{
+    public string? Name { get; set; }
+    public bool Disabled { get; set; }
+
+    public AgentKnowledgeBase()
+    {
+        
+    }
+
+    public AgentKnowledgeBase(string name, bool enabled)
+    {
+        Name = name;
+        Disabled = enabled;
+    }
+
+    public override string ToString()
+    {
+        return Name ?? string.Empty;
+    }
+}
diff --git a/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.UpdateAgent.cs b/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.UpdateAgent.cs
@@ -39,6 +39,7 @@ public async Task UpdateAgent(Agent agent, AgentField updateField)
         record.Responses = agent.Responses ?? [];
         record.Samples = agent.Samples ?? [];
         record.Utilities = agent.Utilities ?? [];
+        record.KnowledgeBases = agent.KnowledgeBases ?? [];
         if (agent.LlmConfig != null && !agent.LlmConfig.IsInherit)
         {
             record.LlmConfig = agent.LlmConfig;
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.KnowledgeBase:
+                    UpdateAgentKnowledgeBases(agent.Id, agent.KnowledgeBases);
+                    break;
                 case AgentField.MaxMessageCount:
                     UpdateAgentMaxMessageCount(agent.Id, agent.MaxMessageCount);
                     break;
@@ -168,6 +171,19 @@ private void UpdateAgentUtilities(string agentId, bool mergeUtility, List<AgentU
             File.WriteAllText(agentFile, json);
         }
 
+        private void UpdateAgentKnowledgeBases(string agentId, List<AgentKnowledgeBase> knowledgeBases)
+        {
+            if (knowledgeBases == null) return;
+
+            var (agent, agentFile) = GetAgentFromFile(agentId);
+            if (agent == null) return;
+
+            agent.KnowledgeBases = knowledgeBases;
+            agent.UpdatedDateTime = DateTime.UtcNow;
+            var json = JsonSerializer.Serialize(agent, _options);
+            File.WriteAllText(agentFile, json);
+        }
+
         private void UpdateAgentRoutingRules(string agentId, List<RoutingRule> rules)
         {
             if (rules == null) return;
@@ -310,6 +326,7 @@ private void UpdateAgentAllFields(Agent inputAgent)
             agent.Type = inputAgent.Type;
             agent.Profiles = inputAgent.Profiles;
             agent.Utilities = inputAgent.Utilities;
+            agent.KnowledgeBases = inputAgent.KnowledgeBases;
             agent.RoutingRules = inputAgent.RoutingRules;
             agent.LlmConfig = inputAgent.LlmConfig;
             agent.MaxMessageCount = inputAgent.MaxMessageCount;
diff --git a/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentUpdateModel.cs b/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentUpdateModel.cs
@@ -39,6 +39,13 @@ public class AgentUpdateModel
     /// </summary>
     public List<AgentUtility>? Utilities { get; set; }
 
+    /// <summary>
+    /// knowledge bases
+    /// </summary>
+    /// 
+    [JsonPropertyName("knowledge_bases")]
+    public List<AgentKnowledgeBase>? KnowledgeBases { get; set; }
+
     /// <summary>
     /// Functions
     /// </summary>
@@ -90,6 +97,7 @@ public Agent ToAgent()
             Functions = Functions ?? new List<FunctionDef>(),
             Responses = Responses ?? new List<AgentResponse>(),
             Utilities = Utilities ?? new List<AgentUtility>(),
+            KnowledgeBases = KnowledgeBases ?? [],
             LlmConfig = LlmConfig
         };
 
diff --git a/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentViewModel.cs b/src/Infrastructure/BotSharp.OpenAPI/ViewModels/Agents/AgentViewModel.cs
@@ -25,6 +25,9 @@ public class AgentViewModel
     public bool MergeUtility { get; set; }
     public List<AgentUtility> Utilities { get; set; }
 
+    [JsonPropertyName("knowledge_bases")]
+    public List<AgentKnowledgeBase> KnowledgeBases { get; set; }
+
     [JsonPropertyName("is_public")]
     public bool IsPublic { get; set; }
 
@@ -75,6 +78,7 @@ public static AgentViewModel FromAgent(Agent agent)
             Responses = agent.Responses,
             Samples = agent.Samples,
             Utilities = agent.Utilities,
+            KnowledgeBases = agent.KnowledgeBases,
             IsPublic= agent.IsPublic,
             Disabled = agent.Disabled,
             MergeUtility = agent.MergeUtility,
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Collections/AgentDocument.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Collections/AgentDocument.cs
@@ -18,6 +18,7 @@ public class AgentDocument : MongoBase
     public List<AgentResponseMongoElement> Responses { get; set; }
     public List<string> Samples { get; set; }
     public List<AgentUtilityMongoElement> Utilities { get; set; }
+    public List<AgentKnowledgeBaseMongoElement> KnowledgeBases { get; set; }
     public List<string> Profiles { get; set; }
     public List<RoutingRuleMongoElement> RoutingRules { get; set; }
     public AgentLlmConfigMongoElement? LlmConfig { get; set; }
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Models/AgentKnowledgeBaseMongoElement.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Models/AgentKnowledgeBaseMongoElement.cs
@@ -0,0 +1,26 @@
+using BotSharp.Abstraction.Agents.Models;
+
+namespace BotSharp.Plugin.MongoStorage.Models;
+
+public class AgentKnowledgeBaseMongoElement
+{
+    public string Name { get; set; }
+    public bool Disabled { get; set; }
+    public static AgentKnowledgeBaseMongoElement ToMongoElement(AgentKnowledgeBase knowledgeBase)
+    {
+        return new AgentKnowledgeBaseMongoElement
+        {
+            Name = knowledgeBase.Name ?? string.Empty,
+            Disabled = knowledgeBase.Disabled,
+        };
+    }
+
+    public static AgentKnowledgeBase ToDomainElement(AgentKnowledgeBaseMongoElement knowledgeBase)
+    {
+        return new AgentKnowledgeBase
+        {
+            Name = knowledgeBase.Name,
+            Disabled = knowledgeBase.Disabled
+        };
+    }
+}
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.Agent.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.Agent.cs
@@ -58,6 +58,9 @@ public void UpdateAgent(Agent agent, AgentField field)
             case AgentField.Utility:
                 UpdateAgentUtilities(agent.Id, agent.MergeUtility, agent.Utilities);
                 break;
+            case AgentField.KnowledgeBase:
+                UpdateAgentKnowledgeBases(agent.Id, agent.KnowledgeBases);
+                break;
             case AgentField.MaxMessageCount:
                 UpdateAgentMaxMessageCount(agent.Id, agent.MaxMessageCount);
                 break;
@@ -239,6 +242,20 @@ private void UpdateAgentUtilities(string agentId, bool mergeUtility, List<AgentU
         _dc.Agents.UpdateOne(filter, update);
     }
 
+    private void UpdateAgentKnowledgeBases(string agentId, List<AgentKnowledgeBase> knowledgeBases)
+    {
+        if (knowledgeBases == null) return;
+
+        var elements = knowledgeBases?.Select(x => AgentKnowledgeBaseMongoElement.ToMongoElement(x))?.ToList() ?? [];
+
+        var filter = Builders<AgentDocument>.Filter.Eq(x => x.Id, agentId);
+        var update = Builders<AgentDocument>.Update
+            .Set(x => x.KnowledgeBases, elements)
+            .Set(x => x.UpdatedTime, DateTime.UtcNow);
+
+        _dc.Agents.UpdateOne(filter, update);
+    }
+
     private void UpdateAgentLlmConfig(string agentId, AgentLlmConfig? config)
     {
         var llmConfig = AgentLlmConfigMongoElement.ToMongoElement(config);
@@ -279,6 +296,7 @@ private void UpdateAgentAllFields(Agent agent)
             .Set(x => x.Responses, agent.Responses.Select(r => AgentResponseMongoElement.ToMongoElement(r)).ToList())
             .Set(x => x.Samples, agent.Samples)
             .Set(x => x.Utilities, agent.Utilities.Select(u => AgentUtilityMongoElement.ToMongoElement(u)).ToList())
+            .Set(x => x.KnowledgeBases, agent.KnowledgeBases.Select(u => AgentKnowledgeBaseMongoElement.ToMongoElement(u)).ToList())
             .Set(x => x.LlmConfig, AgentLlmConfigMongoElement.ToMongoElement(agent.LlmConfig))
             .Set(x => x.IsPublic, agent.IsPublic)
             .Set(x => x.UpdatedTime, DateTime.UtcNow);

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ public enum AgentField`
`18`	`18`	`Sample,`
`19`	`19`	`LlmConfig,`
`20`	`20`	`Utility,`
	`21`	`+ KnowledgeBase,`
`21`	`22`	`MaxMessageCount`
`22`	`23`	`}`
`23`	`24`
Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ public async Task UpdateAgent(Agent agent, AgentField updateField)`
`39`	`39`	`record.Responses = agent.Responses ?? [];`
`40`	`40`	`record.Samples = agent.Samples ?? [];`
`41`	`41`	`record.Utilities = agent.Utilities ?? [];`
	`42`	`+ record.KnowledgeBases = agent.KnowledgeBases ?? [];`
`42`	`43`	`if (agent.LlmConfig != null && !agent.LlmConfig.IsInherit)`
`43`	`44`	`{`
`44`	`45`	`record.LlmConfig = agent.LlmConfig;`