Skip to content

Commit

Permalink
feat: ✨ add cancellation checks across core functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@pelikhan
pelikhan committed Feb 6, 2025
1 parent fcd2cc1 commit 723fa8f
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 7 deletions.
14 changes: 11 additions & 3 deletions packages/core/src/image.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,15 @@ import { TraceOptions } from "./trace"
import { logVerbose } from "./util"
import { deleteUndefinedValues } from "./cleaners"
import pLimit from "p-limit"
import { CancellationOptions, checkCancelled } from "./cancellation"

async function prepare(
url: BufferLike,
options: DefImagesOptions & TraceOptions
options: DefImagesOptions & TraceOptions & CancellationOptions
) {
// Dynamically import the Jimp library and its alignment enums
let {
cancellationToken,
autoCrop,
maxHeight,
maxWidth,
Expand All @@ -28,6 +30,7 @@ async function prepare(
flip,
detail,
} = options
checkCancelled(cancellationToken)

// https://platform.openai.com/docs/guides/vision/calculating-costs#managing-images
// If the URL is a string, resolve it to a data URI
Expand All @@ -50,6 +53,7 @@ async function prepare(
.join(", ")}`
)

checkCancelled(cancellationToken)
// Read the image using Jimp
const { Jimp, HorizontalAlign, VerticalAlign } = await import("jimp")
const img = await Jimp.read(buffer)
Expand Down Expand Up @@ -83,6 +87,8 @@ async function prepare(

if (greyscale) img.greyscale()

checkCancelled(cancellationToken)

// https://platform.openai.com/docs/guides/vision/low-or-high-fidelity-image-understanding#low-or-high-fidelity-image-understanding
if (detail === "low") {
contain(
Expand Down Expand Up @@ -155,23 +161,25 @@ async function encode(
*/
export async function imageEncodeForLLM(
url: BufferLike,
options: DefImagesOptions & TraceOptions
options: DefImagesOptions & TraceOptions & CancellationOptions
) {
const img = await prepare(url, options)
return await encode(img, options)
}

export async function imageTileEncodeForLLM(
urls: BufferLike[],
options: DefImagesOptions & TraceOptions
options: DefImagesOptions & TraceOptions & CancellationOptions
) {
if (urls.length === 0)
throw new Error("image: no images provided for tiling")

const { cancellationToken } = options
const limit = pLimit(4)
const imgs = await Promise.all(
urls.map((url) => limit(() => prepare(url, options)))
)
checkCancelled(cancellationToken)

logVerbose(`image: tiling ${imgs.length} images`)
const imgw = imgs.reduce((acc, img) => Math.max(acc, img.width), 0)
Expand Down
24 changes: 20 additions & 4 deletions packages/core/src/runpromptcontext.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,7 @@ export function createChatGenerationContext(

// Default output processor for the prompt
const defOutputProcessor = (fn: PromptOutputProcessorHandler) => {
checkCancelled(cancellationToken)
if (fn) appendChild(node, createOutputProcessor(fn))
}

Expand All @@ -349,6 +350,7 @@ export function createChatGenerationContext(
fn?: ChatFunctionHandler,
defOptions?: DefToolOptions
) => void = (name, description, parameters, fn, defOptions) => {
checkCancelled(cancellationToken)
if (name === undefined || name === null)
throw new Error("tool name is missing")

Expand Down Expand Up @@ -421,6 +423,7 @@ export function createChatGenerationContext(
) => Promise<void>,
options?: DefAgentOptions
): void => {
checkCancelled(cancellationToken)
const { tools, system, disableMemory, disableMemoryQuery, ...rest } =
options || {}
const memory = !disableMemory
Expand Down Expand Up @@ -463,6 +466,7 @@ export function createChatGenerationContext(
},
async (args) => {
// the LLM automatically adds extract arguments to the context
checkCancelled(cancellationToken)
const { context, ...rest } = args
const { query, ...argsNoQuery } = rest
infoCb?.({
Expand Down Expand Up @@ -537,6 +541,7 @@ export function createChatGenerationContext(
schema: JSONSchema,
defOptions?: DefSchemaOptions
) => {
checkCancelled(cancellationToken)
appendChild(node, createSchemaNode(name, schema, defOptions))

return name
Expand All @@ -548,6 +553,7 @@ export function createChatGenerationContext(
>,
defOptions?: DefImagesOptions
) => {
checkCancelled(cancellationToken)
if (files === undefined || files === null) {
if (defOptions?.ignoreEmpty) return
throw new Error("no images provided")
Expand All @@ -568,6 +574,7 @@ export function createChatGenerationContext(
if (!files.length) return undefined
const encoded = await imageTileEncodeForLLM(files, {
...defOptions,
cancellationToken,
trace,
})
return encoded
Expand All @@ -587,6 +594,7 @@ export function createChatGenerationContext(
(async () => {
const encoded = await imageEncodeForLLM(img, {
...defOptions,
cancellationToken,
trace,
})
return encoded
Expand All @@ -601,6 +609,7 @@ export function createChatGenerationContext(
(async () => {
const encoded = await imageEncodeForLLM(file, {
...defOptions,
cancellationToken,
trace,
})
return {
Expand All @@ -617,6 +626,7 @@ export function createChatGenerationContext(
generator: ChatParticipantHandler,
options?: ChatParticipantOptions
) => {
checkCancelled(cancellationToken)
if (generator)
appendChild(node, createChatParticipant({ generator, options }))
}
Expand All @@ -626,6 +636,7 @@ export function createChatGenerationContext(
description: string,
options?: FileOutputOptions
): void => {
checkCancelled(cancellationToken)
if (pattern)
appendChild(
node,
Expand All @@ -643,6 +654,7 @@ export function createChatGenerationContext(
strings: TemplateStringsArray,
...args: any[]
): RunPromptResultPromiseWithOptions => {
checkCancelled(cancellationToken)
const options: PromptGeneratorOptions = {}
const p: RunPromptResultPromiseWithOptions =
new Promise<RunPromptResult>(async (resolve, reject) => {
Expand All @@ -668,6 +680,7 @@ export function createChatGenerationContext(
audio: string | WorkspaceFile,
options?: TranscriptionOptions
): Promise<TranscriptionResult> => {
checkCancelled(cancellationToken)
const { cache, ...rest } = options || {}
const transcriptionTrace = trace.startTraceDetails("🎤 transcribe")
try {
Expand Down Expand Up @@ -772,6 +785,7 @@ export function createChatGenerationContext(
input: string,
options?: SpeechOptions
): Promise<SpeechResult> => {
checkCancelled(cancellationToken)
const { cache, voice, ...rest } = options || {}
const speechTrace = trace.startTraceDetails("🦜 speak")
try {
Expand Down Expand Up @@ -833,10 +847,16 @@ export function createChatGenerationContext(
}
}

const defFileMerge = (fn: FileMergeHandler) => {
checkCancelled(cancellationToken)
appendChild(node, createFileMerge(fn))
}

const runPrompt = async (
generator: string | PromptGenerator,
runOptions?: PromptGeneratorOptions
): Promise<RunPromptResult> => {
checkCancelled(cancellationToken)
const { label, applyEdits, throwOnError } = runOptions || {}
const runTrace = trace.startTraceDetails(`🎁 run prompt ${label || ""}`)
let messages: ChatCompletionMessageParam[] = []
Expand Down Expand Up @@ -1073,10 +1093,6 @@ export function createChatGenerationContext(
}
}

const defFileMerge = (fn: FileMergeHandler) => {
appendChild(node, createFileMerge(fn))
}

const ctx: RunPromptContextNode = Object.freeze<RunPromptContextNode>({
...turnCtx,
defAgent,
Expand Down

0 comments on commit 723fa8f

Please sign in to comment.