diff --git a/bin/scratch.ts b/bin/scratch.ts index 3441579..7c2b8de 100644 --- a/bin/scratch.ts +++ b/bin/scratch.ts @@ -29,7 +29,8 @@ async function main() { const wikipedia = new WikipediaClient() const res = await wikipedia.getPageSummary({ - title: 'Naruto_(TV_series)' + // title: 'Naruto_(TV_series)' + title: 'SpaceX' }) console.log(JSON.stringify(res, null, 2)) diff --git a/src/ai-function-set.ts b/src/ai-function-set.ts index 50aaec5..a053a56 100644 --- a/src/ai-function-set.ts +++ b/src/ai-function-set.ts @@ -1,6 +1,13 @@ import type * as types from './types.ts' import { AIFunctionsProvider } from './fns.js' +/** + * A set of AI functions intended to make it easier to work with large sets of + * AI functions across different clients. + * + * This class mimics a built-in `Set`, but with additional utility + * methods like `pick`, `omit`, and `map`. + */ export class AIFunctionSet implements Iterable { protected readonly _map: Map diff --git a/src/fns.ts b/src/fns.ts index 1543e15..7f827ef 100644 --- a/src/fns.ts +++ b/src/fns.ts @@ -7,7 +7,7 @@ import { AIFunctionSet } from './ai-function-set.js' import { createAIFunction } from './create-ai-function.js' import { assert } from './utils.js' -export interface Invocable { +export interface PrivateAIFunctionMetadata { name: string description: string inputSchema: z.AnyZodObject @@ -21,7 +21,8 @@ export abstract class AIFunctionsProvider { if (!this._functions) { const metadata = this.constructor[Symbol.metadata] assert(metadata) - const invocables = (metadata?.invocables as Invocable[]) ?? [] + const invocables = + (metadata?.invocables as PrivateAIFunctionMetadata[]) ?? [] // console.log({ metadata, invocables }) const aiFunctions = invocables.map((invocable) => { @@ -71,7 +72,7 @@ export function aiFunction< if (!context.metadata.invocables) { context.metadata.invocables = [] } - ;(context.metadata.invocables as Invocable[]).push({ + ;(context.metadata.invocables as PrivateAIFunctionMetadata[]).push({ name: name ?? methodName, description, inputSchema, diff --git a/src/services/serper-client.ts b/src/services/serper-client.ts index 556dfb8..9fb227b 100644 --- a/src/services/serper-client.ts +++ b/src/services/serper-client.ts @@ -2,7 +2,7 @@ import defaultKy, { type KyInstance } from 'ky' import { z } from 'zod' import { aiFunction, AIFunctionsProvider } from '../fns.js' -import { assert, getEnv } from '../utils.js' +import { assert, getEnv, omit } from '../utils.js' export namespace serper { export const BASE_URL = 'https://google.serper.dev' @@ -13,10 +13,25 @@ export namespace serper { gl: z.string().default('us').optional(), hl: z.string().default('en').optional(), page: z.number().int().positive().default(1).optional(), - num: z.number().int().positive().default(10).optional() + num: z + .number() + .int() + .positive() + .default(10) + .optional() + .describe('number of results to return') }) export type SearchParams = z.infer + export const GeneralSearchSchema = SearchParamsSchema.extend({ + type: z + .enum(['search', 'images', 'videos', 'places', 'news', 'shopping']) + .default('search') + .optional() + .describe('Type of Google search to perform') + }) + export type GeneralSearchParams = z.infer + export interface SearchResponse { searchParameters: SearchParameters & { type: 'search' } organic: Organic[] @@ -233,12 +248,19 @@ export class SerperClient extends AIFunctionsProvider { name: 'serper_google_search', description: 'Uses Google Search to return the most relevant web pages for a given query. Can also be used to find up-to-date news and information about many topics.', - inputSchema: serper.SearchParamsSchema.pick({ - q: true + inputSchema: serper.GeneralSearchSchema.pick({ + q: true, + num: true, + type: true }) }) - async search(queryOrOpts: string | serper.SearchParams) { - return this._fetch('search', queryOrOpts) + async search(queryOrOpts: string | serper.GeneralSearchParams) { + const searchType = + typeof queryOrOpts === 'string' ? 'search' : queryOrOpts.type || 'search' + return this._fetch( + searchType, + typeof queryOrOpts === 'string' ? queryOrOpts : omit(queryOrOpts, 'type') + ) } async searchImages(queryOrOpts: string | serper.SearchParams) { diff --git a/src/services/wikipedia-client.ts b/src/services/wikipedia-client.ts index 7dfd591..e8f0ebe 100644 --- a/src/services/wikipedia-client.ts +++ b/src/services/wikipedia-client.ts @@ -1,10 +1,12 @@ import defaultKy, { type KyInstance } from 'ky' import pThrottle from 'p-throttle' +import { z } from 'zod' +import { aiFunction, AIFunctionsProvider } from '../fns.js' import { assert, getEnv, throttleKy } from '../utils.js' export namespace wikipedia { - // Only allow 200 requests per second + // Only allow 200 requests per second by default. export const throttle = pThrottle({ limit: 200, interval: 1000 @@ -94,7 +96,7 @@ export namespace wikipedia { } } -export class WikipediaClient { +export class WikipediaClient extends AIFunctionsProvider { readonly apiBaseUrl: string readonly apiUserAgent: string readonly ky: KyInstance @@ -114,6 +116,7 @@ export class WikipediaClient { } = {}) { assert(apiBaseUrl, 'WikipediaClient missing required "apiBaseUrl"') assert(apiUserAgent, 'WikipediaClient missing required "apiUserAgent"') + super() this.apiBaseUrl = apiBaseUrl this.apiUserAgent = apiUserAgent @@ -127,6 +130,13 @@ export class WikipediaClient { }) } + @aiFunction({ + name: 'wikipedia_search', + description: 'Searches Wikipedia for pages matching the given query.', + inputSchema: z.object({ + query: z.string().describe('Search query') + }) + }) async search({ query, ...opts }: wikipedia.SearchOptions) { return ( // https://www.mediawiki.org/wiki/API:REST_API @@ -138,12 +148,26 @@ export class WikipediaClient { ) } + @aiFunction({ + name: 'wikipedia_get_page_summary', + description: 'Gets a summary of the given Wikipedia page.', + inputSchema: z.object({ + title: z.string().describe('Wikipedia page title'), + acceptLanguage: z + .string() + .optional() + .default('en-us') + .describe('Locale code for the language to use.') + }) + }) async getPageSummary({ title, acceptLanguage = 'en-us', redirect = true, ...opts }: wikipedia.PageSummaryOptions) { + title = title.trim().replaceAll(' ', '_') + // https://en.wikipedia.org/api/rest_v1/ return this.ky .get(`page/summary/${title}`, {