130 lines
4.7 KiB
TypeScript
130 lines
4.7 KiB
TypeScript
import { Injectable } from '@nestjs/common';
|
|
import { InjectRepository } from '@nestjs/typeorm';
|
|
import { Repository } from 'typeorm';
|
|
import { BaseTenantRepository, TenantContextService } from '@iconsulting/shared';
|
|
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
|
|
extends BaseTenantRepository<TokenUsageEntity, TokenUsageORM>
|
|
implements ITokenUsageRepository
|
|
{
|
|
constructor(
|
|
@InjectRepository(TokenUsageORM) repo: Repository<TokenUsageORM>,
|
|
tenantContext: TenantContextService,
|
|
) {
|
|
super(repo, tenantContext);
|
|
}
|
|
|
|
async save(tokenUsage: TokenUsageEntity): Promise<TokenUsageEntity> {
|
|
const orm = this.toORM(tokenUsage);
|
|
orm.tenantId = this.getTenantId();
|
|
const saved = await this.repo.save(orm);
|
|
return this.toEntity(saved);
|
|
}
|
|
|
|
async findByConversationId(conversationId: string): Promise<TokenUsageEntity[]> {
|
|
const orms = await this.createTenantQueryBuilder('token_usage')
|
|
.andWhere('token_usage.conversation_id = :conversationId', { conversationId })
|
|
.orderBy('token_usage.created_at', 'ASC')
|
|
.getMany();
|
|
return orms.map((orm) => this.toEntity(orm));
|
|
}
|
|
|
|
async findByUserId(userId: string, options?: { limit?: number }): Promise<TokenUsageEntity[]> {
|
|
const queryBuilder = this.createTenantQueryBuilder('token_usage')
|
|
.andWhere('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.createTenantQueryBuilder('token_usage')
|
|
.select('SUM(token_usage.input_tokens)', 'totalInputTokens')
|
|
.addSelect('SUM(token_usage.output_tokens)', 'totalOutputTokens')
|
|
.addSelect('SUM(token_usage.estimated_cost)', 'totalCost')
|
|
.andWhere('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.createTenantQueryBuilder('token_usage')
|
|
.select('SUM(token_usage.input_tokens)', 'totalInputTokens')
|
|
.addSelect('SUM(token_usage.output_tokens)', 'totalOutputTokens')
|
|
.addSelect('SUM(token_usage.estimated_cost)', 'totalCost')
|
|
.andWhere('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.tenantId = this.getTenantId();
|
|
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,
|
|
});
|
|
}
|
|
}
|