diff --git a/lib/Controller/ChattyLLMController.php b/lib/Controller/ChattyLLMController.php index cc750a60..0123d80e 100644 --- a/lib/Controller/ChattyLLMController.php +++ b/lib/Controller/ChattyLLMController.php @@ -581,7 +581,6 @@ public function regenerateForSession(int $sessionId, int $messageId): JSONRespon * @param int $taskId The message generation task ID * @param int $sessionId The chat session ID * @return JSONResponse|JSONResponse|numeric|string>|null}, array{}>|JSONResponse|JSONResponse - * @throws MultipleObjectsReturnedException * @throws \OCP\DB\Exception * * 200: The task was successful, a message has been generated diff --git a/lib/Db/ChattyLLM/MessageMapper.php b/lib/Db/ChattyLLM/MessageMapper.php index 8b44208e..5522f9c6 100644 --- a/lib/Db/ChattyLLM/MessageMapper.php +++ b/lib/Db/ChattyLLM/MessageMapper.php @@ -150,7 +150,6 @@ public function getMessageById(int $sessionId, int $messageId): Message { * @param int $ocpTaskId * @return Message * @throws DoesNotExistException - * @throws MultipleObjectsReturnedException * @throws \OCP\DB\Exception */ public function getMessageByTaskId(int $sessionId, int $ocpTaskId): Message { @@ -158,9 +157,17 @@ public function getMessageByTaskId(int $sessionId, int $ocpTaskId): Message { $qb->select(Message::$columns) ->from($this->getTableName()) ->where($qb->expr()->eq('session_id', $qb->createPositionalParameter($sessionId, IQueryBuilder::PARAM_INT))) - ->andWhere($qb->expr()->eq('ocp_task_id', $qb->createPositionalParameter($ocpTaskId, IQueryBuilder::PARAM_INT))); + ->andWhere($qb->expr()->eq('ocp_task_id', $qb->createPositionalParameter($ocpTaskId, IQueryBuilder::PARAM_INT))) + // concurrent task processing executors may run the same task twice, + // which results in duplicate rows with the same ocp_task_id + ->orderBy('id', 'ASC') + ->setMaxResults(1); - return $this->findEntity($qb); + $entities = $this->findEntities($qb); + if ($entities === []) { + throw new DoesNotExistException('Assistant chat message with session_id ' . $sessionId . ' and ocp_task_id ' . $ocpTaskId . ' does not exist'); + } + return $entities[0]; } /** diff --git a/lib/Listener/ChattyLLMTaskListener.php b/lib/Listener/ChattyLLMTaskListener.php index 388f720d..c1dcc058 100644 --- a/lib/Listener/ChattyLLMTaskListener.php +++ b/lib/Listener/ChattyLLMTaskListener.php @@ -15,6 +15,7 @@ use OCA\Assistant\Db\ChattyLLM\SessionMapper; use OCA\Assistant\Service\NotificationService; use OCA\Assistant\Service\TaskProcessingService; +use OCP\AppFramework\Db\DoesNotExistException; use OCP\EventDispatcher\Event; use OCP\EventDispatcher\IEventListener; use OCP\TaskProcessing\Events\TaskFailedEvent; @@ -76,6 +77,15 @@ public function handle(Event $event): void { if (preg_match('/^chatty-llm:(\d+)/', $customId, $matches)) { $sessionId = (int)$matches[1]; + // concurrent task processing executors can process the same task twice and + // dispatch this event twice, in which case the message was already stored + try { + $this->messageMapper->getMessageByTaskId($sessionId, $task->getId()); + return; + } catch (DoesNotExistException $e) { + // the message was not stored yet, continue with the normal handling + } + $isAgency = class_exists('OCP\\TaskProcessing\\TaskTypes\\ContextAgentInteraction') && $taskTypeId === \OCP\TaskProcessing\TaskTypes\ContextAgentInteraction::ID; $isAudioChat = class_exists('OCP\\TaskProcessing\\TaskTypes\\AudioToAudioChat') diff --git a/tests/unit/Listener/ChattyLLMTaskListenerTest.php b/tests/unit/Listener/ChattyLLMTaskListenerTest.php new file mode 100644 index 00000000..845fb614 --- /dev/null +++ b/tests/unit/Listener/ChattyLLMTaskListenerTest.php @@ -0,0 +1,141 @@ +messageMapper = $this->createMock(MessageMapper::class); + $this->sessionMapper = $this->createMock(SessionMapper::class); + + $session = new Session(); + $session->setAssignmentId(null); + $this->sessionMapper->method('getUserSession')->willReturn($session); + + $this->listener = new ChattyLLMTaskListener( + $this->messageMapper, + $this->sessionMapper, + $this->createMock(TaskProcessingService::class), + $this->createMock(LoggerInterface::class), + $this->createMock(NotificationService::class), + ); + } + + private function makeTask(): Task { + $task = new Task( + TextToTextChat::ID, + [ + 'input' => 'What is a cat?', + 'system_prompt' => '', + 'history' => [], + ], + 'assistant:chatty-llm', + 'admin', + 'chatty-llm:' . self::SESSION_ID, + ); + $task->setId(self::TASK_ID); + $task->setOutput(['output' => 'A cat is an animal.']); + return $task; + } + + private function handleTask(): void { + $event = new TaskSuccessfulEvent($this->makeTask()); + $this->listener->handle($event); + } + + public function testMessageIsStored(): void { + $this->messageMapper->method('getMessageByTaskId') + ->with(self::SESSION_ID, self::TASK_ID) + ->willThrowException(new DoesNotExistException('Message not found')); + $storedMessage = null; + $this->messageMapper->expects($this->once()) + ->method('insert') + ->willReturnCallback(static function (Message $message) use (&$storedMessage): Message { + $storedMessage = $message; + $message->setId(3); + return $message; + }); + $this->sessionMapper->expects($this->once())->method('update'); + + $this->handleTask(); + + $this->assertNotNull($storedMessage); + $this->assertSame(self::SESSION_ID, $storedMessage->getSessionId()); + $this->assertSame(self::TASK_ID, $storedMessage->getOcpTaskId()); + $this->assertSame(Message::ROLE_ASSISTANT, $storedMessage->getRole()); + $this->assertSame('A cat is an animal.', $storedMessage->getContent()); + } + + public function testMessageIsStoredOnlyOnce(): void { + $storedMessage = new Message(); + $storedMessage->setSessionId(self::SESSION_ID); + $storedMessage->setOcpTaskId(self::TASK_ID); + $this->messageMapper->method('getMessageByTaskId') + ->with(self::SESSION_ID, self::TASK_ID) + ->willReturn($storedMessage); + $this->messageMapper->expects($this->never())->method('insert'); + $this->sessionMapper->expects($this->never())->method('update'); + + $this->handleTask(); + } + + public function testTitleTasksAreUntouched(): void { + $task = new Task( + TextToTextChat::ID, + [ + 'input' => 'Generate a title', + 'system_prompt' => '', + 'history' => [], + ], + 'assistant:chatty-llm', + 'admin', + 'chatty-title:' . self::SESSION_ID, + ); + $task->setId(self::TASK_ID); + $task->setOutput(['output' => 'Cats']); + + $event = new TaskSuccessfulEvent($task); + $this->messageMapper->expects($this->never())->method('getMessageByTaskId'); + $this->messageMapper->expects($this->never())->method('insert'); + $this->sessionMapper->expects($this->once()) + ->method('updateSessionTitle') + ->with('admin', self::SESSION_ID, 'Cats'); + + $this->listener->handle($event); + } +}