Merge pull request #54 from Unity-Lab-AI/codex/fix-image-generation-and-json-formatting

Add JSON repair and structured response handling

Author: hackall360

index.html

@@ -478,6 +478,8 @@ <h3 class="modal-title">Voice Chat</h3>
     <script defer src="js/storage/storage.js"></script>
     <script defer src="js/storage/memory-api.js"></script>
 
+    <!-- JSON repair helpers for AI responses -->
+    <script type="module" defer src="js/chat/json-repair.js"></script>
     <!-- chat-core FIRST so PolliLib's default client helpers are available -->
     <script defer src="js/chat/chat-core.js"></script>
 

js/chat/chat-core.js

@@ -421,17 +421,60 @@ document.addEventListener("DOMContentLoaded", () => {
         });
 
     async function handleToolJson(raw, { imageUrls, audioUrls }) {
-        try {
-            const obj = JSON.parse(raw);
+        const obj = (window.repairJson || (() => ({ text: raw })))(raw);
+        let handled = false;
+
+        if (obj.tool) {
             const fn = toolbox.get(obj.tool);
-            if (!fn) return { handled: false, text: raw };
-            const res = await fn(obj);
-            if (res?.imageUrl) imageUrls.push(res.imageUrl);
-            if (res?.audioUrl) audioUrls.push(res.audioUrl);
-            return { handled: true, text: res?.text || '' };
-        } catch {
-            return { handled: false, text: raw };
+            if (fn) {
+                try {
+                    const res = await fn(obj);
+                    if (res?.imageUrl) imageUrls.push(res.imageUrl);
+                    if (res?.audioUrl) audioUrls.push(res.audioUrl);
+                    handled = true;
+                    return { handled: true, text: res?.text || '' };
+                } catch (e) {
+                    console.warn('tool execution failed', e);
+                }
+            }
+        }
+
+        const imgPrompts = obj.image ? [obj.image] : Array.isArray(obj.images) ? obj.images : [];
+        for (const prompt of imgPrompts) {
+            if (!(window.polliLib && window.polliClient) || !prompt) continue;
+            try {
+                const url = window.polliLib.mcp.generateImageUrl(window.polliClient, { prompt });
+                imageUrls.push(url);
+                handled = true;
+            } catch (e) {
+                console.warn('polliLib generateImageUrl failed', e);
+            }
         }
+
+        const audioText = obj.audio || obj.tts;
+        if (audioText && window.polliLib && window.polliClient) {
+            try {
+                const blob = await window.polliLib.tts(audioText, { model: 'openai-audio' }, window.polliClient);
+                const url = URL.createObjectURL(blob);
+                audioUrls.push(url);
+                handled = true;
+            } catch (e) {
+                console.warn('polliLib tts failed', e);
+            }
+        }
+
+        const command = obj.ui || obj.command;
+        if (command) {
+            if (validateUICommand(command)) {
+                try { executeCommand(command); } catch (e) { console.warn('executeCommand failed', e); }
+                handled = true;
+            } else {
+                console.warn('invalid ui command', command);
+            }
+        }
+
+        const text = typeof obj.text === 'string' ? obj.text : raw;
+        return { handled, text };
     }
 
     function handleVoiceCommand(text) {

js/chat/json-repair.js

@@ -0,0 +1,27 @@
+export function repairJson(raw) {
+    if (typeof raw !== 'string') return { text: '' };
+    let text = raw.trim();
+    if (!text) return { text: '' };
+    const attempts = [];
+    attempts.push(text);
+    // Attempt: replace single quotes with double quotes
+    attempts.push(text.replace(/'/g, '"'));
+    // Attempt: quote unquoted keys, replace single quotes, remove trailing commas
+    attempts.push(
+        text
+            .replace(/([,{]\s*)([A-Za-z0-9_]+)\s*:/g, '$1"$2":')
+            .replace(/'/g, '"')
+            .replace(/,\s*([}\]])/g, '$1')
+    );
+    for (const str of attempts) {
+        try {
+            return JSON.parse(str);
+        } catch {}
+    }
+    // Fallback: treat as plain text
+    return { text };
+}
+
+if (typeof window !== 'undefined') {
+    window.repairJson = repairJson;
+}

js/chat/markdown-sanitizer.js

@@ -0,0 +1,16 @@
+export const defaultBlockedFenceTypes = ['image', 'audio', 'ui'];
+
+export function sanitizeMarkdown(content, blocked = defaultBlockedFenceTypes) {
+  if (!content) return '';
+  const pattern = /```(\w+)\n[\s\S]*?```/g;
+  return content.replace(pattern, (match, type) => {
+    return blocked.includes(type.toLowerCase()) ? '' : match;
+  });
+}
+
+if (typeof window !== 'undefined') {
+  window.sanitizeMarkdown = sanitizeMarkdown;
+  if (!window.blockedFenceTypes) {
+    window.blockedFenceTypes = [...defaultBlockedFenceTypes];
+  }
+}

package.json

@@ -3,7 +3,7 @@
   "private": true,
   "type": "module",
   "scripts": {
-    "test": "node tests/pollilib-smoke.mjs && node tests/markdown-sanitization.mjs && node tests/ai-response.mjs && node tests/json-tools.mjs"
+    "test": "node tests/pollilib-smoke.mjs && node tests/markdown-sanitization.mjs && node tests/ai-response.mjs && node tests/json-tools.mjs && node tests/json-repair.mjs"
   },
   "devDependencies": {
     "marked": "^11.2.0"

prompts/ai-instruct.md

@@ -119,22 +119,30 @@ tell me a joke in a calm tone
 
 ## JSON Tools
 
-- As an alternative to fenced blocks, respond with a JSON object to invoke tools.
-- The object must include a `tool` field:
-  - `image` with a `prompt` string to generate an image.
-  - `tts` with a `text` string for text-to-speech.
-  - `ui` with a `command` object that follows `docs/ui-command.schema.json`.
-- Example:
+- As an alternative to fenced blocks, respond with a JSON object.
+- The object may include:
+  - `tool` to invoke a tool (`image`, `tts`, or `ui`).
+  - `text` for plain responses.
+  - `image` or `images` with prompt strings to generate images.
+  - `audio` with text for text-to-speech.
+  - `command`/`ui` objects that follow `docs/ui-command.schema.json`.
+- Examples:
 
 ```json
 {"tool":"image","prompt":"a glowing neon cityscape at night with flying cars"}
 ```
 
 ```json
-{"tool":"ui","command":{"action":"openScreensaver"}}
+{"text":"Hello there"}
+```
+
+```json
+{"images":["a tiny house"],"text":"Here you go"}
 ```
 
+- Always return valid JSON (double quotes, no trailing commas).
 - Do not include extra commentary outside the JSON object.
+- If you must send plain text, wrap it as `{ "text": "..." }`.
 
 ---
 

tests/json-repair.mjs

@@ -0,0 +1,11 @@
+import { strict as assert } from 'node:assert';
+import { repairJson } from '../js/chat/json-repair.js';
+
+const fixed = repairJson("{tool:'image', prompt:'apple',}");
+assert.equal(fixed.tool, 'image');
+assert.equal(fixed.prompt, 'apple');
+
+const plain = repairJson('just some text');
+assert.deepEqual(plain, { text: 'just some text' });
+
+console.log('json-repair test passed');