diff --git a/config.json b/config.json index 45eeb09..5a67b0a 100644 --- a/config.json +++ b/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": {}, diff --git a/src/agent.ts b/src/agent.ts index a26cd12..a31fdc6 100644 --- a/src/agent.ts +++ b/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' diff --git a/src/config.ts b/src/config.ts index f1f2b40..e63c142 100644 --- a/src/config.ts +++ b/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); } diff --git a/src/tools/grounding.ts b/src/tools/grounding.ts new file mode 100644 index 0000000..52161ff --- /dev/null +++ b/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 { + 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} + +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; + } +} \ No newline at end of file