feat: improve robustness of serpapi, diffbot, and add ky rate limiting

examples/llm-with-search.ts

@@ -2,7 +2,7 @@ import { OpenAIClient } from '@agentic/openai-fetch'
 import 'dotenv/config'
 import { z } from 'zod'
 
-import { Agentic, SerpAPITool } from '@/index'
+import { Agentic, DiffbotTool, SerpAPITool } from '@/index'
 
 async function main() {
   const openai = new OpenAIClient({ apiKey: process.env.OPENAI_API_KEY! })
@@ -12,7 +12,7 @@ async function main() {
     .gpt4(
       `Can you summarize the top {{numResults}} results for today's news about {{topic}}?`
     )
-    .tools([new SerpAPITool()])
+    .tools([new SerpAPITool(), new DiffbotTool()])
     .input(
       z.object({
         topic: z.string(),

package.json

@@ -56,6 +56,7 @@
     "normalize-url": "^8.0.0",
     "p-map": "^6.0.0",
     "p-retry": "^5.1.2",
+    "p-throttle": "^5.1.0",
     "p-timeout": "^6.1.2",
     "pino": "^8.14.1",
     "pino-pretty": "^10.0.0",
@@ -104,6 +105,7 @@
   },
   "ava": {
     "snapshotDir": "test/.snapshots",
+    "failFast": true,
     "extensions": {
       "ts": "module"
     },
@@ -126,4 +128,4 @@
     "guardrails",
     "plugins"
   ]
-}
+}

pnpm-lock.yaml

@@ -1,4 +1,4 @@
-lockfileVersion: '6.1'
+lockfileVersion: '6.0'
 
 settings:
   autoInstallPeers: true
@@ -59,6 +59,9 @@ dependencies:
   p-retry:
     specifier: ^5.1.2
     version: 5.1.2
+  p-throttle:
+    specifier: ^5.1.0
+    version: 5.1.0
   p-timeout:
     specifier: ^6.1.2
     version: 6.1.2
@@ -3183,6 +3186,11 @@ packages:
       retry: 0.13.1
     dev: false
 
+  /p-throttle@5.1.0:
+    resolution: {integrity: sha512-+N+s2g01w1Zch4D0K3OpnPDqLOKmLcQ4BvIFq3JC0K29R28vUOjWpO+OJZBNt8X9i3pFCksZJZ0YXkUGjaFE6g==}
+    engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0}
+    dev: false
+
   /p-timeout@5.1.0:
     resolution: {integrity: sha512-auFDyzzzGZZZdHz3BtET9VEz0SE/uMEAx7uWfGPucfzEwwe/xH0iVeZibQmANYE/hp9T2+UUZT5m+BKyrDp3Ew==}
     engines: {node: '>=12'}

src/errors.ts

@@ -1,4 +1,4 @@
-import type { JsonObject } from 'type-fest'
+import type { Jsonifiable } from 'type-fest'
 import type { ZodError } from 'zod'
 import { ValidationError, fromZodError } from 'zod-validation-error'
 
@@ -10,12 +10,12 @@ export type ErrorOptions = {
   cause?: unknown
 
   /** Additional context to be added to the error. */
-  context?: JsonObject
+  context?: Jsonifiable
 }
 
 export class BaseError extends Error {
   status?: number
-  context?: JsonObject
+  context?: Jsonifiable
 
   constructor(message: string, opts: ErrorOptions = {}) {
     if (opts.cause) {

src/human-feedback/feedback.ts

@@ -273,8 +273,8 @@ export abstract class HumanFeedbackMechanism<
 }
 
 export function withHumanFeedback<
-  TInput extends void | types.JsonObject,
+  TInput extends types.TaskInput,
-  TOutput extends types.JsonValue,
+  TOutput extends types.TaskOutput,
   V extends HumanFeedbackType
 >(
   task: BaseTask<TInput, TOutput>,

src/llms/anthropic.ts

@@ -9,8 +9,8 @@ import { BaseChatCompletion } from './chat'
 const defaultStopSequences = [anthropic.HUMAN_PROMPT]
 
 export class AnthropicChatCompletion<
-  TInput extends void | types.JsonObject = any,
+  TInput extends types.TaskInput = any,
-  TOutput extends types.JsonValue = string
+  TOutput extends types.TaskOutput = string
 > extends BaseChatCompletion<
   TInput,
   TOutput,

src/llms/chat.ts

@@ -9,6 +9,7 @@ import * as types from '@/types'
 import { BaseTask } from '@/task'
 import { getCompiledTemplate } from '@/template'
 import {
+  extractFunctionIdentifierFromString,
   extractJSONArrayFromString,
   extractJSONObjectFromString
 } from '@/utils'
@@ -20,8 +21,8 @@ import {
 } from './llm-utils'
 
 export abstract class BaseChatCompletion<
-  TInput extends void | types.JsonObject = void,
+  TInput extends types.TaskInput = void,
-  TOutput extends types.JsonValue = string,
+  TOutput extends types.TaskOutput = string,
   TModelParams extends Record<string, any> = Record<string, any>,
   TChatCompletionResponse extends Record<string, any> = Record<string, any>
 > extends BaseLLM<TInput, TOutput, TModelParams> {
@@ -41,7 +42,7 @@ export abstract class BaseChatCompletion<
   }
 
   // TODO: use polymorphic `this` type to return correct BaseLLM subclass type
-  input<U extends void | types.JsonObject>(
+  input<U extends types.TaskInput>(
     inputSchema: ZodType<U>
   ): BaseChatCompletion<U, TOutput, TModelParams> {
     const refinedInstance = this as unknown as BaseChatCompletion<
@@ -54,7 +55,7 @@ export abstract class BaseChatCompletion<
   }
 
   // TODO: use polymorphic `this` type to return correct BaseLLM subclass type
-  output<U extends types.JsonValue>(
+  output<U extends types.TaskOutput>(
     outputSchema: ZodType<U>
   ): BaseChatCompletion<TInput, U, TModelParams> {
     const refinedInstance = this as unknown as BaseChatCompletion<
@@ -242,9 +243,10 @@ export abstract class BaseChatCompletion<
         `<<< Task createChatCompletion "${this.nameForHuman}"`
       )
       ctx.metadata.completion = completion
+      const message = completion.message
 
-      if (completion.message.function_call) {
+      if (message.function_call) {
-        const functionCall = completion.message.function_call
+        const functionCall = message.function_call
 
         if (!isUsingTools) {
           // TODO: not sure what we should do in this case...
@@ -252,16 +254,31 @@ export abstract class BaseChatCompletion<
           break
         }
 
+        const functionName = extractFunctionIdentifierFromString(
+          functionCall.name
+        )
+
+        if (!functionName) {
+          throw new errors.OutputValidationError(
+            `Unrecognized function call "${functionCall.name}"`
+          )
+        }
+
         const tool = this._tools!.find(
-          (tool) => tool.nameForModel === functionCall.name
+          (tool) => tool.nameForModel === functionName
         )
 
         if (!tool) {
           throw new errors.OutputValidationError(
-            `Function not found "${functionCall.name}"`
+            `Function not found "${functionName}"`
           )
         }
 
+        if (functionName !== functionCall.name) {
+          // fix function name hallucinations
+          functionCall.name = functionName
+        }
+
         let functionArguments: any
         try {
           functionArguments = JSON.parse(jsonrepair(functionCall.arguments))
@@ -281,12 +298,12 @@ export abstract class BaseChatCompletion<
         }
 
         // console.log('>>> sub-task', {
-        //   task: functionCall.name,
+        //   task: functionName,
         //   input: functionArguments
         // })
         this._logger.info(
           {
-            task: functionCall.name,
+            task: functionName,
             input: functionArguments
           },
           `>>> Sub-Task "${tool.nameForHuman}"`
@@ -297,14 +314,14 @@ export abstract class BaseChatCompletion<
 
         this._logger.info(
           {
-            task: functionCall.name,
+            task: functionName,
             input: functionArguments,
             output: toolCallResponse.result
           },
           `<<< Sub-Task "${tool.nameForHuman}"`
         )
         // console.log('<<< sub-task', {
-        //   task: functionCall.name,
+        //   task: functionName,
         //   input: functionArguments,
         //   output: toolCallResponse.result
         // })
@@ -321,7 +338,7 @@ export abstract class BaseChatCompletion<
         messages.push(completion.message as any)
         messages.push({
           role: 'function',
-          name: functionCall.name,
+          name: functionName,
           content: taskCallContent
         })
 

src/llms/llm.ts

@@ -7,8 +7,8 @@ import { Tokenizer, getTokenizerForModel } from '@/tokenizer'
 
 // TODO: TInput should only be allowed to be void or an object
 export abstract class BaseLLM<
-  TInput extends void | types.JsonObject = void,
+  TInput extends types.TaskInput = void,
-  TOutput extends types.JsonValue = string,
+  TOutput extends types.TaskOutput = string,
   TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseTask<TInput, TOutput> {
   protected _inputSchema: ZodType<TInput> | undefined
@@ -38,7 +38,7 @@ export abstract class BaseLLM<
   }
 
   // TODO: use polymorphic `this` type to return correct BaseLLM subclass type
-  input<U extends void | types.JsonObject>(
+  input<U extends types.TaskInput>(
     inputSchema: ZodType<U>
   ): BaseLLM<U, TOutput, TModelParams> {
     const refinedInstance = this as unknown as BaseLLM<U, TOutput, TModelParams>
@@ -47,7 +47,7 @@ export abstract class BaseLLM<
   }
 
   // TODO: use polymorphic `this` type to return correct BaseLLM subclass type
-  output<U extends types.JsonValue>(
+  output<U extends types.TaskOutput>(
     outputSchema: ZodType<U>
   ): BaseLLM<TInput, U, TModelParams> {
     const refinedInstance = this as unknown as BaseLLM<TInput, U, TModelParams>

src/llms/openai.ts

@@ -14,8 +14,8 @@ const openaiModelsSupportingFunctions = new Set([
 ])
 
 export class OpenAIChatCompletion<
-  TInput extends void | types.JsonObject = any,
+  TInput extends types.TaskInput = any,
-  TOutput extends types.JsonValue = string
+  TOutput extends types.TaskOutput = string
 > extends BaseChatCompletion<
   TInput,
   TOutput,

src/services/diffbot.ts

@@ -1,4 +1,7 @@
 import defaultKy from 'ky'
+import pThrottle from 'p-throttle'
+
+import { throttleKy } from '@/utils'
 
 export const DIFFBOT_API_BASE_URL = 'https://api.diffbot.com'
 export const DIFFBOT_KNOWLEDGE_GRAPH_API_BASE_URL = 'https://kg.diffbot.com'
@@ -94,9 +97,17 @@ export interface DiffbotObject {
   categories?: DiffbotCategory[]
   authors: DiffbotAuthor[]
   breadcrumb?: DiffbotBreadcrumb[]
+  items?: DiffbotListItem[]
   meta?: any
 }
 
+interface DiffbotListItem {
+  title: string
+  link: string
+  summary: string
+  image?: string
+}
+
 interface DiffbotAuthor {
   name: string
   link: string
@@ -307,6 +318,12 @@ interface DiffbotSkill {
   diffbotUri: string
 }
 
+const throttle = pThrottle({
+  limit: 5,
+  interval: 1000,
+  strict: true
+})
+
 export class DiffbotClient {
   api: typeof defaultKy
   apiKnowledgeGraph: typeof defaultKy
@@ -336,8 +353,14 @@ export class DiffbotClient {
     this.apiBaseUrl = apiBaseUrl
     this.apiKnowledgeGraphBaseUrl = apiKnowledgeGraphBaseUrl
 
-    this.api = ky.extend({ prefixUrl: apiBaseUrl, timeout: timeoutMs })
+    const throttledKy = throttleKy(ky, throttle)
-    this.apiKnowledgeGraph = ky.extend({
+
+    this.api = throttledKy.extend({
+      prefixUrl: apiBaseUrl,
+      timeout: timeoutMs
+    })
+
+    this.apiKnowledgeGraph = throttledKy.extend({
       prefixUrl: apiKnowledgeGraphBaseUrl,
       timeout: timeoutMs
     })
@@ -364,10 +387,13 @@ export class DiffbotClient {
       }
     }
 
+    console.log(`DiffbotClient._extract: ${endpoint}`, searchParams)
+
     return this.api
       .get(endpoint, {
         searchParams,
-        headers
+        headers,
+        retry: 2
       })
       .json<T>()
   }

src/services/serpapi.ts

@@ -656,6 +656,7 @@ export class SerpAPIClient {
         : queryOrOpts
     const { timeout, ...rest } = this.params
 
+    // console.log(options)
     return this.api
       .get('search', {
         searchParams: {

src/task.ts

@@ -20,8 +20,8 @@ import { defaultIDGeneratorFn, isValidTaskIdentifier } from './utils'
  *    - Invoking sub-agents
  */
 export abstract class BaseTask<
-  TInput extends void | types.JsonObject = void,
+  TInput extends types.TaskInput = void,
-  TOutput extends types.JsonValue = string
+  TOutput extends types.TaskOutput = string
 > {
   protected _agentic: Agentic
   protected _id: string

src/tools/diffbot.ts

@@ -0,0 +1,125 @@
+import { z } from 'zod'
+
+import * as types from '@/types'
+import { DiffbotClient } from '@/services/diffbot'
+import { BaseTask } from '@/task'
+import { omit, pick } from '@/utils'
+
+export const DiffbotInputSchema = z.object({
+  url: z.string().describe('URL of page to scrape')
+})
+export type DiffbotInput = z.infer<typeof DiffbotInputSchema>
+
+export const DiffbotImageSchema = z.object({
+  url: z.string().optional(),
+  naturalWidth: z.number().optional(),
+  naturalHeight: z.number().optional(),
+  width: z.number().optional(),
+  height: z.number().optional(),
+  isCached: z.boolean().optional(),
+  primary: z.boolean().optional()
+})
+
+export const DiffbotListItemSchema = z.object({
+  title: z.string().optional(),
+  link: z.string().optional(),
+  summary: z.string().optional(),
+  image: z.string().optional()
+})
+
+export const DiffbotObjectSchema = z.object({
+  type: z.string().optional(),
+  title: z.string().optional(),
+  siteName: z.string().optional(),
+  author: z.string().optional(),
+  authorUrl: z.string().optional(),
+  pageUrl: z.string().optional(),
+  date: z.string().optional(),
+  estimatedDate: z.string().optional(),
+  humanLanguage: z.string().optional(),
+  text: z.string().describe('core text content of the page').optional(),
+  tags: z.array(z.string()).optional(),
+  images: z.array(DiffbotImageSchema).optional(),
+  items: z.array(DiffbotListItemSchema).optional()
+})
+
+export const DiffbotOutputSchema = z.object({
+  type: z.string().optional(),
+  title: z.string().optional(),
+  objects: z.array(DiffbotObjectSchema).optional()
+})
+export type DiffbotOutput = z.infer<typeof DiffbotOutputSchema>
+
+export class DiffbotTool extends BaseTask<DiffbotInput, DiffbotOutput> {
+  protected _diffbotClient: DiffbotClient
+
+  constructor(
+    opts: {
+      diffbot?: DiffbotClient
+    } & types.BaseTaskOptions = {}
+  ) {
+    super(opts)
+
+    this._diffbotClient =
+      opts.diffbot ?? new DiffbotClient({ ky: opts.agentic?.ky })
+  }
+
+  public override get inputSchema() {
+    return DiffbotInputSchema
+  }
+
+  public override get outputSchema() {
+    return DiffbotOutputSchema
+  }
+
+  public override get nameForModel(): string {
+    return 'scrapeWebPage'
+  }
+
+  public override get nameForHuman(): string {
+    return 'Diffbot Scrape Web Page'
+  }
+
+  public override get descForModel(): string {
+    return 'Scrapes a web page for its content and structured data.'
+  }
+
+  protected override async _call(
+    ctx: types.TaskCallContext<DiffbotInput>
+  ): Promise<DiffbotOutput> {
+    const res = await this._diffbotClient.extractAnalyze({
+      url: ctx.input!.url
+    })
+
+    this._logger.info(res, `Diffbot response for url "${ctx.input!.url}"`)
+    console.log(res)
+
+    const pickedRes = {
+      type: res.type,
+      title: res.title,
+      objects: res.objects.map((obj) => ({
+        ...pick(
+          obj,
+          'type',
+          'siteName',
+          'author',
+          'authorUrl',
+          'pageUrl',
+          'date',
+          'estimatedDate',
+          'humanLanguage',
+          'items',
+          'text'
+        ),
+        tags: obj.tags?.map((tag) => tag.label),
+        images: obj.images?.map((image) => omit(image, 'diffbotUri'))
+      }))
+    }
+
+    this._logger.info(
+      pickedRes,
+      `Diffbot picked response for url "${ctx.input!.url}"`
+    )
+    return this.outputSchema.parse(pickedRes)
+  }
+}

src/tools/index.ts

@@ -1,4 +1,5 @@
 export * from './calculator'
+export * from './diffbot'
 export * from './metaphor'
 export * from './novu'
 export * from './serpapi'

src/tools/serpapi.ts

@@ -11,32 +11,32 @@ export const SerpAPIInputSchema = z.object({
 export type SerpAPIInput = z.infer<typeof SerpAPIInputSchema>
 
 export const SerpAPIOrganicSearchResult = z.object({
-  position: z.number(),
+  position: z.number().optional(),
-  title: z.string(),
+  title: z.string().optional(),
-  link: z.string(),
+  link: z.string().optional(),
-  displayed_link: z.string(),
+  displayed_link: z.string().optional(),
-  snippet: z.string(),
+  snippet: z.string().optional(),
   source: z.string().optional(),
   date: z.string().optional()
 })
 
 export const SerpAPIAnswerBox = z.object({
-  type: z.string(),
+  type: z.string().optional(),
-  title: z.string(),
+  title: z.string().optional(),
-  link: z.string(),
+  link: z.string().optional(),
-  displayed_link: z.string(),
+  displayed_link: z.string().optional(),
-  snippet: z.string()
+  snippet: z.string().optional()
 })
 
 export const SerpAPIKnowledgeGraph = z.object({
-  type: z.string(),
+  type: z.string().optional(),
-  description: z.string()
+  description: z.string().optional()
 })
 
 export const SerpAPIOutputSchema = z.object({
   knowledgeGraph: SerpAPIKnowledgeGraph.optional(),
   answerBox: SerpAPIAnswerBox.optional(),
-  organicResults: z.array(SerpAPIOrganicSearchResult)
+  organicResults: z.array(SerpAPIOrganicSearchResult).optional()
 })
 export type SerpAPIOutput = z.infer<typeof SerpAPIOutputSchema>
 
@@ -82,7 +82,10 @@ export class SerpAPITool extends BaseTask<SerpAPIInput, SerpAPIOutput> {
       num: ctx.input!.numResults
     })
 
-    this._logger.debug(res, `SerpAPI response for query "${ctx.input!.query}"`)
+    this._logger.debug(
+      res,
+      `SerpAPI response for query ${JSON.stringify(ctx.input, null, 2)}"`
+    )
 
     return this.outputSchema.parse({
       knowledgeGraph: res.knowledge_graph,

src/types.ts

@@ -2,7 +2,7 @@ import * as openai from '@agentic/openai-fetch'
 import * as anthropic from '@anthropic-ai/sdk'
 import ky from 'ky'
 import type { Options as RetryOptions } from 'p-retry'
-import type { JsonObject, JsonValue } from 'type-fest'
+import type { JsonObject, Jsonifiable } from 'type-fest'
 import { SafeParseReturnType, ZodType, ZodTypeAny, output, z } from 'zod'
 
 import type { Agentic } from './agentic'
@@ -15,9 +15,16 @@ import type { BaseTask } from './task'
 
 export { anthropic, openai }
 
-export type { JsonObject, JsonValue, Logger }
+export type { Jsonifiable, Logger }
 export type KyInstance = typeof ky
 
+export type JsonifiableObject =
+  | { [Key in string]?: Jsonifiable }
+  | { toJSON: () => Jsonifiable }
+
+export type TaskInput = void | JsonifiableObject
+export type TaskOutput = Jsonifiable
+
 export type ParsedData<T extends ZodTypeAny> = T extends ZodTypeAny
   ? output<T>
   : never
@@ -40,8 +47,8 @@ export interface BaseTaskOptions {
 }
 
 export interface BaseLLMOptions<
-  TInput extends void | JsonObject = void,
+  TInput extends TaskInput = void,
-  TOutput extends JsonValue = string,
+  TOutput extends TaskOutput = string,
   TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseTaskOptions {
   inputSchema?: ZodType<TInput>
@@ -54,8 +61,8 @@ export interface BaseLLMOptions<
 }
 
 export interface LLMOptions<
-  TInput extends void | JsonObject = void,
+  TInput extends TaskInput = void,
-  TOutput extends JsonValue = string,
+  TOutput extends TaskOutput = string,
   TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseLLMOptions<TInput, TOutput, TModelParams> {
   promptTemplate?: string
@@ -67,8 +74,8 @@ export type ChatMessage = openai.ChatMessage
 export type ChatMessageRole = openai.ChatMessageRole
 
 export interface ChatModelOptions<
-  TInput extends void | JsonObject = void,
+  TInput extends TaskInput = void,
-  TOutput extends JsonValue = string,
+  TOutput extends TaskOutput = string,
   TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseLLMOptions<TInput, TOutput, TModelParams> {
   messages: ChatMessage[]
@@ -116,7 +123,7 @@ export interface LLMTaskResponseMetadata<
 }
 
 export interface TaskResponse<
-  TOutput extends JsonValue = string,
+  TOutput extends TaskOutput = string,
   TMetadata extends TaskResponseMetadata = TaskResponseMetadata
 > {
   result: TOutput
@@ -124,7 +131,7 @@ export interface TaskResponse<
 }
 
 export interface TaskCallContext<
-  TInput extends void | JsonObject = void,
+  TInput extends TaskInput = void,
   TMetadata extends TaskResponseMetadata = TaskResponseMetadata
 > {
   input?: TInput

src/utils.ts

@@ -1,4 +1,5 @@
 import { customAlphabet, urlAlphabet } from 'nanoid'
+import type { ThrottledFunction } from 'p-throttle'
 
 import * as types from './types'
 
@@ -43,6 +44,30 @@ export function isValidTaskIdentifier(id: string): boolean {
   return !!id && taskNameRegex.test(id)
 }
 
+export function extractFunctionIdentifierFromString(
+  text: string
+): string | undefined {
+  text = text?.trim()
+
+  if (!text) {
+    return
+  }
+
+  if (isValidTaskIdentifier(text)) {
+    return text
+  }
+
+  const splits = text
+    .split(/[^a-zA-Z0-9_-]/)
+    .map((s) => {
+      s = s.trim()
+      return isValidTaskIdentifier(s) ? s : undefined
+    })
+    .filter(Boolean)
+
+  return splits[splits.length - 1]
+}
+
 /**
  * Chunks a string into an array of chunks.
  *
@@ -81,7 +106,7 @@ export function chunkString(text: string, maxLength: number): string[] {
  * @param json - JSON value to stringify
  * @returns stringified value with all double quotes around object keys removed
  */
-export function stringifyForModel(json: types.JsonValue): string {
+export function stringifyForModel(json: types.TaskOutput): string {
   const UNIQUE_PREFIX = defaultIDGeneratorFn()
   return (
     JSON.stringify(json, replacer)
@@ -112,3 +137,33 @@ export function stringifyForModel(json: types.JsonValue): string {
     return value
   }
 }
+
+export function pick<T extends object, U = T>(obj: T, ...keys: string[]): U {
+  return Object.fromEntries(
+    keys.filter((key) => key in obj).map((key) => [key, obj[key]])
+  ) as U
+}
+
+export function omit<T extends object, U = T>(obj: T, ...keys: string[]): U {
+  return Object.fromEntries<T>(
+    Object.entries(obj).filter(([key]) => !keys.includes(key))
+  ) as U
+}
+
+const noop = () => undefined
+
+/**
+ * Throttles HTTP requests made by a ky instance. Very useful for enforcing rate limits.
+ */
+export function throttleKy(
+  ky: types.KyInstance,
+  throttleFn: <Argument extends readonly unknown[], ReturnValue>(
+    function_: (...args: Argument) => ReturnValue
+  ) => ThrottledFunction<Argument, ReturnValue>
+) {
+  return ky.extend({
+    hooks: {
+      beforeRequest: [throttleFn(noop)]
+    }
+  })
+}

test/_utils.ts

@@ -5,9 +5,10 @@ import 'dotenv/config'
 import hashObject from 'hash-obj'
 import Redis from 'ioredis'
 import Keyv from 'keyv'
-import defaultKy from 'ky'
+import defaultKy, { AfterResponseHook, BeforeRequestHook } from 'ky'
 import pMemoize from 'p-memoize'
 
+import * as types from '@/types'
 import { Agentic } from '@/agentic'
 import { normalizeUrl } from '@/url-utils'
 
@@ -62,6 +63,9 @@ function getCacheKeyForRequest(request: Request): string | null {
   return null
 }
 
+const AGENTIC_TEST_CACHE_HEADER = 'x-agentic-test-cache'
+const AGENTIC_TEST_MOCK_HEADER = 'x-agentic-test-mock'
+
 /**
  * Custom `ky` instance that caches GET JSON requests.
  *
@@ -69,68 +73,148 @@ function getCacheKeyForRequest(request: Request): string | null {
  *  - support non-GET requests
  *  - support non-JSON responses
  */
-export const ky = defaultKy.extend({
+export function createTestKyInstance(
-  hooks: {
+  ky: types.KyInstance = defaultKy
-    beforeRequest: [
+): types.KyInstance {
-      async (request) => {
+  return ky.extend({
-        try {
+    hooks: {
-          // console.log(`beforeRequest ${request.method} ${request.url}`)
+      beforeRequest: [
+        async (request) => {
+          try {
+            const cacheKey = getCacheKeyForRequest(request)
+            // console.log(
+            //   `beforeRequest ${request.method} ${request.url} ⇒ ${cacheKey}`
+            // )
 
-          const cacheKey = getCacheKeyForRequest(request)
+            // console.log({ cacheKey })
-          // console.log({ cacheKey })
+            if (!cacheKey) {
-          if (!cacheKey) {
+              return
-            return
+            }
+
+            if (!(await keyv.has(cacheKey))) {
+              return
+            }
+
+            const cachedResponse = await keyv.get(cacheKey)
+            // console.log({ cachedResponse })
+
+            if (!cachedResponse) {
+              return
+            }
+
+            return new Response(JSON.stringify(cachedResponse), {
+              status: 200,
+              headers: {
+                'Content-Type': 'application/json',
+                [AGENTIC_TEST_CACHE_HEADER]: '1'
+              }
+            })
+          } catch (err) {
+            console.error('ky beforeResponse cache error', err)
           }
-
-          if (!(await keyv.has(cacheKey))) {
-            return
-          }
-
-          const cachedResponse = await keyv.get(cacheKey)
-          // console.log({ cachedResponse })
-
-          if (!cachedResponse) {
-            return
-          }
-
-          return new Response(JSON.stringify(cachedResponse), {
-            status: 200,
-            headers: { 'Content-Type': 'application/json' }
-          })
-        } catch (err) {
-          console.error('ky beforeResponse cache error', err)
         }
-      }
+      ],
-    ],
 
-    afterResponse: [
+      afterResponse: [
-      async (request, _options, response) => {
+        async (request, _options, response) => {
-        try {
+          try {
-          // console.log(
+            if (response.headers.get(AGENTIC_TEST_CACHE_HEADER)) {
-          //   `afterRequest ${request.method} ${request.url} ⇒ ${response.status}`
+              // console.log('cached')
-          // )
+              return
+            }
 
-          if (response.status < 200 || response.status >= 300) {
+            if (response.headers.get(AGENTIC_TEST_MOCK_HEADER)) {
-            return
+              // console.log('mocked')
+              return
+            }
+
+            const contentType = response.headers.get('content-type')
+            // console.log(
+            //   `afterRequest ${request.method} ${request.url} ⇒ ${response.status} ${contentType}`
+            // )
+
+            if (response.status < 200 || response.status >= 300) {
+              return
+            }
+
+            if (contentType !== 'application/json') {
+              return
+            }
+
+            const cacheKey = getCacheKeyForRequest(request)
+            // console.log({ cacheKey })
+            if (!cacheKey) {
+              console.log('222')
+              return
+            }
+
+            const responseBody = await response.json()
+            // console.log({ responseBody })
+
+            await keyv.set(cacheKey, responseBody)
+          } catch (err) {
+            console.error('ky afterResponse cache error', err)
           }
-
-          const cacheKey = getCacheKeyForRequest(request)
-          // console.log({ cacheKey })
-          if (!cacheKey) {
-            return
-          }
-
-          const responseBody = await response.json()
-          // console.log({ responseBody })
-
-          await keyv.set(cacheKey, responseBody)
-        } catch (err) {
-          console.error('ky afterResponse cache error', err)
         }
+      ]
+    }
+  })
+}
+
+function defaultBeforeRequest(request: Request): Response {
+  return new Response(
+    JSON.stringify({
+      url: request.url,
+      normalizedUrl: normalizeUrl(request.url),
+      method: request.method,
+      headers: request.headers
+    }),
+    {
+      status: 200,
+      headers: {
+        'Content-Type': 'application/json',
+        [AGENTIC_TEST_MOCK_HEADER]: '1'
       }
-    ]
+    }
-  }
+  )
-})
+}
+
+export function mockKyInstance(
+  ky: types.KyInstance = defaultKy,
+  {
+    beforeRequest = defaultBeforeRequest,
+    afterResponse = null
+  }: {
+    beforeRequest?: BeforeRequestHook | null
+    afterResponse?: AfterResponseHook | null
+  } = {}
+): types.KyInstance {
+  return ky.extend({
+    hooks: {
+      beforeRequest: beforeRequest === null ? [] : [beforeRequest],
+      afterResponse: afterResponse === null ? [] : [afterResponse]
+    }
+  })
+}
+
+/*
+ * NOTE: ky hooks are appended when doing `ky.extend`, so if you already have a
+ * beforeRequest hook, it will be called before any passed to `ky.extend`.
+ *
+ * For example:
+ *
+ * ```ts
+ * // runs caching first, then mocking
+ * const ky0 = mockKyInstance(createTestKyInstance(ky))
+ *
+ * // runs mocking first, then caching
+ * const ky1 = createTestKyInstance(mockKyInstance(ky))
+ *
+ * // runs throttling first, then mocking
+ * const ky2 = mockKyInstance(throttleKy(ky, throttle))
+ * ```
+ */
+export const ky = createTestKyInstance()
 
 export class OpenAITestClient extends OpenAIClient {
   createChatCompletion = pMemoize(super.createChatCompletion, {

test/services/diffbot.test.ts

@@ -36,7 +36,7 @@ test('Diffbot.extractArticle', async (t) => {
   t.is(result.objects[0].type, 'article')
 })
 
-test.only('Diffbot.knowledgeGraphSearch', async (t) => {
+test('Diffbot.knowledgeGraphSearch', async (t) => {
   if (!process.env.DIFFBOT_API_KEY || isCI) {
     return t.pass()
   }

test/utils.test.ts

@@ -1,15 +1,21 @@
 import test from 'ava'
+import ky from 'ky'
+import pThrottle from 'p-throttle'
 
 import {
   chunkString,
   defaultIDGeneratorFn,
+  extractFunctionIdentifierFromString,
   extractJSONArrayFromString,
   extractJSONObjectFromString,
   isValidTaskIdentifier,
   sleep,
-  stringifyForModel
+  stringifyForModel,
+  throttleKy
 } from '@/utils'
 
+import { mockKyInstance } from './_utils'
+
 test('isValidTaskIdentifier - valid', async (t) => {
   t.true(isValidTaskIdentifier('foo'))
   t.true(isValidTaskIdentifier('foo_bar_179'))
@@ -124,3 +130,49 @@ test('stringifyForModel should stringify objects with null values correctly', (t
 
   t.is(actualOutput, expectedOutput)
 })
+
+test('extractFunctionIdentifierFromString valid', (t) => {
+  t.is(extractFunctionIdentifierFromString('foo'), 'foo')
+  t.is(extractFunctionIdentifierFromString('fooBar_BAZ'), 'fooBar_BAZ')
+  t.is(extractFunctionIdentifierFromString('functions.fooBar'), 'fooBar')
+  t.is(
+    extractFunctionIdentifierFromString('function fooBar1234_'),
+    'fooBar1234_'
+  )
+})
+
+test('extractFunctionIdentifierFromString invalid', (t) => {
+  t.is(extractFunctionIdentifierFromString(''), undefined)
+  t.is(extractFunctionIdentifierFromString('  '), undefined)
+  t.is(extractFunctionIdentifierFromString('.-'), undefined)
+})
+
+test('throttleKy should rate-limit requests to ky properly', async (t) => {
+  t.timeout(30_1000)
+
+  const interval = 1000
+  const throttle = pThrottle({
+    limit: 1,
+    interval,
+    strict: true
+  })
+
+  const ky2 = mockKyInstance(throttleKy(ky, throttle))
+
+  const url = 'https://httpbin.org/get'
+
+  for (let i = 0; i < 10; i++) {
+    const before = Date.now()
+    const res = await ky2.get(url)
+    const after = Date.now()
+
+    const duration = after - before
+    // console.log(duration, res.status)
+    t.is(res.status, 200)
+
+    // leave a bit of wiggle room for the interval
+    if (i > 0) {
+      t.true(duration >= interval - interval / 5)
+    }
+  }
+})