From 0321999511a281555f6a72f4319e1a6ac10be411 Mon Sep 17 00:00:00 2001 From: Acbox Date: Sat, 31 Jan 2026 19:58:45 +0800 Subject: [PATCH] feat(agent): basic subagent function of agent gateway --- agent/src/agent.ts | 84 +++++++++++++++++++++++++++++------ agent/src/prompts/subagent.ts | 21 +++++++++ agent/src/prompts/system.ts | 9 ++++ agent/src/tools/subagent.ts | 61 +++++++++++++++++++++++++ agent/src/types.ts | 7 +++ 5 files changed, 169 insertions(+), 13 deletions(-) create mode 100644 agent/src/prompts/subagent.ts create mode 100644 agent/src/tools/subagent.ts diff --git a/agent/src/agent.ts b/agent/src/agent.ts index e0e547c2..1751e729 100644 --- a/agent/src/agent.ts +++ b/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,16 +44,20 @@ 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, @@ -54,6 +65,19 @@ export const createAgent = ( 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 => { + 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, AgentResult> { messages.push(...input.messages) const user: ModelMessage = { @@ -148,5 +205,6 @@ export const createAgent = ( ask, stream, triggerSchedule, + askAsSubagent, } } \ No newline at end of file diff --git a/agent/src/prompts/subagent.ts b/agent/src/prompts/subagent.ts new file mode 100644 index 00000000..473babce --- /dev/null +++ b/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. +` +} \ No newline at end of file diff --git a/agent/src/prompts/system.ts b/agent/src/prompts/system.ts index 26a7d804..3e938a09 100644 --- a/agent/src/prompts/system.ts +++ b/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() } \ No newline at end of file diff --git a/agent/src/tools/subagent.ts b/agent/src/tools/subagent.ts new file mode 100644 index 00000000..f0f09677 --- /dev/null +++ b/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, + } +} \ No newline at end of file diff --git a/agent/src/types.ts b/agent/src/types.ts index 8d322549..75820d86 100644 --- a/agent/src/types.ts +++ b/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