fix json object generation

src/extract-json.test.ts

@@ -0,0 +1,14 @@
+import { describe, it, expect } from "vitest";
+import { extractJSON } from "./extract-json";
+
+describe("extractJSON", () => {
+  it("correctly extracts JSON part", () => {
+    const text = '```json{"key": "value"}```';
+    expect(extractJSON(text)).toBe(`{"key": "value"}`);
+  });
+
+  it("leaves the text as it was before if no json indicator found", () => {
+    const text = "Just some text";
+    expect(extractJSON(text)).toBe("Just some text");
+  });
+});

src/extract-json.ts

@@ -0,0 +1,8 @@
+export function extractJSON(text: string) {
+  const trimmedText = text.trim();
+  const json = trimmedText.startsWith('```')
+    ? trimmedText.split(/```(?:json)?/)[1]
+    : trimmedText;
+
+  return json;
+}

src/language-model.test.ts

@@ -119,9 +119,36 @@ describe('language-model', () => {
       expect(textPart).toBe('test');
     }
   });
-
+  
   it('should do generate object', async () => {
     const prompt = vi.fn(async (prompt: string) => '{"hello":"world"}');
+    vi.stubGlobal("ai", {
+      canCreateTextSession: vi.fn(async () => "readily"),
+      defaultTextSessionOptions: vi.fn(async () => ({})),
+      createTextSession: vi.fn(async () => ({ prompt })),
+    });
+
+    const { object } = await generateObject({
+      model: new ChromeAIChatLanguageModel('text'),
+      schema: z.object({
+        hello: z.string(),
+      }),
+      prompt: 'test',
+    });
+
+    expect(object).toMatchObject({ hello: "world" });
+  });
+
+  it('should do stream json object when text is in markdown syntax', async () => {
+    const promptStreaming = vi.fn((prompt: string) => {
+      const stream = new ReadableStream<string>({
+        start(controller) {
+          controller.enqueue('```json{"hello":"world"}```');
+          controller.close();
+        },
+      });
+      return stream;
+    });
     vi.stubGlobal('ai', {
       canCreateTextSession: vi.fn(async () => 'readily'),
       textModelInfo: vi.fn(async () => ({})),

src/language-model.ts

@@ -16,7 +16,8 @@ import {
 } from '@ai-sdk/provider';
 import { ChromeAISession, ChromeAISessionOptions } from './global';
 import createDebug from 'debug';
-import { objectStartSequence, objectStopSequence, StreamAI } from './stream-ai';
+import { StreamAI } from './stream-ai';
+import { extractJSON } from './extract-json';
 
 const debug = createDebug('chromeai');
 
@@ -164,15 +165,12 @@ export class ChromeAIChatLanguageModel implements LanguageModelV1 {
 
     const session = await this.getSession();
     const message = this.formatMessages(options);
-    let text = await session.prompt(message);
 
-    if (options.mode.type === 'object-json') {
-      text = text.replace(new RegExp('^' + objectStartSequence, 'ig'), '');
-      text = text.replace(new RegExp(objectStopSequence + '$', 'ig'), '');
-    }
-
-    debug('generate result:', text);
+    const rawText = await session.prompt(message);
+    const text =
+      options.mode.type === "object-json" ? extractJSON(rawText) : rawText;
 
+    debug("generate result:", text);
     return {
       text,
       finishReason: 'stop',

src/stream-ai.ts

@@ -3,6 +3,7 @@ import {
   LanguageModelV1StreamPart,
 } from '@ai-sdk/provider';
 import createDebug from 'debug';
+import { extractJSON } from './extract-json';
 
 const debug = createDebug('chromeai');
 
@@ -32,21 +33,10 @@ export class StreamAI extends TransformStream<
         });
       },
       transform: (chunk, controller) => {
-        if (options.mode.type === 'object-json') {
-          transforming =
-            chunk.startsWith(objectStartSequence) &&
-            !chunk.endsWith(objectStopSequence);
-          chunk = chunk.replace(
-            new RegExp('^' + objectStartSequence, 'ig'),
-            ''
-          );
-          chunk = chunk.replace(new RegExp(objectStopSequence + '$', 'ig'), '');
-        } else {
-          transforming = true;
-        }
-        const textDelta = chunk.replace(buffer, ''); // See: https://github.com/jeasonstudio/chrome-ai/issues/11
-        if (transforming) controller.enqueue({ type: 'text-delta', textDelta });
-        buffer = chunk;
+        const jsonPart = extractJSON(chunk);
+        const textDelta = jsonPart.replace(textTemp, "");
+        textTemp += textDelta;
+        controller.enqueue({ type: 'text-delta', textDelta });
       },
       flush: (controller) => {
         controller.enqueue({