diff --git a/lib/AppInfo/Application.php b/lib/AppInfo/Application.php index d894e987..c66749bc 100644 --- a/lib/AppInfo/Application.php +++ b/lib/AppInfo/Application.php @@ -8,6 +8,7 @@ namespace OCA\OpenAi\AppInfo; use OCA\OpenAi\Capabilities; +use OCA\OpenAi\Notification\Notifier; use OCA\OpenAi\OldProcessing\Translation\TranslationProvider as OldTranslationProvider; use OCA\OpenAi\TaskProcessing\AudioToAudioChatProvider; use OCA\OpenAi\TaskProcessing\AudioToTextProvider; @@ -152,6 +153,7 @@ public function register(IRegistrationContext $context): void { } $context->registerCapability(Capabilities::class); + $context->registerNotifierService(Notifier::class); } public function boot(IBootContext $context): void { diff --git a/lib/Notification/Notifier.php b/lib/Notification/Notifier.php new file mode 100644 index 00000000..0c1d8949 --- /dev/null +++ b/lib/Notification/Notifier.php @@ -0,0 +1,86 @@ +factory->get(Application::APP_ID)->t('OpenAI Integration'); + } + + public function prepare(INotification $notification, string $languageCode): INotification { + if ($notification->getApp() !== Application::APP_ID) { + // Not my app => throw + throw new InvalidArgumentException(); + } + if ($notification->getSubject() !== 'quota_exceeded') { + // Not a valid subject => throw + throw new InvalidArgumentException(); + } + + $l = $this->factory->get(Application::APP_ID, $languageCode); + + $params = $notification->getSubjectParameters(); + + $subject = $l->t('Quota exceeded'); + $content = ''; + switch ($params['type']) { + case Application::QUOTA_TYPE_TEXT: + $content = $l->t('Text generation quota exceeded'); + break; + case Application::QUOTA_TYPE_IMAGE: + $content = $l->t('Image generation quota exceeded'); + break; + case Application::QUOTA_TYPE_TRANSCRIPTION: + $content = $l->t('Audio transcription quota exceeded'); + break; + case Application::QUOTA_TYPE_SPEECH: + $content = $l->t('Speech generation quota exceeded'); + break; + } + + $link = $this->url->getWebroot() . '/settings/user/ai'; + $iconUrl = $this->url->getAbsoluteURL($this->url->imagePath(Application::APP_ID, 'app-dark.svg')); + + $notification + ->setParsedSubject($subject) + ->setParsedMessage($content) + ->setLink($link) + ->setIcon($iconUrl); + + $actionLabel = $params['actionLabel'] ?? $l->t('View quota'); + $action = $notification->createAction(); + $action->setLabel($actionLabel) + ->setParsedLabel($actionLabel) + ->setLink($notification->getLink(), IAction::TYPE_WEB) + ->setPrimary(true); + + $notification->addParsedAction($action); + + return $notification; + } +} diff --git a/lib/Service/OpenAiAPIService.php b/lib/Service/OpenAiAPIService.php index 0f25548b..2cc8d2f8 100644 --- a/lib/Service/OpenAiAPIService.php +++ b/lib/Service/OpenAiAPIService.php @@ -7,6 +7,7 @@ namespace OCA\OpenAi\Service; +use DateTime; use Exception; use GuzzleHttp\Exception\ClientException; use GuzzleHttp\Exception\ServerException; @@ -25,6 +26,7 @@ use OCP\ICacheFactory; use OCP\IL10N; use OCP\Lock\LockedException; +use OCP\Notification\IManager as INotificationManager; use OCP\TaskProcessing\ShapeEnumValue; use Psr\Log\LoggerInterface; use RuntimeException; @@ -46,6 +48,7 @@ public function __construct( private ICacheFactory $cacheFactory, private QuotaUsageMapper $quotaUsageMapper, private OpenAiSettingsService $openAiSettingsService, + private INotificationManager $notificationManager, IClientService $clientService, ) { $this->client = $clientService->newClient(); @@ -252,8 +255,21 @@ public function isQuotaExceeded(?string $userId, int $type): bool { $this->logger->warning('Could not retrieve quota usage for user: ' . $userId . ' and quota type: ' . $type . '. Error: ' . $e->getMessage()); throw new Exception('Could not retrieve quota usage.', Http::STATUS_INTERNAL_SERVER_ERROR); } - - return $quotaUsage >= $quota; + if ($quotaUsage < $quota) { + return false; + } + $cache = $this->cacheFactory->createLocal(Application::APP_ID); + if ($cache->get('quota_exceeded_' . $userId . '_' . $type) === null) { + $notification = $this->notificationManager->createNotification(); + $notification->setApp(Application::APP_ID) + ->setUser($userId) + ->setDateTime(new DateTime()) + ->setObject('quota_exceeded', (string)$type) + ->setSubject('quota_exceeded', ['type' => $type]); + $this->notificationManager->notify($notification); + $cache->set('quota_exceeded_' . $userId . '_' . $type, true, 3600); + } + return true; } /** diff --git a/tests/unit/Providers/OpenAiProviderTest.php b/tests/unit/Providers/OpenAiProviderTest.php index f746d860..bdec3f6e 100644 --- a/tests/unit/Providers/OpenAiProviderTest.php +++ b/tests/unit/Providers/OpenAiProviderTest.php @@ -56,7 +56,7 @@ public static function setUpBeforeClass(): void { parent::setUpBeforeClass(); $backend = new Dummy(); $backend->createUser(self::TEST_USER1, self::TEST_USER1); - \OC::$server->get(\OCP\IUserManager::class)->registerBackend($backend); + \OCP\Server::get(\OCP\IUserManager::class)->registerBackend($backend); } protected function setUp(): void { @@ -64,11 +64,11 @@ protected function setUp(): void { $this->loginAsUser(self::TEST_USER1); - $this->openAiSettingsService = \OC::$server->get(OpenAiSettingsService::class); + $this->openAiSettingsService = \OCP\Server::get(OpenAiSettingsService::class); - $this->chunkService = \OC::$server->get(ChunkService::class); + $this->chunkService = \OCP\Server::get(ChunkService::class); - $this->quotaUsageMapper = \OC::$server->get(QuotaUsageMapper::class); + $this->quotaUsageMapper = \OCP\Server::get(QuotaUsageMapper::class); // We'll hijack the client service and subsequently iClient to return a mock response from the OpenAI API $clientService = $this->createMock(IClientService::class); @@ -76,12 +76,13 @@ protected function setUp(): void { $clientService->method('newClient')->willReturn($this->iClient); $this->openAiApiService = new OpenAiAPIService( - \OC::$server->get(\Psr\Log\LoggerInterface::class), + \OCP\Server::get(\Psr\Log\LoggerInterface::class), $this->createMock(\OCP\IL10N::class), - \OC::$server->get(IAppConfig::class), - \OC::$server->get(ICacheFactory::class), - \OC::$server->get(QuotaUsageMapper::class), + \OCP\Server::get(IAppConfig::class), + \OCP\Server::get(ICacheFactory::class), + \OCP\Server::get(QuotaUsageMapper::class), $this->openAiSettingsService, + $this->createMock(\OCP\Notification\IManager::class), $clientService, ); @@ -90,7 +91,7 @@ protected function setUp(): void { public static function tearDownAfterClass(): void { // Delete quota usage for test user - $quotaUsageMapper = \OC::$server->get(QuotaUsageMapper::class); + $quotaUsageMapper = \OCP\Server::get(QuotaUsageMapper::class); try { $quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER1); } catch (\OCP\Db\Exception|\RuntimeException|\Exception|\Throwable $e) { @@ -99,7 +100,7 @@ public static function tearDownAfterClass(): void { $backend = new \Test\Util\User\Dummy(); $backend->deleteUser(self::TEST_USER1); - \OC::$server->get(\OCP\IUserManager::class)->removeBackend($backend); + \OCP\Server::get(\OCP\IUserManager::class)->removeBackend($backend); parent::tearDownAfterClass(); } @@ -107,7 +108,7 @@ public static function tearDownAfterClass(): void { public function testFreePromptProvider(): void { $freePromptProvider = new TextToTextProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), self::TEST_USER1, @@ -170,7 +171,7 @@ public function testFreePromptProvider(): void { public function testEmojiProvider(): void { $emojiProvider = new EmojiProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), self::TEST_USER1, @@ -235,7 +236,7 @@ public function testEmojiProvider(): void { public function testHeadlineProvider(): void { $headlineProvider = new HeadlineProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), self::TEST_USER1, @@ -299,7 +300,7 @@ public function testHeadlineProvider(): void { public function testChangeToneProvider(): void { $changeToneProvider = new ChangeToneProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), $this->chunkService, @@ -366,7 +367,7 @@ public function testChangeToneProvider(): void { public function testSummaryProvider(): void { $summaryProvider = new SummaryProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), $this->chunkService, @@ -435,7 +436,7 @@ public function testSummaryProvider(): void { public function testProofreadProvider(): void { $proofreadProvider = new ProofreadProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), $this->chunkService, @@ -503,10 +504,10 @@ public function testProofreadProvider(): void { public function testTranslationProvider(): void { $translationProvider = new TranslateProvider( $this->openAiApiService, - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), $this->openAiSettingsService, $this->createMock(\OCP\IL10N::class), - \OC::$server->get(\OCP\L10N\IFactory::class), + \OCP\Server::get(\OCP\L10N\IFactory::class), $this->createMock(\OCP\ICacheFactory::class), $this->createMock(\Psr\Log\LoggerInterface::class), $this->chunkService, @@ -576,7 +577,7 @@ public function testTextToSpeechProvider(): void { $this->openAiApiService, $this->createMock(\OCP\IL10N::class), $this->createMock(\Psr\Log\LoggerInterface::class), - \OC::$server->get(IAppConfig::class), + \OCP\Server::get(IAppConfig::class), self::TEST_USER1, ); diff --git a/tests/unit/Quota/QuotaTest.php b/tests/unit/Quota/QuotaTest.php new file mode 100644 index 00000000..9545e4d3 --- /dev/null +++ b/tests/unit/Quota/QuotaTest.php @@ -0,0 +1,123 @@ +createUser(self::TEST_USER1, self::TEST_USER1); + \OCP\Server::get(IUserManager::class)->registerBackend($backend); + } + + protected function setUp(): void { + parent::setUp(); + + $this->loginAsUser(self::TEST_USER1); + + $this->openAiSettingsService = \OCP\Server::get(OpenAiSettingsService::class); + + $this->quotaUsageMapper = \OCP\Server::get(QuotaUsageMapper::class); + + $this->notificationManager = $this->createMock(IManager::class); + + $this->cacheFactory = $this->createMock(ICacheFactory::class); + + + $this->openAiApiService = new OpenAiAPIService( + \OCP\Server::get(LoggerInterface::class), + $this->createMock(IL10N::class), + \OCP\Server::get(IAppConfig::class), + $this->cacheFactory, + \OCP\Server::get(QuotaUsageMapper::class), + $this->openAiSettingsService, + $this->notificationManager, + \OCP\Server::get(IClientService::class), + ); + } + + public static function tearDownAfterClass(): void { + // Delete quota usage for test user + $quotaUsageMapper = \OCP\Server::get(QuotaUsageMapper::class); + try { + $quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER1); + } catch (\OCP\Db\Exception|RuntimeException|Exception|Throwable $e) { + // Ignore + } + + $backend = new Dummy(); + $backend->deleteUser(self::TEST_USER1); + \OCP\Server::get(IUserManager::class)->removeBackend($backend); + + parent::tearDownAfterClass(); + } + + public function testNotification(): void { + $this->openAiSettingsService->setQuotas([1, 1, 1, 1]); + $cache = $this->createMock(ICache::class); + $this->cacheFactory->method('createLocal')->willReturn($cache); + $key = 'quota_exceeded_' . self::TEST_USER1 . '_' . Application::QUOTA_TYPE_TEXT; + + $cache->expects($this->any())->method('get')->with($key)->willReturn($this->onConsecutiveCalls(null, true, true)); + $cache->expects($this->once())->method('set')->with($key, true, 3600); + + $this->notificationManager->expects($this->once())->method('notify'); + $this->assertFalse($this->openAiApiService->isQuotaExceeded(self::TEST_USER1, Application::QUOTA_TYPE_TEXT)); + $this->quotaUsageMapper->createQuotaUsage(self::TEST_USER1, Application::QUOTA_TYPE_TEXT, 100); + + // Send notification + $this->assertTrue($this->openAiApiService->isQuotaExceeded(self::TEST_USER1, Application::QUOTA_TYPE_TEXT)); + // Try again to make sure a notification is only sent once + $this->assertTrue($this->openAiApiService->isQuotaExceeded(self::TEST_USER1, Application::QUOTA_TYPE_TEXT)); + // Clear quota usage + $this->quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER1); + } + +}