protofy agent network: inherit protocol information from parents to c…

…hildren
Protofy-xyz · Nov 17, 2024 · e2e6096 · e2e6096
1 parent fca5a12
commit e2e6096
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 51 deletions.
diff --git 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
@@ -5,6 +5,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema';
 let userSchema: z.ZodObject<any>;
 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' });