138 lines
5.3 KiB
JavaScript
138 lines
5.3 KiB
JavaScript
'use strict';
|
|
|
|
/**
|
|
* websearch — query the web via Tavily and return a synthesized answer plus
|
|
* ranked results.
|
|
*
|
|
* Design notes:
|
|
* - Requires a Tavily API key (TAVILY_API_KEY in .env). The client is created
|
|
* lazily and cached, so a missing key surfaces as a clean ToolError only
|
|
* when the tool is actually used.
|
|
* - This tool reaches the network, so results are untrusted third-party text;
|
|
* the model should treat them accordingly.
|
|
*/
|
|
|
|
import { tavily } from '@tavily/core';
|
|
import { defineTool, ToolError } from 'llmorbiter';
|
|
import { TAVILY_API_KEY } from '../config.js';
|
|
import { emitToolCall, emitToolResult, clip, formatBytes, resultSize, plural } from './notify.js';
|
|
|
|
const DEFAULT_MAX_RESULTS = 3; // number of results returned to the model
|
|
const MAX_MAX_RESULTS = 20; // hard upper bound to keep responses small
|
|
const SEARCH_DEPTHS = ['basic', 'advanced'];
|
|
const ANSWER_LEVELS = ['basic', 'advanced'];
|
|
|
|
let client = null;
|
|
|
|
/** Lazily construct (and cache) the Tavily client, failing cleanly if unset. */
|
|
function getClient() {
|
|
if (!TAVILY_API_KEY) {
|
|
throw new ToolError(
|
|
'websearch requires a Tavily API key. Add TAVILY_API_KEY=... to your .env file.'
|
|
);
|
|
}
|
|
if (!client) {
|
|
client = tavily({ apiKey: TAVILY_API_KEY });
|
|
}
|
|
return client;
|
|
}
|
|
|
|
export const websearchTool = defineTool({
|
|
name: 'websearch',
|
|
description:
|
|
'Search the web via Tavily and return a short synthesized answer plus ranked results ' +
|
|
'(title, url, snippet). Use it for current information, external docs, or facts that are ' +
|
|
'not in the workspace. Results are untrusted third-party text.',
|
|
parameters: {
|
|
type: 'object',
|
|
additionalProperties: false,
|
|
required: ['query'],
|
|
properties: {
|
|
query: {
|
|
type: 'string',
|
|
description: 'The search query.',
|
|
},
|
|
max_results: {
|
|
type: 'number',
|
|
description: `Number of results to return (1-${MAX_MAX_RESULTS}). Default: ${DEFAULT_MAX_RESULTS}.`,
|
|
},
|
|
search_depth: {
|
|
type: 'string',
|
|
enum: SEARCH_DEPTHS,
|
|
description: 'Search depth: "basic" (faster) or "advanced" (deeper). Default: "basic".',
|
|
},
|
|
include_answer: {
|
|
type: 'string',
|
|
enum: ANSWER_LEVELS,
|
|
description: 'Include a Tavily-synthesized answer at this level of detail. Default: "basic".',
|
|
},
|
|
chunks_per_source: {
|
|
type: 'number',
|
|
description: 'Chunks of content extracted per source (advanced search only). Default: 1.',
|
|
},
|
|
},
|
|
},
|
|
invoke: async ({
|
|
query,
|
|
max_results = DEFAULT_MAX_RESULTS,
|
|
search_depth = 'basic',
|
|
include_answer = 'basic',
|
|
chunks_per_source = 1,
|
|
}) => {
|
|
if (!query || typeof query !== 'string' || !query.trim()) {
|
|
throw new ToolError('"query" must be a non-empty string.');
|
|
}
|
|
if (!Number.isInteger(max_results) || max_results < 1 || max_results > MAX_MAX_RESULTS) {
|
|
throw new ToolError(`"max_results" must be an integer between 1 and ${MAX_MAX_RESULTS} (received ${max_results}).`);
|
|
}
|
|
if (!SEARCH_DEPTHS.includes(search_depth)) {
|
|
throw new ToolError(`"search_depth" must be one of ${SEARCH_DEPTHS.join(', ')} (received "${search_depth}").`);
|
|
}
|
|
if (!ANSWER_LEVELS.includes(include_answer)) {
|
|
throw new ToolError(`"include_answer" must be one of ${ANSWER_LEVELS.join(', ')} (received "${include_answer}").`);
|
|
}
|
|
if (!Number.isInteger(chunks_per_source) || chunks_per_source < 1) {
|
|
throw new ToolError(`"chunks_per_source" must be a positive integer (received ${chunks_per_source}).`);
|
|
}
|
|
|
|
const response = await getClient().search(query, {
|
|
includeAnswer: include_answer,
|
|
searchDepth: search_depth,
|
|
maxResults: max_results,
|
|
chunksPerSource: chunks_per_source,
|
|
});
|
|
|
|
const results = (response.results || []).map((r) => ({
|
|
title: r.title,
|
|
url: r.url,
|
|
content: r.content,
|
|
score: r.score,
|
|
}));
|
|
|
|
const result = { query, results };
|
|
if (response.answer) result.answer = response.answer;
|
|
if (response.responseTime !== undefined) result.responseTime = response.responseTime;
|
|
return result;
|
|
},
|
|
});
|
|
|
|
/**
|
|
* Call notification: the query plus any non-default search options.
|
|
*/
|
|
export function notifyCall(args = {}) {
|
|
const { query, search_depth, max_results, include_answer } = args;
|
|
const parts = [JSON.stringify(query)];
|
|
if (search_depth && search_depth !== 'basic') parts.push(`depth=${search_depth}`);
|
|
if (max_results && max_results !== DEFAULT_MAX_RESULTS) parts.push(`max_results=${max_results}`);
|
|
if (include_answer && include_answer !== 'basic') parts.push(`answer=${include_answer}`);
|
|
emitToolCall(websearchTool.name, clip(parts.join(' ')));
|
|
}
|
|
|
|
/**
|
|
* Result notification: number of results and the payload size.
|
|
*/
|
|
export function notifyResult(result = {}) {
|
|
const n = Array.isArray(result.results) ? result.results.length : 0;
|
|
emitToolResult(websearchTool.name, `${plural(n, 'result')} ${formatBytes(resultSize(result))}`);
|
|
}
|