From 23867c0b6e3f227a351053e949f72465fd83412e Mon Sep 17 00:00:00 2001 From: Aitor <1726644+aaitor@users.noreply.github.com> Date: Thu, 27 Jun 2024 15:34:24 +0200 Subject: [PATCH] feat: adding nvm cost --- .env.local.sh | 6 +++++ src/api/agent/agent.controller.ts | 34 ++++++++++++++++++++++--- src/api/agent/agent.service.ts | 11 +++++++- src/common/utils/ipfs-helper.ts | 3 ++- src/configuration.ts | 7 ++++- src/database/1718898482722-AgentBase.ts | 12 +++++++++ src/database/entities/step.entity.ts | 3 +++ src/database/entities/task.entity.ts | 3 +++ src/processor/processor.controller.ts | 9 ++++--- 9 files changed, 78 insertions(+), 10 deletions(-) diff --git a/.env.local.sh b/.env.local.sh index 4134d54..9bd86d6 100644 --- a/.env.local.sh +++ b/.env.local.sh @@ -4,3 +4,9 @@ export API_AUTH_TOKEN="12345" export NVM_ENVIRONMENT="appStaging" export NVM_AGENT_DID="did:nv:1234" export NVM_SUBSCRIPTION_DID="did:nv:aabb" + +export DATABASE_NAME="/tmp/agent-youtube.db" + +export IPFS_GATEWAY="https://ipfs.infura.io:5001" +export IPFS_PROJECT_ID="" +export IPFS_PROJECT_SECRET="" \ No newline at end of file diff --git a/src/api/agent/agent.controller.ts b/src/api/agent/agent.controller.ts index 434570a..2780a8b 100644 --- a/src/api/agent/agent.controller.ts +++ b/src/api/agent/agent.controller.ts @@ -4,6 +4,7 @@ import { Get, Param, Post, + Res, UseGuards, UsePipes, ValidationPipe @@ -17,13 +18,21 @@ import { } from '@nestjs/swagger' import { CreateTaskDto } from './dto/create-task-dto' import { AuthorizationGuard } from '../../common/guards/auth.guard' +import { ConfigService } from '@nestjs/config' +import { Response } from 'express' @ApiTags('Agent') @Controller() @UseGuards(AuthorizationGuard) @ApiBearerAuth('Authorization') export class AgentController { - constructor(private readonly agentService: AgentService) {} + static readonly DEFAULT_COST_CREDITS = '1' + static readonly NVM_COST_HEADER = 'NVMCreditsConsumed' + + constructor( + private readonly agentService: AgentService, + private readonly configService: ConfigService + ) {} @Post() @ApiOperation({ @@ -43,8 +52,19 @@ export class AgentController { description: 'Unauthorized' }) @UsePipes(new ValidationPipe()) - async createAgentTask(@Body() agentTaskDto: CreateTaskDto) { - return this.agentService.createAgentTask(agentTaskDto) + async createAgentTask( + @Body() agentTaskDto: CreateTaskDto, + @Res({ passthrough: true }) response: Response + ) { + const task = await this.agentService.createAgentTask(agentTaskDto) + const createTaskCost = this.configService.get('nvm.credits.createTask') + if (createTaskCost > 0) { + response.setHeader( + this.configService.get('nvm.credits.header'), + createTaskCost + ) + } + return task } @Get(':task_id') @@ -56,8 +76,14 @@ export class AgentController { status: 200, description: 'Return the task details' }) - async getAgentTask(@Param('task_id') taskId: string) { + async getAgentTask( + @Param('task_id') taskId: string, + @Res({ passthrough: true }) response: Response + ) { const fullTask = await this.agentService.getTaskById(taskId) + const taskCost = + fullTask.task.cost || this.configService.get('nvm.credits.getTask') + response.setHeader(this.configService.get('nvm.credits.header'), taskCost) return { task: fullTask.task, steps: fullTask.steps diff --git a/src/api/agent/agent.service.ts b/src/api/agent/agent.service.ts index 251a915..8231bc7 100644 --- a/src/api/agent/agent.service.ts +++ b/src/api/agent/agent.service.ts @@ -59,13 +59,20 @@ export class AgentService { } updateStep(stepId: string, stepUpdated: Partial) { - Logger.log(`Updating step ${stepId}, with ${JSON.stringify(stepUpdated)}`) + Logger.log(`Updating step ${stepId}, with ${stepUpdated.step_status}`) return this.stepEntity.update( { step_id: stepId }, { ...stepUpdated, updated_at: new Date() } ) } + async getTotalCostForTask(taskId: string) { + const totalCost = await AppDataSource.query( + `SELECT SUM(steps.cost) as total_cost FROM steps WHERE steps.task_id = \'${taskId}\'` + ) + return totalCost[0].total_cost + } + async completeTasksWhenStepsAreDone() { const tasks = await AppDataSource.query( `SELECT tasks.task_id as task_id, steps.output as output, steps.output_artifacts as output_artifacts, steps.output_additional as output_additional FROM tasks, steps WHERE tasks.task_id = steps.task_id AND steps.is_last = true AND steps.step_status = \'COMPLETED\' AND tasks.task_status = \'PENDING\' ` @@ -76,6 +83,7 @@ export class AgentService { { task_id: task.task_id }, { task_status: ExecutionStatus.COMPLETED, + cost: await this.getTotalCostForTask(task.task_id), output: task.output, output_artifacts: task.output_artifacts, output_additional: task.output_additional, @@ -100,6 +108,7 @@ export class AgentService { { task_id: task.task_id }, { task_status: ExecutionStatus.FAILED, + cost: await this.getTotalCostForTask(task.task_id), output: task.output, output_artifacts: task.output_artifacts, output_additional: task.output_additional, diff --git a/src/common/utils/ipfs-helper.ts b/src/common/utils/ipfs-helper.ts index fd7a792..66324ba 100644 --- a/src/common/utils/ipfs-helper.ts +++ b/src/common/utils/ipfs-helper.ts @@ -1,4 +1,5 @@ -import IpfsHttpClientLite from 'ipfs-http-client-lite' +/* eslint-disable @typescript-eslint/no-var-requires */ +const IpfsHttpClientLite = require('ipfs-http-client-lite') export interface IpfsConnectionParameters { gatewayUrl: string diff --git a/src/configuration.ts b/src/configuration.ts index c271202..7c6e9a9 100644 --- a/src/configuration.ts +++ b/src/configuration.ts @@ -13,7 +13,12 @@ export const appConfig = { nvm: { environment: process.env.NVM_ENVIRONMENT, agentDid: process.env.NVM_AGENT_DID, - subscriptionDid: process.env.NVM_SUBSCRIPTION_DID + subscriptionDid: process.env.NVM_SUBSCRIPTION_DID, + credits: { + header: process.env.NVM_CREDITS_HEADER || 'NVMCreditsConsumed', + createTask: Number(process.env.NVM_CREDITS_CREATE_TASK) || 1, + getTask: Number(process.env.NVM_CREDITS_GET_TASK) || 1 + } }, logger: { level: process.env.LOG_LEVEL || ['error', 'log', 'warn'] diff --git a/src/database/1718898482722-AgentBase.ts b/src/database/1718898482722-AgentBase.ts index d1522a1..0919648 100644 --- a/src/database/1718898482722-AgentBase.ts +++ b/src/database/1718898482722-AgentBase.ts @@ -62,6 +62,12 @@ export class AgentBase1718898482722 implements MigrationInterface { type: 'text', isNullable: true }, + { + name: 'cost', + type: 'integer', + default: 0, + isNullable: true + }, ...baseSchema ] }), @@ -137,6 +143,12 @@ export class AgentBase1718898482722 implements MigrationInterface { type: 'text', isNullable: true }, + { + name: 'cost', + type: 'integer', + default: 0, + isNullable: true + }, ...baseSchema ] }), diff --git a/src/database/entities/step.entity.ts b/src/database/entities/step.entity.ts index b2a1285..7aab33d 100644 --- a/src/database/entities/step.entity.ts +++ b/src/database/entities/step.entity.ts @@ -17,6 +17,9 @@ export class StepEntity extends BaseEntity { @Column('integer') retries: number + @Column('integer') + cost: number + @Column('boolean') is_last: boolean diff --git a/src/database/entities/task.entity.ts b/src/database/entities/task.entity.ts index d4ab59b..c06c9d6 100644 --- a/src/database/entities/task.entity.ts +++ b/src/database/entities/task.entity.ts @@ -18,6 +18,9 @@ export class TaskEntity extends BaseEntity { @Column('varchar') name: string + @Column('integer') + cost: number + @Column('text') input_query: string diff --git a/src/processor/processor.controller.ts b/src/processor/processor.controller.ts index 1f7fa1e..09363ba 100644 --- a/src/processor/processor.controller.ts +++ b/src/processor/processor.controller.ts @@ -56,13 +56,15 @@ export class ProcessorController implements OnModuleInit { Logger.log(`[${task.task_id}] Simulate a step that is still in progress`) await this.agentService.updateStep(step.step_id, { step_status: ExecutionStatus.IN_PROGRESS, - retries: step.retries + 1 + retries: step.retries + 1, + cost: 0 }) } else if (randomIndex === 1) { Logger.log(`[${task.task_id}] Simulate a step that is failing`) await this.agentService.updateStep(step.step_id, { step_status: ExecutionStatus.FAILED, - output: `{"result": 500, "message": "Failed because ${randomIndex} is 1"}` + output: `{"result": 500, "message": "Failed because ${randomIndex} is 1"}`, + cost: 0 }) } else { // Simulate a successful step @@ -70,7 +72,8 @@ export class ProcessorController implements OnModuleInit { await this.agentService.updateStep(step.step_id, { step_status: ExecutionStatus.COMPLETED, is_last: true, - output: `{"result": 200, "message": "${randomIndex}"}` + output: `{"result": 200, "message": "${randomIndex}"}`, + cost: randomIndex + 1 }) } }