feat: create agent, tool primitives and update example

example/surprise_trip.ts

@@ -2,6 +2,7 @@
  * Example borrowed from CrewAI.
  */
 import { Agent, Team } from '@dead-simple-ai-agent/framework'
+import { tool } from '@dead-simple-ai-agent/framework/tool'
 import { WikipediaQueryRun } from '@langchain/community/tools/wikipedia_query_run'
 import { z } from 'zod'
 
@@ -10,15 +11,6 @@ const wikipedia = new WikipediaQueryRun({
   maxDocContentLength: 4000,
 })
 
-const wikipediaTool = {
-  name: 'wikipedia',
-  description: 'Tool for querying Wikipedia',
-  parameters: z.object({
-    query: z.string().describe('The query to search Wikipedia with'),
-  }),
-  function: ({ query }) => wikipedia.invoke(query),
-}
-
 const personalizedActivityPlanner = new Agent({
   role: 'Activity Planner',
   description: `
@@ -35,7 +27,15 @@ const landmarkScout = new Agent({
     You are skilled at researching and finding interesting landmarks at the destination.
     Your goal is to find historical landmarks, museums, and other interesting places.
   `,
-  tools: [wikipediaTool],
+  tools: {
+    wikipedia: tool({
+      description: 'Tool for querying Wikipedia',
+      parameters: z.object({
+        query: z.string().describe('The query to search Wikipedia with'),
+      }),
+      execute: ({ query }) => wikipedia.invoke(query),
+    }),
+  },
 })
 
 const restaurantScout = new Agent({

packages/framework/package.json

@@ -6,6 +6,9 @@
   "exports": {
     ".": {
       "bun": "./src/index.ts"
+    },
+    "./tool": {
+      "bun": "./src/tool.ts"
     }
   },
   "type": "module",

packages/framework/src/agent.ts

@@ -0,0 +1,26 @@
+import { Tool } from './tool.js'
+import { RequiredOptionals } from './types.js'
+
+export type AgentOptions = {
+  role: string
+  description: string
+  tools?: {
+    [key: string]: Tool
+  }
+  model?: string
+}
+
+const defaults: RequiredOptionals<AgentOptions> = {
+  tools: {},
+  model: 'gpt-4o',
+}
+
+/**
+ * Helper utility to create an agent with defaults.
+ */
+export const agent = (options: AgentOptions) => {
+  return {
+    ...defaults,
+    ...options,
+  }
+}

packages/framework/src/index.ts

@@ -2,7 +2,9 @@ import s from 'dedent'
 import OpenAI from 'openai'
 import { zodFunction, zodResponseFormat } from 'openai/helpers/zod'
 import type { ChatCompletionMessageParam } from 'openai/resources/chat/completions'
-import { z } from 'zod'
+import { z, ZodType, ZodTypeAny } from 'zod'
+
+import { Tool } from './tool.js'
 
 // tbd: abstract this away or not? most APIs are OpenAI compatible
 const openai = new OpenAI()
@@ -26,17 +28,14 @@ class CLIProtocol implements Protocol {
 }
 
 // tbd: make this Vercel AI SDK compatible, as it has nicest interface
-type ToolDefinition<T extends z.ZodObject<{}>> = {
-  name: string
-  description: string
-  parameters: T
-  function: (parameters: z.infer<T>) => Promise<string>
+type Toolkit = {
+  [key: string]: Tool
 }
 
 interface AgentConfig {
   role: string
   description: string
-  tools?: ToolDefinition<any>[]
+  tools?: Toolkit
   model?: string
   protocol?: Protocol
 }
@@ -46,14 +45,14 @@ export class Agent {
   role: string
 
   private description: string
-  private tools: ToolDefinition<any>[]
+  private tools: Toolkit
   private model: string
   private protocol: Protocol
 
   constructor({
     role,
     description,
-    tools = [],
+    tools = {},
     model = 'gpt-4o',
     protocol = new CLIProtocol(),
   }: AgentConfig) {
@@ -65,43 +64,15 @@ export class Agent {
   }
 
   async executeTask(messages: Message[], agents: Agent[]): Promise<string> {
-    // tbd: after we implememt this, it keeps delegating
-    // const tools = [
-    //   {
-    //     name: 'ask_a_question',
-    //     description: s`
-    //       Ask a specific question to one of the following agents:
-    //       <agents>
-    //         ${agents.map((agent, index) => `<agent index="${index}">${agent.role}</agent>`)}
-    //       </agents>
-    //     `,
-    //     parameters: z.object({
-    //       agent: z.number().describe('The index of the agent to ask the question'),
-    //       question: z.string().describe('The question you want the agent to answer'),
-    //       context: z.string().describe('All context necessary to execute the task'),
-    //     }),
-    //     function: async ({ agent, question, context }) => {
-    //       console.log('ask_a_question', agent, question, context)
-    //       const selectedAgent = agents[agent]
-    //       if (!selectedAgent) {
-    //         throw new Error('Invalid agent')
-    //       }
-    //       return selectedAgent.executeTask(
-    //         [
-    //           {
-    //             role: 'user',
-    //             content: context,
-    //           },
-    //           {
-    //             role: 'user',
-    //             content: question,
-    //           },
-    //         ],
-    //         agents
-    //       )
-    //     },
-    //   },
-    // ]
+    const tools = Object.entries(this.tools).map(([name, tool]) =>
+      zodFunction({
+        name,
+        parameters: tool.parameters,
+        function: tool.execute,
+        description: tool.description,
+      })
+    )
+
     const response = await openai.beta.chat.completions.parse({
       model: this.model,
       // tbd: verify the prompt
@@ -121,10 +92,7 @@ export class Agent {
         },
         ...messages,
       ],
-      // tbd: add other tools
-      // tbd: should we include agent description in the prompt too? we need to list responsibilities
-      // but keep context window in mind
-      tools: this.tools.length > 0 ? this.tools.map(zodFunction) : undefined,
+      tools: tools.length > 0 ? tools : undefined,
       response_format: zodResponseFormat(
         z.object({
           response: z.discriminatedUnion('kind', [
@@ -151,12 +119,12 @@ export class Agent {
             throw new Error('Tool call is not a function')
           }
 
-          const tool = this.tools.find((t) => t.name === toolCall.function.name)
+          const tool = this.tools[toolCall.function.name]
           if (!tool) {
             throw new Error(`Unknown tool: ${toolCall.function.name}`)
           }
 
-          const content = await tool.function(toolCall.function.parsed_arguments)
+          const content = await tool.execute(toolCall.function.parsed_arguments)
           return {
             role: 'tool' as const,
             tool_call_id: toolCall.id,

packages/framework/src/tool.ts

@@ -0,0 +1,9 @@
+import z, { ZodTypeAny } from 'zod'
+
+export type Tool<P extends ZodTypeAny = any> = {
+  description: string
+  parameters: P
+  execute: (parameters: z.infer<P>) => Promise<string>
+}
+
+export const tool = <P extends ZodTypeAny>(tool: Tool<P>): Tool<P> => tool

packages/framework/src/types.ts

@@ -0,0 +1,16 @@
+/**
+ * Utility type to get optional keys from T.
+ */
+export type OptionalKeys<T> = {
+  [K in keyof T]-?: undefined extends T[K] ? K : never
+}[keyof T]
+
+/**
+ * Utility type to get optional properties from T.
+ */
+export type OptionalProperties<T> = Pick<T, OptionalKeys<T>>
+
+/**
+ * Utility type to make optional properties required (only includes optional props).
+ */
+export type RequiredOptionals<T> = Required<OptionalProperties<T>>