Handle structured OpenAI stream content (#84)

diff --git a/src/masking/extractors/openai.test.ts b/src/masking/extractors/openai.test.ts
index e1078c688cda14f8c1e28109fcb43a35868de1ba..2674784dc931dc5ad919f3c30896b2d18bb4df2d 100644 (file)
--- 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 b03212be5e8d7cb6005116718ff46282a735e0ed..50052117a77cf8265ff8d8e799a05d3fcc986891 100644 (file)
--- 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<OpenAIRequest, OpenAIResponse> =
         ...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 797e738db004c096d9e6f9e2d71f9caea7666712..c586147eda3b755c5f18cfbfd12248314ddd6727 100644 (file)
--- 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 b807a8697369e1922acf7b5d91119e2865bcf0b0..c382ddc57734106109ab76a3024df5e56d51dce7 100644 (file)
--- 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`));
                   }