Move LLM work to background

Author: santiagolizardo

src/Cli/Commands/TaskProcessorCommand.php

@@ -1,18 +1,23 @@
-<?php declare(strict_types=1);
+<?php
+
+declare(strict_types=1);
 
 namespace Reconmap\Cli\Commands;
 
 use Reconmap\Services\QueueProcessor;
+use Reconmap\Tasks\AiTaskProcessor;
 use Reconmap\Tasks\TaskResultProcessor;
 use Symfony\Component\Console\Command\Command;
 use Symfony\Component\Console\Input\InputInterface;
 use Symfony\Component\Console\Output\OutputInterface;
 
 class TaskProcessorCommand extends Command
 {
-    public function __construct(private readonly QueueProcessor      $queueProcessor,
-                                private readonly TaskResultProcessor $taskResultProcessor)
-    {
+    public function __construct(
+        private readonly QueueProcessor      $queueProcessor,
+        private readonly TaskResultProcessor $taskResultProcessor,
+        private readonly AiTaskProcessor $aiTaskProcessor,
+    ) {
         parent::__construct();
     }
 
@@ -24,6 +29,7 @@ protected function configure(): void
 
     protected function execute(InputInterface $input, OutputInterface $output): int
     {
-        return $this->queueProcessor->run($this->taskResultProcessor, 'tasks:queue');
+        return $this->queueProcessor->run($this->taskResultProcessor, 'tasks:queue') +
+            $this->queueProcessor->run($this->aiTaskProcessor, 'llm:queue');
     }
 }

src/Controllers/Agents/BootAgentController.php

@@ -1,8 +1,11 @@
-<?php declare(strict_types=1);
+<?php
+
+declare(strict_types=1);
 
 namespace Reconmap\Controllers\Agents;
 
 use OpenApi\Attributes as OpenApi;
+use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Reconmap\Controllers\Controller;
 use Reconmap\Http\Docs\Default200OkResponse;
@@ -15,14 +18,16 @@
 class BootAgentController extends Controller
 {
 
-    public function __construct(private readonly AgentRepository $repository)
-    {
-    }
+    public function __construct(private readonly AgentRepository $repository) {}
 
-    public function __invoke(ServerRequestInterface $request): array
+    public function __invoke(ServerRequestInterface $request): ResponseInterface
     {
         $info = $this->getJsonBodyDecoded($request);
 
-        return $this->repository->updateLastBootAt($info);
+        if ($this->repository->updateLastBootAt($info)) {
+            return $this->createNoContentResponse();
+        }
+
+        return $this->createBadRequestResponse();
     }
 }

src/Controllers/Agents/PingAgentController.php

@@ -1,8 +1,11 @@
-<?php declare(strict_types=1);
+<?php
+
+declare(strict_types=1);
 
 namespace Reconmap\Controllers\Agents;
 
 use OpenApi\Attributes as OpenApi;
+use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Reconmap\Controllers\Controller;
 use Reconmap\Http\Docs\Default200OkResponse;
@@ -15,12 +18,14 @@
 class PingAgentController extends Controller
 {
 
-    public function __construct(private readonly AgentRepository $repository)
-    {
-    }
+    public function __construct(private readonly AgentRepository $repository) {}
 
-    public function __invoke(ServerRequestInterface $request): array
+    public function __invoke(ServerRequestInterface $request): ResponseInterface
     {
-        return $this->repository->updateLastPingAt();
+        if ($this->repository->updateLastPingAt()) {
+            return $this->createNoContentResponse();
+        }
+
+        return $this->createBadRequestResponse();
     }
 }

src/Controllers/Vulnerabilities/GenerateVulnerabilityRemediationController.php

@@ -1,65 +1,35 @@
-<?php declare(strict_types=1);
+<?php
+
+declare(strict_types=1);
 
 namespace Reconmap\Controllers\Vulnerabilities;
 
-use LLPhant\Chat\OllamaChat;
-use LLPhant\OllamaConfig;
 use Psr\Http\Message\ResponseInterface;
 use Psr\Http\Message\ServerRequestInterface;
 use Reconmap\Controllers\Controller;
 use Reconmap\Models\AuditActions\AuditActions;
-use Reconmap\Models\Notification;
-use Reconmap\Repositories\NotificationsRepository;
-use Reconmap\Repositories\VulnerabilityRepository;
 use Reconmap\Services\ActivityPublisherService;
 use Reconmap\Services\ApplicationConfig;
 use Reconmap\Services\RedisServer;
 
 class GenerateVulnerabilityRemediationController extends Controller
 {
-    public function __construct(private readonly VulnerabilityRepository  $repository,
-                                private readonly ActivityPublisherService $activityPublisherService,
-                                private readonly NotificationsRepository  $notificationsRepository,
-                                private readonly RedisServer              $redisServer,
-                                private readonly ApplicationConfig        $config,
-    )
-    {
-    }
+    public function __construct(
+        private readonly ActivityPublisherService $activityPublisherService,
+        private readonly RedisServer              $redisServer,
+        private readonly ApplicationConfig        $config,
+    ) {}
 
     public function __invoke(ServerRequestInterface $request, array $args): ResponseInterface
     {
         $vulnerabilityId = intval($args['vulnerabilityId']);
 
-        $appConfig = $this->config->getSettings('integrations');
-
-        $config = new OllamaConfig();
-        $config->url = $appConfig['ollama']['url'];
-        $config->model = $appConfig['ollama']['model'];
-        $config->stream = false;
-
-        $vulnerability = $this->repository->findById($vulnerabilityId);
-
-        $chat = new OllamaChat($config);
-        $chat->setSystemMessage('You are a vulnerability and pentesting expert system');
-        $remediation = $chat->generateText('Write instructions on how to remediate this vulnerability: ' . $vulnerability['summary']);
-
-        $success = $this->repository->updateById($vulnerabilityId, ['remediation' => $remediation]);
-
         $loggedInUserId = $request->getAttribute('userId');
 
-        $notification = new Notification();
-        $notification->toUserId = $loggedInUserId;
-        $notification->title = "AI job completed";
-        $notification->content = "The vulnerability remediation instructions have now been generated";
-        $this->notificationsRepository->insert($notification);
-        $this->redisServer->lPush("notifications:queue", json_encode(['type' => 'message']));
-
-        if ($success) {
-            $this->auditAction($loggedInUserId, $vulnerabilityId);
-            return $this->createNoContentResponse();
-        }
+        $this->redisServer->lPush("llm:queue", json_encode(['vulnerabilityId' => $vulnerabilityId, 'loggedInUserId' => $loggedInUserId]));
 
-        return $this->createBadRequestResponse();
+        $this->auditAction($loggedInUserId, $vulnerabilityId);
+        return $this->createNoContentResponse();
     }
 
     private function auditAction(int $loggedInUserId, int $vulnerabilityId): void

src/Tasks/AiTaskProcessor.php

@@ -0,0 +1,56 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Reconmap\Tasks;
+
+use LLPhant\Chat\OllamaChat;
+use LLPhant\OllamaConfig;
+use Psr\Log\LoggerInterface;
+use Reconmap\Models\Notification;
+use Reconmap\Repositories\NotificationsRepository;
+use Reconmap\Repositories\VulnerabilityRepository;
+use Reconmap\Services\ApplicationConfig;
+use Reconmap\Services\RedisServer;
+
+readonly class AiTaskProcessor implements ItemProcessor
+{
+    public function __construct(
+        private ApplicationConfig $config,
+        private LoggerInterface   $logger,
+        private readonly VulnerabilityRepository  $repository,
+        private readonly NotificationsRepository  $notificationsRepository,
+        private readonly RedisServer $redisServer,
+    ) {}
+
+    public function process(object $message): void
+    {
+        $this->logger->debug('Processing LLM job');
+
+        try {
+            $appConfig = $this->config->getSettings('integrations');
+
+            $config = new OllamaConfig();
+            $config->url = $appConfig['ollama']['url'];
+            $config->model = $appConfig['ollama']['model'];
+            $config->stream = false;
+
+            $vulnerability = $this->repository->findById($message->vulnerabilityId);
+            $this->logger->debug($vulnerability['summary']);
+            $chat = new OllamaChat($config);
+            $chat->setSystemMessage('You are a vulnerability and pentesting expert system');
+            $remediation = $chat->generateText('Write instructions on how to remediate this vulnerability: ' . $vulnerability['summary']);
+
+            $success = $this->repository->updateById($message->vulnerabilityId, ['remediation' => $remediation]);
+
+            $notification = new Notification();
+            $notification->toUserId = $message->loggedInUserId;
+            $notification->title = "AI job completed";
+            $notification->content = "The vulnerability remediation instructions have now been generated";
+            $this->notificationsRepository->insert($notification);
+            $this->redisServer->lPush("notifications:queue", json_encode(['type' => 'message']));
+        } catch (\Exception $e) {
+            $this->logger->error($e->getMessage());
+        }
+    }
+}