generator-begcode/dist/agent/functions/WebSearch.js

Spring Boot + Angular/React/Vue in one handy generator

import axios from 'axios';
import { AgentOutputType, ArrayRecombiner, ChatMessageBuilder, Prompt, Rag, TextChunker, trimText, } from '../agent-core/index.js';
import { FUNCTION_CALL_FAILED, FUNCTION_CALL_SUCCESS_CONTENT } from '../agents/Scripter/utils.js';
import { LlmAgentFunctionBase, processWebpage, searchOnGoogle } from './utils/index.js';
export class WebSearchFunction extends LlmAgentFunctionBase {
    constructor(llm, tokenizer) {
        super(llm, tokenizer);
    }
    name = 'web_search';
    description = 'Searches the web for multiple queries.';
    parameters = {
        type: 'object',
        properties: {
            queries: {
                type: 'array',
                items: {
                    type: 'string',
                },
                description: 'Queries with all details to search the web for in parallel',
            },
        },
        required: ['queries'],
        additionalProperties: false,
    };
    buildExecutor(agent) {
        return async (params) => {
            const searches = [];
            for (const query of params.queries) {
                searches.push(this.runQuery(agent, query, JSON.stringify({ queries: [query] })));
            }
            const results = await Promise.all(searches);
            return {
                outputs: results.flatMap(x => x.outputs),
                messages: results.flatMap(x => x.messages),
            };
        };
    }
    async runQuery({ context }, query, rawParams) {
        try {
            let googleResults;
            if (typeof window === 'object') {
                const results = await axios.get(`/api/search?query=${query}`);
                googleResults = results.data.googleResults;
            }
            else {
                if (!context.env.SERP_API_KEY) {
                    throw new Error('SERP_API_KEY environment variable is required to use the websearch plugin. See env.template for help');
                }
                googleResults = await searchOnGoogle(query, context.env.SERP_API_KEY);
            }
            const urls = googleResults.map(x => x.url).filter(u => !u.endsWith('.pdf'));
            const searchMatches = await this.searchInPages({
                urls,
                query,
                context,
            });
            const analyzeChunkMatchesPrompt = new Prompt()
                .text(`
          I will give you chunks of text from different webpages.

          I want to extract ${query}. Keep in mind some information may not be properly formatted.
          Do your best to extract as much information as you can.

          Prioritize decimal precision. Aim for answers with 3 decimal places, if possible; if not settle for 2 decimal places.
          Only take 1 decimal or rounded numbers when absolutely necessary.

          Chunks: ${searchMatches.join('\n------------\n')}
        `)
                .line('Specify if the information is incomplete but still return it')
                .toString();
            const analysisFromChunks = await this.askLlm(analyzeChunkMatchesPrompt, {
                maxResponseTokens: 200,
            });
            return this.onSuccess({ queries: [query] }, analysisFromChunks, rawParams, context.variables);
        }
        catch (err) {
            return this.onError({ queries: [query] }, err.toString(), rawParams, context.variables);
        }
    }
    onSuccess(params, result, rawParams, variables) {
        return {
            outputs: [
                {
                    type: AgentOutputType.Success,
                    title: `Web search for '${params.queries}'`,
                    content: FUNCTION_CALL_SUCCESS_CONTENT(this.name, params, `Found the following result for the web search: '${params.queries}'` +
                        '\n--------------\n' +
                        `${result}\n` +
                        '\n--------------\n'),
                },
            ],
            messages: [
                ChatMessageBuilder.functionCall(this.name, rawParams),
                ChatMessageBuilder.functionCallResult(this.name, `Found the following result for the web search: '${params.queries}'\`\`\`
${result}\n\`\`\``),
            ],
        };
    }
    onError(params, error, rawParams, variables) {
        return {
            outputs: [
                {
                    type: AgentOutputType.Error,
                    title: `Web search for '${params.queries}'`,
                    content: FUNCTION_CALL_FAILED(params, this.name, error),
                },
            ],
            messages: [
                ChatMessageBuilder.functionCall(this.name, rawParams),
                ChatMessageBuilder.functionCallResult(this.name, `Error web searching for '${params.queries}'\n\`\`\`
${trimText(error, 300)}\n\`\`\``),
            ],
        };
    }
    async searchInPages(params) {
        const urlsContents = await Promise.all(params.urls.map(async (url) => {
            try {
                let response;
                if (typeof window === 'object') {
                    const result = await axios.get(`/api/process-web-page?url=${url}`);
                    response = result.data.text;
                }
                else {
                    response = await processWebpage(url);
                }
                return {
                    url,
                    response,
                };
            }
            catch (e) {
                return {
                    url,
                    response: '',
                };
            }
        }));
        const webpagesChunks = urlsContents.map(webpageContent => ({
            url: webpageContent.url,
            chunks: TextChunker.fixedCharacterLength(webpageContent.response, { chunkLength: 550, overlap: 110 }),
        }));
        const results = await Promise.all(webpagesChunks.map(async (webpageChunks) => {
            const items = webpageChunks.chunks.map(chunk => ({
                chunk,
                url: webpageChunks.url,
            }));
            const matches = await (await Rag.standard(params.context))
                .addItems(items)
                .selector(x => x.chunk)
                .query(params.query)
                .recombine(ArrayRecombiner.standard({
                limit: 4,
                unique: true,
            }));
            return matches;
        }));
        const otherResults = await (await Rag.standard(params.context))
            .addItems(results.flat())
            .selector(x => x.chunk)
            .query(params.query)
            .recombine(ArrayRecombiner.standard({
            limit: 10,
            unique: true,
            sort: 'index',
        }));
        return otherResults.map(x => x.chunk);
    }
}