125 lines
4.4 KiB
TypeScript
125 lines
4.4 KiB
TypeScript
import { Injectable } from '@nestjs/common';
|
|
import { InjectRepository } from '@nestjs/typeorm';
|
|
import { Repository } from 'typeorm';
|
|
import { TokenUsageORM } from '../../../infrastructure/database/postgres/entities/token-usage.orm';
|
|
import { ITokenUsageRepository } from '../../../domain/repositories/token-usage.repository.interface';
|
|
import { TokenUsageEntity } from '../../../domain/entities/token-usage.entity';
|
|
|
|
@Injectable()
|
|
export class TokenUsagePostgresRepository implements ITokenUsageRepository {
|
|
constructor(
|
|
@InjectRepository(TokenUsageORM)
|
|
private readonly repo: Repository<TokenUsageORM>,
|
|
) {}
|
|
|
|
async save(tokenUsage: TokenUsageEntity): Promise<TokenUsageEntity> {
|
|
const orm = this.toORM(tokenUsage);
|
|
const saved = await this.repo.save(orm);
|
|
return this.toEntity(saved);
|
|
}
|
|
|
|
async findByConversationId(conversationId: string): Promise<TokenUsageEntity[]> {
|
|
const orms = await this.repo.find({
|
|
where: { conversationId },
|
|
order: { createdAt: 'ASC' },
|
|
});
|
|
return orms.map((orm) => this.toEntity(orm));
|
|
}
|
|
|
|
async findByUserId(userId: string, options?: { limit?: number }): Promise<TokenUsageEntity[]> {
|
|
const queryBuilder = this.repo
|
|
.createQueryBuilder('token_usage')
|
|
.where('token_usage.user_id = :userId', { userId })
|
|
.orderBy('token_usage.created_at', 'DESC');
|
|
|
|
if (options?.limit) {
|
|
queryBuilder.limit(options.limit);
|
|
}
|
|
|
|
const orms = await queryBuilder.getMany();
|
|
return orms.map((orm) => this.toEntity(orm));
|
|
}
|
|
|
|
async sumByConversationId(conversationId: string): Promise<{
|
|
totalInputTokens: number;
|
|
totalOutputTokens: number;
|
|
totalCost: number;
|
|
}> {
|
|
const result = await this.repo
|
|
.createQueryBuilder('token_usage')
|
|
.select('SUM(token_usage.input_tokens)', 'totalInputTokens')
|
|
.addSelect('SUM(token_usage.output_tokens)', 'totalOutputTokens')
|
|
.addSelect('SUM(token_usage.estimated_cost)', 'totalCost')
|
|
.where('token_usage.conversation_id = :conversationId', { conversationId })
|
|
.getRawOne();
|
|
|
|
return {
|
|
totalInputTokens: parseInt(result?.totalInputTokens || '0', 10),
|
|
totalOutputTokens: parseInt(result?.totalOutputTokens || '0', 10),
|
|
totalCost: parseFloat(result?.totalCost || '0'),
|
|
};
|
|
}
|
|
|
|
async sumByUserId(userId: string): Promise<{
|
|
totalInputTokens: number;
|
|
totalOutputTokens: number;
|
|
totalCost: number;
|
|
}> {
|
|
const result = await this.repo
|
|
.createQueryBuilder('token_usage')
|
|
.select('SUM(token_usage.input_tokens)', 'totalInputTokens')
|
|
.addSelect('SUM(token_usage.output_tokens)', 'totalOutputTokens')
|
|
.addSelect('SUM(token_usage.estimated_cost)', 'totalCost')
|
|
.where('token_usage.user_id = :userId', { userId })
|
|
.getRawOne();
|
|
|
|
return {
|
|
totalInputTokens: parseInt(result?.totalInputTokens || '0', 10),
|
|
totalOutputTokens: parseInt(result?.totalOutputTokens || '0', 10),
|
|
totalCost: parseFloat(result?.totalCost || '0'),
|
|
};
|
|
}
|
|
|
|
private toORM(entity: TokenUsageEntity): TokenUsageORM {
|
|
const orm = new TokenUsageORM();
|
|
orm.id = entity.id;
|
|
orm.userId = entity.userId;
|
|
orm.conversationId = entity.conversationId;
|
|
orm.messageId = entity.messageId;
|
|
orm.model = entity.model;
|
|
orm.inputTokens = entity.inputTokens;
|
|
orm.outputTokens = entity.outputTokens;
|
|
orm.cacheCreationTokens = entity.cacheCreationTokens;
|
|
orm.cacheReadTokens = entity.cacheReadTokens;
|
|
orm.totalTokens = entity.totalTokens;
|
|
orm.estimatedCost = entity.estimatedCost;
|
|
orm.intentType = entity.intentType;
|
|
orm.toolCalls = entity.toolCalls;
|
|
orm.responseLength = entity.responseLength;
|
|
orm.latencyMs = entity.latencyMs;
|
|
orm.createdAt = entity.createdAt;
|
|
return orm;
|
|
}
|
|
|
|
private toEntity(orm: TokenUsageORM): TokenUsageEntity {
|
|
return TokenUsageEntity.fromPersistence({
|
|
id: orm.id,
|
|
userId: orm.userId,
|
|
conversationId: orm.conversationId,
|
|
messageId: orm.messageId,
|
|
model: orm.model,
|
|
inputTokens: orm.inputTokens,
|
|
outputTokens: orm.outputTokens,
|
|
cacheCreationTokens: orm.cacheCreationTokens,
|
|
cacheReadTokens: orm.cacheReadTokens,
|
|
totalTokens: orm.totalTokens,
|
|
estimatedCost: Number(orm.estimatedCost),
|
|
intentType: orm.intentType,
|
|
toolCalls: orm.toolCalls,
|
|
responseLength: orm.responseLength,
|
|
latencyMs: orm.latencyMs,
|
|
createdAt: orm.createdAt,
|
|
});
|
|
}
|
|
}
|