Merge pull request #58 from FormulaMonks/add/kurt-cache

feat: add `kurt-cache` package with `KurtCache` adapter

examples/basic/package.json

@@ -24,6 +24,7 @@
   },
   "dependencies": {
     "@formula-monks/kurt": "workspace:*",
+    "@formula-monks/kurt-cache": "workspace:*",
     "@formula-monks/kurt-open-ai": "workspace:*",
     "@formula-monks/kurt-vertex-ai": "workspace:*",
     "@google-cloud/vertexai": "1.1.0",

examples/basic/src/cache.ts

@@ -0,0 +1,30 @@
+import { Kurt } from "@formula-monks/kurt"
+import { KurtCache } from "@formula-monks/kurt-cache"
+import { KurtOpenAI } from "@formula-monks/kurt-open-ai"
+import OpenAI from "openai"
+import { z } from "zod"
+
+const cacheAdapter = new KurtCache(
+  // This is the directory where the cache will be stored.
+  `${__dirname}/.kurt-cache`,
+  // This is the cache prefix. It should identify this adapter configuration,
+  // If the prefix changes, prior matching cache entries will no longer match.
+  "openai-gpt-3.5-turbo-0125",
+  // This function will only be run the first time we encounter a cache miss.
+  () =>
+    new KurtOpenAI({
+      openAI: new OpenAI(),
+      model: "gpt-3.5-turbo-0125",
+    })
+)
+
+const kurt = new Kurt(cacheAdapter)
+
+const schema = z.object({ say: z.string().describe("A single word to say") })
+const stream1 = kurt.generateStructuredData({ prompt: "Say hello!", schema })
+const stream2 = kurt.generateStructuredData({ prompt: "Say hello!", schema })
+const stream3 = kurt.generateStructuredData({ prompt: "Say hi!", schema })
+
+console.log((await stream1.result).data) // (cache miss on first run)
+console.log((await stream2.result).data) // (always cached; identical to prior)
+console.log((await stream3.result).data) // (cache miss on first run)

packages/kurt-cache/.gitignore

@@ -0,0 +1,2 @@
+// Ignore the test cache directory (but NOT the test-retain cache directory)
+.kurt-cache/test

packages/kurt-cache/.kurt-cache/test-retain/stub-ad557ba1818e8013f9e2fbc9598c034e263c96c5fe7edd491e75be8ce450f5c9.yaml

@@ -0,0 +1,13 @@
+messages:
+  - role: user
+    text: Was this cached?
+sampling:
+  maxOutputTokens: 4096
+  temperature: 0.5
+  topP: 0.95
+tools: {}
+response:
+  - chunk: This was cached
+  - chunk: " on disk"
+  - finished: true
+    text: This was cached on disk

packages/kurt-cache/README.md

@@ -0,0 +1,16 @@
+# Kurt Adapter for Caching
+
+[Kurt](https://github.com/FormulaMonks/kurt) is a TypeScript library by [Formula.Monks](https://www.formula.co/) that wraps AI SDKs, making it easy to build structured LLM-based applications (RAG, agents, etc) that work with any LLM that supports structured output (via function calling features).
+
+This package implements an adapter for Kurt that caches responses to disk. This is most useful for testing and development, for:
+- ensuring determinism of the test data and code paths
+- avoiding unnecessary AI usage costs for repetitive requests
+- allowing for running existing tests without requiring an API key for the AI service
+
+The cache entries are YAML files, which can be easily inspected and modified for your test cases, as well as checked into your code repository for meaningful code review.
+
+[Read here for more information about Kurt](https://github.com/FormulaMonks/kurt/blob/main/README.md).
+
+## Examples
+
+[This example code](../../examples/basic/src/openai.ts) shows how to set up and use Kurt with OpenAI.

packages/kurt-cache/jest.config.js

@@ -0,0 +1,5 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: "ts-jest",
+  testEnvironment: "node",
+}

packages/kurt-cache/package.json

@@ -0,0 +1,44 @@
+{
+  "name": "@formula-monks/kurt-cache",
+  "description": "Caching plugin for Kurt - A wrapper for AI SDKs, for building LLM-agnostic structured AI applications",
+  "license": "MIT",
+  "version": "1.0.0",
+  "homepage": "https://github.com/FormulaMonks/kurt",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/FormulaMonks/kurt.git",
+    "directory": "packages/kurt-cache"
+  },
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "files": ["dist"],
+  "scripts": {
+    "test": "jest",
+    "build": "tsc --build",
+    "prepack": "pnpm run build",
+    "format": "pnpm biome format --write .",
+    "lint": "pnpm biome lint --apply .",
+    "check": "pnpm biome check .",
+    "prepublish": "../../scripts/interpolate-example-code.sh README.md",
+    "release": "pnpm exec semantic-release"
+  },
+  "release": {
+    "branches": ["main"],
+    "extends": "semantic-release-monorepo"
+  },
+  "dependencies": {
+    "@formula-monks/kurt": "^1.4.0",
+    "yaml": "^2.4.5",
+    "zod-to-json-schema": "^3.23.3"
+  },
+  "devDependencies": {
+    "@jest/globals": "^29.7.0",
+    "@types/node": "^18.19.32",
+    "jest": "^29.7.0",
+    "semantic-release": "^23.0.8",
+    "semantic-release-monorepo": "^8.0.2",
+    "ts-jest": "^29.1.2",
+    "type-fest": "^4.30.0",
+    "typescript": "^5.4.5"
+  }
+}

packages/kurt-cache/spec/KurtCache.spec.ts

@@ -0,0 +1,214 @@
+import { describe, test, expect } from "@jest/globals"
+import {
+  existsSync,
+  readFileSync,
+  readdirSync,
+  rmSync,
+  rmdirSync,
+} from "node:fs"
+import { randomBytes } from "node:crypto"
+import {
+  Kurt,
+  type KurtAdapterV1,
+  type KurtMessage,
+  type KurtSamplingOptions,
+  type KurtSchema,
+  type KurtSchemaInner,
+  type KurtSchemaInnerMap,
+  type KurtSchemaMap,
+  type KurtSchemaMapSingleResult,
+  type KurtSchemaResult,
+  type KurtStreamEvent,
+} from "@formula-monks/kurt"
+import { KurtCache } from "../src"
+import { zodToJsonSchema } from "zod-to-json-schema"
+import type { NonEmptyTuple } from "type-fest"
+
+// Define the cache directories that will be used by this test.
+const cacheDir = `${__dirname}/../.kurt-cache/test`
+const cacheDirRetain = `${cacheDir}-retain`
+
+// A convenience function to make the test cases succinct one-liners.
+const gen = async (kurt: Kurt, prompt: string) =>
+  (await kurt.generateNaturalLanguage({ prompt }).result).text
+
+describe("KurtCache", () => {
+  test("when cache misses, runs the adapter setup fn just once", async () => {
+    // Remove the cache dir first, to demonstrate that the cache dir will be
+    // automatically created by the KurtCache adapter.
+    if (existsSync(cacheDir)) rmdirSync(cacheDir, { recursive: true })
+
+    // Use a random string to ensure an initial cache miss for this test run.
+    const random = randomBytes(8).toString("hex")
+
+    // Use the cache adapter configured appropriately to find the cache entry.
+    let adapterFnCallCount = 0
+    const kurt = new Kurt(
+      new KurtCache(cacheDir, "stub", () => {
+        adapterFnCallCount++
+        return new StubAdapter([
+          ["World ", random],
+          ["bar ", random],
+          ["never ", random],
+        ])
+      })
+    )
+
+    // Expect the canned responses from the stub adapter to be returned,
+    // with cache misses on the first call for each prompt, but with
+    // cache hits on subsequent calls using a prior prompt.
+    expect(await gen(kurt, `Hello ${random}`)).toEqual(`World ${random}`)
+    expect(await gen(kurt, `foo ${random}`)).toEqual(`bar ${random}`)
+    expect(await gen(kurt, `Hello ${random}`)).toEqual(`World ${random}`)
+    expect(await gen(kurt, `foo ${random}`)).toEqual(`bar ${random}`)
+    expect(await gen(kurt, `foo ${random}`)).toEqual(`bar ${random}`)
+
+    // Expect that the adapter setup function was called just once.
+    expect(adapterFnCallCount).toEqual(1)
+
+    // Expect that the cache dir contains exactly two files.
+    expect(readdirSync(cacheDir)).toHaveLength(2)
+  })
+
+  test("when cache hits, works without running the adapter fn", async () => {
+    // We compare with a hard-coded hash here to test that the hash function is
+    // stable/deterministic across library versions of KurtCache.
+    //
+    // If you find yourself needing to change this hash value, it means
+    // that you're breaking all existing cache entries, which is a breaking
+    // change for users of KurtCache who rely on it for their test suites.
+    const hash =
+      "ad557ba1818e8013f9e2fbc9598c034e263c96c5fe7edd491e75be8ce450f5c9"
+    const filePath = `${cacheDirRetain}/stub-${hash}.yaml`
+
+    // Assert that the cache file entry already exists (it has been
+    // committed into the repo and retained there)
+    const cached = readFileSync(filePath, "utf8")
+    expect(cached).toContain("text: This was cached on disk")
+
+    // Use the cache adapter configured appropriately to find the cache entry.
+    let adapterFnCallCount = 0
+    const kurt = new Kurt(
+      new KurtCache(cacheDirRetain, "stub", () => {
+        adapterFnCallCount++
+        return new StubAdapter([["This was cached", " on disk"]])
+      })
+    )
+
+    // Expect the cache hit to return the result text from the file.
+    expect(await gen(kurt, "Was this cached?")).toEqual(
+      "This was cached on disk"
+    )
+
+    // Expect that the adapter setup function was never called.
+    expect(adapterFnCallCount).toEqual(0)
+
+    // Expect that the cache file entry was not modified.
+    expect(readFileSync(filePath, "utf8")).toEqual(cached)
+
+    // Delete the cache file and prove that it regenerates exactly the same.
+    rmSync(filePath)
+    expect(await gen(kurt, "Was this cached?")).toEqual(
+      "This was cached on disk"
+    )
+    expect(adapterFnCallCount).toEqual(1)
+    expect(readFileSync(filePath, "utf8")).toEqual(cached)
+  })
+})
+
+class StubAdapter
+  implements
+    KurtAdapterV1<{
+      rawMessage: { bytes: string }
+      rawSchema: ReturnType<typeof zodToJsonSchema>
+      rawTool: {
+        name: string
+        desc: string
+        schema: ReturnType<typeof zodToJsonSchema>
+      }
+      rawEvent: { bytes: string }
+    }>
+{
+  constructor(readonly cannedResponses: NonEmptyTuple<NonEmptyTuple<string>>) {}
+
+  /**
+   * Rotate through the canned responses, returning the next one each time.
+   */
+  private nextCannedResponseIndex = 0
+  private nextCannedResponse() {
+    let cannedResponse = this.cannedResponses[this.nextCannedResponseIndex]
+    if (cannedResponse === undefined) {
+      cannedResponse = this.cannedResponses[0]
+      this.nextCannedResponseIndex = 0
+    }
+
+    this.nextCannedResponseIndex++
+    return cannedResponse
+  }
+
+  kurtAdapterVersion = "v1" as const
+
+  transformToRawMessages(messages: KurtMessage[]) {
+    return messages.map((message) => ({
+      bytes: message.text ?? JSON.stringify(message),
+    }))
+  }
+
+  transformToRawSchema<I extends KurtSchemaInner>(schema: KurtSchema<I>) {
+    return zodToJsonSchema(schema)
+  }
+
+  transformToRawTool(tool: {
+    name: string
+    description: string
+    parameters: ReturnType<typeof zodToJsonSchema>
+  }) {
+    return {
+      name: tool.name,
+      desc: tool.description,
+      schema: tool.parameters,
+    }
+  }
+
+  async *generateRawEvents(options: {
+    messages: { bytes: string }[]
+    sampling: Required<KurtSamplingOptions>
+    tools: {
+      [key: string]: {
+        name: string
+        desc: string
+        schema: ReturnType<typeof zodToJsonSchema>
+      }
+    }
+    forceTool?: string
+  }) {
+    for (const bytes of this.nextCannedResponse()) {
+      yield { bytes }
+    }
+  }
+
+  async *transformNaturalLanguageFromRawEvents(
+    rawEvents: AsyncIterable<{ bytes: string }>
+  ): AsyncIterable<KurtStreamEvent<undefined>> {
+    let text = ""
+    for await (const { bytes } of rawEvents) {
+      text += bytes
+      yield { chunk: bytes }
+    }
+    yield { finished: true, text, data: undefined }
+  }
+
+  transformStructuredDataFromRawEvents<I extends KurtSchemaInner>(
+    schema: KurtSchema<I>,
+    rawEvents: AsyncIterable<{ bytes: string }>
+  ): AsyncIterable<KurtStreamEvent<KurtSchemaResult<I>>> {
+    throw new Error("Not implemented because tests here don't use it")
+  }
+
+  transformWithOptionalToolsFromRawEvents<I extends KurtSchemaInnerMap>(
+    tools: KurtSchemaMap<I>,
+    rawEvents: AsyncIterable<{ bytes: string }>
+  ): AsyncIterable<KurtStreamEvent<KurtSchemaMapSingleResult<I> | undefined>> {
+    throw new Error("Not implemented because tests here don't use it")
+  }
+}