diff --git a/appinfo/info.xml b/appinfo/info.xml index 440f869e..3fe1f4d9 100644 --- a/appinfo/info.xml +++ b/appinfo/info.xml @@ -101,7 +101,7 @@ Negative: Learn more about the Nextcloud Ethical AI Rating [in our blog](https://nextcloud.com/blog/nextcloud-ethical-ai-rating/). ]]> - 3.7.1 + 3.8.0 agpl Julien Veyssier OpenAi diff --git a/appinfo/routes.php b/appinfo/routes.php index ce8e4355..beed6668 100644 --- a/appinfo/routes.php +++ b/appinfo/routes.php @@ -16,5 +16,9 @@ ['name' => 'openAiAPI#getModels', 'url' => '/models', 'verb' => 'GET'], ['name' => 'openAiAPI#getUserQuotaInfo', 'url' => '/quota-info', 'verb' => 'GET'], ['name' => 'openAiAPI#getAdminQuotaInfo', 'url' => '/admin-quota-info', 'verb' => 'GET'], + + ['name' => 'quotaRule#addRule', 'url' => '/quota/rule', 'verb' => 'POST'], + ['name' => 'quotaRule#updateRule', 'url' => '/quota/rule', 'verb' => 'PUT'], + ['name' => 'quotaRule#deleteRule', 'url' => '/quota/rule', 'verb' => 'DELETE'], ], ]; diff --git a/lib/AppInfo/Application.php b/lib/AppInfo/Application.php index 9087c24d..5774e419 100644 --- a/lib/AppInfo/Application.php +++ b/lib/AppInfo/Application.php @@ -77,6 +77,7 @@ class Application extends App implements IBootstrap { ]; public const MODELS_CACHE_KEY = 'models'; + public const QUOTA_RULES_CACHE_PREFIX = 'quota_rules'; public const MODELS_CACHE_TTL = 60 * 30; private IAppConfig $appConfig; diff --git a/lib/Controller/QuotaRuleController.php b/lib/Controller/QuotaRuleController.php new file mode 100644 index 00000000..d830c924 --- /dev/null +++ b/lib/Controller/QuotaRuleController.php @@ -0,0 +1,82 @@ +quotaRuleService->addRule(); + return new DataResponse($result); + } catch (Exception $e) { + return new DataResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST); + } + } + + /** + * PUT /rule + * @param int $id + * @param array $rule expects: type, amount, priority, pool, entities[] + * @return DataResponse + */ + public function updateRule(int $id, array $rule): DataResponse { + if (!isset($rule['type']) || !is_int($rule['type'])) { + return new DataResponse(['error' => 'Missing or invalid type'], Http::STATUS_BAD_REQUEST); + } + if (!isset($rule['amount']) || !is_int($rule['amount'])) { + return new DataResponse(['error' => 'Missing or invalid amount'], Http::STATUS_BAD_REQUEST); + } + if (!isset($rule['priority']) || !is_int($rule['priority'])) { + return new DataResponse(['error' => 'Missing or invalid priority'], Http::STATUS_BAD_REQUEST); + } + if (!isset($rule['pool']) || !is_bool($rule['pool'])) { + return new DataResponse(['error' => 'Missing or invalid pool value'], Http::STATUS_BAD_REQUEST); + } + if (!isset($rule['entities']) || !is_array($rule['entities'])) { + return new DataResponse(['error' => 'Missing or invalid entities'], Http::STATUS_BAD_REQUEST); + } + try { + $result = $this->quotaRuleService->updateRule($id, $rule); + return new DataResponse($result); + } catch (Exception $e) { + return new DataResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST); + } + } + + /** + * DELETE /rule + * @param int $id + * @return DataResponse + */ + public function deleteRule(int $id): DataResponse { + try { + $this->quotaRuleService->deleteRule($id); + return new DataResponse(''); + } catch (Exception $e) { + return new DataResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST); + } + } +} diff --git a/lib/Db/QuotaRule.php b/lib/Db/QuotaRule.php new file mode 100644 index 00000000..584a7b03 --- /dev/null +++ b/lib/Db/QuotaRule.php @@ -0,0 +1,54 @@ +addType('type', Types::INTEGER); + $this->addType('amount', Types::INTEGER); + $this->addType('priority', Types::INTEGER); + $this->addType('pool', Types::INTEGER); + } + + #[ReturnTypeWillChange] + public function jsonSerialize() { + return [ + 'id' => $this->getId(), + 'type' => $this->getType(), + 'amount' => $this->getAmount(), + 'priority' => $this->getPriority(), + 'pool' => $this->getPool() + ]; + } +} diff --git a/lib/Db/QuotaRuleMapper.php b/lib/Db/QuotaRuleMapper.php new file mode 100644 index 00000000..883e62b1 --- /dev/null +++ b/lib/Db/QuotaRuleMapper.php @@ -0,0 +1,131 @@ + + */ +class QuotaRuleMapper extends QBMapper { + public function __construct( + IDBConnection $db, + ) { + parent::__construct($db, 'openai_quota_rule', QuotaRule::class); + } + + /** + * @return array + * @throws Exception + */ + public function getRules(): array { + $qb = $this->db->getQueryBuilder(); + + $qb->select('*') + ->from($this->getTableName()); + + return $this->findEntities($qb); + } + + /** + * @param int $quotaType + * @param string $userId + * @param array $groups + * @return QuotaRule + * @throws DoesNotExistException + * @throws Exception + * @throws MultipleObjectsReturnedException + */ + public function getRule(int $quotaType, string $userId, array $groups): QuotaRule { + $qb = $this->db->getQueryBuilder(); + + $qb->select('r.*') + ->from($this->getTableName(), 'r') + ->leftJoin('r', 'openai_quota_user', 'u', 'r.id = u.rule_id') + ->where( + $qb->expr()->eq('r.type', $qb->createNamedParameter($quotaType, IQueryBuilder::PARAM_INT)) + )->andWhere( + $qb->expr()->orX( + $qb->expr()->andX( + $qb->expr()->eq('u.entity_type', $qb->createNamedParameter(EntityType::USER->value, IQueryBuilder::PARAM_INT)), + $qb->expr()->eq('u.entity_id', $qb->createNamedParameter($userId, IQueryBuilder::PARAM_STR)) + ), + $qb->expr()->andX( + $qb->expr()->eq('u.entity_type', $qb->createNamedParameter(EntityType::GROUP->value, IQueryBuilder::PARAM_INT)), + $qb->expr()->in('u.entity_id', $qb->createNamedParameter($groups, IQueryBuilder::PARAM_STR_ARRAY)) + ), + + ) + )->orderBy('r.priority', 'ASC') + ->setMaxResults(1); + /** @var QuotaRule $entity */ + $entity = $this->findEntity($qb); + return $entity; + } + /** + * @param int $quotaType + * @param int $amount + * @param int $priority + * @param int $pool + * @return int + * @throws Exception + */ + public function addRule(int $quotaType, int $amount, int $priority, int $pool): int { + $qb = $this->db->getQueryBuilder(); + + $qb->insert($this->getTableName()) + ->values( + [ + 'type' => $qb->createNamedParameter($quotaType, IQueryBuilder::PARAM_INT), + 'amount' => $qb->createNamedParameter($amount, IQueryBuilder::PARAM_INT), + 'priority' => $qb->createNamedParameter($priority, IQueryBuilder::PARAM_INT), + 'pool' => $qb->createNamedParameter($pool, IQueryBuilder::PARAM_INT) + ] + ); + $qb->executeStatement(); + return $qb->getLastInsertId(); + } + /** + * @param int $id + * @param int $quotaType + * @param int $amount + * @param int $priority + * @param int $pool + * @return void + * @throws Exception + */ + public function updateRule(int $id, int $quotaType, int $amount, int $priority, int $pool): void { + $qb = $this->db->getQueryBuilder(); + + $qb->update($this->getTableName()) + ->set('type', $qb->createNamedParameter($quotaType, IQueryBuilder::PARAM_INT)) + ->set('amount', $qb->createNamedParameter($amount, IQueryBuilder::PARAM_INT)) + ->set('priority', $qb->createNamedParameter($priority, IQueryBuilder::PARAM_INT)) + ->set('pool', $qb->createNamedParameter($pool, IQueryBuilder::PARAM_INT)) + ->where( + $qb->expr()->eq('id', $qb->createNamedParameter($id, IQueryBuilder::PARAM_INT)) + ); + $qb->executeStatement(); + } + /** + * @param int $id + * @throws Exception + */ + public function deleteRule(int $id): void { + $qb = $this->db->getQueryBuilder(); + + $qb->delete($this->getTableName()) + ->where($qb->expr()->eq('id', $qb->createNamedParameter($id, IQueryBuilder::PARAM_INT))); + $qb->executeStatement(); + } +} diff --git a/lib/Db/QuotaUsage.php b/lib/Db/QuotaUsage.php index 80a9cbe7..64680b50 100644 --- a/lib/Db/QuotaUsage.php +++ b/lib/Db/QuotaUsage.php @@ -9,7 +9,10 @@ namespace OCA\OpenAi\Db; +use JsonSerializable; use OCP\AppFramework\Db\Entity; +use OCP\DB\Types; +use ReturnTypeWillChange; /** * @method string getUserId() @@ -20,8 +23,10 @@ * @method void setUnits(int $units) * @method int getTimestamp() * @method void setTimestamp(int $timestamp) + * @method int getPool() + * @method void setPool(int $pool) */ -class QuotaUsage extends Entity implements \JsonSerializable { +class QuotaUsage extends Entity implements JsonSerializable { /** @var string */ protected $userId; /** @var int */ @@ -30,22 +35,26 @@ class QuotaUsage extends Entity implements \JsonSerializable { protected $units; /** @var int */ protected $timestamp; + /** @var int */ + protected $pool; public function __construct() { - $this->addType('user_id', 'string'); - $this->addType('type', 'integer'); - $this->addType('units', 'integer'); - $this->addType('timestamp', 'integer'); + $this->addType('user_id', Types::STRING); + $this->addType('type', Types::INTEGER); + $this->addType('units', Types::INTEGER); + $this->addType('timestamp', Types::INTEGER); + $this->addType('pool', Types::INTEGER); } - #[\ReturnTypeWillChange] + #[ReturnTypeWillChange] public function jsonSerialize() { return [ - 'id' => $this->id, - 'user_id' => $this->userId, - 'type' => $this->type, - 'units' => $this->units, - 'timestamp' => $this->timestamp, + 'id' => $this->getId(), + 'user_id' => $this->getUserId(), + 'type' => $this->getType(), + 'units' => $this->getUnits(), + 'timestamp' => $this->getTimestamp(), + 'pool' => $this->getPool() ]; } } diff --git a/lib/Db/QuotaUsageMapper.php b/lib/Db/QuotaUsageMapper.php index 7203cf2d..926815c6 100644 --- a/lib/Db/QuotaUsageMapper.php +++ b/lib/Db/QuotaUsageMapper.php @@ -17,6 +17,7 @@ use OCP\DB\Exception; use OCP\DB\QueryBuilder\IQueryBuilder; use OCP\IDBConnection; +use RuntimeException; /** * @extends QBMapper @@ -101,13 +102,14 @@ public function getQuotaUnitsInTimePeriod(int $type, int $periodStart): int { * @param string $userId * @param int $type Type of the quota * @param int $periodStart Start time of quota + * @param int|null $pool * @return int * @throws DoesNotExistException * @throws Exception * @throws MultipleObjectsReturnedException - * @throws \RuntimeException + * @throws RuntimeException */ - public function getQuotaUnitsOfUserInTimePeriod(string $userId, int $type, int $periodStart): int { + public function getQuotaUnitsOfUserInTimePeriod(string $userId, int $type, int $periodStart, ?int $pool = null): int { $qb = $this->db->getQueryBuilder(); // Get the sum of the units used in the time period @@ -116,12 +118,18 @@ public function getQuotaUnitsOfUserInTimePeriod(string $userId, int $type, int $ ->where( $qb->expr()->eq('type', $qb->createNamedParameter($type, IQueryBuilder::PARAM_INT)) ) - ->andWhere( - $qb->expr()->eq('user_id', $qb->createNamedParameter($userId, IQueryBuilder::PARAM_STR)) - ) ->andWhere( $qb->expr()->gt('timestamp', $qb->createNamedParameter($periodStart, IQueryBuilder::PARAM_INT)) ); + if ($pool === null) { + $qb->andWhere( + $qb->expr()->eq('user_id', $qb->createNamedParameter($userId, IQueryBuilder::PARAM_STR)) + ); + } else { + $qb->andWhere( + $qb->expr()->eq('pool', $qb->createNamedParameter($pool, IQueryBuilder::PARAM_INT)) + ); + } // Execute the query and return the result $result = (int)$qb->executeQuery()->fetchOne(); @@ -181,15 +189,17 @@ public function getQuotaUnitsOfUser(string $userId, int $type): int { * @param string $userId * @param int $type * @param int $units + * @param int $pool * @return QuotaUsage * @throws Exception */ - public function createQuotaUsage(string $userId, int $type, int $units): QuotaUsage { + public function createQuotaUsage(string $userId, int $type, int $units, int $pool = -1): QuotaUsage { $quotaUsage = new QuotaUsage(); $quotaUsage->setUserId($userId); $quotaUsage->setType($type); $quotaUsage->setUnits($units); + $quotaUsage->setPool($pool); $quotaUsage->setTimestamp((new DateTime())->getTimestamp()); $insertedQuotaUsage = $this->insert($quotaUsage); diff --git a/lib/Db/QuotaUser.php b/lib/Db/QuotaUser.php new file mode 100644 index 00000000..e3ba1242 --- /dev/null +++ b/lib/Db/QuotaUser.php @@ -0,0 +1,46 @@ +addType('rule_id', Types::INTEGER); + $this->addType('entity_type', Types::INTEGER); + $this->addType('entity_id', Types::STRING); + } + + #[ReturnTypeWillChange] + public function jsonSerialize() { + return [ + 'id' => $this->getId(), + 'rule_id' => $this->getRuleId(), + 'entity_type' => $this->getEntityType(), + 'entity_id' => $this->getEntityId() + ]; + } +} diff --git a/lib/Db/QuotaUserMapper.php b/lib/Db/QuotaUserMapper.php new file mode 100644 index 00000000..858b0622 --- /dev/null +++ b/lib/Db/QuotaUserMapper.php @@ -0,0 +1,93 @@ + + */ +class QuotaUserMapper extends QBMapper { + public function __construct(IDBConnection $db) { + parent::__construct($db, 'openai_quota_user', QuotaUser::class); + } + /** + * @param int $ruleId + * @return array + * @throws Exception + */ + public function getUsers(int $ruleId): array { + $qb = $this->db->getQueryBuilder(); + + $qb->select('*') + ->from($this->getTableName()) + ->where($qb->expr()->eq('rule_id', $qb->createNamedParameter($ruleId, IQueryBuilder::PARAM_INT))); + return $this->findEntities($qb); + } + /** + * @param int $ruleId + * @param array $users + * @throws Exception + */ + public function setUsers(int $ruleId, array $users): void { + $qb = $this->db->getQueryBuilder(); + + $qb->select('*') + ->from($this->getTableName()) + ->where($qb->expr()->eq('rule_id', $qb->createNamedParameter($ruleId, IQueryBuilder::PARAM_INT))); + $this->db->beginTransaction(); + try { + $oldUsers = $this->findEntities($qb); + $oldUsersById = array_reduce($oldUsers, function (array $carry, QuotaUser $oldUser) { + $carry[$oldUser->getEntityType() . '-' . $oldUser->getEntityId()] = $oldUser; + return $carry; + }, []); + $usersById = []; + // Add users that are in the new list but not in the old list + foreach ($users as $user) { + if (!isset($oldUsersById[$user['entity_type'] . '-' . $user['entity_id']])) { + $newUser = new QuotaUser(); + $newUser->setRuleId($ruleId); + $newUser->setEntityType($user['entity_type']); + $newUser->setEntityId($user['entity_id']); + $this->insert($newUser); + } + $usersById[$user['entity_type'] . '-' . $user['entity_id']] = $user; + } + // Delete users that are not in the new list but are in the old list + foreach ($oldUsers as $oldUser) { + if (!isset($usersById[$oldUser->getEntityType() . '-' . $oldUser->getEntityId()])) { + $this->delete($oldUser); + } + } + $this->db->commit(); + } catch (\Throwable $e) { + $this->db->rollBack(); + throw $e; + } + } + + /** + * @param int $ruleId + * @throws Exception + */ + public function deleteByRuleId(int $ruleId): void { + $qb = $this->db->getQueryBuilder(); + $qb->delete($this->getTableName()) + ->where($qb->expr()->eq('rule_id', $qb->createNamedParameter($ruleId, IQueryBuilder::PARAM_INT))); + $qb->executeStatement(); + } +} diff --git a/lib/Migration/Version030800Date20250812122830.php b/lib/Migration/Version030800Date20250812122830.php new file mode 100644 index 00000000..8d79ff1c --- /dev/null +++ b/lib/Migration/Version030800Date20250812122830.php @@ -0,0 +1,88 @@ +hasTable('openai_quota_rule')) { + $table = $schema->createTable('openai_quota_rule'); + $table->addColumn('id', Types::BIGINT, [ + 'autoincrement' => true, + 'notnull' => true, + ]); + $table->addColumn('type', Types::INTEGER, [ + 'notnull' => true, + ]); + $table->addColumn('amount', Types::BIGINT, [ + 'notnull' => true + ]); + $table->addColumn('priority', Types::INTEGER, [ + 'notnull' => true + ]); + $table->addColumn('pool', Types::INTEGER, [ + 'notnull' => true + ]); + $table->setPrimaryKey(['id']); + $table->addIndex(['type'], 'oai_rule_type'); + } + if (!$schema->hasTable('openai_quota_user')) { + $table = $schema->createTable('openai_quota_user'); + $table->addColumn('id', Types::BIGINT, [ + 'autoincrement' => true, + 'notnull' => true, + ]); + $table->addColumn('rule_id', Types::BIGINT, [ + 'notnull' => true, + ]); + $table->addColumn('entity_type', Types::INTEGER, [ + 'notnull' => true, + ]); + $table->addColumn('entity_id', Types::STRING, [ + 'notnull' => true, + 'length' => 300, + ]); + $table->setPrimaryKey(['id']); + $table->addIndex(['rule_id'], 'oai_rule_id'); + $table->addIndex(['entity_id', 'entity_type'], 'oai_user_id_type'); + } + if ($schema->hasTable('openai_quota_usage')) { + $table = $schema->getTable('openai_quota_usage'); + if (!$table->hasColumn('pool')) { + $table->addColumn('pool', Types::BIGINT, [ + 'notnull' => true, + 'default' => -1 + ]); + $table->addIndex(['pool'], 'oai_usage_pool'); + } + } + + return $schema; + } +} diff --git a/lib/Service/OpenAiAPIService.php b/lib/Service/OpenAiAPIService.php index 283eef99..c08cf038 100644 --- a/lib/Service/OpenAiAPIService.php +++ b/lib/Service/OpenAiAPIService.php @@ -49,11 +49,23 @@ public function __construct( private QuotaUsageMapper $quotaUsageMapper, private OpenAiSettingsService $openAiSettingsService, private INotificationManager $notificationManager, + private QuotaRuleService $quotaRuleService, IClientService $clientService, ) { $this->client = $clientService->newClient(); } + /** + * @param string $userId + * @param int $type + * @param int $usage + * @throws Exception If there is an error creating the quota usage. + */ + public function createQuotaUsage(string $userId, int $type, int $usage) { + $rule = $this->quotaRuleService->getRule($type, $userId); + $this->quotaUsageMapper->createQuotaUsage($userId, $type, $usage, $rule['pool'] ? $rule['id'] : -1); + } + /** * @return bool */ @@ -238,9 +250,9 @@ public function isQuotaExceeded(?string $userId, int $type): bool { // User has specified own OpenAI API key, no quota limit: return false; } - - // Get quota limits - $quota = $this->openAiSettingsService->getQuotas()[$type]; + $rule = $this->quotaRuleService->getRule($type, $userId); + $quota = $rule['amount']; + $pool = $rule['pool'] ? $rule['id'] : null; if ($quota === 0) { // Unlimited quota: @@ -250,7 +262,7 @@ public function isQuotaExceeded(?string $userId, int $type): bool { $quotaStart = $this->openAiSettingsService->getQuotaStart(); try { - $quotaUsage = $this->quotaUsageMapper->getQuotaUnitsOfUserInTimePeriod($userId, $type, $quotaStart); + $quotaUsage = $this->quotaUsageMapper->getQuotaUnitsOfUserInTimePeriod($userId, $type, $quotaStart, $pool); } catch (DoesNotExistException|MultipleObjectsReturnedException|DBException|RuntimeException $e) { $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); @@ -319,7 +331,7 @@ public function translatedQuotaUnit(int $type): string { */ public function getUserQuotaInfo(string $userId): array { // Get quota limits (if the user has specified an own OpenAI API key, no quota limit, just supply default values as fillers) - $quotas = $this->hasOwnOpenAiApiKey($userId) ? Application::DEFAULT_QUOTAS : $this->openAiSettingsService->getQuotas(); + $ownApikey = $this->hasOwnOpenAiApiKey($userId); // Get quota period $quotaPeriod = $this->openAiSettingsService->getQuotaPeriod(); $quotaStart = $this->openAiSettingsService->getQuotaStart(); @@ -334,7 +346,15 @@ public function getUserQuotaInfo(string $userId): array { $this->logger->warning('Could not retrieve quota usage for user: ' . $userId . ' and quota type: ' . $quotaType . '. Error: ' . $e->getMessage(), ['app' => Application::APP_ID]); throw new Exception($this->l10n->t('Unknown error while retrieving quota usage.'), Http::STATUS_INTERNAL_SERVER_ERROR); } - $quotaInfo[$quotaType]['limit'] = intval($quotas[$quotaType]); + if ($ownApikey) { + $quotaInfo[$quotaType]['limit'] = Application::DEFAULT_QUOTAS[$quotaType]; + } else { + $rule = $this->quotaRuleService->getRule($quotaType, $userId); + $quotaInfo[$quotaType]['limit'] = $rule['amount']; + if ($rule['pool']) { + $quotaInfo[$quotaType]['used_pool'] = $this->quotaUsageMapper->getQuotaUnitsOfUserInTimePeriod($userId, $quotaType, $quotaStart, $rule['id']); + } + } $quotaInfo[$quotaType]['unit'] = $this->translatedQuotaUnit($quotaType); } @@ -426,7 +446,7 @@ public function createCompletion( if (isset($response['usage'], $response['usage']['total_tokens'])) { $usage = $response['usage']['total_tokens']; try { - $this->quotaUsageMapper->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TEXT, $usage); + $this->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TEXT, $usage); } catch (DBException $e) { $this->logger->warning('Could not create quota usage for user: ' . $userId . ' and quota type: ' . Application::QUOTA_TYPE_TEXT . '. Error: ' . $e->getMessage(), ['app' => Application::APP_ID]); } @@ -597,7 +617,7 @@ public function createChatCompletion( if (isset($response['usage'], $response['usage']['total_tokens'])) { $usage = $response['usage']['total_tokens']; try { - $this->quotaUsageMapper->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TEXT, $usage); + $this->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TEXT, $usage); } catch (DBException $e) { $this->logger->warning('Could not create quota usage for user: ' . $userId . ' and quota type: ' . Application::QUOTA_TYPE_TEXT . '. Error: ' . $e->getMessage(), ['app' => Application::APP_ID]); } @@ -745,7 +765,7 @@ public function transcribe( $audioDuration = intval(round(floatval(array_pop($response['segments'])['end']))); try { - $this->quotaUsageMapper->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TRANSCRIPTION, $audioDuration); + $this->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_TRANSCRIPTION, $audioDuration); } catch (DBException $e) { $this->logger->warning('Could not create quota usage for user: ' . $userId . ' and quota type: ' . Application::QUOTA_TYPE_TRANSCRIPTION . '. Error: ' . $e->getMessage(), ['app' => Application::APP_ID]); } @@ -784,7 +804,7 @@ public function requestImageCreation( } else { try { - $this->quotaUsageMapper->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_IMAGE, $n); + $this->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_IMAGE, $n); } catch (DBException $e) { $this->logger->warning('Could not create quota usage for user: ' . $userId . ' and quota type: ' . Application::QUOTA_TYPE_IMAGE . '. Error: ' . $e->getMessage(), ['app' => Application::APP_ID]); } @@ -849,7 +869,7 @@ public function requestSpeechCreation( try { $charCount = mb_strlen($prompt); - $this->quotaUsageMapper->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_SPEECH, $charCount); + $this->createQuotaUsage($userId ?? '', Application::QUOTA_TYPE_SPEECH, $charCount); } catch (DBException $e) { $this->logger->warning('Could not create quota usage for user: ' . $userId . ' and quota type: ' . Application::QUOTA_TYPE_SPEECH . '. Error: ' . $e->getMessage()); } diff --git a/lib/Service/OpenAiSettingsService.php b/lib/Service/OpenAiSettingsService.php index ec4ff028..ad2f14a7 100644 --- a/lib/Service/OpenAiSettingsService.php +++ b/lib/Service/OpenAiSettingsService.php @@ -518,6 +518,8 @@ public function setQuotas(array $quotas): void { } $this->appConfig->setValueString(Application::APP_ID, 'quotas', json_encode($quotas, JSON_THROW_ON_ERROR), lazy: true); + $cache = $this->cacheFactory->createDistributed(Application::APP_ID); + $cache->clear(Application::QUOTA_RULES_CACHE_PREFIX); } /** @@ -681,7 +683,9 @@ public function setQuotaPeriod(array $quotaPeriod): void { if (!isset($quotaPeriod['length']) || !is_int($quotaPeriod['length'])) { throw new Exception('Invalid quota period length'); } - $quotaPeriod['length'] = max(1, $quotaPeriod['length']); + if ($quotaPeriod['length'] < 1) { + throw new Exception('Invalid quota period length'); + } if (!isset($quotaPeriod['unit']) || !is_string($quotaPeriod['unit'])) { throw new Exception('Invalid quota period unit'); } @@ -690,8 +694,12 @@ public function setQuotaPeriod(array $quotaPeriod): void { if (!isset($quotaPeriod['day']) || !is_int($quotaPeriod['day'])) { throw new Exception('Invalid quota period day'); } - $quotaPeriod['day'] = max(1, $quotaPeriod['day']); - $quotaPeriod['day'] = min($quotaPeriod['day'], 28); + if ($quotaPeriod['day'] < 1) { + throw new Exception('Invalid quota period day'); + } + if ($quotaPeriod['day'] > 28) { + throw new Exception('Invalid quota period day'); + } } elseif ($quotaPeriod['unit'] !== 'day') { throw new Exception('Invalid quota period unit'); } diff --git a/lib/Service/QuotaRuleService.php b/lib/Service/QuotaRuleService.php new file mode 100644 index 00000000..67f6f163 --- /dev/null +++ b/lib/Service/QuotaRuleService.php @@ -0,0 +1,181 @@ +cacheFactory->createDistributed(Application::APP_ID); + $cacheKey = Application::QUOTA_RULES_CACHE_PREFIX . $quotaType . '-' . $userId; + $rule = $cache->get($cacheKey); + if ($rule === null) { + $user = $this->userManager->get($userId); + $groups = $this->groupManager->getUserGroupIds($user); + try { + $rule = $this->quotaRuleMapper->getRule($quotaType, $userId, $groups)->jsonSerialize(); + } catch (DoesNotExistException|MultipleObjectsReturnedException) { + $rule = [ + 'amount' => $this->openAiSettingsService->getQuotas()[$quotaType], + 'pool' => false, + 'id' => null, + ]; + } + $cache->set($cacheKey, $rule); + } + return $rule; + } + + /** + * Clears the cache for the quota rules + */ + public function clearCache(): void { + $cache = $this->cacheFactory->createDistributed(Application::APP_ID); + $cache->clear(Application::QUOTA_RULES_CACHE_PREFIX); + } + + /** + * @return array + * @throws Exception + */ + public function getRules(): array { + $rules = $this->quotaRuleMapper->getRules(); + $userManager = $this->userManager; + return array_map(function ($rule) use ($userManager) { + $entities = $this->quotaUserMapper->getUsers($rule->getId()); + $result = $rule->jsonSerialize(); + $result['entities'] = array_map(static function ($u) use ($userManager) { + $displayName = $u->getEntityId(); + if ($u->getEntityType() === EntityType::USER->value) { + $user = $userManager->get($u->getEntityId()); + $displayName = $user->getDisplayName(); + } + return [ + 'display_name' => $displayName, + 'entity_type' => $u->getEntityType(), + 'entity_id' => $u->getEntityId(), + 'id' => $u->getEntityType() . '-' . $u->getEntityId(), + ]; + }, $entities); + $result['pool'] = $result['pool'] === 1; + return $result; + }, $rules); + } + + /** + * @return array created rule with entities + * @throws Exception + */ + public function addRule(): array { + $id = $this->quotaRuleMapper->addRule(0, 0, 0, 0); + $this->clearCache(); + return [ + 'id' => $id, + 'type' => 0, + 'amount' => 0, + 'priority' => 0, + 'pool' => false, + 'entities' => [], + ]; + } + + /** + * @param int $id + * @param array $rule + * @return array updated rule with entities + * @throws Exception + */ + public function updateRule(int $id, array $rule): array { + $this->validateRuleBasics($rule); + $this->validateEntities($rule['entities']); + $this->quotaRuleMapper->updateRule($id, $rule['type'], $rule['amount'], $rule['priority'], $rule['pool']); + $this->quotaUserMapper->setUsers($id, $rule['entities']); + $rule['id'] = $id; + $this->clearCache(); + return $rule; + } + + /** + * @param int $id + * @throws Exception + */ + public function deleteRule(int $id): void { + $this->quotaUserMapper->deleteByRuleId($id); + $this->quotaRuleMapper->deleteRule($id); + $this->clearCache(); + } + + /** + * Validate the basic parts of a quota rule: type and amount + * + * @param array $rule with keys 'type' and 'amount' + * + * @throws Exception if the type is invalid or the amount is less than 0 + */ + private function validateRuleBasics(array $rule): void { + $validTypes = [ + Application::QUOTA_TYPE_TEXT, + Application::QUOTA_TYPE_IMAGE, + Application::QUOTA_TYPE_TRANSCRIPTION, + Application::QUOTA_TYPE_SPEECH, + ]; + if (!in_array($rule['type'], $validTypes, true)) { + throw new Exception('Invalid quota type'); + } + if ($rule['amount'] < 0) { + throw new Exception('Amount must be >= 0'); + } + } + + /** + * Validate the entities of a quota rule + * + * @param array $entities contains each entity as an array with keys 'entity_type' and 'entity_id' + * + * @throws Exception if an entity is invalid + */ + private function validateEntities(array $entities) { + foreach ($entities as $e) { + if (!is_array($e)) { + $this->logger->warning('Invalid entity', $e); + throw new Exception('Invalid entity'); + } + if (!isset($e['entity_type'], $e['entity_id']) || EntityType::tryFrom($e['entity_type']) === null || !is_string($e['entity_id'])) { + $this->logger->warning('Invalid entity', $e); + throw new Exception('Invalid entity'); + } + } + } +} diff --git a/lib/Settings/Admin.php b/lib/Settings/Admin.php index 6f1a89c5..0655bb50 100644 --- a/lib/Settings/Admin.php +++ b/lib/Settings/Admin.php @@ -9,6 +9,7 @@ use OCA\OpenAi\AppInfo\Application; use OCA\OpenAi\Service\OpenAiSettingsService; +use OCA\OpenAi\Service\QuotaRuleService; use OCP\App\IAppManager; use OCP\AppFramework\Http\TemplateResponse; use OCP\AppFramework\Services\IInitialState; @@ -18,6 +19,7 @@ class Admin implements ISettings { public function __construct( private IInitialState $initialStateService, private OpenAiSettingsService $openAiSettingsService, + private QuotaRuleService $quotaRuleService, private IAppManager $appManager, ) { } @@ -32,6 +34,8 @@ public function getForm(): TemplateResponse { $isAssistantEnabled = $this->appManager->isEnabledForUser('assistant'); $adminConfig['assistant_enabled'] = $isAssistantEnabled; $this->initialStateService->provideInitialState('admin-config', $adminConfig); + $rules = $this->quotaRuleService->getRules(); + $this->initialStateService->provideInitialState('rules', $rules); return new TemplateResponse(Application::APP_ID, 'adminSettings'); } diff --git a/src/components/AdminSettings.vue b/src/components/AdminSettings.vue index be5cfc79..ae5e455a 100644 --- a/src/components/AdminSettings.vue +++ b/src/components/AdminSettings.vue @@ -495,7 +495,7 @@ input-id="openai-tts-voices-select" @click="onInput()" />
-

+

{{ t('integration_openai', 'Usage limits') }}

@@ -508,7 +508,7 @@ {{ t('integration_openai', 'Usage quotas per time period') }} - {{ t('integration_openai', 'A per-user quota for each quota type can be set. If the user has not provided their own API key, this quota will be enforced.') }} + {{ t('integration_openai', 'A per-user quota for each quota type can be set. If the user has not provided their own API key, and a rule is not specified for this user or any of their groups, this quota will be enforced.') }} {{ t('integration_openai', '"0" means unlimited usage for a particular quota type.') }} @@ -546,6 +546,8 @@ +

{{ t('integration_openai', 'Quota Rules') }}

+

@@ -593,6 +595,7 @@ import EarthIcon from 'vue-material-design-icons/Earth.vue' import HelpCircleOutlineIcon from 'vue-material-design-icons/HelpCircleOutline.vue' import KeyOutlineIcon from 'vue-material-design-icons/KeyOutline.vue' import TimerAlertOutlineIcon from 'vue-material-design-icons/TimerAlertOutline.vue' +import QuotaRules from './Rules/QuotaRules.vue' import OpenAiIcon from './icons/OpenAiIcon.vue' @@ -631,6 +634,7 @@ export default { NcTextField, NcInputField, NcNoteCard, + QuotaRules, }, data() { @@ -980,4 +984,7 @@ export default { margin: 0 !important; } } +.notecard { + max-width: 900px; +} diff --git a/src/components/PersonalSettings.vue b/src/components/PersonalSettings.vue index 6fd4c7e2..95a86f75 100644 --- a/src/components/PersonalSettings.vue +++ b/src/components/PersonalSettings.vue @@ -87,6 +87,9 @@

{{ t('integration_openai', 'Usage quota info') }}

+ + {{ t('integration_openai', 'If you see a shared quota usage of 50% and a usage of 10% that means that you have used 10% of the total shared quota, and the sum of all other users affected by this quota is 40%.') }} + @@ -95,6 +98,9 @@ {{ t('integration_openai', 'Quota type') }} + @@ -106,6 +112,12 @@ + +
{{ t('integration_openai', 'Usage') }} + {{ t('integration_openai', 'Shared Usage') }} +
{{ quota.used + ' ' + quota.unit }} + {{ quota.limit > 0 ? Math.round(quota.used_pool / quota.limit * 100) + ' %' : quota.used_pool + ' ' + quota.unit }} + + {{ t('integration_openai', 'Not Shared') }} +
@@ -175,6 +187,9 @@ export default { : n('integration_openai', 'The quota is kept over a floating period of the last %n day', 'The quota is kept over a floating period of the last %n days', this.quotaInfo.period.length) }, + poolUsed() { + return this.quotaInfo !== null && this.quotaInfo.quota_usage.some((quota) => quota.used_pool) + }, }, watch: {}, diff --git a/src/components/QuotaPeriodPicker.vue b/src/components/QuotaPeriodPicker.vue index 12db6fc3..ca3b6053 100644 --- a/src/components/QuotaPeriodPicker.vue +++ b/src/components/QuotaPeriodPicker.vue @@ -16,9 +16,10 @@ + @update:model-value="$event >= 1 && update('length', $event)" /> + @update:model-value="$event >= 1 && $event <= 28 && update('day', $event)" />
{{ resetText }} @@ -114,6 +117,9 @@ export default { display: flex; align-items: center; gap: 5px; + .select { + margin-bottom: 0; + } } .container { diff --git a/src/components/Rules/MultiselectWho.vue b/src/components/Rules/MultiselectWho.vue new file mode 100644 index 00000000..59a57564 --- /dev/null +++ b/src/components/Rules/MultiselectWho.vue @@ -0,0 +1,244 @@ + + + + + + diff --git a/src/components/Rules/QuotaRules.vue b/src/components/Rules/QuotaRules.vue new file mode 100644 index 00000000..9eb772db --- /dev/null +++ b/src/components/Rules/QuotaRules.vue @@ -0,0 +1,101 @@ + + + + + + diff --git a/src/components/Rules/Rule.vue b/src/components/Rules/Rule.vue new file mode 100644 index 00000000..e29afed7 --- /dev/null +++ b/src/components/Rules/Rule.vue @@ -0,0 +1,186 @@ + + + + + + diff --git a/tests/unit/Providers/OpenAiProviderTest.php b/tests/unit/Providers/OpenAiProviderTest.php index bdec3f6e..87d76fd0 100644 --- a/tests/unit/Providers/OpenAiProviderTest.php +++ b/tests/unit/Providers/OpenAiProviderTest.php @@ -17,6 +17,7 @@ use OCA\OpenAi\Service\ChunkService; use OCA\OpenAi\Service\OpenAiAPIService; use OCA\OpenAi\Service\OpenAiSettingsService; +use OCA\OpenAi\Service\QuotaRuleService; use OCA\OpenAi\TaskProcessing\ChangeToneProvider; use OCA\OpenAi\TaskProcessing\EmojiProvider; use OCA\OpenAi\TaskProcessing\HeadlineProvider; @@ -83,6 +84,7 @@ protected function setUp(): void { \OCP\Server::get(QuotaUsageMapper::class), $this->openAiSettingsService, $this->createMock(\OCP\Notification\IManager::class), + \OCP\Server::get(QuotaRuleService::class), $clientService, ); diff --git a/tests/unit/Quota/QuotaTest.php b/tests/unit/Quota/QuotaTest.php index 9545e4d3..dff21fc6 100644 --- a/tests/unit/Quota/QuotaTest.php +++ b/tests/unit/Quota/QuotaTest.php @@ -1,22 +1,23 @@ createUser(self::TEST_USER1, self::TEST_USER1); + $backend->createUser(self::TEST_USER2, self::TEST_USER2); \OCP\Server::get(IUserManager::class)->registerBackend($backend); } @@ -70,6 +74,8 @@ protected function setUp(): void { $this->cacheFactory = $this->createMock(ICacheFactory::class); + $this->quotaRuleService = \OCP\Server::get(QuotaRuleService::class); + $this->openAiApiService = new OpenAiAPIService( \OCP\Server::get(LoggerInterface::class), @@ -79,6 +85,7 @@ protected function setUp(): void { \OCP\Server::get(QuotaUsageMapper::class), $this->openAiSettingsService, $this->notificationManager, + \OCP\Server::get(QuotaRuleService::class), \OCP\Server::get(IClientService::class), ); } @@ -91,14 +98,18 @@ public static function tearDownAfterClass(): void { } catch (\OCP\Db\Exception|RuntimeException|Exception|Throwable $e) { // Ignore } + $rules = \OCP\Server::get(QuotaRuleService::class)->getRules(); + foreach ($rules as $rule) { + \OCP\Server::get(QuotaRuleService::class)->deleteRule($rule['id']); + } $backend = new Dummy(); $backend->deleteUser(self::TEST_USER1); + $backend->deleteUser(self::TEST_USER2); \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); @@ -119,5 +130,34 @@ public function testNotification(): void { // Clear quota usage $this->quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER1); } + public function testQuotaPool(): void { + // Create a quota rule for both test users as a pool + $this->openAiSettingsService->setQuotas([1000, 1, 1, 1]); + $rule = $this->quotaRuleService->addRule(); + $rule['type'] = Application::QUOTA_TYPE_TEXT; + $rule['amount'] = 10; + $rule['pool'] = true; + $rule['entities'] = [ + [ + 'entity_id' => self::TEST_USER1, + 'entity_type' => EntityType::USER->value, + ], + [ + 'entity_id' => self::TEST_USER2, + 'entity_type' => EntityType::USER->value, + ] + ]; + $this->quotaRuleService->updateRule($rule['id'], $rule); + + $this->assertFalse($this->openAiApiService->isQuotaExceeded(self::TEST_USER1, Application::QUOTA_TYPE_TEXT)); + $this->quotaUsageMapper->createQuotaUsage(self::TEST_USER1, Application::QUOTA_TYPE_TEXT, 100, $rule['id']); + + $this->assertTrue($this->openAiApiService->isQuotaExceeded(self::TEST_USER1, Application::QUOTA_TYPE_TEXT)); + // Check other user + $this->assertTrue($this->openAiApiService->isQuotaExceeded(self::TEST_USER2, Application::QUOTA_TYPE_TEXT)); + // Clear quota usage + $this->quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER1); + $this->quotaUsageMapper->deleteUserQuotaUsages(self::TEST_USER2); + } }