🌀

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))
 

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<AIFunction>`, but with additional utility
+ * methods like `pick`, `omit`, and `map`.
+ */
 export class AIFunctionSet implements Iterable<types.AIFunction> {
   protected readonly _map: Map<string, types.AIFunction>
 

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,

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<typeof SearchParamsSchema>
 
+  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<typeof GeneralSearchSchema>
+
   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({
+    inputSchema: serper.GeneralSearchSchema.pick({
-      q: true
+      q: true,
+      num: true,
+      type: true
     })
   })
-  async search(queryOrOpts: string | serper.SearchParams) {
+  async search(queryOrOpts: string | serper.GeneralSearchParams) {
-    return this._fetch<serper.SearchResponse>('search', queryOrOpts)
+    const searchType =
+      typeof queryOrOpts === 'string' ? 'search' : queryOrOpts.type || 'search'
+    return this._fetch<serper.SearchResponse>(
+      searchType,
+      typeof queryOrOpts === 'string' ? queryOrOpts : omit(queryOrOpts, 'type')
+    )
   }
 
   async searchImages(queryOrOpts: string | serper.SearchParams) {

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}`, {