feat: add google search grounding

config.json

@@ -32,6 +32,7 @@
         "maxTokens": 8000
       },
       "tools": {
+        "search-grounding": { "temperature": 0 },
         "dedup": { "temperature": 0.1 },
         "evaluator": {},
         "errorAnalyzer": {},
@@ -47,6 +48,7 @@
         "maxTokens": 8000
       },
       "tools": {
+        "search-grounding": { "temperature": 0 },
         "dedup": { "temperature": 0.1 },
         "evaluator": {},
         "errorAnalyzer": {},

src/agent.ts

@@ -1,6 +1,6 @@
 import {z} from 'zod';
 import {generateObject} from 'ai';
-import {getModel, getMaxTokens, SEARCH_PROVIDER, STEP_SLEEP} from "./config";
+import {getModel, getMaxTokens, SEARCH_PROVIDER, STEP_SLEEP, LLM_PROVIDER} from "./config";
 import {readUrl} from "./tools/read";
 import {handleGenerateObjectError} from './utils/error-handling';
 import fs from 'fs/promises';
@@ -15,6 +15,7 @@ import {ActionTracker} from "./utils/action-tracker";
 import {StepAction, AnswerAction} from "./types";
 import {TrackerContext} from "./types";
 import {search} from "./tools/jina-search";
+import {grounding} from "./tools/grounding";
 
 async function sleep(ms: number) {
   const seconds = Math.ceil(ms / 1000);
@@ -504,15 +505,20 @@ But then you realized you have asked them before. You decided to to think out of
       keywordsQueries = dedupedQueries;
 
       if (keywordsQueries.length > 0) {
+        let googleGrounded = '';
         const searchResults = [];
         for (const query of keywordsQueries) {
           console.log(`Search query: ${query}`);
 
           let results;
+
           switch (SEARCH_PROVIDER) {
             case 'jina':
               // use jinaSearch
               results = {results: (await search(query, context.tokenTracker)).response?.data || []};
+              if (LLM_PROVIDER === 'gemini') {
+                googleGrounded = await grounding(query, context.tokenTracker);
+              }
               break;
             case 'duck':
               results = await duckSearch(query, {safeSearch: SafeSearchType.STRICT});
@@ -551,7 +557,7 @@ But then you realized you have asked them before. You decided to to think out of
 
         allKnowledge.push({
           question: `What do Internet say about ${thisStep.searchQuery}?`,
-          answer: removeHTMLtags(searchResults.map(r => r.results.map(r => r.description).join('; ')).join('; ')),
+          answer: googleGrounded + removeHTMLtags(searchResults.map(r => r.results.map(r => r.description).join('; ')).join('; ')),
           // flatten into one url list, and take unique urls
           references: searchResults.map(r => r.results.map(r => r.url)).flat().filter((v, i, a) => a.indexOf(v) === i),
           type: 'side-info'

src/config.ts

@@ -126,6 +126,9 @@ export function getModel(toolName: ToolName) {
     throw new Error('GEMINI_API_KEY not found');
   }
 
+  if (toolName === 'search-grounding') {
+    return createGoogleGenerativeAI({ apiKey: GEMINI_API_KEY })(config.model, { useSearchGrounding: true });
+  }
   return createGoogleGenerativeAI({ apiKey: GEMINI_API_KEY })(config.model);
 }
 

src/tools/grounding.ts

@@ -0,0 +1,38 @@
+import { generateText } from 'ai';
+import {getModel} from "../config";
+import { GoogleGenerativeAIProviderMetadata } from '@ai-sdk/google';
+import {TokenTracker} from "../utils/token-tracker";
+
+const model = getModel('search-grounding')
+
+export async function grounding(query: string, tracker?: TokenTracker): Promise<string> {
+  try {
+    const { text, experimental_providerMetadata, usage } = await generateText({
+      model,
+      prompt:
+      `Current date is ${new Date().toISOString()}. Find the latest answer to the following question: 
+<query>
+${query}
+</query>      
+Must include the date and time of the latest answer.`,
+    });
+
+    const metadata = experimental_providerMetadata?.google as
+  | GoogleGenerativeAIProviderMetadata
+  | undefined;
+    const groundingMetadata = metadata?.groundingMetadata;
+
+    // Extract and concatenate all groundingSupport text into a single line
+    const groundedText = groundingMetadata?.groundingSupports
+      ?.map(support => support.segment.text)
+      .join(' ') || '';
+
+    (tracker || new TokenTracker()).trackUsage('grounding', usage.totalTokens);
+    console.log('Grounding:', {text, groundedText});
+    return text + '|' + groundedText;
+
+  } catch (error) {
+    console.error('Error in search:', error);
+    throw error;
+  }
+}