refactor: move agent evaluation to background jobs for non-blocking execution (#16234)

2025-12-03 10:25:47 +05:30
parent 05b30554c3
commit f248b3f7f4
11 changed files with 299 additions and 118 deletions
@@ -19,65 +19,65 @@ export class AgentMessagePartDTO {
  @Field()
  type: string;

-  @Field({ nullable: true })
-  textContent?: string;
+  @Field(() => String, { nullable: true })
+  textContent: string | null;

-  @Field({ nullable: true })
-  reasoningContent?: string;
+  @Field(() => String, { nullable: true })
+  reasoningContent: string | null;

-  @Field({ nullable: true })
-  toolName?: string;
+  @Field(() => String, { nullable: true })
+  toolName: string | null;

-  @Field({ nullable: true })
-  toolCallId?: string;
+  @Field(() => String, { nullable: true })
+  toolCallId: string | null;

  @Field(() => GraphQLJSON, { nullable: true })
-  toolInput?: Record<string, unknown>;
+  toolInput: unknown | null;

  @Field(() => GraphQLJSON, { nullable: true })
-  toolOutput?: Record<string, unknown>;
+  toolOutput: unknown | null;

-  @Field({ nullable: true })
-  state?: string;
+  @Field(() => String, { nullable: true })
+  state: string | null;

-  @Field({ nullable: true })
-  errorMessage?: string;
+  @Field(() => String, { nullable: true })
+  errorMessage: string | null;

  @Field(() => GraphQLJSON, { nullable: true })
-  errorDetails?: Record<string, unknown>;
+  errorDetails: Record<string, unknown> | null;

-  @Field({ nullable: true })
-  sourceUrlSourceId?: string;
+  @Field(() => String, { nullable: true })
+  sourceUrlSourceId: string | null;

-  @Field({ nullable: true })
-  sourceUrlUrl?: string;
+  @Field(() => String, { nullable: true })
+  sourceUrlUrl: string | null;

-  @Field({ nullable: true })
-  sourceUrlTitle?: string;
+  @Field(() => String, { nullable: true })
+  sourceUrlTitle: string | null;

-  @Field({ nullable: true })
-  sourceDocumentSourceId?: string;
+  @Field(() => String, { nullable: true })
+  sourceDocumentSourceId: string | null;

-  @Field({ nullable: true })
-  sourceDocumentMediaType?: string;
+  @Field(() => String, { nullable: true })
+  sourceDocumentMediaType: string | null;

-  @Field({ nullable: true })
-  sourceDocumentTitle?: string;
+  @Field(() => String, { nullable: true })
+  sourceDocumentTitle: string | null;

-  @Field({ nullable: true })
-  sourceDocumentFilename?: string;
+  @Field(() => String, { nullable: true })
+  sourceDocumentFilename: string | null;

-  @Field({ nullable: true })
-  fileMediaType?: string;
+  @Field(() => String, { nullable: true })
+  fileMediaType: string | null;

-  @Field({ nullable: true })
-  fileFilename?: string;
+  @Field(() => String, { nullable: true })
+  fileFilename: string | null;

-  @Field({ nullable: true })
-  fileUrl?: string;
+  @Field(() => String, { nullable: true })
+  fileUrl: string | null;

  @Field(() => GraphQLJSON, { nullable: true })
-  providerMetadata?: Record<string, unknown>;
+  providerMetadata: Record<string, unknown> | null;

  @IsDateString()
  @Field()
@@ -6,9 +6,12 @@ import { AiAgentModule } from 'src/engine/metadata-modules/ai/ai-agent/ai-agent.
 import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity';
 import { AiChatModule } from 'src/engine/metadata-modules/ai/ai-chat/ai-chat.module';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
+import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
 import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module';
 import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';

+import { EvaluateAgentTurnJob } from './jobs/evaluate-agent-turn.job';
+import { RunEvaluationInputJob } from './jobs/run-evaluation-input.job';
 import { AgentTurnEvaluationEntity } from './entities/agent-turn-evaluation.entity';
 import { AgentTurnResolver } from './resolvers/agent-turn.resolver';
 import { AgentTurnGraderService } from './services/agent-turn-grader.service';
@@ -17,6 +20,7 @@ import { AgentTurnGraderService } from './services/agent-turn-grader.service';
  imports: [
    TypeOrmModule.forFeature([
      AgentTurnEvaluationEntity,
+      AgentTurnEntity,
      AgentEntity,
      AgentChatThreadEntity,
    ]),
@@ -26,7 +30,12 @@ import { AgentTurnGraderService } from './services/agent-turn-grader.service';
    AiModelsModule,
    PermissionsModule,
  ],
-  providers: [AgentTurnGraderService, AgentTurnResolver],
+  providers: [
+    AgentTurnGraderService,
+    AgentTurnResolver,
+    EvaluateAgentTurnJob,
+    RunEvaluationInputJob,
+  ],
  exports: [AgentTurnGraderService],
 })
 export class AiAgentMonitorModule {}
@@ -15,8 +15,8 @@ export class AgentTurnEvaluationDTO {
  @Field(() => Int)
  score: number;

-  @Field({ nullable: true })
-  comment?: string;
+  @Field(() => String, { nullable: true })
+  comment: string | null;

  @IsDateString()
  @Field()
@@ -0,0 +1,35 @@
+import { Logger } from '@nestjs/common';
+
+import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
+import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
+import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
+import { AgentTurnGraderService } from 'src/engine/metadata-modules/ai/ai-agent-monitor/services/agent-turn-grader.service';
+
+export type EvaluateAgentTurnJobData = {
+  turnId: string;
+  workspaceId: string;
+};
+
+@Processor(MessageQueue.aiQueue)
+export class EvaluateAgentTurnJob {
+  private readonly logger = new Logger(EvaluateAgentTurnJob.name);
+
+  constructor(private readonly graderService: AgentTurnGraderService) {}
+
+  @Process(EvaluateAgentTurnJob.name)
+  async handle(data: EvaluateAgentTurnJobData): Promise<void> {
+    if (!data.turnId) {
+      throw new Error('Turn ID is required');
+    }
+
+    if (!data.workspaceId) {
+      throw new Error('Workspace ID is required');
+    }
+
+    const evaluation = await this.graderService.evaluateTurn(data.turnId);
+
+    this.logger.log(
+      `Evaluation completed for turn ${data.turnId}: score=${evaluation.score}`,
+    );
+  }
+}
@@ -0,0 +1,85 @@
+import { Logger } from '@nestjs/common';
+import { InjectRepository } from '@nestjs/typeorm';
+
+import { Repository } from 'typeorm';
+
+import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decorators/message-queue.decorator';
+import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
+import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
+import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
+import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service';
+import { AgentAsyncExecutorService } from 'src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service';
+import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity';
+import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
+
+import { EvaluateAgentTurnJob } from './evaluate-agent-turn.job';
+
+export type RunEvaluationInputJobData = {
+  turnId: string;
+  threadId: string;
+  agentId: string;
+  input: string;
+  workspaceId: string;
+};
+
+@Processor(MessageQueue.aiQueue)
+export class RunEvaluationInputJob {
+  private readonly logger = new Logger(RunEvaluationInputJob.name);
+
+  constructor(
+    @InjectRepository(AgentEntity)
+    private readonly agentRepository: Repository<AgentEntity>,
+    private readonly agentChatService: AgentChatService,
+    private readonly aiAgentExecutorService: AgentAsyncExecutorService,
+    @InjectMessageQueue(MessageQueue.aiQueue)
+    private readonly messageQueueService: MessageQueueService,
+  ) {}
+
+  @Process(RunEvaluationInputJob.name)
+  async handle(data: RunEvaluationInputJobData): Promise<void> {
+    await this.agentChatService.addMessage({
+      threadId: data.threadId,
+      turnId: data.turnId,
+      uiMessage: {
+        role: 'user',
+        parts: [{ type: 'text', text: data.input }],
+      },
+    });
+
+    const agent = await this.agentRepository.findOne({
+      where: { id: data.agentId },
+    });
+
+    if (!agent) {
+      throw new Error(`Agent ${data.agentId} not found`);
+    }
+
+    const executionResult = await this.aiAgentExecutorService.executeAgent({
+      agent,
+      userPrompt: data.input,
+    });
+
+    await this.agentChatService.addMessage({
+      threadId: data.threadId,
+      turnId: data.turnId,
+      agentId: agent.id,
+      uiMessage: {
+        role: 'assistant',
+        parts: [
+          {
+            type: 'text',
+            text: JSON.stringify(executionResult.result) || '',
+          },
+        ],
+      },
+    });
+
+    await this.messageQueueService.add<{
+      turnId: string;
+      workspaceId: string;
+    }>(EvaluateAgentTurnJob.name, {
+      turnId: data.turnId,
+      workspaceId: data.workspaceId,
+    });
+  }
+}
@@ -1,10 +1,15 @@
-import { UseGuards } from '@nestjs/common';
+import { Logger, UseGuards } from '@nestjs/common';
 import { Args, Mutation, Query, Resolver } from '@nestjs/graphql';
 import { InjectRepository } from '@nestjs/typeorm';

+import { msg } from '@lingui/core/macro';
 import { Repository } from 'typeorm';

 import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
+import { NotFoundError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
+import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decorators/message-queue.decorator';
+import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
+import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service';
 import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
 import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
@@ -12,27 +17,25 @@ import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.g
 import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
 import { AgentTurnDTO } from 'src/engine/metadata-modules/ai/ai-agent-execution/dtos/agent-turn.dto';
 import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
-import { AgentAsyncExecutorService } from 'src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service';
 import { AgentTurnEvaluationDTO } from 'src/engine/metadata-modules/ai/ai-agent-monitor/dtos/agent-turn-evaluation.dto';
+import { RunEvaluationInputJob } from 'src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job';
 import { AgentTurnGraderService } from 'src/engine/metadata-modules/ai/ai-agent-monitor/services/agent-turn-grader.service';
-import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
-import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
 import { PermissionFlagType } from 'src/engine/metadata-modules/permissions/constants/permission-flag-type.constants';

@UseGuards(WorkspaceAuthGuard, SettingsPermissionGuard(PermissionFlagType.AI))
@Resolver()
 export class AgentTurnResolver {
+  private readonly logger = new Logger(AgentTurnResolver.name);
+
  constructor(
    @InjectRepository(AgentTurnEntity)
    private readonly turnRepository: Repository<AgentTurnEntity>,
    @InjectRepository(AgentChatThreadEntity)
    private readonly threadRepository: Repository<AgentChatThreadEntity>,
-    @InjectRepository(AgentEntity)
-    private readonly agentRepository: Repository<AgentEntity>,
+    @InjectMessageQueue(MessageQueue.aiQueue)
+    private readonly messageQueueService: MessageQueueService,
    private readonly graderService: AgentTurnGraderService,
-    private readonly aiAgentExecutorService: AgentAsyncExecutorService,
-    private readonly agentChatService: AgentChatService,
  ) {}

  @Query(() => [AgentTurnDTO])
@@ -45,7 +48,7 @@ export class AgentTurnResolver {
      order: { createdAt: 'DESC' },
    });

-    return turns as unknown as AgentTurnDTO[];
+    return turns;
  }

  @Mutation(() => AgentTurnEvaluationDTO)
@@ -54,7 +57,7 @@ export class AgentTurnResolver {
  ): Promise<AgentTurnEvaluationDTO> {
    const evaluation = await this.graderService.evaluateTurn(turnId);

-    return evaluation as unknown as AgentTurnEvaluationDTO;
+    return evaluation;
  }

  @Mutation(() => AgentTurnDTO)
@@ -76,50 +79,31 @@ export class AgentTurnResolver {
    });
    const savedTurn = await this.turnRepository.save(turn);

-    const agent = await this.agentRepository.findOne({
-      where: { id: agentId },
-    });
-
-    await this.agentChatService.addMessage({
-      threadId: savedThread.id,
+    this.messageQueueService.add<{
+      turnId: string;
+      threadId: string;
+      agentId: string;
+      input: string;
+      workspaceId: string;
+    }>(RunEvaluationInputJob.name, {
      turnId: savedTurn.id,
-      uiMessage: {
-        role: 'user',
-        parts: [{ type: 'text', text: input }],
-      },
-    });
-
-    const executionResult = await this.aiAgentExecutorService.executeAgent({
-      agent,
-      userPrompt: input,
-    });
-
-    await this.agentChatService.addMessage({
      threadId: savedThread.id,
-      turnId: savedTurn.id,
-      agentId: agent?.id,
-      uiMessage: {
-        role: 'assistant',
-        parts: [
-          {
-            type: 'text',
-            text: JSON.stringify(executionResult.result) || '',
-          },
-        ],
-      },
+      agentId,
+      input,
+      workspaceId: workspace.id,
    });

-    await this.graderService.evaluateTurn(savedTurn.id);
-
-    const turnWithEvaluations = await this.turnRepository.findOne({
+    const turnWithRelations = await this.turnRepository.findOne({
      where: { id: savedTurn.id },
      relations: ['evaluations', 'messages', 'messages.parts'],
    });

-    if (!turnWithEvaluations) {
-      throw new Error('Turn not found after execution');
+    if (!turnWithRelations) {
+      throw new NotFoundError('Turn not found after creation', {
+        userFriendlyMessage: msg`Failed to create evaluation. Please try again.`,
+      });
    }

-    return turnWithEvaluations as unknown as AgentTurnDTO;
+    return turnWithRelations;
  }
 }