🔍 feat: Add Google Search Tool for Assistants (#1994)

This commit is contained in:
Walber Cardoso 2024-03-07 12:49:48 -03:00 committed by GitHub
parent a5f4292d2d
commit d4fe8fc82d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 66 additions and 122 deletions

View file

@ -1,121 +0,0 @@
const { google } = require('googleapis');
const { Tool } = require('langchain/tools');
const { logger } = require('~/config');
/**
* Represents a tool that allows an agent to use the Google Custom Search API.
* @extends Tool
*/
class GoogleSearchAPI extends Tool {
constructor(fields = {}) {
super();
this.cx = fields.GOOGLE_CSE_ID || this.getCx();
this.apiKey = fields.GOOGLE_API_KEY || this.getApiKey();
this.customSearch = undefined;
}
/**
* The name of the tool.
* @type {string}
*/
name = 'google';
/**
* A description for the agent to use
* @type {string}
*/
description =
'Use the \'google\' tool to retrieve internet search results relevant to your input. The results will return links and snippets of text from the webpages';
description_for_model =
'Use the \'google\' tool to retrieve internet search results relevant to your input. The results will return links and snippets of text from the webpages';
getCx() {
const cx = process.env.GOOGLE_CSE_ID || '';
if (!cx) {
throw new Error('Missing GOOGLE_CSE_ID environment variable.');
}
return cx;
}
getApiKey() {
const apiKey = process.env.GOOGLE_API_KEY || '';
if (!apiKey) {
throw new Error('Missing GOOGLE_API_KEY environment variable.');
}
return apiKey;
}
getCustomSearch() {
if (!this.customSearch) {
const version = 'v1';
this.customSearch = google.customsearch(version);
}
return this.customSearch;
}
resultsToReadableFormat(results) {
let output = 'Results:\n';
results.forEach((resultObj, index) => {
output += `Title: ${resultObj.title}\n`;
output += `Link: ${resultObj.link}\n`;
if (resultObj.snippet) {
output += `Snippet: ${resultObj.snippet}\n`;
}
if (index < results.length - 1) {
output += '\n';
}
});
return output;
}
/**
* Calls the tool with the provided input and returns a promise that resolves with a response from the Google Custom Search API.
* @param {string} input - The input to provide to the API.
* @returns {Promise<String>} A promise that resolves with a response from the Google Custom Search API.
*/
async _call(input) {
try {
const metadataResults = [];
const response = await this.getCustomSearch().cse.list({
q: input,
cx: this.cx,
auth: this.apiKey,
num: 5, // Limit the number of results to 5
});
// return response.data;
// logger.debug(response.data);
if (!response.data.items || response.data.items.length === 0) {
return this.resultsToReadableFormat([
{ title: 'No good Google Search Result was found', link: '' },
]);
}
// const results = response.items.slice(0, numResults);
const results = response.data.items;
for (const result of results) {
const metadataResult = {
title: result.title || '',
link: result.link || '',
};
if (result.snippet) {
metadataResult.snippet = result.snippet;
}
metadataResults.push(metadataResult);
}
return this.resultsToReadableFormat(metadataResults);
} catch (error) {
logger.error('[GoogleSearchAPI]', error);
// throw error;
return 'There was an error searching Google.';
}
}
}
module.exports = GoogleSearchAPI;

View file

@ -1,7 +1,6 @@
const availableTools = require('./manifest.json');
// Basic Tools
const CodeBrew = require('./CodeBrew');
const GoogleSearchAPI = require('./GoogleSearch');
const WolframAlphaAPI = require('./Wolfram');
const AzureAiSearch = require('./AzureAiSearch');
const OpenAICreateImage = require('./DALL-E');
@ -16,6 +15,7 @@ const CodeSherpa = require('./structured/CodeSherpa');
const StructuredSD = require('./structured/StableDiffusion');
const StructuredACS = require('./structured/AzureAISearch');
const CodeSherpaTools = require('./structured/CodeSherpaTools');
const GoogleSearchAPI = require('./structured/GoogleSearch');
const StructuredWolfram = require('./structured/Wolfram');
const TavilySearchResults = require('./structured/TavilySearchResults');
const TraversaalSearch = require('./structured/TraversaalSearch');

View file

@ -0,0 +1,65 @@
const { z } = require('zod');
const { Tool } = require('@langchain/core/tools');
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
class GoogleSearchResults extends Tool {
static lc_name() {
return 'GoogleSearchResults';
}
constructor(fields = {}) {
super(fields);
this.envVarApiKey = 'GOOGLE_API_KEY';
this.envVarSearchEngineId = 'GOOGLE_CSE_ID';
this.override = fields.override ?? false;
this.apiKey = fields.apiKey ?? getEnvironmentVariable(this.envVarApiKey);
this.searchEngineId =
fields.searchEngineId ?? getEnvironmentVariable(this.envVarSearchEngineId);
this.kwargs = fields?.kwargs ?? {};
this.name = 'google';
this.description =
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.';
this.schema = z.object({
query: z.string().min(1).describe('The search query string.'),
max_results: z
.number()
.min(1)
.max(10)
.optional()
.describe('The maximum number of search results to return. Defaults to 10.'),
// Note: Google API has its own parameters for search customization, adjust as needed.
});
}
async _call(input) {
const validationResult = this.schema.safeParse(input);
if (!validationResult.success) {
throw new Error(`Validation failed: ${JSON.stringify(validationResult.error.issues)}`);
}
const { query, max_results = 5 } = validationResult.data;
const response = await fetch(
`https://www.googleapis.com/customsearch/v1?key=${this.apiKey}&cx=${
this.searchEngineId
}&q=${encodeURIComponent(query)}&num=${max_results}`,
{
method: 'GET',
headers: {
'Content-Type': 'application/json',
},
},
);
const json = await response.json();
if (!response.ok) {
throw new Error(`Request failed with status ${response.status}: ${json.error.message}`);
}
return JSON.stringify(json);
}
}
module.exports = GoogleSearchResults;