Introduce agent hints to reduce context bloat (#15763)

Introducing a new pattern that should reduce token consumption by 90% for the most common use-cases
2025-11-13 16:15:01 +01:00
parent 261cba40ef
commit cd5d0f9eac
10 changed files with 286 additions and 155 deletions
@@ -26,6 +26,7 @@ import { AGENT_CONFIG } from 'src/engine/metadata-modules/agent/constants/agent-
 import { AGENT_SYSTEM_PROMPTS } from 'src/engine/metadata-modules/agent/constants/agent-system-prompts.const';
 import { AgentActorContextService } from 'src/engine/metadata-modules/agent/services/agent-actor-context.service';
 import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
+import { type ToolHints } from 'src/engine/metadata-modules/ai-router/types/tool-hints.interface';
 import { getObjectMetadataMapItemByNameSingular } from 'src/engine/metadata-modules/utils/get-object-metadata-map-item-by-name-singular.util';
 import { WorkspacePermissionsCacheService } from 'src/engine/metadata-modules/workspace-permissions-cache/workspace-permissions-cache.service';
 import { TwentyORMGlobalManager } from 'src/engine/twenty-orm/twenty-orm-global.manager';
@@ -43,6 +44,16 @@ export interface AgentExecutionResult {
  usage: LanguageModelUsage;
 }

+export interface StreamChatResponseResult {
+  stream: ReturnType<typeof streamText>;
+  timings: {
+    contextBuildTimeMs: number;
+    toolGenerationTimeMs: number;
+    aiRequestPrepTimeMs: number;
+    toolCount: number;
+  };
+}
+
@Injectable()
 export class AgentExecutionService implements AgentExecutionContext {
  private readonly logger = new Logger(AgentExecutionService.name);
@@ -68,6 +79,7 @@ export class AgentExecutionService implements AgentExecutionContext {
    roleIds,
    excludeHandoffTools = false,
    userWorkspaceId,
+    toolHints,
  }: {
    system: string;
    agent: AgentEntity | null;
@@ -76,6 +88,7 @@ export class AgentExecutionService implements AgentExecutionContext {
    roleIds?: string[];
    excludeHandoffTools?: boolean;
    userWorkspaceId?: string;
+    toolHints?: ToolHints;
  }) {
    try {
      if (agent) {
@@ -98,6 +111,7 @@ export class AgentExecutionService implements AgentExecutionContext {
            actorContext,
            roleIds,
            userWorkspaceId,
+            toolHints,
          );

        let handoffTools = {};
@@ -169,6 +183,9 @@ export class AgentExecutionService implements AgentExecutionContext {
    }
  }

+  // Fetches and formats record data to provide context for AI agents
+  // Respects permissions and field restrictions based on user role
+  // Returns a JSON string with record data and workspace URLs
  async getContextForSystemPrompt(
    workspace: WorkspaceEntity,
    recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType,
@@ -272,12 +289,14 @@ export class AgentExecutionService implements AgentExecutionContext {
    agentId,
    messages,
    recordIdsByObjectMetadataNameSingular,
+    toolHints,
  }: {
    workspace: WorkspaceEntity;
    userWorkspaceId: string;
    agentId: string;
    messages: UIMessage<unknown, UIDataTypes, UITools>[];
    recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
+    toolHints?: ToolHints;
  }): Promise<{
    stream: ReturnType<typeof streamText>;
    timings: {
@@ -335,6 +354,7 @@ export class AgentExecutionService implements AgentExecutionContext {
        actorContext,
        roleIds: [roleId, ...(agent?.roleId ? [agent?.roleId] : [])],
        userWorkspaceId,
+        toolHints,
      });

      const aiRequestPrepTime = Date.now() - aiRequestPrepStart;
@@ -27,6 +27,12 @@ import {
 import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
 import { AiRouterService } from 'src/engine/metadata-modules/ai-router/ai-router.service';

+export type TokenUsage = {
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+};
+
 export type StreamAgentChatOptions = {
  threadId: string;
  userWorkspaceId: string;
@@ -87,17 +93,18 @@ export class AgentStreamingService {
          });

          const routingStart = Date.now();
+          const includeDebugInfo = true;
          const routeResult = await this.aiRouterService.routeMessage(
            {
              messages,
              workspaceId: workspace.id,
              routerModel: workspace.routerModel,
            },
-            true,
+            includeDebugInfo,
          );

          const routingTime = Date.now() - routingStart;
-          const { agent, debugInfo } = routeResult;
+          const { agent, debugInfo, toolHints } = routeResult;

          if (!agent) {
            writer.write({
@@ -154,6 +161,7 @@ export class AgentStreamingService {
            userWorkspaceId,
            messages,
            recordIdsByObjectMetadataNameSingular,
+            toolHints,
          });

          const routedStatusPart = {
@@ -199,41 +207,7 @@ export class AgentStreamingService {
                  part.type.startsWith('tool-'),
                ).length;

-                let tokenUsage: {
-                  promptTokens: number;
-                  completionTokens: number;
-                  totalTokens: number;
-                } | null = null;
-
-                try {
-                  const usage = await result.usage;
-
-                  const usageWithTokens = usage as {
-                    inputTokens?: number;
-                    outputTokens?: number;
-                    promptTokens?: number;
-                    completionTokens?: number;
-                    totalTokens?: number;
-                  };
-
-                  tokenUsage = {
-                    promptTokens:
-                      usageWithTokens.inputTokens ??
-                      usageWithTokens.promptTokens ??
-                      0,
-                    completionTokens:
-                      usageWithTokens.outputTokens ??
-                      usageWithTokens.completionTokens ??
-                      0,
-                    totalTokens: usageWithTokens.totalTokens ?? 0,
-                  };
-
-                  this.logger.log(
-                    `Agent execution usage: ${tokenUsage.promptTokens} prompt + ${tokenUsage.completionTokens} completion = ${tokenUsage.totalTokens} total tokens`,
-                  );
-                } catch (error) {
-                  this.logger.warn('Failed to get token usage:', error);
-                }
+                const tokenUsage = await this.extractTokenUsage(result.usage);

                const agentExecutionTime = Date.now() - agentExecutionStart;

@@ -325,8 +299,45 @@ export class AgentStreamingService {

      pipeUIMessageStreamToResponse({ stream, response });
    } catch (error) {
-      this.logger.error(error.message);
+      this.logger.error(
+        'Failed to stream agent chat:',
+        error instanceof Error ? error.message : String(error),
+      );
      response.end();
    }
  }
+
+  private async extractTokenUsage(
+    usagePromise: Promise<unknown>,
+  ): Promise<TokenUsage | null> {
+    try {
+      const usage = await usagePromise;
+
+      const usageWithTokens = usage as {
+        inputTokens?: number;
+        outputTokens?: number;
+        promptTokens?: number;
+        completionTokens?: number;
+        totalTokens?: number;
+      };
+
+      const tokenUsage = {
+        promptTokens:
+          usageWithTokens.inputTokens ?? usageWithTokens.promptTokens ?? 0,
+        completionTokens:
+          usageWithTokens.outputTokens ?? usageWithTokens.completionTokens ?? 0,
+        totalTokens: usageWithTokens.totalTokens ?? 0,
+      };
+
+      this.logger.log(
+        `Agent execution usage: ${tokenUsage.promptTokens} prompt + ${tokenUsage.completionTokens} completion = ${tokenUsage.totalTokens} total tokens`,
+      );
+
+      return tokenUsage;
+    } catch (error) {
+      this.logger.warn('Failed to get token usage:', error);
+
+      return null;
+    }
+  }
 }
@@ -9,6 +9,7 @@ import { ToolAdapterService } from 'src/engine/core-modules/ai/services/tool-ada
 import { ToolService } from 'src/engine/core-modules/ai/services/tool.service';
 import { SearchArticlesTool } from 'src/engine/core-modules/tool/tools/search-articles-tool/search-articles-tool';
 import { AgentEntity } from 'src/engine/metadata-modules/agent/agent.entity';
+import { type ToolHints } from 'src/engine/metadata-modules/ai-router/types/tool-hints.interface';
 import { PermissionFlagType } from 'src/engine/metadata-modules/permissions/constants/permission-flag-type.constants';
 import { PermissionsService } from 'src/engine/metadata-modules/permissions/permissions.service';
 import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
@@ -37,6 +38,7 @@ export class AgentToolGeneratorService {
    actorContext?: ActorMetadata,
    roleIds?: string[],
    userWorkspaceId?: string,
+    toolHints?: ToolHints,
  ): Promise<ToolSet> {
    let tools: ToolSet = {};

@@ -78,6 +80,7 @@ export class AgentToolGeneratorService {
        workspaceId,
        actorContext,
        userWorkspaceId,
+        toolHints,
      );

      tools = { ...tools, ...databaseTools };
@@ -19,6 +19,8 @@ import { ObjectMetadataService } from 'src/engine/metadata-modules/object-metada
 import { DATA_MANIPULATOR_AGENT } from 'src/engine/workspace-manager/workspace-sync-metadata/standard-agents/agents/data-manipulator-agent';
 import { HELPER_AGENT } from 'src/engine/workspace-manager/workspace-sync-metadata/standard-agents/agents/helper-agent';

+import { type ToolHints } from './types/tool-hints.interface';
+
 export interface AiRouterContext {
  messages: UIMessage<unknown, UIDataTypes, UITools>[];
  workspaceId: string;
@@ -27,6 +29,7 @@ export interface AiRouterContext {

 export interface AiRouterResult {
  agent: AgentEntity | null;
+  toolHints?: ToolHints;
  debugInfo?: {
    availableAgents: Array<{ id: string; label: string }>;
    routerModel: string;
@@ -47,6 +50,9 @@ export class AiRouterService {
    private readonly objectMetadataService: ObjectMetadataService,
  ) {}

+  // Routes a user message to the most appropriate agent
+  // Uses AI to analyze the conversation and select the best agent
+  // Returns the selected agent along with tool hints for optimization
  async routeMessage(
    context: AiRouterContext,
    includeDebugInfo = false,
@@ -105,18 +111,42 @@ export class AiRouterService {
        currentMessage,
      );

+      const agentIds = availableAgents.map((agent) => agent.id);
+
+      if (agentIds.length === 0) {
+        throw new Error('No agent IDs available for routing schema');
+      }
+
      const routerDecisionSchema = z.object({
        agentId: z
-          .enum(availableAgents.map((agent) => agent.id))
+          .enum([agentIds[0], ...agentIds.slice(1)])
          .describe('The ID of the most suitable agent to handle this message'),
+        toolHints: z
+          .object({
+            relevantObjects: z
+              .array(z.string())
+              .optional()
+              .describe(
+                'Names of the specific objects mentioned in the query (e.g., "person", "company")',
+              ),
+            operations: z
+              .array(z.enum(['find', 'create', 'update', 'delete']))
+              .optional()
+              .describe(
+                'Specific operations needed: find (search/query), create (new records), update (modify), delete (remove)',
+              ),
+          })
+          .optional(),
      });

+      const ROUTER_TEMPERATURE = 0.1; // Low temperature for deterministic routing
+
      const result = await generateObject({
        model,
        system: systemPrompt,
        prompt: userPrompt,
        schema: routerDecisionSchema,
-        temperature: 0.1,
+        temperature: ROUTER_TEMPERATURE,
        experimental_telemetry: AI_TELEMETRY_CONFIG,
      });

@@ -148,7 +178,11 @@ export class AiRouterService {
        }
      }

-      return { agent: selectedAgent ?? null, debugInfo };
+      return {
+        agent: selectedAgent ?? null,
+        toolHints: result.object.toolHints,
+        debugInfo,
+      };
    } catch (error) {
      this.logger.error(
        'Routing to agent failed, falling back to Helper agent:',
@@ -260,7 +294,25 @@ ${workspaceObjectsList}`;
 Available agents:
 ${agentDescriptions}

-Your task is to analyze the user's message and conversation history, then select the most appropriate agent to handle it. Choose the agent whose description and capabilities best match the user's request.`;
+Your task is to:
+1. Select the most appropriate agent
+2. Identify specific objects mentioned in the query (if any)
+3. Determine which operations are needed
+
+For toolHints:
+- relevantObjects: Extract object names the user is asking about (e.g., if asking about "companies and people", return ["company", "person"])
+- operations: Array of needed operations from: ["find", "create", "update", "delete"]
+  - "find": for searching, querying, or reading data
+  - "create": for creating new records
+  - "update": for modifying existing records
+  - "delete": for removing records
+
+Examples:
+- "Show me all companies" → operations: ["find"]
+- "Create a task for John" → operations: ["create"]
+- "Update the company name" → operations: ["find", "update"]
+
+This helps optimize the agent's tool context by only loading relevant tools.`;
  }

  private buildRouterUserPrompt(
@@ -0,0 +1,8 @@
+export type ToolOperation = 'find' | 'create' | 'update' | 'delete';
+
+export interface ToolHints {
+  // Object names (singular or plural) that are relevant to the query
+  relevantObjects?: string[];
+  // Specific CRUD operations needed for the query
+  operations?: ToolOperation[];
+}