TeamsAIssistantPlanner.cs

using Microsoft.Bot.Builder;
using Microsoft.KernelMemory;
using Microsoft.Teams.AI;
using Microsoft.Teams.AI.AI;
using Microsoft.Teams.AI.AI.Planners;
using Microsoft.Teams.AI.AI.Planners.Experimental;
using TeamsAIssistant.DataSources;
using TeamsAIssistant.Extensions;
using TeamsAIssistant.Services;
using TeamsAIssistant.State;

namespace TeamsAIssistant.Planner
{
    public class TeamsAIssistantPlanner(
        ActionPlannerOptions<TeamsAIssistantState> actionPlannerOptions,
        KernelMemoryData kernelMemoryData,
        TeamsAdapter teamsAdapter,
        AssistantService assistantService,
        AssistantsPlannerOptions assistantsPlannerOptions, ILoggerFactory logger) : IPlanner<TeamsAIssistantState>
    {
        private readonly ActionPlanner<TeamsAIssistantState> actionPlanner = new(actionPlannerOptions, logger);
        private readonly AssistantsPlanner<TeamsAIssistantState> assistantsPlanner = new(assistantsPlannerOptions, logger, teamsAdapter.HttpClientFactory.CreateClient("AssistantAI"));

        public async Task<Plan> BeginTaskAsync(ITurnContext turnContext, TeamsAIssistantState turnState,
            AI<TeamsAIssistantState> ai, CancellationToken cancellationToken = default)
        {
            var citations = await EnsureContext(turnContext, turnState, cancellationToken)!;
            var plan = await assistantsPlanner.BeginTaskAsync(turnContext, turnState, ai, cancellationToken);

            return plan.AddCitations(citations, turnState.MaxCitations);
        }

        public async Task<Plan> ContinueTaskAsync(ITurnContext turnContext, TeamsAIssistantState turnState,
            AI<TeamsAIssistantState> ai, CancellationToken cancellationToken = default)
        {
            var citations = await EnsureContext(turnContext, turnState, cancellationToken)!;
            var plan = await assistantsPlanner.ContinueTaskAsync(turnContext, turnState, ai, cancellationToken);

            return plan.AddCitations(citations, turnState.MaxCitations);
        }

        private async Task<IEnumerable<Citation>?>? EnsureContext(ITurnContext turnContext,
            TeamsAIssistantState turnState,
            CancellationToken cancellationToken = default)
        {
            if (turnState.HasIndexes())
            {
                var query = await CreateContextQuery(turnContext, turnState, cancellationToken);

                var (context, citations) = await kernelMemoryData.RenderDataAsync(query,
                    turnState,
                    actionPlanner.Options.Tokenizer,
                    turnState.ContextLength ?? Constants.AIConstants.DefaultContextTokenLength);

                var contextString = $"Context:\n{context}";

                if (turnState.AdditionalInstructionsContext == true)
                {
                    turnState.Temp.AdditionalInstructions += $"\n\n{contextString}";
                }
                else
                {
                    turnState.Temp.Input = $"{contextString}\n\nUser: {turnState.Temp.Input}";
                }

                return citations;
            }

            return null;
        }

        private async Task<string> CreateContextQuery(ITurnContext turnContext, TeamsAIssistantState turnState, CancellationToken cancellationToken = default)
        {
            var query = $"{turnState.AdditionalInstructions}\n\n{turnState.Temp.Input}";

            if (turnState.ThreadId != null)
            {
                var lastMessages = await assistantService.GetLastMessages(turnState.ThreadId, 10);
                var messages = lastMessages.Reverse();
                var promptTemplate = actionPlannerOptions.Prompts.GetPrompt("Chat");

                string messageHistory = string.Join("\n\n", messages?.Select(t => t.ToEmbeddingsSearch())!);
                turnState.ThreadMessageHistory = messageHistory;

                var result = await actionPlanner.CompletePromptAsync(turnContext, turnState, promptTemplate, null, cancellationToken);
                query = result.Message?.Content ?? query;

            }

            return query;
        }
    }
}