PaperlessAI/Services/DocumentProcessingService.cs

using System;
using System.Diagnostics.CodeAnalysis;
using PaperlessAI.Context;
using PaperlessAI.Enums;
using PaperlessAI.Models;
using PaperlessAI.Models.LMStudio;
using PaperlessAI.Repositories;

namespace PaperlessAI.Services;

public class DocumentProcessingService(IServiceScopeFactory scopeFactory, ILogger<DocumentProcessingService> logger)
{
    private Queue<PaperlessDocument> paperlessDocumentsQueue = new Queue<PaperlessDocument>();
    private bool isWorking = false;

    public async Task Enqueue(List<PaperlessDocument> paperlessDocuments)
    {
        if(paperlessDocumentsQueue.Count != 0)
            return;

        foreach(var doc in paperlessDocuments)
        {
            await Enqueue(doc);
        }
    }

    public async Task Enqueue(PaperlessDocument paperlessDocument)
    {
        if(paperlessDocumentsQueue.Contains(paperlessDocument, new PaperlessDocumentSameDocument()))
            return;

        await using var scope = scopeFactory.CreateAsyncScope();
        var repository = scope.ServiceProvider.GetRequiredService<PaperlessDocumentRepository>();

        paperlessDocumentsQueue.Enqueue(paperlessDocument);

        paperlessDocument.SetDocumentState(EDocumentState.Enqueued);
        await repository.Update(paperlessDocument);

        logger.LogInformation($"Enqueued document: {paperlessDocument}");
    }

    public async Task Process()
    {
        if(isWorking)
            return;

        if(paperlessDocumentsQueue.Count == 0)
            return;

        if(paperlessDocumentsQueue.TryDequeue(out PaperlessDocument? docToProcess))
        {
            logger.LogInformation($"Dequeued: {docToProcess}");
            var llmResponse = await ProcessDocument(docToProcess);
            await UpdateDocumentData(docToProcess, llmResponse);
            logger.LogInformation($"Document processed: {docToProcess} | Success: {llmResponse?.Success ?? false}");
        }

        if(paperlessDocumentsQueue.Count == 0)
        {
            await UnloadModel();
        }
    }

    private async Task UpdateDocumentData(PaperlessDocument document, LmStudioLlmResponse llmResponse)
    {
        await using var scope = scopeFactory.CreateAsyncScope();
        var repository = scope.ServiceProvider.GetRequiredService<PaperlessDocumentRepository>();

        var isSameId = document.DocumentId == llmResponse?.DocumentId;
        if(!llmResponse!.IsValid || !isSameId)
        {
            document.SetDocumentState(EDocumentState.Failed);
            document.SetLastLlmMessage("Invalid LLM Response.");
        }
        else
        {
            document.SetDocumentState(llmResponse.Success ? EDocumentState.Success : EDocumentState.Failed);
            document.SetLastLlmMessage(llmResponse.Message);
        }

        await repository.Update(document);
    }

    private async Task UnloadModel()
    {
        try
        {
            isWorking = true;
            await using var scope = scopeFactory.CreateAsyncScope();
            var agentService = scope.ServiceProvider.GetRequiredService<LLMAgentService>();

            logger.LogInformation($"Unloading Model");
            await agentService.UnloadModel();
        }
        catch (Exception e)
        {
            logger.LogError(e.ToString());
            throw;
        }
        finally
        {
            isWorking = false;
        }
        
    }

    private async Task<LmStudioLlmResponse> ProcessDocument(PaperlessDocument docToProcess)
    {
        var llmResponse = new LmStudioLlmResponse();
        try
        {
            isWorking = true;

            await using var scope = scopeFactory.CreateAsyncScope();
            var agentService = scope.ServiceProvider.GetRequiredService<LLMAgentService>();
            var repository = scope.ServiceProvider.GetRequiredService<PaperlessDocumentRepository>();

            logger.LogInformation($"Processing: {docToProcess}");

            docToProcess.DocumentState = EDocumentState.Started;
            await repository.Update(docToProcess);
            
            llmResponse = await agentService.ProcessDocumentWithAI(docToProcess);
        }
        catch (Exception e)
        {
            logger.LogError(e.ToString());
            throw;
        }
        finally
        {
            isWorking = false;
        }
        return llmResponse ?? new LmStudioLlmResponse();
    }
}

class PaperlessDocumentSameDocument : EqualityComparer<PaperlessDocument>
{
    public override bool Equals(PaperlessDocument? x, PaperlessDocument? y)
    {
        if (x == null && y == null)
            return true;
        else if (x == null || y == null)
            return false;

        return x.Id == y.Id || x.DocumentId == y.DocumentId;
    }

    public override int GetHashCode([DisallowNull] PaperlessDocument obj)
    {
        var hcode = obj.Id + obj.DocumentId;
        return hcode.GetHashCode();
    }
}