BotSharp/src/Plugins/BotSharp.Plugin.FileReader/Hooks/FileHandlerHook.cs

70 lines
2.2 KiB
C#
Raw Normal View History

2024-07-16 17:08:37 +00:00
using BotSharp.Abstraction.Agents.Enums;
using BotSharp.Abstraction.Agents.Settings;
using BotSharp.Abstraction.Functions.Models;
using BotSharp.Abstraction.Repositories;
namespace BotSharp.Plugin.FileHandler.Hooks;
public class FileHandlerHook : AgentHookBase, IAgentHook
{
2024-07-16 19:13:29 +00:00
private const string READ_IMAGE_FN = "read_image";
private const string READ_PDF_FN = "read_pdf";
2024-07-16 17:08:37 +00:00
public override string SelfId => string.Empty;
public FileHandlerHook(IServiceProvider services, AgentSettings settings) : base(services, settings)
{
}
public override void OnAgentLoaded(Agent agent)
{
var conv = _services.GetRequiredService<IConversationService>();
var isConvMode = conv.IsConversationMode();
if (isConvMode)
{
2024-07-16 19:13:29 +00:00
AddUtility(agent, UtilityName.ImageReader, READ_IMAGE_FN);
AddUtility(agent, UtilityName.PdfReader, READ_PDF_FN);
2024-07-16 17:08:37 +00:00
}
base.OnAgentLoaded(agent);
}
private bool IsEnableUtility(Agent agent, string utility)
{
return !agent.Utilities.IsNullOrEmpty() && agent.Utilities.Contains(utility);
}
private (string, FunctionDef?) GetPromptAndFunction(string functionName)
{
var db = _services.GetRequiredService<IBotSharpRepository>();
var agent = db.GetAgent(BuiltInAgentId.UtilityAssistant);
var prompt = agent?.Templates?.FirstOrDefault(x => x.Name.IsEqualTo($"{functionName}.fn"))?.Content ?? string.Empty;
var loadAttachmentFn = agent?.Functions?.FirstOrDefault(x => x.Name.IsEqualTo(functionName));
return (prompt, loadAttachmentFn);
}
private void AddUtility(Agent agent, string utility, string functionName)
{
if (!IsEnableUtility(agent, utility)) return;
var (prompt, fn) = GetPromptAndFunction(functionName);
if (fn != null)
{
if (!string.IsNullOrWhiteSpace(prompt))
{
agent.Instruction += $"\r\n\r\n{prompt}\r\n\r\n";
}
if (agent.Functions == null)
{
agent.Functions = new List<FunctionDef> { fn };
}
else
{
agent.Functions.Add(fn);
}
}
}
}