Luna/Luna.Core/Session/SessionManager.cs

using System.Collections.Concurrent;
using System.Globalization;
using System.Runtime.CompilerServices;
using System.Text;
using Luna.Agents.Abstractions;
using Luna.Core.Abstractions;
using Luna.Core.Services;
using Luna.Memory;
using Luna.Shared;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Options;
using SessionOptions = Luna.Configuration.SessionOptions;

namespace Luna.Core.Session;

public class SessionManager(
    [FromKeyedServices("Core")] IAgent coreAgent,
    [FromKeyedServices("Librarian")] IAgent librarianAgent,
    IOptions<SessionOptions> sessionOptions,
    IChatStreamUpdateBuilder streamUpdateBuilder,
    IMemoryStore memoryStore) : ISessionManager
{
    private readonly Dictionary<string, Session> sessions = new();
    private readonly ConcurrentDictionary<string, string> connectedClients = new();
    
    private const string CompactionPrompt =
        "Summarize the following conversation history for context preservation. Keep it short (max 12 bullet points).";
    
    private SessionOptions Options => sessionOptions.Value;
    
    public async IAsyncEnumerable<ChatStreamUpdate> RouteMessagesAsync(
        string content,
        string conversationId,
        string connectionId,
        [EnumeratorCancellation] CancellationToken ct = default)
    {
        if (!connectedClients.ContainsKey(connectionId))
            connectedClients[connectionId] = conversationId;
        
        if (!sessions.ContainsKey(conversationId))
        {
            var memory = new ChatMessage(ChatRole.System, await memoryStore.GetMemoriesAsync());
            sessions[conversationId] = new Session
            {
                SessionId = conversationId,
                Messages = string.IsNullOrWhiteSpace(memory.Text)
                    ? []
                    : [memory]
            };
        }
        
        var session = sessions[conversationId];

        var coreOptions = coreAgent.Options;
        if (session.TokenAmount > coreOptions.MaxContextTokens * Options.ContextTokenThreshold)
            await CompactSessionAsync(conversationId, Options.RetainedMessagesAfterCompacting, ct);
        
        session.Messages.Add(new ChatMessage(ChatRole.User, content));

        var messageId = Guid.NewGuid().ToString();
        var responseBuilder = new StringBuilder();
        await foreach (var update in coreAgent.ProcessStreamingAsync(session.Messages, ct))
        {
            if (update.Contents is { Count: > 0 })
            {
                foreach (var chunk in update.Contents)
                {
                    var streamUpdate = streamUpdateBuilder.FromAiContent(chunk, messageId);
                    responseBuilder.Append(streamUpdate.Content);
                    yield return streamUpdate;
                }
            }
            else if (update.Text is { Length: > 0 } text)
            {
                var streamUpdate = streamUpdateBuilder.FromString(text, messageId);
                responseBuilder.Append(streamUpdate.Content);
                yield return streamUpdate;
            }
        }
        yield return streamUpdateBuilder.CreateCompleteUpdate(messageId);
        
        session.Messages.Add(new ChatMessage(ChatRole.Assistant, responseBuilder.ToString()));
    }
    
    public async Task ClientDisconnectedAsync(string connectionId)
    {
        if (!connectedClients.TryRemove(connectionId, out var conversationId))
            return;
        
        await SaveSessionLogAsync(conversationId);
    }
    
    private async Task SaveSessionLogAsync(string conversationId)
    {
        var session = sessions[conversationId];
        var logBuilder = new StringBuilder();

        foreach (var message in session.Messages
            .Where(m => m.Role != ChatRole.System || m.Role != ChatRole.Tool))
        {
            logBuilder.AppendLine($"[{message.Role.Value}]:");
            logBuilder.AppendLine(message.Text);
            logBuilder.AppendLine();
        }
        
        await memoryStore.AddMemoryAsync(logBuilder.ToString());
        
        sessions.Remove(conversationId);
    }
    
    private async Task CompactSessionAsync(string conversationId, int retainMessages = 0, CancellationToken ct = default)
    {
        var session = sessions[conversationId];
        var retainedMessages = session.Messages.Count > retainMessages 
            ? session.Messages.TakeLast(retainMessages)
            : [];
        var messagesToCompact = session.Messages.Count > retainMessages
            ? session.Messages.SkipLast(retainMessages)
            : session.Messages;

        var compactionMessage = BuildCompactionMessage(messagesToCompact.ToList());
        var response = await librarianAgent.ProcessAsync([compactionMessage], ct);
        
        var memoryBuilder = new StringBuilder();
        memoryBuilder.AppendLine("<---- MEMORY BEGIN ---->");
        memoryBuilder.AppendLine($"[Meta] Conversation Recorded at: {DateTime.Now.ToString(CultureInfo.CurrentCulture)}");
        foreach (var message in response)
        {
            if (message.Role != ChatRole.System)
            {
                memoryBuilder.AppendLine(message.Text);
            }
        }
        memoryBuilder.AppendLine("<---- MEMORY END ---->");
        
        session.Messages = [
            new ChatMessage(ChatRole.Assistant, memoryBuilder.ToString()),
        ];
        
        session.Messages.AddRange(retainedMessages);
    }
    
    private ChatMessage BuildCompactionMessage(IReadOnlyList<ChatMessage> messages)
    {
        var builder = new StringBuilder();
        builder.AppendLine(CompactionPrompt);
        builder.AppendLine();
        builder.AppendLine("<---- CONVERSATION BEGIN ---->");
        builder.AppendLine();
        foreach (var message in messages)
        {
            if (message.Text.Contains("<---- MEMORY BEGIN ---->"))
                continue;
            
            builder.AppendLine($"{message.Role}: {message.Text}");
        }
        builder.AppendLine();
        builder.AppendLine("<---- CONVERSATION END ---->");
        return new ChatMessage(ChatRole.User, builder.ToString());
    }
}