feat: web_search tool with brave search api

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 = () => {

agent/src/config.ts

@@ -8,6 +8,10 @@ type AgentGatewayConfig = {
   },
   'server': {
     addr?: string
+  },
+  'brave': {
+    api_key?: string
+    base_url?: string
   }
 }
 

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({

agent/src/tools/index.ts

@@ -0,0 +1,3 @@
+export { getWebTools } from './web'
+export { getScheduleTools } from './schedule'
+

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,
+  }
+}

config.toml.example

@@ -44,4 +44,8 @@ timeout_seconds = 10
 ## Agent Gateway
 [agent_gateway]
 host = "127.0.0.1"
-port = 8081
+port = 8081
+
+[brave]
+api_key = ""
+base_url = "https://api.search.brave.com/res/v1/"