From 5d54d949168e959189358a392d6d9f0c00438c85 Mon Sep 17 00:00:00 2001 From: Joe McIlvain Date: Mon, 6 May 2024 09:54:10 -0700 Subject: [PATCH] refactor: Shrink interface of KurtEvent and rename as KurtStream As discussed with @InfraK, we want to shrink the interface a bit because it will lessen the impact of future changes to the final event structure. In doing so, we also decided to rename some things, such that the word `result` is only used for the final event of the stream, rather than the entire stream itself. BREAKING CHANGE: `finalText` and `finalData` have been removed, and some types have been renamed --- README.md | 8 +- ...{KurtResult.spec.ts => KurtStream.spec.ts} | 94 ++++++++----------- src/Kurt.ts | 6 +- src/KurtOpenAI.ts | 18 ++-- src/{KurtResult.ts => KurtStream.ts} | 52 ++++------ src/KurtVertexAI.ts | 14 +-- 6 files changed, 82 insertions(+), 110 deletions(-) rename spec/{KurtResult.spec.ts => KurtStream.spec.ts} (67%) rename src/{KurtResult.ts => KurtStream.ts} (79%) diff --git a/README.md b/README.md index b43dde7..867f16d 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,9 @@ for await (const event of stream) { // data: undefined, // } -await stream.finalText // "Hello! How can I assist you today?" +const { text } = await stream.result +console.log(text) +// "Hello! How can I assist you today?" ``` ## Generate Structured Data Output @@ -91,5 +93,7 @@ for await (const event of stream) { // { chunk: '"}' } // { finished: true, text: '{"say":"hello"}', data: { say: "hello" } } -await stream.finalData // { say: "hello" } +const { data } = await stream.result +console.log(data) +// { say: "hello" } ``` diff --git a/spec/KurtResult.spec.ts b/spec/KurtStream.spec.ts similarity index 67% rename from spec/KurtResult.spec.ts rename to spec/KurtStream.spec.ts index 6279e5e..af5515f 100644 --- a/spec/KurtResult.spec.ts +++ b/spec/KurtStream.spec.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from "@jest/globals" import { z } from "zod" -import { KurtResult, type KurtResultEvent } from "../src/KurtResult" +import { KurtStream, type KurtStreamEvent } from "../src/KurtStream" function kurtSayHelloEvents() { return [ @@ -17,19 +17,19 @@ function kurtSayHelloEvents() { ] } -function kurtSayHelloFinalEvent() { +function kurtSayHelloResult() { const events = kurtSayHelloEvents() // biome-ignore lint/style/noNonNullAssertion: Pending explanation return events[events.length - 1]! } -function kurtResultSayHello( +function kurtStreamSayHello( opts: { errorBeforeFinish?: boolean } = {} ) { const schema = z.object({ say: z.string() }) - return new KurtResult<(typeof schema)["shape"]>( + return new KurtStream<(typeof schema)["shape"]>( (async function* gen() { const events = kurtSayHelloEvents() @@ -41,7 +41,7 @@ function kurtResultSayHello( if (opts.errorBeforeFinish && event.finished) throw new Error("Whoops!") // Send the event. - yield event as KurtResultEvent<(typeof schema)["shape"]> + yield event as KurtStreamEvent<(typeof schema)["shape"]> } })() ) @@ -59,85 +59,73 @@ async function expectThrow(message: string, fn: () => Promise) { } describe("KurtResult", () => { - test("with an await for final text", async () => { - const result = kurtResultSayHello() - - expect(await result.finalText).toEqual(kurtSayHelloFinalEvent().text) - }) - - test("with an await for final data", async () => { - const result = kurtResultSayHello() - - expect(await result.finalData).toEqual(kurtSayHelloFinalEvent().data) - }) - test("with an await for final event", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() - expect(await result.finalEvent).toEqual(kurtSayHelloFinalEvent()) + expect(await stream.result).toEqual(kurtSayHelloResult()) }) test("with an await for final event catching an error", async () => { - const result = kurtResultSayHello({ errorBeforeFinish: true }) + const stream = kurtStreamSayHello({ errorBeforeFinish: true }) - expectThrow("Whoops!", async () => await result.finalEvent) + expectThrow("Whoops!", async () => await stream.result) }) test("with one event listener", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() const events: unknown[] = [] - for await (const event of result) events.push(event) + for await (const event of stream) events.push(event) expect(events).toEqual(kurtSayHelloEvents()) }) test("with one event listener, catching an error", async () => { - const result = kurtResultSayHello({ errorBeforeFinish: true }) + const stream = kurtStreamSayHello({ errorBeforeFinish: true }) const events: unknown[] = [] await expectThrow("Whoops!", async () => { - for await (const event of result) events.push(event) + for await (const event of stream) events.push(event) }) expect(events).toEqual(kurtSayHelloEvents().slice(0, -1)) }) test("with final awaits before listener", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() - expect(await result.finalEvent).toEqual(kurtSayHelloFinalEvent()) - expect(await result.finalText).toEqual(kurtSayHelloFinalEvent().text) - expect(await result.finalData).toEqual(kurtSayHelloFinalEvent().data) + expect(await stream.result).toEqual(kurtSayHelloResult()) + expect(await stream.result).toEqual(kurtSayHelloResult()) + expect(await stream.result).toEqual(kurtSayHelloResult()) const events: unknown[] = [] - for await (const event of result) events.push(event) + for await (const event of stream) events.push(event) expect(events).toEqual(kurtSayHelloEvents()) }) test("with final awaits after listener", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() const events: unknown[] = [] - for await (const event of result) events.push(event) + for await (const event of stream) events.push(event) - expect(await result.finalEvent).toEqual(kurtSayHelloFinalEvent()) - expect(await result.finalText).toEqual(kurtSayHelloFinalEvent().text) - expect(await result.finalData).toEqual(kurtSayHelloFinalEvent().data) + expect(await stream.result).toEqual(kurtSayHelloResult()) + expect(await stream.result).toEqual(kurtSayHelloResult()) + expect(await stream.result).toEqual(kurtSayHelloResult()) expect(events).toEqual(kurtSayHelloEvents()) }) test("with three listeners, each after the last one finished", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() const events1: unknown[] = [] const events2: unknown[] = [] const events3: unknown[] = [] - for await (const event of result) events1.push(event) - for await (const event of result) events2.push(event) - for await (const event of result) events3.push(event) + for await (const event of stream) events1.push(event) + for await (const event of stream) events2.push(event) + for await (const event of stream) events3.push(event) expect(events1).toEqual(kurtSayHelloEvents()) expect(events2).toEqual(kurtSayHelloEvents()) @@ -145,19 +133,19 @@ describe("KurtResult", () => { }) test("with three listeners, each catching an error", async () => { - const result = kurtResultSayHello({ errorBeforeFinish: true }) + const stream = kurtStreamSayHello({ errorBeforeFinish: true }) const events1: unknown[] = [] const events2: unknown[] = [] const events3: unknown[] = [] await expectThrow("Whoops!", async () => { - for await (const event of result) events1.push(event) + for await (const event of stream) events1.push(event) }) await expectThrow("Whoops!", async () => { - for await (const event of result) events2.push(event) + for await (const event of stream) events2.push(event) }) await expectThrow("Whoops!", async () => { - for await (const event of result) events3.push(event) + for await (const event of stream) events3.push(event) }) expect(events1).toEqual(kurtSayHelloEvents().slice(0, -1)) @@ -166,14 +154,14 @@ describe("KurtResult", () => { }) test("with many listeners, interleaved with one another", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() const events: unknown[][] = [] const listeners: Promise[] = [] async function listen(spawnMoreListeners = false) { events.push([]) const listenerIndex = events.length - 1 - for await (const event of result) { + for await (const event of stream) { events[listenerIndex]?.push(event) if (spawnMoreListeners) listeners.push(listen()) } @@ -193,7 +181,7 @@ describe("KurtResult", () => { }) test("with many listeners, interleaved with final event awaits", async () => { - const result = kurtResultSayHello() + const stream = kurtStreamSayHello() const events: unknown[][] = [] const listeners: Promise[] = [] @@ -201,15 +189,13 @@ describe("KurtResult", () => { async function listen(spawnMoreListeners = false) { events.push([]) const listenerIndex = events.length - 1 - for await (const event of result) { + for await (const event of stream) { events[listenerIndex]?.push(event) if (spawnMoreListeners) { listeners.push(listen()) awaits.push( (async () => - expect(await result.finalEvent).toEqual( - kurtSayHelloFinalEvent() - ))() + expect(await stream.result).toEqual(kurtSayHelloResult()))() ) } } @@ -230,7 +216,7 @@ describe("KurtResult", () => { }) test("with many listeners/awaits, interleaved, with error", async () => { - const result = kurtResultSayHello({ errorBeforeFinish: true }) + const stream = kurtStreamSayHello({ errorBeforeFinish: true }) const events: unknown[][] = [] const listeners: Promise[] = [] @@ -238,15 +224,13 @@ describe("KurtResult", () => { async function listen(spawnMoreListeners = false) { events.push([]) const listenerIndex = events.length - 1 - for await (const event of result) { + for await (const event of stream) { events[listenerIndex]?.push(event) if (spawnMoreListeners) { const listener = listen() listeners.push(listener) errors.push(expectThrow("Whoops!", async () => await listener)) - errors.push( - expectThrow("Whoops!", async () => await result.finalEvent) - ) + errors.push(expectThrow("Whoops!", async () => await stream.result)) } } } diff --git a/src/Kurt.ts b/src/Kurt.ts index fbe22a3..98d677e 100644 --- a/src/Kurt.ts +++ b/src/Kurt.ts @@ -1,14 +1,14 @@ -import type { KurtResult } from "./KurtResult" +import type { KurtStream } from "./KurtStream" import type { KurtSchema, KurtSchemaInner } from "./KurtSchema" export interface Kurt { generateNaturalLanguage( options: KurtGenerateNaturalLanguageOptions - ): KurtResult + ): KurtStream generateStructuredData( options: KurtGenerateStructuredDataOptions - ): KurtResult + ): KurtStream } export interface KurtMessage { diff --git a/src/KurtOpenAI.ts b/src/KurtOpenAI.ts index 16df941..b48c88d 100644 --- a/src/KurtOpenAI.ts +++ b/src/KurtOpenAI.ts @@ -6,7 +6,7 @@ import type { KurtGenerateStructuredDataOptions, KurtMessage, } from "./Kurt" -import { KurtResult, type KurtResultEvent } from "./KurtResult" +import { KurtStream, type KurtStreamEvent } from "./KurtStream" import type { KurtSchemaInner, KurtSchemaInnerMaybe, @@ -37,7 +37,7 @@ export class KurtOpenAI implements Kurt { generateNaturalLanguage( options: KurtGenerateNaturalLanguageOptions - ): KurtResult { + ): KurtStream { return this.handleStream( undefined, this.options.openAI.chat.completions.create({ @@ -50,7 +50,7 @@ export class KurtOpenAI implements Kurt { generateStructuredData( options: KurtGenerateStructuredDataOptions - ): KurtResult { + ): KurtStream { const schema = options.schema return this.handleStream( @@ -80,7 +80,7 @@ export class KurtOpenAI implements Kurt { private handleStream( schema: KurtSchemaMaybe, response: OpenAIResponse - ): KurtResult { + ): KurtStream { async function* generator() { const stream = await response const chunks: string[] = [] @@ -91,13 +91,13 @@ export class KurtOpenAI implements Kurt { const textChunk = choice.delta.content if (textChunk) { - yield { chunk: textChunk } as KurtResultEvent + yield { chunk: textChunk } as KurtStreamEvent chunks.push(textChunk) } const dataChunk = choice.delta.tool_calls?.at(0)?.function?.arguments if (dataChunk) { - yield { chunk: dataChunk } as KurtResultEvent + yield { chunk: dataChunk } as KurtStreamEvent chunks.push(dataChunk) } @@ -111,19 +111,19 @@ export class KurtOpenAI implements Kurt { finished: true, text, data, - } as KurtResultEvent + } as KurtStreamEvent } else { yield { finished: true, text, data: undefined, - } as KurtResultEvent + } as KurtStreamEvent } } } } - return new KurtResult(generator()) + return new KurtStream(generator()) } private toOpenAIMessages = ({ diff --git a/src/KurtResult.ts b/src/KurtStream.ts similarity index 79% rename from src/KurtResult.ts rename to src/KurtStream.ts index b3bd496..07517ce 100644 --- a/src/KurtResult.ts +++ b/src/KurtStream.ts @@ -1,18 +1,18 @@ import type { Promisable } from "type-fest" import type { KurtSchemaInnerMaybe, KurtSchemaResultMaybe } from "./KurtSchema" -export type KurtResultEventChunk = { chunk: string } -export type KurtResultEventFinal = { +export type KurtStreamEventChunk = { chunk: string } +export type KurtResult = { finished: true text: string data: KurtSchemaResultMaybe } -export type KurtResultEvent = - | KurtResultEventChunk - | KurtResultEventFinal +export type KurtStreamEvent = + | KurtStreamEventChunk + | KurtResult type _AdditionalListener = ( - event: KurtResultEvent | { uncaughtError: unknown } + event: KurtStreamEvent | { uncaughtError: unknown } ) => void // This class represents the result of a call to an LLM. @@ -25,41 +25,25 @@ type _AdditionalListener = ( // each listener will see exactly the same stream of events, regardless // of when each one started listening. // -// It also exposes a few convenience getters for callers who are only -// interested in the final result event, or the text/data from that event. -export class KurtResult - implements AsyncIterable> +// It also exposes a `result` convenience getter for callers who are only +// interested in the final result event. +export class KurtStream + implements AsyncIterable> { private started = false private finished = false - private seenEvents: KurtResultEvent[] = [] + private seenEvents: KurtStreamEvent[] = [] private finalError?: { uncaughtError: unknown } private additionalListeners = new Set<_AdditionalListener>() // Create a new result stream, from the given underlying stream generator. - constructor(private gen: AsyncGenerator>) {} + constructor(private gen: AsyncGenerator>) {} // Get the final event from the end of the result stream, when it is ready. - get finalEvent(): Promise> { + get result(): Promise> { return toFinal(this) } - // Get the text from the end of the result stream, when it is ready. - get finalText(): Promise { - return this._finalText() - } - private async _finalText() { - return (await this.finalEvent).text - } - - // Get the data from the end of the result stream, when it is ready. - get finalData(): Promise> { - return this._finalData() - } - private async _finalData() { - return (await this.finalEvent).data - } - // Get each event in the stream (each yielded from this `AsyncGenerator`). async *[Symbol.asyncIterator]() { // If some other caller has already started iterating on this stream, @@ -123,10 +107,10 @@ export class KurtResult // To make this generator work, we need to set up a replaceable promise // that will receive the next event (or error) via the listener callback. - let nextEventResolve: (value: Promisable>) => void + let nextEventResolve: (value: Promisable>) => void let nextEventReject: (reason?: unknown) => void const createNextEventPromise = () => { - return new Promise>((resolve, reject) => { + return new Promise>((resolve, reject) => { nextEventResolve = resolve nextEventReject = reject }) @@ -162,9 +146,9 @@ export class KurtResult } async function toFinal( - result: KurtResult -): Promise> { - for await (const event of result) { + stream: KurtStream +): Promise> { + for await (const event of stream) { if ("finished" in event) { return event } diff --git a/src/KurtVertexAI.ts b/src/KurtVertexAI.ts index 873b61f..1eee83b 100644 --- a/src/KurtVertexAI.ts +++ b/src/KurtVertexAI.ts @@ -8,7 +8,7 @@ import type { KurtGenerateStructuredDataOptions, KurtMessage, } from "./Kurt" -import { KurtResult, type KurtResultEvent } from "./KurtResult" +import { KurtStream, type KurtStreamEvent } from "./KurtStream" import type { KurtSchema, KurtSchemaInner, @@ -38,7 +38,7 @@ export class KurtVertexAI implements Kurt { generateNaturalLanguage( options: KurtGenerateNaturalLanguageOptions - ): KurtResult { + ): KurtStream { const llm = this.options.vertexAI.getGenerativeModel({ model: this.options.model, }) as VertexAIGenerativeModel @@ -53,7 +53,7 @@ export class KurtVertexAI implements Kurt { generateStructuredData( options: KurtGenerateStructuredDataOptions - ): KurtResult { + ): KurtStream { const schema = options.schema const llm = this.options.vertexAI.getGenerativeModel({ @@ -83,7 +83,7 @@ export class KurtVertexAI implements Kurt { private handleStream( schema: KurtSchemaMaybe, response: VertexAIResponse - ): KurtResult { + ): KurtStream { async function* generator() { const { stream } = await response const chunks: string[] = [] @@ -114,7 +114,7 @@ export class KurtVertexAI implements Kurt { finished: true, text, data, - } as KurtResultEvent + } as KurtStreamEvent } else { const text = chunks.join("") const data = undefined @@ -122,14 +122,14 @@ export class KurtVertexAI implements Kurt { finished: true, text, data, - } as KurtResultEvent + } as KurtStreamEvent } } } } } - return new KurtResult(generator()) + return new KurtStream(generator()) } private toVertexAIMessages = ({