feat(usage): add token usage service and IPC router endpoints
This commit is contained in:
1 parent
2a5e54184f
commit
ce21283170
3 files changed
+101
No files matched your search
@@ -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,
|
||||
});
|
||||
@@ -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);
|
||||
}),
|
||||
});
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user