feat: add Brave Search integration (#4)

* feat: add Brave Search integration

- Add Brave Search implementation
- Configure Brave API key in config
- Make search provider configurable with Brave as default

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

* chore: first commit

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Han Xiao <han.xiao@jina.ai>

package.json

@@ -16,8 +16,11 @@
   "description": "",
   "dependencies": {
     "@google/generative-ai": "^0.21.0",
+    "@types/node-fetch": "^2.6.12",
+    "axios": "^1.7.9",
     "dotenv": "^16.4.7",
     "duck-duck-scrape": "^2.2.7",
+    "node-fetch": "^3.3.2",
     "undici": "^7.3.0"
   },
   "devDependencies": {

src/agent.ts

@@ -1,13 +1,14 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
 import { readUrl } from "./tools/read";
 import fs from 'fs/promises';
-import { SafeSearchType, search } from "duck-duck-scrape";
+import { SafeSearchType, search as duckSearch } from "duck-duck-scrape";
+import { braveSearch } from "./tools/brave-search";
 import { rewriteQuery } from "./tools/query-rewriter";
 import { dedupQueries } from "./tools/dedup";
 import { evaluateAnswer } from "./tools/evaluator";
 import { StepData } from "./tools/getURLIndex";
 import { analyzeSteps } from "./tools/error-analyzer";
-import { GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME } from "./config";
+import { GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME, BRAVE_API_KEY, SEARCH_PROVIDER } from "./config";
 import { tokenTracker } from "./utils/token-tracker";
 
 async function sleep(ms: number) {
@@ -260,7 +261,6 @@ function removeAllLineBreaks(text: string) {
 }
 
 async function getResponse(question: string, tokenBudget: number = 1000000, maxBadAttempts: number = 3) {
-  let totalTokens = 0;
   let step = 0;
   let totalStep = 0;
   let badAttempts = 0;
@@ -271,7 +271,7 @@ async function getResponse(question: string, tokenBudget: number = 1000000, maxB
   const badContext = [];
   let diaryContext = [];
   const allURLs: Record<string, string> = {};
-  while (totalTokens < tokenBudget) {
+  while (tokenTracker.getTotalUsage() < tokenBudget) {
     // add 1s delay to avoid rate limiting
     await sleep(1000);
     step++;
@@ -462,9 +462,21 @@ But then you realized you have asked them before. You decided to to think out of
           const searchResults = [];
           for (const query of keywordsQueries) {
             console.log(`Search query: ${query}`);
-            const results = await search(query, {
-              safeSearch: SafeSearchType.STRICT
-            });
+            let results;
+            if (SEARCH_PROVIDER === 'duck') {
+              results = await duckSearch(query, {
+                safeSearch: SafeSearchType.STRICT
+              });
+            } else {
+              const { response } = await braveSearch(query);
+              results = {
+                results: response.web.results.map(r => ({
+                  title: r.title,
+                  url: r.url,
+                  description: r.description
+                }))
+              };
+            }
             const minResults = results.results.map(r => ({
               title: r.title,
               url: r.url,

src/config.ts

@@ -16,6 +16,9 @@ if (process.env.https_proxy) {
 
 export const GEMINI_API_KEY = process.env.GEMINI_API_KEY as string;
 export const JINA_API_KEY = process.env.JINA_API_KEY as string;
+export const BRAVE_API_KEY = process.env.BRAVE_API_KEY as string;
+export const SEARCH_PROVIDER = BRAVE_API_KEY ? 'brave' : 'duck';
+
 export const MODEL_NAME = 'gemini-1.5-flash';
 
 if (!GEMINI_API_KEY) throw new Error("GEMINI_API_KEY not found");

src/test-duck.ts

@@ -1,16 +1,37 @@
-import {search, SafeSearchType} from 'duck-duck-scrape';
+import axios from 'axios';  // You'll need to npm install axios first
+
+async function testRequest(): Promise<any> {
+  console.log('Starting test request...');
 
-const query = process.argv[2] || "jina ai";
-async function runTest() {
   try {
-    const results = await search(query, {
-      safeSearch: SafeSearchType.STRICT
+    const response = await axios.get('https://jsonplaceholder.typicode.com/posts/1', {
+      timeout: 10000,
+      headers: {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+      }
     });
-    console.log('Search results:', results);
+
+    console.log('Response status:', response.status);
+    console.log('Request completed');
+    return response.data;
+
   } catch (error) {
-    console.error('Test failed:', error);
+    if (axios.isAxiosError(error)) {
+      console.error('Axios error:', {
+        message: error.message,
+        code: error.code,
+        status: error.response?.status
+      });
+    } else {
+      console.error('Unexpected error:', error);
+    }
+    throw error;
   }
 }
 
-runTest();
-
+// Test
+console.log('Before test request');
+testRequest()
+  .then(result => console.log('Success:', result))
+  .catch(error => console.error('Error:', error));
+console.log('After test request');

src/tools/brave-search.ts

@@ -0,0 +1,36 @@
+import axios from 'axios';
+import {BRAVE_API_KEY} from "../config";
+
+interface BraveSearchResponse {
+  web: {
+    results: Array<{
+      title: string;
+      description: string;
+      url: string;
+    }>;
+  };
+}
+
+export async function braveSearch(query: string): Promise<{ response: BraveSearchResponse }> {
+  const response = await axios.get<BraveSearchResponse>('https://api.search.brave.com/res/v1/web/search', {
+    params: {
+      q: query,
+      count: 5,
+      safesearch: 'off'
+    },
+    headers: {
+      'Accept': 'application/json',
+      'X-Subscription-Token': BRAVE_API_KEY
+    },
+    timeout: 10000
+  });
+
+  // Keep the same console.log as the original code
+  console.log('Brave Search:', response.data.web.results.map(item => ({
+    title: item.title,
+    url: item.url
+  })));
+
+  // Maintain the same return structure as the original code
+  return { response: response.data };
+}