diff --git a/packages/protofy/src/Agent.ts b/packages/protofy/src/Agent.ts index 2b4ca4c73..56ef6d370 100644 --- a/packages/protofy/src/Agent.ts +++ b/packages/protofy/src/Agent.ts @@ -42,11 +42,13 @@ export class Agent { children: Agent[]; input: AgentInterface | undefined; output: AgentInterface | undefined; - constructor(data: AgentData, agents: Agent[] = []) { + parent: Agent | undefined; + constructor(data: AgentData, children: Agent[] = [], parent: Agent = undefined, ) { this.data = data; - this.children = agents; + this.children = children; this.input = data.input && new AgentInputInterface(data.input, this); this.output = data.output && new AgentOutputInterface(data.output, this); + this.parent = parent; } getName() { @@ -66,6 +68,12 @@ export class Agent { } getProtocol() { + if(this.parent) { + return { + ...this.parent.getProtocol(), + ...this.data.protocol + } + } return this.data.protocol } @@ -86,6 +94,7 @@ export class Agent { } addChildren(agents: Agent[]) { + agents.forEach(agent => agent.setParent(this)); this.children.push(...agents); } @@ -94,12 +103,17 @@ export class Agent { } addChild(agent: Agent) { + agent.setParent(this); this.children.push(agent); } getChild(id: string) { return this.children && this.children.find(agent => agent.data.id === id); } + + setParent(agent: Agent) { + this.parent = agent; + } } export class AgentInterface { diff --git a/packages/protofy/tests/agent.test.ts b/packages/protofy/tests/agent.test.ts index 58d20908e..780e316b8 100644 --- a/packages/protofy/tests/agent.test.ts +++ b/packages/protofy/tests/agent.test.ts @@ -5,6 +5,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema'; let userSchema: z.ZodObject; let returnSchema: z.ZodString; let agent: Agent; +let childAgent: Agent; describe('Agents basic behavior', () => { beforeEach(() => { @@ -17,13 +18,20 @@ describe('Agents basic behavior', () => { returnSchema = z.string() + childAgent = new Agent({ + id: 'childAgent', + name: 'childAgent', + description: 'Child agent', + tags: ['child'] + }); + agent = new Agent({ id: 'getDisplayInfo', name: 'getDisplayInfo', description: 'Get display info of a user', tags: ['user', 'display'], protocol: { - type: 'function' + type: 'function', }, input: { shape: zodToJsonSchema(userSchema, "user"), @@ -49,64 +57,16 @@ describe('Agents basic behavior', () => { }); it('Should be able to add children to the agent', () => { - const childAgent = new Agent({ - id: 'childAgent', - name: 'childAgent', - description: 'Child agent', - tags: ['child'], - protocol: { - type: 'function' - }, - input: { - shape: zodToJsonSchema(userSchema, "user") - }, - output: { - shape: zodToJsonSchema(returnSchema, "displayInfo") - } - }) - agent.addChildren([childAgent]); expect(agent.getChildren()).toEqual([childAgent]); }); it('Should be able to add a single child to the agent', () => { - const childAgent = new Agent({ - id: 'childAgent', - name: 'childAgent', - description: 'Child agent', - tags: ['child'], - protocol: { - type: 'function' - }, - input: { - shape: zodToJsonSchema(userSchema, "user") - }, - output: { - shape: zodToJsonSchema(returnSchema, "displayInfo") - } - }) - agent.addChild(childAgent); expect(agent.getChildren()).toEqual([childAgent]); }); it('Should be able to get a child by id', () => { - const childAgent = new Agent({ - id: 'childAgent', - name: 'childAgent', - description: 'Child agent', - tags: ['child'], - protocol: { - type: 'function' - }, - input: { - shape: zodToJsonSchema(userSchema, "user") - }, - output: { - shape: zodToJsonSchema(returnSchema, "displayInfo") - } - }) - agent.addChild(childAgent); expect(agent.getChild('childAgent')).toEqual(childAgent); }); @@ -115,6 +75,13 @@ describe('Agents basic behavior', () => { expect(agent.getChild('childAgent')).toBeUndefined(); }); + it('Should return combined agent protocol from the parent', () => { + agent.addChild(childAgent); + expect(agent.getChild('childAgent').getProtocol()).toEqual({ + type: 'function' + }); + }); + it('Should combine agent protocol definition with input and output protocol definition, to reduce verbosity', () => { expect(agent.getInputProtocol()).toEqual({ type: 'function', encoder: 'object' }); expect(agent.getOutputProtocol()).toEqual({ type: 'function' });