OpenChat/OpenChatApi/Services/MongoDBService .cs

using Microsoft.Extensions.Options;
using MongoDB.Driver;
using OpenChatApi.Models;

namespace OpenChatApi.Services
{
    public class MongoDBService : IMongoDBService
    {
        private readonly IMongoCollection<MessageDocument> _messagesCollection;

        public MongoDBService(IOptions<MongoDBSettings> settings)
        {
            var client = new MongoClient(settings.Value.ConnectionString);
            var database = client.GetDatabase("ChatMessages");
            _messagesCollection = database.GetCollection<MessageDocument>("Messages");
        }

        public async Task<MessageDocument> GetUserMessages(string userId)
        {
            var messages = await _messagesCollection
                .Find(x => x.UserId == userId)
                .FirstOrDefaultAsync();

            return messages ?? new MessageDocument { UserId = userId };
        }
        public async Task UpdateSystemPrompt(string userId, string systemPrompt)
        {
            var update = Builders<MessageDocument>.Update
                .Set(x => x.SystemPrompt, systemPrompt);

            await _messagesCollection.UpdateOneAsync(
                x => x.UserId == userId,
                update,
                new UpdateOptions { IsUpsert = true }
            );
        }

        public async Task UpdateMessages(string userId, string userMessage, string assistantResponse)
        {
            var document = await GetUserMessages(userId);

            // Add new messages
            document.Messages.Add(new Message { Type = "user", Content = userMessage });
            document.Messages.Add(new Message { Type = "assistant", Content = assistantResponse });

            // Keep only last 5 message pairs
            if (document.Messages.Count > 10)
            {
                document.Messages = document.Messages.Skip(2).ToList();
            }

            document.LastMessage = DateTime.Now;

            // Upsert document
            await _messagesCollection.ReplaceOneAsync(
                x => x.UserId == userId,
                document,
                new ReplaceOptions { IsUpsert = true }
            );
        }
    }
}