@bastani/pi-ai 0.9.43 → 0.9.44-alpha.2

package/dist/api/openai-completions.js

@@ -4,7 +4,6 @@ import { AssistantMessageEventStream } from "../utils/event-stream.js";
 import { shortHash } from "../utils/hash.js";
 import { headersToRecord } from "../utils/headers.js";
 import { getIdentityHeaders } from "../utils/identity.js";
-import { parseStreamingJson, parseToolCallArguments } from "../utils/json-parse.js";
 import { applyOpenRouterReportedCost, isOpenRouterModel, } from "../utils/openrouter-cost.js";
 import { getProviderEnvValue } from "../utils/provider-env.js";
 import { retryProviderRequest } from "../utils/provider-retry.js";
@@ -12,9 +11,10 @@ import { admitOpenAIRequests, getFetchAdmission } from "../utils/request-admissi
 import { getProviderRetryAfterMs } from "../utils/retry-hints.js";
 import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
 import { createStreamDeadline, withStreamDeadline } from "../utils/stream-deadline.js";
+import { ToolCallArgumentsStream } from "../utils/streaming-json.js";
 import { getSystemMessageText, renderSystemMessageUpdate } from "../utils/text.js";
 import { getDeclaredTools, resolveTranscript, resolveTranscriptTools, } from "../utils/transcript.js";
-import { appendGrammarToolInputJsonDelta, createGrammarToolInputProperties, getGrammarToolInput, getJsonSchemaToolParameters, resolveGrammarConstrainedSampling, resolveJsonSchemaStrictSampling, } from "./constrained-sampling.js";
+import { appendGrammarToolInputJsonDelta, closeGrammarToolInputJson, createGrammarToolInputProperties, getGrammarToolInput, getJsonSchemaToolParameters, resolveGrammarConstrainedSampling, resolveJsonSchemaStrictSampling, } from "./constrained-sampling.js";
 import { buildCopilotDynamicHeaders, hasCopilotVisionInput, preserveCopilotIntegrationHeader, } from "./github-copilot-headers.js";
 import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.js";
 import { RetryHintOpenAI } from "./openai-retry-hint-client.js";
@@ -222,23 +222,14 @@ export const stream = (model, context, options) => {
             let hasFinishReason = false;
             const toolCallBlocksByIndex = new Map();
             const toolCallBlocksById = new Map();
+            const toolCallStates = new Map();
             const blocks = output.content;
             const getContentIndex = (block) => blocks.indexOf(block);
-            const getCustomToolCallInput = (block) => {
-                const property = block.customInput?.property;
-                if (property === undefined)
-                    return "";
-                const value = block.arguments[property];
-                return typeof value === "string" ? value : "";
-            };
-            const appendCustomToolCallInput = (block, nextInput, close) => {
-                const customInput = block.customInput;
-                if (!customInput)
-                    return undefined;
-                const delta = appendGrammarToolInputJsonDelta(customInput.jsonBuffer, customInput.property, nextInput, close);
-                block.arguments = { [customInput.property]: nextInput };
-                return delta;
-            };
+            const customInput = (property) => ({
+                kind: "custom",
+                property,
+                buffer: { input: "", started: false, closed: false },
+            });
             const finishBlock = (block) => {
                 const contentIndex = getContentIndex(block);
                 if (contentIndex === -1) {
@@ -262,8 +253,10 @@ export const stream = (model, context, options) => {
                     });
                 }
                 else if (block.type === "toolCall") {
-                    if (block.customInput) {
-                        const delta = appendCustomToolCallInput(block, getCustomToolCallInput(block), true);
+                    const input = toolCallStates.get(block).input;
+                    toolCallStates.delete(block);
+                    if (input.kind === "custom") {
+                        const delta = closeGrammarToolInputJson(input.buffer, input.property, undefined);
                         if (delta !== undefined) {
                             stream.push({
                                 type: "toolcall_delta",
@@ -274,13 +267,8 @@ export const stream = (model, context, options) => {
                         }
                     }
                     else {
-                        Object.assign(block, parseToolCallArguments(block.partialArgs));
+                        Object.assign(block, input.arguments.finish());
                     }
-                    // Finalize in-place and strip the scratch buffers so replay only
-                    // carries parsed arguments.
-                    delete block.partialArgs;
-                    delete block.customInput;
-                    delete block.streamIndex;
                     stream.push({
                         type: "toolcall_end",
                         contentIndex,
@@ -320,18 +308,18 @@ export const stream = (model, context, options) => {
                     // Note: the "input" fallback here should/must not be taken.  in case the LLM makes up
                     // a tool we don't knwo about, we at least have a place to stash our stuff.
                     const customInputProperty = toolCall.custom && !toolCall.function ? (grammarToolInputProperties.get(name) ?? "input") : undefined;
-                    const hasCustomInput = customInputProperty !== undefined;
                     block = {
                         type: "toolCall",
                         id: toolCall.id || "",
                         name,
-                        arguments: hasCustomInput ? { [customInputProperty]: "" } : {},
-                        partialArgs: hasCustomInput ? undefined : "",
-                        customInput: hasCustomInput
-                            ? { property: customInputProperty, jsonBuffer: { input: "", started: false, closed: false } }
-                            : undefined,
-                        streamIndex,
+                        arguments: customInputProperty !== undefined ? { [customInputProperty]: "" } : {},
                     };
+                    toolCallStates.set(block, {
+                        streamIndex,
+                        input: customInputProperty !== undefined
+                            ? customInput(customInputProperty)
+                            : { kind: "json", arguments: new ToolCallArgumentsStream() },
+                    });
                     if (streamIndex !== undefined) {
                         toolCallBlocksByIndex.set(streamIndex, block);
                     }
@@ -345,8 +333,9 @@ export const stream = (model, context, options) => {
                         partial: output,
                     });
                 }
-                if (streamIndex !== undefined && block.streamIndex === undefined) {
-                    block.streamIndex = streamIndex;
+                const state = toolCallStates.get(block);
+                if (streamIndex !== undefined && state.streamIndex === undefined) {
+                    state.streamIndex = streamIndex;
                     toolCallBlocksByIndex.set(streamIndex, block);
                 }
                 if (toolCall.id) {
@@ -355,16 +344,12 @@ export const stream = (model, context, options) => {
                 if (!block.name && name) {
                     block.name = name;
                 }
-                if (toolCall.custom && !toolCall.function && !block.customInput) {
+                if (toolCall.custom && !toolCall.function && state.input.kind !== "custom") {
                     const customInputProperty = grammarToolInputProperties.get(block.name) ?? "input";
                     block.arguments = { [customInputProperty]: "" };
-                    block.customInput = {
-                        property: customInputProperty,
-                        jsonBuffer: { input: "", started: false, closed: false },
-                    };
-                    delete block.partialArgs;
+                    state.input = customInput(customInputProperty);
                 }
-                return block;
+                return { block, state };
             };
             for await (const chunk of withStreamDeadline(openaiStream, streamDeadline.deadlineMs, streamDeadline.abort)) {
                 await options?.onProviderStreamEvent?.(chunk, model);
@@ -454,7 +439,7 @@ export const stream = (model, context, options) => {
                     }
                     if (choice?.delta?.tool_calls) {
                         for (const toolCall of choice.delta.tool_calls) {
-                            const block = ensureToolCallBlock(toolCall);
+                            const { block, state } = ensureToolCallBlock(toolCall);
                             if (!block.id && toolCall.id) {
                                 block.id = toolCall.id;
                                 toolCallBlocksById.set(toolCall.id, block);
@@ -464,14 +449,15 @@ export const stream = (model, context, options) => {
                                 block.name = name;
                             }
                             let delta = "";
-                            if (toolCall.function?.arguments) {
+                            const input = state.input;
+                            if (toolCall.function?.arguments && input.kind === "json") {
                                 delta = toolCall.function.arguments;
-                                block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments;
-                                block.arguments = parseStreamingJson(block.partialArgs);
+                                block.arguments = input.arguments.append(delta);
                             }
-                            else if (toolCall.custom?.input) {
-                                const nextInput = getCustomToolCallInput(block) + toolCall.custom.input;
-                                delta = appendCustomToolCallInput(block, nextInput, false) ?? "";
+                            else if (toolCall.custom?.input && input.kind === "custom") {
+                                delta =
+                                    appendGrammarToolInputJsonDelta(input.buffer, input.property, toolCall.custom.input) ?? "";
+                                block.arguments = { [input.property]: input.buffer.input };
                             }
                             stream.push({
                                 type: "toolcall_delta",
@@ -509,11 +495,6 @@ export const stream = (model, context, options) => {
                 if (block.type === "thinking") {
                     applyStreamedReasoningDetails(block);
                 }
-                delete block.index;
-                // Streaming scratch buffers are only used during parsing; never persist them.
-                delete block.partialArgs;
-                delete block.customInput;
-                delete block.streamIndex;
             }
             output.stopReason = options?.signal?.aborted ? "aborted" : "error";
             output.errorMessage = formatProviderError(normalizeProviderError(error));