From: thstyl2000 Date: Tue, 2 Jun 2026 15:22:32 +0000 (+0200) Subject: Handle structured OpenAI stream content (#84) X-Git-Tag: v0.3.6~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=2b0a278e11fe02e7a69c1fb544b6988a975fb475;p=sgasser-llm-shield.git Handle structured OpenAI stream content (#84) --- diff --git a/src/masking/extractors/openai.test.ts b/src/masking/extractors/openai.test.ts index e1078c6..2674784 100644 --- a/src/masking/extractors/openai.test.ts +++ b/src/masking/extractors/openai.test.ts @@ -298,6 +298,45 @@ describe("OpenAI Text Extractor", () => { expect(result.choices[0].message.content).toBeNull(); }); + + test("unmasks text parts inside structured response content arrays", () => { + const response: OpenAIResponse = { + id: "test-id", + object: "chat.completion", + created: 123456, + model: "gpt-4", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: [ + { type: "reference", reference_ids: ["ref"] }, + { type: "text", text: "Hello [[PERSON_1]]" }, + // biome-ignore lint/suspicious/noExplicitAny: testing structured content preservation + ] as any, + }, + finish_reason: "stop", + }, + ], + }; + + const context: PlaceholderContext = { + mapping: { "[[PERSON_1]]": "John" }, + reverseMapping: { John: "[[PERSON_1]]" }, + counters: { PERSON: 1 }, + }; + + const result = openaiExtractor.unmaskResponse(response, context); + const content = result.choices[0].message.content as Array<{ + type: string; + text?: string; + reference_ids?: string[]; + }>; + + expect(content[0]).toEqual({ type: "reference", reference_ids: ["ref"] }); + expect(content[1]).toEqual({ type: "text", text: "Hello John" }); + }); }); describe("unknown field preservation", () => { diff --git a/src/masking/extractors/openai.ts b/src/masking/extractors/openai.ts index b03212b..5005211 100644 --- a/src/masking/extractors/openai.ts +++ b/src/masking/extractors/openai.ts @@ -14,6 +14,31 @@ import type { OpenAIRequest, OpenAIResponse } from "../../providers/openai/types import type { OpenAIContentPart } from "../../utils/content"; import type { MaskedSpan, RequestExtractor, TextSpan } from "../types"; +function unmaskContent( + content: OpenAIResponse["choices"][number]["message"]["content"], + context: PlaceholderContext, + formatValue?: (original: string) => string, +) { + if (typeof content === "string") { + return restorePlaceholders(content, context, formatValue); + } + + if (Array.isArray(content)) { + return content.map((part: OpenAIContentPart) => { + if (part.type === "text" && typeof part.text === "string") { + return { + ...part, + text: restorePlaceholders(part.text, context, formatValue), + }; + } + + return part; + }); + } + + return content; +} + /** * OpenAI request extractor * @@ -102,10 +127,7 @@ export const openaiExtractor: RequestExtractor = ...choice, message: { ...choice.message, - content: - typeof choice.message.content === "string" - ? restorePlaceholders(choice.message.content, context, formatValue) - : choice.message.content, + content: unmaskContent(choice.message.content, context, formatValue), }, })), }; diff --git a/src/providers/openai/stream-transformer.test.ts b/src/providers/openai/stream-transformer.test.ts index 797e738..c586147 100644 --- a/src/providers/openai/stream-transformer.test.ts +++ b/src/providers/openai/stream-transformer.test.ts @@ -151,4 +151,22 @@ describe("createUnmaskingStream", () => { expect(result).toContain("not-json"); }); + + test("preserves structured content arrays and only unmasks text parts", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "John"; + + const sseData = + 'data: {"choices":[{"delta":{"content":[{"type":"reference","reference_ids":["ref"]},{"type":"text","text":"Hello [[PERSON_1]]"}]}}]}\n\n'; + const source = createSSEStream([sseData]); + + const unmaskedStream = createUnmaskingStream(source, context, defaultConfig); + const result = await consumeStream(unmaskedStream); + + expect(result).not.toContain("[object Object]"); + expect(result).toContain('"type":"reference"'); + expect(result).toContain('"reference_ids":["ref"]'); + expect(result).toContain('"type":"text"'); + expect(result).toContain('"text":"Hello John"'); + }); }); diff --git a/src/providers/openai/stream-transformer.ts b/src/providers/openai/stream-transformer.ts index b807a86..c382ddc 100644 --- a/src/providers/openai/stream-transformer.ts +++ b/src/providers/openai/stream-transformer.ts @@ -2,6 +2,43 @@ import type { MaskingConfig } from "../../config"; import type { PlaceholderContext } from "../../masking/context"; import { flushMaskingBuffer, unmaskStreamChunk } from "../../pii/mask"; import { flushSecretsMaskingBuffer, unmaskSecretsStreamChunk } from "../../secrets/mask"; +import type { OpenAIContentPart } from "../../utils/content"; + +function unmaskTextContent( + text: string, + piiBuffer: string, + piiContext: PlaceholderContext | undefined, + config: MaskingConfig, + secretsBuffer: string, + secretsContext?: PlaceholderContext, +): { text: string; piiBuffer: string; secretsBuffer: string } { + let processedText = text; + let nextPiiBuffer = piiBuffer; + let nextSecretsBuffer = secretsBuffer; + + if (piiContext) { + const { output, remainingBuffer } = unmaskStreamChunk( + nextPiiBuffer, + processedText, + piiContext, + config, + ); + nextPiiBuffer = remainingBuffer; + processedText = output; + } + + if (secretsContext && processedText) { + const { output, remainingBuffer } = unmaskSecretsStreamChunk( + nextSecretsBuffer, + processedText, + secretsContext, + ); + nextSecretsBuffer = remainingBuffer; + processedText = output; + } + + return { text: processedText, piiBuffer: nextPiiBuffer, secretsBuffer: nextSecretsBuffer }; +} /** * Creates a transform stream that unmasks SSE content @@ -81,36 +118,49 @@ export function createUnmaskingStream( try { const parsed = JSON.parse(data); - const content = parsed.choices?.[0]?.delta?.content || ""; + const content = parsed.choices?.[0]?.delta?.content; - if (content) { - let processedContent = content; + if (typeof content === "string") { + const unmasked = unmaskTextContent( + content, + piiBuffer, + piiContext, + config, + secretsBuffer, + secretsContext, + ); + piiBuffer = unmasked.piiBuffer; + secretsBuffer = unmasked.secretsBuffer; - // First unmask PII if context provided - if (piiContext) { - const { output, remainingBuffer } = unmaskStreamChunk( + if (unmasked.text) { + parsed.choices[0].delta.content = unmasked.text; + controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); + } + } else if (Array.isArray(content)) { + const processedContent = content.flatMap((part: OpenAIContentPart) => { + if (part.type !== "text" || typeof part.text !== "string") { + return [part]; + } + + const unmasked = unmaskTextContent( + part.text, piiBuffer, - processedContent, piiContext, config, - ); - piiBuffer = remainingBuffer; - processedContent = output; - } - - // Then unmask secrets if context provided - if (secretsContext && processedContent) { - const { output, remainingBuffer } = unmaskSecretsStreamChunk( secretsBuffer, - processedContent, secretsContext, ); - secretsBuffer = remainingBuffer; - processedContent = output; - } + piiBuffer = unmasked.piiBuffer; + secretsBuffer = unmasked.secretsBuffer; + + if (!unmasked.text) { + return []; + } + + return [{ ...part, text: unmasked.text }]; + }); - if (processedContent) { - // Update the parsed object with processed content + if (processedContent.length > 0) { parsed.choices[0].delta.content = processedContent; controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); }