1
0

TextProcessingTest.php 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352
  1. <?php
  2. /**
  3. * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors
  4. * SPDX-License-Identifier: AGPL-3.0-or-later
  5. */
  6. namespace Test\TextProcessing;
  7. use OC\AppFramework\Bootstrap\Coordinator;
  8. use OC\AppFramework\Bootstrap\RegistrationContext;
  9. use OC\AppFramework\Bootstrap\ServiceRegistration;
  10. use OC\EventDispatcher\EventDispatcher;
  11. use OC\TextProcessing\Db\Task as DbTask;
  12. use OC\TextProcessing\Db\TaskMapper;
  13. use OC\TextProcessing\Manager;
  14. use OC\TextProcessing\RemoveOldTasksBackgroundJob;
  15. use OC\TextProcessing\TaskBackgroundJob;
  16. use OCP\AppFramework\Db\DoesNotExistException;
  17. use OCP\AppFramework\Utility\ITimeFactory;
  18. use OCP\BackgroundJob\IJobList;
  19. use OCP\Common\Exception\NotFoundException;
  20. use OCP\EventDispatcher\IEventDispatcher;
  21. use OCP\IConfig;
  22. use OCP\IServerContainer;
  23. use OCP\PreConditionNotMetException;
  24. use OCP\TextProcessing\Events\TaskFailedEvent;
  25. use OCP\TextProcessing\Events\TaskSuccessfulEvent;
  26. use OCP\TextProcessing\FreePromptTaskType;
  27. use OCP\TextProcessing\IManager;
  28. use OCP\TextProcessing\IProvider;
  29. use OCP\TextProcessing\SummaryTaskType;
  30. use OCP\TextProcessing\Task;
  31. use OCP\TextProcessing\TopicsTaskType;
  32. use PHPUnit\Framework\Constraint\IsInstanceOf;
  33. use Psr\Log\LoggerInterface;
  34. use Test\BackgroundJob\DummyJobList;
  35. class SuccessfulSummaryProvider implements IProvider {
  36. public bool $ran = false;
  37. public function getName(): string {
  38. return 'TEST Vanilla LLM Provider';
  39. }
  40. public function process(string $prompt): string {
  41. $this->ran = true;
  42. return $prompt . ' Summarize';
  43. }
  44. public function getTaskType(): string {
  45. return SummaryTaskType::class;
  46. }
  47. }
  48. class FailingSummaryProvider implements IProvider {
  49. public bool $ran = false;
  50. public function getName(): string {
  51. return 'TEST Vanilla LLM Provider';
  52. }
  53. public function process(string $prompt): string {
  54. $this->ran = true;
  55. throw new \Exception('ERROR');
  56. }
  57. public function getTaskType(): string {
  58. return SummaryTaskType::class;
  59. }
  60. }
  61. class FreePromptProvider implements IProvider {
  62. public bool $ran = false;
  63. public function getName(): string {
  64. return 'TEST Free Prompt Provider';
  65. }
  66. public function process(string $prompt): string {
  67. $this->ran = true;
  68. return $prompt . ' Free Prompt';
  69. }
  70. public function getTaskType(): string {
  71. return FreePromptTaskType::class;
  72. }
  73. }
  74. class TextProcessingTest extends \Test\TestCase {
  75. private IManager $manager;
  76. private Coordinator $coordinator;
  77. private array $providers;
  78. private IServerContainer $serverContainer;
  79. private IEventDispatcher $eventDispatcher;
  80. private RegistrationContext $registrationContext;
  81. private \DateTimeImmutable $currentTime;
  82. private TaskMapper $taskMapper;
  83. private array $tasksDb;
  84. private IJobList $jobList;
  85. protected function setUp(): void {
  86. parent::setUp();
  87. $this->providers = [
  88. SuccessfulSummaryProvider::class => new SuccessfulSummaryProvider(),
  89. FailingSummaryProvider::class => new FailingSummaryProvider(),
  90. FreePromptProvider::class => new FreePromptProvider(),
  91. ];
  92. $this->serverContainer = $this->createMock(IServerContainer::class);
  93. $this->serverContainer->expects($this->any())->method('get')->willReturnCallback(function ($class) {
  94. return $this->providers[$class];
  95. });
  96. $this->eventDispatcher = new EventDispatcher(
  97. new \Symfony\Component\EventDispatcher\EventDispatcher(),
  98. $this->serverContainer,
  99. \OC::$server->get(LoggerInterface::class),
  100. );
  101. $this->registrationContext = $this->createMock(RegistrationContext::class);
  102. $this->coordinator = $this->createMock(Coordinator::class);
  103. $this->coordinator->expects($this->any())->method('getRegistrationContext')->willReturn($this->registrationContext);
  104. $this->currentTime = new \DateTimeImmutable('now');
  105. $this->taskMapper = $this->createMock(TaskMapper::class);
  106. $this->tasksDb = [];
  107. $this->taskMapper
  108. ->expects($this->any())
  109. ->method('insert')
  110. ->willReturnCallback(function (DbTask $task) {
  111. $task->setId(count($this->tasksDb) ? max(array_keys($this->tasksDb)) : 1);
  112. $task->setLastUpdated($this->currentTime->getTimestamp());
  113. $this->tasksDb[$task->getId()] = $task->toRow();
  114. return $task;
  115. });
  116. $this->taskMapper
  117. ->expects($this->any())
  118. ->method('update')
  119. ->willReturnCallback(function (DbTask $task) {
  120. $task->setLastUpdated($this->currentTime->getTimestamp());
  121. $this->tasksDb[$task->getId()] = $task->toRow();
  122. return $task;
  123. });
  124. $this->taskMapper
  125. ->expects($this->any())
  126. ->method('find')
  127. ->willReturnCallback(function (int $id) {
  128. if (!isset($this->tasksDb[$id])) {
  129. throw new DoesNotExistException('Could not find it');
  130. }
  131. return DbTask::fromRow($this->tasksDb[$id]);
  132. });
  133. $this->taskMapper
  134. ->expects($this->any())
  135. ->method('deleteOlderThan')
  136. ->willReturnCallback(function (int $timeout) {
  137. $this->tasksDb = array_filter($this->tasksDb, function (array $task) use ($timeout) {
  138. return $task['last_updated'] >= $this->currentTime->getTimestamp() - $timeout;
  139. });
  140. });
  141. $this->jobList = $this->createPartialMock(DummyJobList::class, ['add']);
  142. $this->jobList->expects($this->any())->method('add')->willReturnCallback(function () {
  143. });
  144. $config = $this->createMock(IConfig::class);
  145. $config->method('getAppValue')
  146. ->with('core', 'ai.textprocessing_provider_preferences', '')
  147. ->willReturn('');
  148. $this->manager = new Manager(
  149. $this->serverContainer,
  150. $this->coordinator,
  151. \OC::$server->get(LoggerInterface::class),
  152. $this->jobList,
  153. $this->taskMapper,
  154. $config
  155. );
  156. }
  157. public function testShouldNotHaveAnyProviders() {
  158. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([]);
  159. $this->assertCount(0, $this->manager->getAvailableTaskTypes());
  160. $this->assertFalse($this->manager->hasProviders());
  161. $this->expectException(PreConditionNotMetException::class);
  162. $this->manager->runTask(new \OCP\TextProcessing\Task(FreePromptTaskType::class, 'Hello', 'test', null));
  163. }
  164. public function testProviderShouldBeRegisteredAndRun() {
  165. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([
  166. new ServiceRegistration('test', SuccessfulSummaryProvider::class)
  167. ]);
  168. $this->assertCount(1, $this->manager->getAvailableTaskTypes());
  169. $this->assertTrue($this->manager->hasProviders());
  170. $this->assertEquals('Hello Summarize', $this->manager->runTask(new Task(SummaryTaskType::class, 'Hello', 'test', null)));
  171. // Summaries are not implemented by the vanilla provider, only free prompt
  172. $this->expectException(PreConditionNotMetException::class);
  173. $this->manager->runTask(new Task(FreePromptTaskType::class, 'Hello', 'test', null));
  174. }
  175. public function testProviderShouldBeRegisteredAndScheduled() {
  176. // register provider
  177. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([
  178. new ServiceRegistration('test', SuccessfulSummaryProvider::class)
  179. ]);
  180. $this->assertCount(1, $this->manager->getAvailableTaskTypes());
  181. $this->assertTrue($this->manager->hasProviders());
  182. // create task object
  183. $task = new Task(SummaryTaskType::class, 'Hello', 'test', null);
  184. $this->assertNull($task->getId());
  185. $this->assertNull($task->getOutput());
  186. // schedule works
  187. $this->assertEquals(Task::STATUS_UNKNOWN, $task->getStatus());
  188. $this->manager->scheduleTask($task);
  189. // Task object is up-to-date
  190. $this->assertNotNull($task->getId());
  191. $this->assertNull($task->getOutput());
  192. $this->assertEquals(Task::STATUS_SCHEDULED, $task->getStatus());
  193. // Task object retrieved from db is up-to-date
  194. $task2 = $this->manager->getTask($task->getId());
  195. $this->assertEquals($task->getId(), $task2->getId());
  196. $this->assertEquals('Hello', $task2->getInput());
  197. $this->assertNull($task2->getOutput());
  198. $this->assertEquals(Task::STATUS_SCHEDULED, $task2->getStatus());
  199. $this->eventDispatcher = $this->createMock(IEventDispatcher::class);
  200. $this->eventDispatcher->expects($this->once())->method('dispatchTyped')->with(new IsInstanceOf(TaskSuccessfulEvent::class));
  201. // run background job
  202. $bgJob = new TaskBackgroundJob(
  203. \OC::$server->get(ITimeFactory::class),
  204. $this->manager,
  205. $this->eventDispatcher,
  206. );
  207. $bgJob->setArgument(['taskId' => $task->getId()]);
  208. $bgJob->start($this->jobList);
  209. $provider = $this->providers[SuccessfulSummaryProvider::class];
  210. $this->assertTrue($provider->ran);
  211. // Task object retrieved from db is up-to-date
  212. $task3 = $this->manager->getTask($task->getId());
  213. $this->assertEquals($task->getId(), $task3->getId());
  214. $this->assertEquals('Hello', $task3->getInput());
  215. $this->assertEquals('Hello Summarize', $task3->getOutput());
  216. $this->assertEquals(Task::STATUS_SUCCESSFUL, $task3->getStatus());
  217. }
  218. public function testMultipleProvidersShouldBeRegisteredAndRunCorrectly() {
  219. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([
  220. new ServiceRegistration('test', SuccessfulSummaryProvider::class),
  221. new ServiceRegistration('test', FreePromptProvider::class),
  222. ]);
  223. $this->assertCount(2, $this->manager->getAvailableTaskTypes());
  224. $this->assertTrue($this->manager->hasProviders());
  225. // Try free prompt again
  226. $this->assertEquals('Hello Free Prompt', $this->manager->runTask(new Task(FreePromptTaskType::class, 'Hello', 'test', null)));
  227. // Try summary task
  228. $this->assertEquals('Hello Summarize', $this->manager->runTask(new Task(SummaryTaskType::class, 'Hello', 'test', null)));
  229. // Topics are not implemented by both the vanilla provider and the full provider
  230. $this->expectException(PreConditionNotMetException::class);
  231. $this->manager->runTask(new Task(TopicsTaskType::class, 'Hello', 'test', null));
  232. }
  233. public function testNonexistentTask() {
  234. $this->expectException(NotFoundException::class);
  235. $this->manager->getTask(2147483646);
  236. }
  237. public function testTaskFailure() {
  238. // register provider
  239. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([
  240. new ServiceRegistration('test', FailingSummaryProvider::class),
  241. ]);
  242. $this->assertCount(1, $this->manager->getAvailableTaskTypes());
  243. $this->assertTrue($this->manager->hasProviders());
  244. // create task object
  245. $task = new Task(SummaryTaskType::class, 'Hello', 'test', null);
  246. $this->assertNull($task->getId());
  247. $this->assertNull($task->getOutput());
  248. // schedule works
  249. $this->assertEquals(Task::STATUS_UNKNOWN, $task->getStatus());
  250. $this->manager->scheduleTask($task);
  251. // Task object is up-to-date
  252. $this->assertNotNull($task->getId());
  253. $this->assertNull($task->getOutput());
  254. $this->assertEquals(Task::STATUS_SCHEDULED, $task->getStatus());
  255. // Task object retrieved from db is up-to-date
  256. $task2 = $this->manager->getTask($task->getId());
  257. $this->assertEquals($task->getId(), $task2->getId());
  258. $this->assertEquals('Hello', $task2->getInput());
  259. $this->assertNull($task2->getOutput());
  260. $this->assertEquals(Task::STATUS_SCHEDULED, $task2->getStatus());
  261. $this->eventDispatcher = $this->createMock(IEventDispatcher::class);
  262. $this->eventDispatcher->expects($this->once())->method('dispatchTyped')->with(new IsInstanceOf(TaskFailedEvent::class));
  263. // run background job
  264. $bgJob = new TaskBackgroundJob(
  265. \OC::$server->get(ITimeFactory::class),
  266. $this->manager,
  267. $this->eventDispatcher,
  268. );
  269. $bgJob->setArgument(['taskId' => $task->getId()]);
  270. $bgJob->start($this->jobList);
  271. $provider = $this->providers[FailingSummaryProvider::class];
  272. $this->assertTrue($provider->ran);
  273. // Task object retrieved from db is up-to-date
  274. $task3 = $this->manager->getTask($task->getId());
  275. $this->assertEquals($task->getId(), $task3->getId());
  276. $this->assertEquals('Hello', $task3->getInput());
  277. $this->assertNull($task3->getOutput());
  278. $this->assertEquals(Task::STATUS_FAILED, $task3->getStatus());
  279. }
  280. public function testOldTasksShouldBeCleanedUp() {
  281. $this->registrationContext->expects($this->any())->method('getTextProcessingProviders')->willReturn([
  282. new ServiceRegistration('test', SuccessfulSummaryProvider::class)
  283. ]);
  284. $this->assertCount(1, $this->manager->getAvailableTaskTypes());
  285. $this->assertTrue($this->manager->hasProviders());
  286. $task = new Task(SummaryTaskType::class, 'Hello', 'test', null);
  287. $this->assertEquals('Hello Summarize', $this->manager->runTask($task));
  288. $this->currentTime = $this->currentTime->add(new \DateInterval('P1Y'));
  289. // run background job
  290. $bgJob = new RemoveOldTasksBackgroundJob(
  291. \OC::$server->get(ITimeFactory::class),
  292. $this->taskMapper,
  293. \OC::$server->get(LoggerInterface::class),
  294. );
  295. $bgJob->setArgument([]);
  296. $bgJob->start($this->jobList);
  297. $this->expectException(NotFoundException::class);
  298. $this->manager->getTask($task->getId());
  299. }
  300. }