feat(agent): basic subagent function of agent gateway

Author: sheepbox8646

agent/src/agent.ts

@@ -1,24 +1,31 @@
 import { generateText, ModelMessage, stepCountIs, streamText, TextStreamPart, ToolSet } from 'ai'
 import { createChatGateway } from './gateway'
-import { ClientType, Schedule } from './types'
+import { BaseModelConfig, Schedule } from './types'
 import { system, schedule } from './prompts'
 import { AuthFetcher } from './index'
 import { getScheduleTools } from './tools/schedule'
 import { getWebTools } from './tools/web'
+import { subagentSystem } from './prompts/subagent'
+import { getSubagentTools } from './tools/subagent'
 
-export interface AgentParams {
-  apiKey: string
-  baseUrl: string
-  model: string
-  clientType: ClientType
+export enum AgentAction {
+  WebSearch = 'web_search',
+  Message = 'message',
+  Subagent = 'subagent',
+  Schedule = 'schedule',
+  Skill = 'skill',
+}
+
+export interface AgentParams extends BaseModelConfig {
   locale?: Intl.LocalesArgument
   language?: string
   maxSteps?: number
-  maxContextLoadTime: number
+  maxContextLoadTime?: number
   platforms?: string[]
   currentPlatform?: string
   braveApiKey?: string
   braveBaseUrl?: string
+  allowed?: AgentAction[]
 }
 
 export interface AgentInput {
@@ -37,23 +44,40 @@ export const createAgent = (
   const gateway = createChatGateway(params.clientType)
   const messages: ModelMessage[] = []
 
+  const allowedActions = params.allowed
+    ?? Object.values(AgentAction)
+
   const maxSteps = params.maxSteps ?? 50
 
   const getTools = () => {
-    const scheduleTools = getScheduleTools({ fetch: fetcher })
-    const tools: ToolSet = {
-      ...scheduleTools,
+    const tools: ToolSet = {}
+
+    if (allowedActions.includes(AgentAction.Schedule)) {
+      const scheduleTools = getScheduleTools({ fetch: fetcher })
+      Object.assign(tools, scheduleTools)
     }
 
-    // Add web search tools if Brave API key is provided
-    if (params.braveApiKey) {
+    if (params.braveApiKey && allowedActions.includes(AgentAction.WebSearch)) {
       const webTools = getWebTools({
         braveApiKey: params.braveApiKey,
         braveBaseUrl: params.braveBaseUrl,
       })
       Object.assign(tools, webTools)
     }
 
+    if (allowedActions.includes(AgentAction.Subagent)) {
+      const subagentTools = getSubagentTools({
+        fetch: fetcher,
+        apiKey: params.apiKey,
+        baseUrl: params.baseUrl,
+        model: params.model,
+        clientType: params.clientType,
+        braveApiKey: params.braveApiKey,
+        braveBaseUrl: params.braveBaseUrl,
+      })
+      Object.assign(tools, subagentTools)
+    }
+    
     return tools
   }
 
@@ -62,7 +86,7 @@ export const createAgent = (
       date: new Date(),
       locale: params.locale,
       language: params.language,
-      maxContextLoadTime: params.maxContextLoadTime,
+      maxContextLoadTime: params.maxContextLoadTime ?? 1550,
       platforms: params.platforms ?? [],
       currentPlatform: params.currentPlatform,
     })
@@ -90,6 +114,39 @@ export const createAgent = (
     }
   }
 
+  const askAsSubagent = async (
+    input: AgentInput,
+    options: {
+      name: string
+      description?: string
+    }
+  ): Promise<AgentResult> => {
+    messages.push(...input.messages)
+    const user: ModelMessage = {
+      role: 'user',
+      content: input.query,
+    }
+    messages.push(user)
+    const { response } = await generateText({
+      model: gateway({
+        apiKey: params.apiKey,
+        baseURL: params.baseUrl,
+      })(params.model),
+      system: subagentSystem({ date: new Date(), name: options.name, description: options.description }),
+      stopWhen: stepCountIs(maxSteps),
+      messages,
+      prepareStep: () => {
+        return {
+          system: subagentSystem({ date: new Date(), name: options.name, description: options.description }),
+        }
+      },
+      tools: getTools(),
+    })
+    return {
+      messages: [user, ...response.messages],
+    }
+  }
+
   async function* stream(input: AgentInput): AsyncGenerator<TextStreamPart<ToolSet>, AgentResult> {
     messages.push(...input.messages)
     const user: ModelMessage = {
@@ -148,5 +205,6 @@ export const createAgent = (
     ask,
     stream,
     triggerSchedule,
+    askAsSubagent,
   }
 }

agent/src/prompts/subagent.ts

@@ -0,0 +1,21 @@
+import { time } from './shared'
+
+export interface SubagentParams {
+  date: Date
+  name: string
+  description?: string
+}
+
+export const subagentSystem = ({ date, name, description }: SubagentParams) => {
+  return `
+---
+${time({ date })}
+name: ${name}
+description: ${description}
+---
+
+You are a subagent, which is a specialized assistant for a specific task.
+
+Your task is communicated with the master agent to complete a task.
+`
+}

agent/src/prompts/system.ts

@@ -46,5 +46,14 @@ Your abilities:
   + The ${quote('message')} is the message to send.
   + IF: the problem is initiated by a user, regardless of the platform the user is using, the content should be directly output in the content.
   + IF: the issue is initiated by a non-user (such as a scheduled task reminder), then it should be sent using the appropriate tools on the platform specified in the requirements.
+
+**Subagent**
+When a task is large, you can create a Subagent to help you complete some tasks in order to save your own context.
+
+- You can use ${quote('create_subagent')} to create a new subagent.
+- You can use ${quote('query_subagent')} to ask a subagent to complete a task.
+  + The ${quote('name')} is the name of the subagent to ask.
+  + The ${quote('query')} is the prompt to ask the subagent to complete the task.
+Before asking a subagent, you should first create a subagent if it does not exist.
   `.trim()
 }

agent/src/tools/subagent.ts

@@ -0,0 +1,61 @@
+import { tool } from 'ai'
+import { z } from 'zod'
+import { AgentAction, createAgent } from '../agent'
+import { BaseModelConfig } from '../types'
+import { AuthFetcher } from '..'
+
+export interface SubagentToolParams extends BaseModelConfig {
+  fetch: AuthFetcher
+  braveApiKey?: string
+  braveBaseUrl?: string
+}
+
+export const getSubagentTools = ({ fetch, apiKey, baseUrl, model, clientType, braveApiKey, braveBaseUrl }: SubagentToolParams) => {
+  const createSubagent = tool({
+    description: 'Create a new subagent',
+    inputSchema: z.object({
+      name: z.string(),
+      description: z.string(),
+    }),
+    execute: async ({ name, description }) => {
+      return {
+        success: true,
+        message: 'Subagent created successfully',
+      }
+    },
+  })
+
+  const querySubagent = tool({
+    description: 'Query a subagent',
+    inputSchema: z.object({
+      name: z.string(),
+      query: z.string().describe('The prompt to ask the subagent to do.'),
+    }),
+    execute: async ({ name, query }) => {
+      const { askAsSubagent } = createAgent({
+        apiKey,
+        baseUrl,
+        model,
+        clientType,
+        braveApiKey,
+        braveBaseUrl,
+        allowed: [
+          AgentAction.WebSearch,
+        ]
+      })
+      const result = await askAsSubagent({
+        messages: [],
+        query,
+      }, { name })
+      return {
+        success: true,
+        result: result.messages[result.messages.length - 1].content,
+      }
+    },
+  })
+
+  return {
+    'create_subagent': createSubagent,
+    'query_subagent': querySubagent,
+  }
+}

agent/src/types.ts

@@ -4,6 +4,13 @@ export enum ClientType {
   GOOGLE = 'google',
 }
 
+export interface BaseModelConfig {
+  apiKey: string
+  baseUrl: string
+  model: string
+  clientType: ClientType
+}
+
 export interface Schedule {
   id: string
   name: string