diff --git a/src/ipc/router.ts b/src/ipc/router.ts index 1a28e1d..5511835 100644 --- a/src/ipc/router.ts +++ b/src/ipc/router.ts @@ -6,6 +6,7 @@ import { accountRouter } from './account/router'; import { cloudRouter } from './cloud/router'; import { configRouter } from './config/router'; import { gatewayRouter } from './gateway/router'; +import { usageRouter } from './usage/router'; import { os } from '@orpc/server'; import { z } from 'zod'; @@ -53,5 +54,6 @@ export const router = os.use(logMiddleware).router({ cloud: cloudRouter, config: configRouter, gateway: gatewayRouter, + usage: usageRouter, system: systemHandler, }); diff --git a/src/ipc/usage/router.ts b/src/ipc/usage/router.ts new file mode 100644 index 0000000..fe3488c --- /dev/null +++ b/src/ipc/usage/router.ts @@ -0,0 +1,70 @@ +import { z } from 'zod'; +import { os } from '@orpc/server'; +import { TokenUsageRepo } from '../database/usageHandler'; +import { + HourlyUsageSchema, + DailyUsageSchema, + WeeklyUsageSchema, + MonthlyUsageSchema, + ModelUsageSchema, +} from '../../types/db'; + +const RecordUsageInputSchema = z.object({ + accountId: z.string(), + model: z.string(), + promptTokens: z.number().int().min(0), + completionTokens: z.number().int().min(0), + totalTokens: z.number().int().min(0), + timestamp: z.number().int(), + requestType: z.string().optional(), +}); + +const UsageQueryInputSchema = z.object({ + accountId: z.string().optional(), + start: z.number().int().optional(), + end: z.number().int().optional(), +}); + +export const usageRouter = os.router({ + recordUsage: os + .input(RecordUsageInputSchema) + .output(z.void()) + .handler(async ({ input }) => { + TokenUsageRepo.recordUsage(input); + }), + + getUsageByHour: os + .input(UsageQueryInputSchema) + .output(z.array(HourlyUsageSchema)) + .handler(async ({ input }) => { + return TokenUsageRepo.getUsageByHour(input.accountId, input.start, input.end); + }), + + getUsageByDay: os + .input(UsageQueryInputSchema) + .output(z.array(DailyUsageSchema)) + .handler(async ({ input }) => { + return TokenUsageRepo.getUsageByDay(input.accountId, input.start, input.end); + }), + + getUsageByWeek: os + .input(UsageQueryInputSchema) + .output(z.array(WeeklyUsageSchema)) + .handler(async ({ input }) => { + return TokenUsageRepo.getUsageByWeek(input.accountId, input.start, input.end); + }), + + getUsageByMonth: os + .input(UsageQueryInputSchema) + .output(z.array(MonthlyUsageSchema)) + .handler(async ({ input }) => { + return TokenUsageRepo.getUsageByMonth(input.accountId, input.start, input.end); + }), + + getUsageByModel: os + .input(UsageQueryInputSchema) + .output(z.array(ModelUsageSchema)) + .handler(async ({ input }) => { + return TokenUsageRepo.getUsageByModel(input.accountId, input.start, input.end); + }), +}); diff --git a/src/server/modules/usage/token-usage.service.ts b/src/server/modules/usage/token-usage.service.ts new file mode 100644 index 0000000..42c6637 --- /dev/null +++ b/src/server/modules/usage/token-usage.service.ts @@ -0,0 +1,29 @@ +import { Injectable } from '@nestjs/common'; +import { TokenUsageRepo, RecordUsageParams } from '../../../ipc/database/usageHandler'; + +@Injectable() +export class TokenUsageService { + recordUsage(params: RecordUsageParams): void { + TokenUsageRepo.recordUsage(params); + } + + getUsageByHour(accountId?: string, start?: number, end?: number) { + return TokenUsageRepo.getUsageByHour(accountId, start, end); + } + + getUsageByDay(accountId?: string, start?: number, end?: number) { + return TokenUsageRepo.getUsageByDay(accountId, start, end); + } + + getUsageByWeek(accountId?: string, start?: number, end?: number) { + return TokenUsageRepo.getUsageByWeek(accountId, start, end); + } + + getUsageByMonth(accountId?: string, start?: number, end?: number) { + return TokenUsageRepo.getUsageByMonth(accountId, start, end); + } + + getUsageByModel(accountId?: string, start?: number, end?: number) { + return TokenUsageRepo.getUsageByModel(accountId, start, end); + } +}