From 8eb232c40cc25143bc52782951a1c0712385ddc2 Mon Sep 17 00:00:00 2001 From: Acbox Date: Sat, 31 Jan 2026 18:14:16 +0800 Subject: [PATCH] feat: web_search tool with brave search api --- agent/src/agent.ts | 16 ++++++- agent/src/config.ts | 4 ++ agent/src/modules/chat.ts | 9 ++++ agent/src/tools/index.ts | 3 ++ agent/src/tools/web.ts | 93 +++++++++++++++++++++++++++++++++++++++ config.toml.example | 6 ++- 6 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 agent/src/tools/web.ts diff --git a/agent/src/agent.ts b/agent/src/agent.ts index 7bd03645..e0e547c2 100644 --- a/agent/src/agent.ts +++ b/agent/src/agent.ts @@ -4,6 +4,7 @@ import { ClientType, Schedule } from './types' import { system, schedule } from './prompts' import { AuthFetcher } from './index' import { getScheduleTools } from './tools/schedule' +import { getWebTools } from './tools/web' export interface AgentParams { apiKey: string @@ -16,6 +17,8 @@ export interface AgentParams { maxContextLoadTime: number platforms?: string[] currentPlatform?: string + braveApiKey?: string + braveBaseUrl?: string } export interface AgentInput { @@ -38,9 +41,20 @@ export const createAgent = ( const getTools = () => { const scheduleTools = getScheduleTools({ fetch: fetcher }) - return { + const tools: ToolSet = { ...scheduleTools, } + + // Add web search tools if Brave API key is provided + if (params.braveApiKey) { + const webTools = getWebTools({ + braveApiKey: params.braveApiKey, + braveBaseUrl: params.braveBaseUrl, + }) + Object.assign(tools, webTools) + } + + return tools } const generateSystem = () => { diff --git a/agent/src/config.ts b/agent/src/config.ts index 1b88f1ba..9f3ca291 100644 --- a/agent/src/config.ts +++ b/agent/src/config.ts @@ -8,6 +8,10 @@ type AgentGatewayConfig = { }, 'server': { addr?: string + }, + 'brave': { + api_key?: string + base_url?: string } } diff --git a/agent/src/modules/chat.ts b/agent/src/modules/chat.ts index ae851f37..5aaadef4 100644 --- a/agent/src/modules/chat.ts +++ b/agent/src/modules/chat.ts @@ -5,6 +5,7 @@ import { createAuthFetcher } from '../index' import { ClientType } from '../types' import { ModelMessage } from 'ai' import { bearerMiddleware } from '../middlewares/bearer' +import { loadConfig } from '../config' const ChatBody = z.object({ apiKey: z.string().min(1, 'API key is required'), @@ -37,6 +38,8 @@ const ScheduleBody = z.object({ }), }).and(ChatBody) +const config = loadConfig('../config.toml') + export const chatModule = new Elysia({ prefix: '/chat' }) .use(bearerMiddleware) .post('/', async ({ body, bearer }) => { @@ -51,6 +54,8 @@ export const chatModule = new Elysia({ prefix: '/chat' }) maxContextLoadTime: body.maxContextLoadTime, platforms: body.platforms, currentPlatform: body.currentPlatform, + braveApiKey: config.brave?.api_key, + braveBaseUrl: config.brave?.base_url, }, createAuthFetcher(bearer)) try { const result = await ask({ @@ -91,6 +96,8 @@ export const chatModule = new Elysia({ prefix: '/chat' }) maxContextLoadTime: body.maxContextLoadTime, platforms: body.platforms, currentPlatform: body.currentPlatform, + braveApiKey: config.brave?.api_key, + braveBaseUrl: config.brave?.base_url, }, createAuthFetcher(bearer)) try { const streanGenerator = stream({ @@ -142,6 +149,8 @@ export const chatModule = new Elysia({ prefix: '/chat' }) maxContextLoadTime: body.maxContextLoadTime, platforms: body.platforms, currentPlatform: body.currentPlatform, + braveApiKey: config.brave?.api_key, + braveBaseUrl: config.brave?.base_url, }, createAuthFetcher(bearer)) try { return await triggerSchedule({ diff --git a/agent/src/tools/index.ts b/agent/src/tools/index.ts index e69de29b..b97b16a8 100644 --- a/agent/src/tools/index.ts +++ b/agent/src/tools/index.ts @@ -0,0 +1,3 @@ +export { getWebTools } from './web' +export { getScheduleTools } from './schedule' + diff --git a/agent/src/tools/web.ts b/agent/src/tools/web.ts new file mode 100644 index 00000000..80edc3dc --- /dev/null +++ b/agent/src/tools/web.ts @@ -0,0 +1,93 @@ +import { tool } from 'ai' +import { z } from 'zod' + +interface WebToolParams { + braveApiKey: string + braveBaseUrl?: string +} + +interface BraveSearchResult { + type: string + title: string + url: string + description?: string + age?: string +} + +interface BraveSearchResponse { + web?: { + results: BraveSearchResult[] + } +} + +export const getWebTools = ({ braveApiKey, braveBaseUrl = 'https://api.search.brave.com/res/v1/' }: WebToolParams) => { + const webSearch = tool({ + description: 'Search the web for information using Brave Search API. Use this when you need current information, facts, news, or any web content.', + inputSchema: z.object({ + query: z.string().describe('The search query to look up on the web'), + count: z.number().optional().describe('Number of results to return (default: 10, max: 20)'), + }), + execute: async ({ query, count = 10 }) => { + try { + const url = new URL('web/search', braveBaseUrl) + url.searchParams.append('q', query) + url.searchParams.append('count', Math.min(count, 20).toString()) + + const response = await fetch(url.toString(), { + method: 'GET', + headers: { + 'Accept': 'application/json', + 'Accept-Encoding': 'gzip', + 'X-Subscription-Token': braveApiKey, + }, + }) + + if (!response.ok) { + const errorText = await response.text() + console.error('[Web Search] error', { + type: 'web_search', + query, + count, + status: response.status, + statusText: response.statusText, + error: errorText, + }) + throw new Error(`Brave Search API error: ${response.status} ${response.statusText}`) + } + + const data: BraveSearchResponse = await response.json() + + const results = data.web?.results || [] + + if (results.length === 0) { + return { + success: false, + message: 'No results found for the query', + query, + } + } + + return { + success: true, + query, + results: results.map((result) => ({ + title: result.title, + url: result.url, + description: result.description, + age: result.age, + })), + } + } catch (error) { + return { + success: false, + error: error instanceof Error ? error.message : 'Unknown error occurred', + query, + } + } + }, + }) + + return { + 'web_search': webSearch, + } +} \ No newline at end of file diff --git a/config.toml.example b/config.toml.example index cec66d69..8b6754b0 100644 --- a/config.toml.example +++ b/config.toml.example @@ -44,4 +44,8 @@ timeout_seconds = 10 ## Agent Gateway [agent_gateway] host = "127.0.0.1" -port = 8081 \ No newline at end of file +port = 8081 + +[brave] +api_key = "" +base_url = "https://api.search.brave.com/res/v1/" \ No newline at end of file