Skip to content

Commit 9260711

Browse files
lukasdotcomlukasdotcom
authored
Merge pull request #257 from nextcloud/feat/usage-export
Feat: quota report download
2 parents 83686a3 + 417ddc1 commit 9260711

8 files changed

Lines changed: 205 additions & 9 deletions

File tree

‎appinfo/routes.php‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,5 +20,6 @@
2020
['name' => 'quotaRule#addRule', 'url' => '/quota/rule', 'verb' => 'POST'],
2121
['name' => 'quotaRule#updateRule', 'url' => '/quota/rule', 'verb' => 'PUT'],
2222
['name' => 'quotaRule#deleteRule', 'url' => '/quota/rule', 'verb' => 'DELETE'],
23+
['name' => 'quotaRule#getQuotaUsage', 'url' => '/quota/download-usage', 'verb' => 'GET'],
2324
],
2425
];

‎lib/Controller/QuotaRuleController.php‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@
1111
use OCA\OpenAi\Service\QuotaRuleService;
1212
use OCP\AppFramework\Controller;
1313
use OCP\AppFramework\Http;
14+
use OCP\AppFramework\Http\Attribute\NoCSRFRequired;
1415
use OCP\AppFramework\Http\DataResponse;
16+
use OCP\AppFramework\Http\Response;
17+
use OCP\AppFramework\Http\TextPlainResponse;
1518
use OCP\IRequest;
1619

1720
class QuotaRuleController extends Controller {
@@ -79,4 +82,35 @@ public function deleteRule(int $id): DataResponse {
7982
return new DataResponse(['error' => $e->getMessage()], Http::STATUS_BAD_REQUEST);
8083
}
8184
}
85+
/**
86+
* Gets the quota usage between two dates
87+
* @param int $startDate
88+
* @param int $endDate
89+
* @param int $type
90+
* @return Http\StreamResponse|TextPlainResponse
91+
*/
92+
#[NoCSRFRequired]
93+
public function getQuotaUsage(int $startDate, int $endDate, int $type): Response {
94+
try {
95+
$result = $this->quotaRuleService->getQuotaUsage($startDate, $endDate, $type);
96+
$csv = fopen('php://memory', 'w');
97+
try {
98+
foreach ($result as $row) {
99+
fputcsv($csv, $row);
100+
}
101+
rewind($csv);
102+
103+
$response = new Http\StreamResponse($csv, Http::STATUS_OK);
104+
$response->setHeaders([
105+
'Content-Type' => 'text/csv',
106+
'Content-Disposition' => 'attachment; filename="quota_usage.csv"'
107+
]);
108+
return $response;
109+
} catch (Exception $e) {
110+
return new TextPlainResponse('Failed to get quota usage:' . $e->getMessage(), Http::STATUS_INTERNAL_SERVER_ERROR);
111+
}
112+
} catch (Exception $e) {
113+
return new TextPlainResponse('Failed to get quota usage:' . $e->getMessage(), Http::STATUS_NOT_FOUND);
114+
}
115+
}
82116
}

‎lib/Cron/CleanupQuotaDb.php‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,15 @@ public function __construct(
3030
protected function run($argument) {
3131
$this->logger->debug('Run cleanup job for OpenAI quota db');
3232
$quota = $this->openAiSettingsService->getQuotaPeriod();
33-
$days = $quota['length'];
33+
$quotaDays = $quota['length'];
3434
if ($quota['unit'] == 'month') {
35-
$days *= 30;
35+
$quotaDays *= 30;
3636
}
37+
$days = $this->openAiSettingsService->getUsageStorageTime();
3738
$this->quotaUsageMapper->cleanupQuotaUsages(
3839
// The mimimum period is limited to DEFAULT_QUOTA_PERIOD to not lose
3940
// the stored quota usage data below this limit.
40-
max($days, Application::DEFAULT_QUOTA_PERIOD)
41+
max($quotaDays, $days, Application::DEFAULT_QUOTA_PERIOD)
4142
);
4243

4344
}

‎lib/Db/QuotaUsageMapper.php‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ public function getQuotaUnitsInTimePeriod(int $type, int $periodStart): int {
8282
$qb = $this->db->getQueryBuilder();
8383

8484
// Get the sum of the units used in the time period
85-
$qb->select($qb->createFunction('SUM(units)'))
85+
$qb->select($qb->func()->sum('units'))
8686
->from($this->getTableName())
8787
->where(
8888
$qb->expr()->eq('type', $qb->createNamedParameter($type, IQueryBuilder::PARAM_INT))
@@ -113,7 +113,7 @@ public function getQuotaUnitsOfUserInTimePeriod(string $userId, int $type, int $
113113
$qb = $this->db->getQueryBuilder();
114114

115115
// Get the sum of the units used in the time period
116-
$qb->select($qb->createFunction('SUM(units)'))
116+
$qb->select($qb->func()->sum('units'))
117117
->from($this->getTableName())
118118
->where(
119119
$qb->expr()->eq('type', $qb->createNamedParameter($type, IQueryBuilder::PARAM_INT))
@@ -170,7 +170,7 @@ public function getQuotaUsagesOfUser(string $userId, int $type): array {
170170
public function getQuotaUnitsOfUser(string $userId, int $type): int {
171171
$qb = $this->db->getQueryBuilder();
172172

173-
$qb->select($qb->createFunction('SUM(units)'))
173+
$qb->select($qb->func()->sum('units'))
174174
->from($this->getTableName())
175175
->where(
176176
$qb->expr()->eq('user_id', $qb->createNamedParameter($userId, IQueryBuilder::PARAM_STR))
@@ -274,4 +274,51 @@ public function cleanupQuotaUsages(int $timePeriod): void {
274274
);
275275
$qb->executeStatement();
276276
}
277+
278+
/**
279+
* Gets quota usage of all users
280+
* @param int $startTime
281+
* @param int $endTime
282+
* @return array
283+
* @throws Exception
284+
* @throws RuntimeException
285+
*/
286+
public function getUsersQuotaUsage(int $startTime, int $endTime, $type): array {
287+
$qb = $this->db->getQueryBuilder();
288+
289+
$qb->select('user_id')
290+
->selectAlias($qb->func()->sum('units'), 'usage')
291+
->from($this->getTableName())
292+
->where($qb->expr()->gte('timestamp', $qb->createNamedParameter($startTime, IQueryBuilder::PARAM_INT)))
293+
->andWhere($qb->expr()->lte('timestamp', $qb->createNamedParameter($endTime, IQueryBuilder::PARAM_INT)))
294+
->andWhere($qb->expr()->eq('type', $qb->createNamedParameter($type, IQueryBuilder::PARAM_INT)))
295+
->groupBy('user_id')
296+
->orderBy('usage', 'DESC');
297+
298+
return $qb->executeQuery()->fetchAll();
299+
}
300+
/**
301+
* Gets quota usage of all pools
302+
* @param int $startTime
303+
* @param int $endTime
304+
* @param int $type
305+
* @return array
306+
* @throws Exception
307+
* @throws RuntimeException
308+
*/
309+
public function getPoolsQuotaUsage(int $startTime, int $endTime, int $type): array {
310+
$qb = $this->db->getQueryBuilder();
311+
312+
$qb->select('pool')
313+
->selectAlias($qb->func()->sum('units'), 'usage')
314+
->from($this->getTableName())
315+
->where($qb->expr()->neq('pool', $qb->createNamedParameter(-1, IQueryBuilder::PARAM_INT)))
316+
->andWhere($qb->expr()->gte('timestamp', $qb->createNamedParameter($startTime, IQueryBuilder::PARAM_INT)))
317+
->andWhere($qb->expr()->lte('timestamp', $qb->createNamedParameter($endTime, IQueryBuilder::PARAM_INT)))
318+
->andWhere($qb->expr()->eq('type', $qb->createNamedParameter($type, IQueryBuilder::PARAM_INT)))
319+
->groupBy('type', 'pool')
320+
->orderBy('usage', 'DESC');
321+
322+
return $qb->executeQuery()->fetchAll();
323+
}
277324
}

‎lib/Service/OpenAiSettingsService.php‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ class OpenAiSettingsService {
3737
'llm_extra_params' => 'string',
3838
'quota_period' => 'array',
3939
'quotas' => 'array',
40+
'usage_storage_time' => 'integer',
4041
'translation_provider_enabled' => 'boolean',
4142
'llm_provider_enabled' => 'boolean',
4243
't2i_provider_enabled' => 'boolean',
@@ -119,7 +120,9 @@ public function getQuotaEnd(): int {
119120
$startDate = new DateTime('2000-01-' . $quotaPeriod['day']);
120121
$months = $startDate->diff($periodEnd)->m + $startDate->diff($periodEnd)->y * 12;
121122
$remainder = $months % $quotaPeriod['length'];
122-
$periodEnd = $periodEnd->add(new DateInterval('P' . $quotaPeriod['length'] - $remainder . 'M'));
123+
if ($remainder != 0) {
124+
$periodEnd = $periodEnd->add(new DateInterval('P' . $quotaPeriod['length'] - $remainder . 'M'));
125+
}
123126
}
124127
}
125128
return $periodEnd->getTimestamp();
@@ -308,6 +311,10 @@ public function getQuotas(): array {
308311
return $quotas;
309312
}
310313

314+
public function getUsageStorageTime() : int {
315+
return $this->appConfig->getValueInt(Application::APP_ID, 'usage_storage_time', Application::DEFAULT_QUOTA_PERIOD, lazy: true);
316+
}
317+
311318
/**
312319
* @return boolean
313320
*/
@@ -395,6 +402,7 @@ public function getAdminConfig(): array {
395402
// Updated to get quota period
396403
'quotas' => $this->getQuotas(),
397404
// Get quotas from the config value and return it
405+
'usage_storage_time' => $this->getUsageStorageTime(),
398406
'translation_provider_enabled' => $this->getTranslationProviderEnabled(),
399407
'llm_provider_enabled' => $this->getLlmProviderEnabled(),
400408
't2i_provider_enabled' => $this->getT2iProviderEnabled(),
@@ -522,6 +530,15 @@ public function setQuotas(array $quotas): void {
522530
$cache->clear(Application::QUOTA_RULES_CACHE_PREFIX);
523531
}
524532

533+
/**
534+
* @param int $usageStorageTime
535+
* @return void
536+
*/
537+
public function setUsageStorageTime(int $usageStorageTime): void {
538+
$usageStorageTime = max(1, $usageStorageTime);
539+
$this->appConfig->setValueInt(Application::APP_ID, 'usage_storage_time', $usageStorageTime, lazy: true);
540+
}
541+
525542
/**
526543
* @param string $apiKey
527544
* @return void
@@ -831,6 +848,9 @@ public function setAdminConfig(array $adminConfig): void {
831848
if (isset($adminConfig['quotas'])) {
832849
$this->setQuotas($adminConfig['quotas']);
833850
}
851+
if (isset($adminConfig['usage_storage_time'])) {
852+
$this->setUsageStorageTime(intval($adminConfig['usage_storage_time']));
853+
}
834854
if (isset($adminConfig['use_max_completion_tokens_param'])) {
835855
$this->setUseMaxCompletionParam($adminConfig['use_max_completion_tokens_param']);
836856
}

‎lib/Service/QuotaRuleService.php‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,13 @@
1111
use OCA\OpenAi\AppInfo\Application;
1212
use OCA\OpenAi\Db\EntityType;
1313
use OCA\OpenAi\Db\QuotaRuleMapper;
14+
use OCA\OpenAi\Db\QuotaUsageMapper;
1415
use OCA\OpenAi\Db\QuotaUserMapper;
1516
use OCP\AppFramework\Db\DoesNotExistException;
1617
use OCP\AppFramework\Db\MultipleObjectsReturnedException;
1718
use OCP\ICacheFactory;
1819
use OCP\IGroupManager;
20+
use OCP\IL10N;
1921
use OCP\IUserManager;
2022
use Psr\Log\LoggerInterface;
2123

@@ -27,6 +29,8 @@ public function __construct(
2729
private IGroupManager $groupManager,
2830
private ICacheFactory $cacheFactory,
2931
private IUserManager $userManager,
32+
private QuotaUsageMapper $quotaUsageMapper,
33+
private IL10N $l10n,
3034
private LoggerInterface $logger,
3135
) {
3236
}
@@ -178,4 +182,29 @@ private function validateEntities(array $entities) {
178182
}
179183
}
180184
}
185+
public function getQuotaUsage(int $startDate, int $endDate, int $type): array {
186+
$data = [[$this->l10n->t('Name'), $this->l10n->t('Usage')]];
187+
$users = $this->quotaUsageMapper->getUsersQuotaUsage($startDate, $endDate, $type);
188+
$pools = $this->quotaUsageMapper->getPoolsQuotaUsage($startDate, $endDate, $type);
189+
$usersIdx = 0;
190+
$poolsIdx = 0;
191+
while ($usersIdx < count($users) && $poolsIdx < count($pools)) {
192+
if ($users[$usersIdx]['usage'] > $pools[$poolsIdx]['usage']) {
193+
$data[] = [$users[$usersIdx]['user_id'], $users[$usersIdx]['usage']];
194+
$usersIdx++;
195+
} else {
196+
$data[] = [$this->l10n->t('Quota pool for rule %d', $pools[$poolsIdx]['pool']), $pools[$poolsIdx]['usage']];
197+
$poolsIdx++;
198+
}
199+
}
200+
while ($usersIdx < count($users)) {
201+
$data[] = [$users[$usersIdx]['user_id'], $users[$usersIdx]['usage']];
202+
$usersIdx++;
203+
}
204+
while ($poolsIdx < count($pools)) {
205+
$data[] = [$this->l10n->t('Quota pool for rule %d', $pools[$poolsIdx]['pool']), $pools[$poolsIdx]['usage']];
206+
$poolsIdx++;
207+
}
208+
return $data;
209+
}
181210
}

‎lib/Settings/Admin.php‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ public function getForm(): TemplateResponse {
3333
$adminConfig['basic_password'] = $adminConfig['basic_password'] === '' ? '' : 'dummyPassword';
3434
$isAssistantEnabled = $this->appManager->isEnabledForUser('assistant');
3535
$adminConfig['assistant_enabled'] = $isAssistantEnabled;
36+
$adminConfig['quota_start_date'] = $this->openAiSettingsService->getQuotaStart();
37+
$adminConfig['quota_end_date'] = $this->openAiSettingsService->getQuotaEnd();
3638
$this->initialStateService->provideInitialState('admin-config', $adminConfig);
3739
$rules = $this->quotaRuleService->getRules();
3840
$this->initialStateService->provideInitialState('rules', $rules);

0 commit comments

Comments
 (0)