diff --git a/agents-run-api/src/__tests__/agents/generateTaskHandler.test.ts b/agents-run-api/src/__tests__/agents/generateTaskHandler.test.ts index 7caf9057f..02d5a4978 100644 --- a/agents-run-api/src/__tests__/agents/generateTaskHandler.test.ts +++ b/agents-run-api/src/__tests__/agents/generateTaskHandler.test.ts @@ -21,6 +21,7 @@ const { getArtifactComponentsForAgentMock, getExternalAgentsForSubAgentMock, getTeamAgentsForSubAgentMock, + getAgentWithDefaultSubAgentMock, getProjectMock, dbResultToMcpToolMock, } = vi.hoisted(() => { @@ -168,6 +169,7 @@ const { data: [ { id: 'team-relation-1', + targetAgentId: 'team-agent-1', targetAgent: { id: 'team-agent-1', name: 'Team Agent 1', @@ -180,6 +182,27 @@ const { }) ); + const getAgentWithDefaultSubAgentMock = vi.fn(() => + vi.fn().mockResolvedValue({ + id: 'team-agent-1', + name: 'Team Agent 1', + description: 'A team agent for delegation', + defaultSubAgent: { + id: 'team-agent-1-default', + name: 'Team Agent 1 Default', + description: 'Default sub agent for team agent 1', + prompt: 'You are a helpful team agent', + conversationHistoryConfig: { + mode: 'full', + limit: 10, + }, + models: null, + createdAt: '2024-01-01T00:00:00Z', + updatedAt: '2024-01-01T00:00:00Z', + }, + }) + ); + const getProjectMock = vi.fn(() => vi.fn().mockResolvedValue({ id: 'test-project', @@ -208,6 +231,7 @@ const { getArtifactComponentsForAgentMock, getExternalAgentsForSubAgentMock, getTeamAgentsForSubAgentMock, + getAgentWithDefaultSubAgentMock, getProjectMock, dbResultToMcpToolMock, }; @@ -221,6 +245,7 @@ vi.mock('@inkeep/agents-core', () => ({ getAgentAgent: getAgentAgentMock, getAgentAgentById: getAgentAgentByIdMock, getTeamAgentsForSubAgent: getTeamAgentsForSubAgentMock, + getAgentWithDefaultSubAgent: getAgentWithDefaultSubAgentMock, getTracer: vi.fn().mockReturnValue({ startSpan: vi.fn().mockReturnValue({ setAttributes: vi.fn(), @@ -714,6 +739,127 @@ describe('generateTaskHandler', () => { }, }); }); + + it('should enhance team relations with default sub agent data', async () => { + // Clear all mocks first + vi.clearAllMocks(); + + // Mock the initial team relations call (first call) + getTeamAgentsForSubAgentMock.mockReturnValueOnce( + vi.fn().mockResolvedValue({ + data: [ + { + id: 'team-relation-1', + targetAgentId: 'team-agent-1', + targetAgent: { + id: 'team-agent-1', + name: 'Team Agent 1', + description: 'A team agent for delegation', + }, + headers: { 'X-Custom-Header': 'team-value' }, + }, + ], + pagination: { page: 1, limit: 10, total: 1, pages: 1 }, + }) + ); + + // Mock the default sub agent for the team agent + getAgentWithDefaultSubAgentMock.mockReturnValueOnce( + vi.fn().mockResolvedValue({ + id: 'team-agent-1', + name: 'Team Agent 1', + description: 'A team agent for delegation', + defaultSubAgent: { + id: 'team-agent-1-default', + name: 'Team Agent 1 Default', + description: 'Default sub agent for team agent 1', + prompt: 'You are a helpful team agent', + conversationHistoryConfig: { + mode: 'full', + limit: 10, + }, + models: null, + createdAt: '2024-01-01T00:00:00Z', + updatedAt: '2024-01-01T00:00:00Z', + }, + }) + ); + + // Mock related agents for the default sub agent (second call) + getRelatedAgentsForAgentMock.mockReturnValueOnce( + vi.fn().mockResolvedValue({ + data: [ + { + id: 'related-agent-1', + name: 'Related Agent 1', + description: 'A related agent', + relationType: 'delegate', + }, + ], + }) + ); + + // Mock external agents for the default sub agent (second call) + getExternalAgentsForSubAgentMock.mockReturnValueOnce( + vi.fn().mockResolvedValue({ + data: [ + { + id: 'external-relation-1', + externalAgent: { + id: 'external-agent-1', + name: 'External Agent 1', + description: 'An external agent', + }, + }, + ], + }) + ); + + // Mock team agents for the default sub agent (second call) + getTeamAgentsForSubAgentMock.mockReturnValueOnce( + vi.fn().mockResolvedValue({ + data: [ + { + id: 'team-relation-2', + targetAgentId: 'team-agent-2', + targetAgent: { + id: 'team-agent-2', + name: 'Team Agent 2', + description: 'Another team agent', + }, + }, + ], + }) + ); + + const taskHandler = createTaskHandler(mockConfig); + + const task: A2ATask = { + id: 'task-123', + input: { + parts: [{ kind: 'text', text: 'Test with enhanced team relations' }], + }, + }; + + await taskHandler(task); + + // Verify that getAgentWithDefaultSubAgent was called for the team agent + expect(getAgentWithDefaultSubAgentMock).toHaveBeenCalledWith(expect.anything()); + + // Verify that the Agent constructor received enhanced team relations + expect(lastAgentConstructorArgs).toBeDefined(); + expect(lastAgentConstructorArgs.delegateRelations).toBeDefined(); + + const teamDelegateRelation = lastAgentConstructorArgs.delegateRelations.find( + (rel: any) => rel.type === 'team' + ); + expect(teamDelegateRelation).toBeDefined(); + expect(teamDelegateRelation.config.id).toBe('team-agent-1'); + expect(teamDelegateRelation.config.name).toBe('Team Agent 1'); + // The description should be enhanced with related agents information + expect(teamDelegateRelation.config.description).toContain('A team agent for delegation'); + expect(teamDelegateRelation.config.description).toContain('Can delegate to:'); + }); }); describe('createTaskHandlerConfig', () => { diff --git a/agents-run-api/src/agents/generateTaskHandler.ts b/agents-run-api/src/agents/generateTaskHandler.ts index 2673094dd..86e11aa9a 100644 --- a/agents-run-api/src/agents/generateTaskHandler.ts +++ b/agents-run-api/src/agents/generateTaskHandler.ts @@ -4,6 +4,7 @@ import { dbResultToMcpTool, generateId, getAgentById, + getAgentWithDefaultSubAgent, getArtifactComponentsForAgent, getDataComponentsForAgent, getExternalAgentsForSubAgent, @@ -17,7 +18,7 @@ import { TaskState, } from '@inkeep/agents-core'; import type { A2ATask, A2ATaskResult } from '../a2a/types'; -import { generateDescriptionWithTransfers } from '../data/agents'; +import { generateDescriptionWithRelationData } from '../data/agents'; import dbClient from '../data/db/dbClient'; import { getLogger } from '../logger'; import { agentSessionManager } from '../services/AgentSession'; @@ -158,11 +159,20 @@ export const createTaskHandler = ( subAgentId: relation.id, }, }); + const relatedAgentTeamAgentRelations = await getTeamAgentsForSubAgent(dbClient)({ + scopes: { + tenantId: config.tenantId, + projectId: config.projectId, + agentId: config.agentId, + subAgentId: relation.id, + }, + }); - const enhancedDescription = generateDescriptionWithTransfers( + const enhancedDescription = generateDescriptionWithRelationData( relation.description || '', relatedAgentRelations.data, - relatedAgentExternalAgentRelations.data + relatedAgentExternalAgentRelations.data, + relatedAgentTeamAgentRelations.data ); return { ...relation, description: enhancedDescription }; } @@ -173,6 +183,78 @@ export const createTaskHandler = ( }) ); + const enhancedTeamRelations = await Promise.all( + teamRelations.data.map(async (relation) => { + try { + // Get the default sub agent for the team agent + const teamAgentWithDefault = await getAgentWithDefaultSubAgent(dbClient)({ + scopes: { + tenantId: config.tenantId, + projectId: config.projectId, + agentId: relation.targetAgentId, + }, + }); + + if (teamAgentWithDefault?.defaultSubAgent) { + const defaultSubAgent = teamAgentWithDefault.defaultSubAgent; + + // Get related agents for the default sub agent + const relatedAgentRelations = await getRelatedAgentsForAgent(dbClient)({ + scopes: { + tenantId: config.tenantId, + projectId: config.projectId, + agentId: relation.targetAgentId, + }, + subAgentId: defaultSubAgent.id, + }); + + // Get external agents for the default sub agent + const relatedAgentExternalAgentRelations = await getExternalAgentsForSubAgent( + dbClient + )({ + scopes: { + tenantId: config.tenantId, + projectId: config.projectId, + agentId: relation.targetAgentId, + subAgentId: defaultSubAgent.id, + }, + }); + + // Get team agents for the default sub agent + const relatedAgentTeamAgentRelations = await getTeamAgentsForSubAgent(dbClient)({ + scopes: { + tenantId: config.tenantId, + projectId: config.projectId, + agentId: relation.targetAgentId, + subAgentId: defaultSubAgent.id, + }, + }); + + const enhancedDescription = generateDescriptionWithRelationData( + teamAgentWithDefault.description || '', + relatedAgentRelations.data, + relatedAgentExternalAgentRelations.data, + relatedAgentTeamAgentRelations.data + ); + + return { + ...relation, + targetAgent: { + ...relation.targetAgent, + description: enhancedDescription, + }, + }; + } + } catch (error) { + logger.warn( + { targetAgentId: relation.targetAgentId, error }, + 'Failed to enhance team agent description' + ); + } + return relation; + }) + ); + const prompt = 'prompt' in config.agentSchema ? config.agentSchema.prompt : ''; const models = 'models' in config.agentSchema ? config.agentSchema.models : undefined; const stopWhen = 'stopWhen' in config.agentSchema ? config.agentSchema.stopWhen : undefined; @@ -261,7 +343,7 @@ export const createTaskHandler = ( relationType: 'delegate', }, })), - ...teamRelations.data.map((relation) => ({ + ...enhancedTeamRelations.map((relation) => ({ type: 'team' as const, config: { id: relation.targetAgent.id, diff --git a/agents-run-api/src/data/agents.ts b/agents-run-api/src/data/agents.ts index e1a406edf..87a3ad348 100644 --- a/agents-run-api/src/data/agents.ts +++ b/agents-run-api/src/data/agents.ts @@ -61,10 +61,11 @@ export function createAgentCard({ * @param internalRelations - Pre-computed internal relations * @param externalRelations - Pre-computed external relations */ -export function generateDescriptionWithTransfers( +export function generateDescriptionWithRelationData( baseDescription: string, internalRelations: any[], - externalRelations: any[] + externalRelations: any[], + teamRelations: any[] ): string { // Filter relations by type const transfers = [...internalRelations.filter((rel) => rel.relationType === 'transfer')]; @@ -72,6 +73,7 @@ export function generateDescriptionWithTransfers( const delegates = [ ...internalRelations.filter((rel) => rel.relationType === 'delegate'), ...externalRelations.map((data) => data.externalAgent), + ...teamRelations.map((data) => data.targetAgent), ]; // If no relations, return base description