Skip to content

Commit b88fc93

Browse files
julien-ncjulien-nc
authored
Merge pull request #156 from nextcloud/fix/noid/chat-polling-switch-session
Fix polling new chat message when switching sessions
2 parents 38c1277 + 23eaccd commit b88fc93

File tree

8 files changed

+267
-36
lines changed

8 files changed

+267
-36
lines changed

‎appinfo/routes.php‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
['name' => 'chattyLLM#getMessages', 'url' => '/chat/messages', 'verb' => 'GET'],
2525
['name' => 'chattyLLM#generateForSession', 'url' => '/chat/generate', 'verb' => 'GET'],
2626
['name' => 'chattyLLM#regenerateForSession', 'url' => '/chat/regenerate', 'verb' => 'GET'],
27+
['name' => 'chattyLLM#checkSession', 'url' => '/chat/check_session', 'verb' => 'GET'],
2728
['name' => 'chattyLLM#checkMessageGenerationTask', 'url' => '/chat/check_generation', 'verb' => 'GET'],
2829
['name' => 'chattyLLM#generateTitle', 'url' => '/chat/generate_title', 'verb' => 'GET'],
2930
['name' => 'chattyLLM#checkTitleGenerationTask', 'url' => '/chat/check_title_generation', 'verb' => 'GET'],

‎lib/AppInfo/Application.php‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
use OCA\Assistant\Capabilities;
66
use OCA\Assistant\Listener\BeforeTemplateRenderedListener;
7+
use OCA\Assistant\Listener\ChattyLLMTaskListener;
78
use OCA\Assistant\Listener\CSPListener;
89
use OCA\Assistant\Listener\FreePrompt\FreePromptReferenceListener;
910
use OCA\Assistant\Listener\SpeechToText\SpeechToTextReferenceListener;
@@ -55,6 +56,7 @@ public function register(IRegistrationContext $context): void {
5556

5657
$context->registerEventListener(TaskSuccessfulEvent::class, TaskSuccessfulListener::class);
5758
$context->registerEventListener(TaskFailedEvent::class, TaskFailedListener::class);
59+
$context->registerEventListener(TaskSuccessfulEvent::class, ChattyLLMTaskListener::class);
5860

5961
$context->registerNotifierService(Notifier::class);
6062

‎lib/Controller/ChattyLLMController.php‎

Lines changed: 82 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,11 @@ public function generateForSession(int $sessionId): JSONResponse {
297297
. PHP_EOL
298298
. 'assistant: ';
299299

300-
$taskId = $this->scheduleLLMTask($stichedPrompt);
300+
try {
301+
$taskId = $this->scheduleLLMTask($stichedPrompt, $sessionId);
302+
} catch (\Exception $e) {
303+
return new JSONResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST);
304+
}
301305

302306
return new JSONResponse(['taskId' => $taskId]);
303307
}
@@ -374,7 +378,7 @@ public function checkMessageGenerationTask(int $taskId, int $sessionId): JSONRes
374378
$message->setRole('assistant');
375379
$message->setContent(trim($task->getOutput()['output'] ?? ''));
376380
$message->setTimestamp(time());
377-
$this->messageMapper->insert($message);
381+
// do not insert here, it is done by the listener
378382
return new JSONResponse($message);
379383
} catch (\OCP\DB\Exception $e) {
380384
$this->logger->warning('Failed to add a chat message into DB', ['exception' => $e]);
@@ -388,6 +392,56 @@ public function checkMessageGenerationTask(int $taskId, int $sessionId): JSONRes
388392
return new JSONResponse(['error' => 'unknown_error', 'task_status' => $task->getstatus()], Http::STATUS_BAD_REQUEST);
389393
}
390394

395+
/**
396+
* Check the status of a session
397+
*
398+
* Used by the frontend to determine if it should poll a generation task status.
399+
*
400+
* @param int $sessionId
401+
* @return JSONResponse
402+
* @throws \JsonException
403+
* @throws \OCP\DB\Exception
404+
*/
405+
#[NoAdminRequired]
406+
public function checkSession(int $sessionId): JSONResponse {
407+
if ($this->userId === null) {
408+
return new JSONResponse(['error' => $this->l10n->t('User not logged in')], Http::STATUS_UNAUTHORIZED);
409+
}
410+
411+
$sessionExists = $this->sessionMapper->exists($this->userId, $sessionId);
412+
if (!$sessionExists) {
413+
return new JSONResponse(['error' => $this->l10n->t('Session not found')], Http::STATUS_NOT_FOUND);
414+
}
415+
416+
try {
417+
$messageTasks = $this->taskProcessingManager->getUserTasksByApp($this->userId, Application::APP_ID . ':chatty-llm', 'chatty-llm:' . $sessionId);
418+
$titleTasks = $this->taskProcessingManager->getUserTasksByApp($this->userId, Application::APP_ID . ':chatty-llm', 'chatty-title:' . $sessionId);
419+
} catch (\OCP\TaskProcessing\Exception\Exception $e) {
420+
return new JSONResponse(['error' => 'task_query_failed'], Http::STATUS_BAD_REQUEST);
421+
}
422+
$messageTasks = array_filter($messageTasks, static function (Task $task) {
423+
return $task->getStatus() === Task::STATUS_RUNNING || $task->getStatus() === Task::STATUS_SCHEDULED;
424+
});
425+
$titleTasks = array_filter($titleTasks, static function (Task $task) {
426+
return $task->getStatus() === Task::STATUS_RUNNING || $task->getStatus() === Task::STATUS_SCHEDULED;
427+
});
428+
$session = $this->sessionMapper->getUserSession($this->userId, $sessionId);
429+
$responseData = [
430+
'messageTaskId' => null,
431+
'titleTaskId' => null,
432+
'sessionTitle' => $session->getTitle(),
433+
];
434+
if (!empty($messageTasks)) {
435+
$task = array_pop($messageTasks);
436+
$responseData['messageTaskId'] = $task->getId();
437+
}
438+
if (!empty($titleTasks)) {
439+
$task = array_pop($titleTasks);
440+
$responseData['titleTaskId'] = $task->getId();
441+
}
442+
return new JSONResponse($responseData);
443+
}
444+
391445
/**
392446
* Schedule a task to generate a title for the chat session
393447
*
@@ -430,7 +484,11 @@ public function generateTitle(int $sessionId): JSONResponse {
430484
. PHP_EOL . PHP_EOL
431485
. $userInstructions;
432486

433-
$taskId = $this->scheduleLLMTask($stichedPrompt);
487+
try {
488+
$taskId = $this->scheduleLLMTask($stichedPrompt, $sessionId, false);
489+
} catch (\Exception $e) {
490+
return new JSONResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST);
491+
}
434492
return new JSONResponse(['taskId' => $taskId]);
435493
} catch (\OCP\DB\Exception $e) {
436494
$this->logger->warning('Failed to generate a title for the chat session', ['exception' => $e]);
@@ -475,8 +533,7 @@ public function checkTitleGenerationTask(int $taskId, int $sessionId): JSONRespo
475533
$title = str_replace('"', '', $title);
476534
$title = explode(PHP_EOL, $title)[0];
477535
$title = trim($title);
478-
479-
$this->sessionMapper->updateSessionTitle($this->userId, $sessionId, $title);
536+
// do not write the title here since it's done in the listener
480537

481538
return new JSONResponse(['result' => $title]);
482539
} catch (\OCP\DB\Exception $e) {
@@ -525,14 +582,32 @@ private function getStichedMessages(int $sessionId): string {
525582
* Schedule the LLM task
526583
*
527584
* @param string $content
585+
* @param int $sessionId
586+
* @param bool $isMessage
528587
* @return int|null
529588
* @throws Exception
530589
* @throws PreConditionNotMetException
531590
* @throws UnauthorizedException
532591
* @throws ValidationException
592+
* @throws \JsonException
533593
*/
534-
private function scheduleLLMTask(string $content): ?int {
535-
$task = new Task(TextToText::ID, ['input' => $content], Application::APP_ID . ':chatty-llm', $this->userId);
594+
private function scheduleLLMTask(string $content, int $sessionId, bool $isMessage = true): ?int {
595+
$customId = ($isMessage
596+
? 'chatty-llm:'
597+
: 'chatty-title:') . $sessionId;
598+
try {
599+
$tasks = $this->taskProcessingManager->getUserTasksByApp($this->userId, Application::APP_ID . ':chatty-llm', $customId);
600+
} catch (\OCP\TaskProcessing\Exception\Exception $e) {
601+
throw new \Exception('task_query_failed');
602+
}
603+
$tasks = array_filter($tasks, static function (Task $task) {
604+
return $task->getStatus() === Task::STATUS_RUNNING || $task->getStatus() === Task::STATUS_SCHEDULED;
605+
});
606+
// prevent scheduling multiple llm tasks simultaneously for one session
607+
if (!empty($tasks)) {
608+
throw new \Exception('session_already_thinking');
609+
}
610+
$task = new Task(TextToText::ID, ['input' => $content], Application::APP_ID . ':chatty-llm', $this->userId, $customId);
536611
$this->taskProcessingManager->scheduleTask($task);
537612
return $task->getId();
538613
}

‎lib/Db/ChattyLLM/Session.php‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
/**
3232
* @method \string getUserId()
3333
* @method \void setUserId(string $userId)
34-
* @method \?string getTitle()
34+
* @method \string|null getTitle()
3535
* @method \void setTitle(?string $title)
3636
* @method \int|null getTimestamp()
3737
* @method \void setTimestamp(?int $timestamp)

‎lib/Db/ChattyLLM/SessionMapper.php‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@
2525

2626
namespace OCA\Assistant\Db\ChattyLLM;
2727

28+
use OCP\AppFramework\Db\DoesNotExistException;
29+
use OCP\AppFramework\Db\MultipleObjectsReturnedException;
2830
use OCP\AppFramework\Db\QBMapper;
31+
use OCP\DB\Exception;
2932
use OCP\DB\QueryBuilder\IQueryBuilder;
3033
use OCP\IDBConnection;
3134

@@ -59,6 +62,24 @@ public function exists(string $userId, int $sessionId): bool {
5962
}
6063
}
6164

65+
/**
66+
* @param string $userId
67+
* @param int $sessionId
68+
* @return Session
69+
* @throws DoesNotExistException
70+
* @throws MultipleObjectsReturnedException
71+
* @throws Exception
72+
*/
73+
public function getUserSession(string $userId, int $sessionId): Session {
74+
$qb = $this->db->getQueryBuilder();
75+
$qb->select('id', 'title', 'timestamp')
76+
->from($this->getTableName())
77+
->where($qb->expr()->eq('id', $qb->createPositionalParameter($sessionId, IQueryBuilder::PARAM_INT)))
78+
->andWhere($qb->expr()->eq('user_id', $qb->createPositionalParameter($userId, IQueryBuilder::PARAM_STR)));
79+
80+
return $this->findEntity($qb);
81+
}
82+
6283
/**
6384
* @param string $userId
6485
* @return array

‎lib/Listener/ChattyLLMTaskListener.php‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
<?php
2+
3+
declare(strict_types=1);
4+
5+
namespace OCA\Assistant\Listener;
6+
7+
use OCA\Assistant\AppInfo\Application;
8+
use OCA\Assistant\Db\ChattyLLM\Message;
9+
use OCA\Assistant\Db\ChattyLLM\MessageMapper;
10+
use OCA\Assistant\Db\ChattyLLM\SessionMapper;
11+
use OCP\EventDispatcher\Event;
12+
use OCP\EventDispatcher\IEventListener;
13+
use OCP\TaskProcessing\Events\TaskSuccessfulEvent;
14+
use Psr\Log\LoggerInterface;
15+
16+
/**
17+
* @template-implements IEventListener<TaskSuccessfulEvent>
18+
*/
19+
class ChattyLLMTaskListener implements IEventListener {
20+
21+
public function __construct(
22+
private MessageMapper $messageMapper,
23+
private SessionMapper $sessionMapper,
24+
private LoggerInterface $logger,
25+
) {
26+
}
27+
28+
public function handle(Event $event): void {
29+
if (!($event instanceof TaskSuccessfulEvent)) {
30+
return;
31+
}
32+
33+
$task = $event->getTask();
34+
$customId = $task->getCustomId();
35+
$appId = $task->getAppId();
36+
37+
if ($customId === null || $appId !== (Application::APP_ID . ':chatty-llm')) {
38+
return;
39+
}
40+
41+
// title generation
42+
if (preg_match('/^chatty-title:(\d+)$/', $customId, $matches)) {
43+
$sessionId = (int)$matches[1];
44+
$title = trim($task->getOutput()['output'] ?? '');
45+
$this->sessionMapper->updateSessionTitle($task->getUserId(), $sessionId, $title);
46+
}
47+
48+
// message generation
49+
if (preg_match('/^chatty-llm:(\d+)$/', $customId, $matches)) {
50+
$sessionId = (int)$matches[1];
51+
52+
$message = new Message();
53+
$message->setSessionId($sessionId);
54+
$message->setRole('assistant');
55+
$message->setContent(trim($task->getOutput()['output'] ?? ''));
56+
$message->setTimestamp(time());
57+
try {
58+
$this->messageMapper->insert($message);
59+
} catch (\OCP\DB\Exception $e) {
60+
$this->logger->error('Message insertion error in chattyllm task listener', ['exception' => $e]);
61+
}
62+
}
63+
}
64+
}

0 commit comments

Comments
 (0)