fix: script returns empty when model response lacks code block markers

package.json

@@ -1,5 +1,5 @@
 {
-  "name": "@builder.io/ai-shell",
+  "name": "@jadenoliver/ai-shell",
   "description": "A CLI that converts natural language to shell commands.",
   "version": "1.0.12",
   "type": "module",

src/helpers/completion.ts

@@ -13,7 +13,6 @@ import type { AxiosError } from 'axios';
 import { streamToString } from './stream-to-string';
 import './replace-all-polyfill';
 import i18n from './i18n';
-import { stripRegexPatterns } from './strip-regex-patterns';
 import readline from 'readline';
 
 const explainInSecondRequest = true;
@@ -25,9 +24,6 @@ function getOpenAi(key: string, apiEndpoint: string) {
   return openAi;
 }
 
-// Openai outputs markdown format for code blocks. It oftne uses
-// a github style like: "```bash"
-const shellCodeExclusions = [/```[a-zA-Z]*\n/gi, /```[a-zA-Z]*/gi, '\n'];
 
 export async function getScriptAndInfo({
   prompt,
@@ -50,8 +46,7 @@ export async function getScriptAndInfo({
   });
   const iterableStream = streamToIterable(stream);
   return {
-    readScript: readData(iterableStream, ...shellCodeExclusions),
-    readInfo: readData(iterableStream, ...shellCodeExclusions),
+    readScript: readData(iterableStream),
   };
 }
 
@@ -182,9 +177,23 @@ export async function getRevision({
   });
   const iterableStream = streamToIterable(stream);
   return {
-    readScript: readData(iterableStream, ...shellCodeExclusions),
+    readScript: readData(iterableStream),
   };
 }
+export function extractCommand(raw: string): string {
+  // Case 1: Standard code block with newline after language identifier
+  const match = raw.match(/```[a-zA-Z]*\n([\s\S]*?)```/);
+  if (match) {
+    return match[1].trim();
+  }
+  // Case 2: Code block without newline (```command```)
+  const noNewlineMatch = raw.match(/```([\s\S]*?)```/);
+  if (noNewlineMatch) {
+    return noNewlineMatch[1].trim();
+  }
+  // Case 3: No code blocks, return as-is
+  return raw.trim();
+}
 
 export const readData =
   (
@@ -194,13 +203,10 @@ export const readData =
   (writer: (data: string) => void): Promise<string> =>
     new Promise(async (resolve) => {
       let stopTextStream = false;
-      let data = '';
+      let raw = '';
       let content = '';
-      let dataStart = false;
-      let buffer = ''; // This buffer will temporarily hold incoming data only for detecting the start
 
-      const [excludedPrefix] = excluded;
-      const stopTextStreamKeys = ['q', 'escape']; //Group of keys that stop the text stream
+      const stopTextStreamKeys = ['q', 'escape'];
 
       const rl = readline.createInterface({
         input: process.stdin,
@@ -217,33 +223,15 @@ export const readData =
         const payloads = chunk.toString().split('\n\n');
         for (const payload of payloads) {
           if (payload.includes('[DONE]') || stopTextStream) {
-            dataStart = false;
-            resolve(data);
+            resolve(raw.trim());
             return;
           }
 
           if (payload.startsWith('data:')) {
             content = parseContent(payload);
-            // Use buffer only for start detection
-            if (!dataStart) {
-              // Append content to the buffer
-              buffer += content;
-              if (buffer.match(excludedPrefix ?? '')) {
-                dataStart = true;
-                // Clear the buffer once it has served its purpose
-                buffer = '';
-                if (excludedPrefix) break;
-              }
-            }
-
-            if (dataStart && content) {
-              const contentWithoutExcluded = stripRegexPatterns(
-                content,
-                excluded
-              );
-
-              data += contentWithoutExcluded;
-              writer(contentWithoutExcluded);
+            if (content) {
+              raw += content;
+              writer(content);
             }
           }
         }
@@ -255,11 +243,11 @@ export const readData =
           const delta = JSON.parse(data.trim());
           return delta.choices?.[0]?.delta?.content ?? '';
         } catch (error) {
-          return `Error with JSON.parse and ${payload}.\n${error}`;
+          return '';
         }
       }
 
-      resolve(data);
+      resolve(raw.trim());
     });
 
 function getExplanationPrompt(script: string) {
@@ -288,7 +276,7 @@ function getOperationSystemDetails() {
   return os.name();
 }
 const generationDetails = dedent`
-    Only reply with the single line command surrounded by three backticks. It must be able to be directly run in the target shell. Do not include any other text.
+    Only reply with the single line command. It must be able to be directly run in the target shell. Do not include any other text, explanations, or markdown formatting.
 
     Make sure the command runs on ${getOperationSystemDetails()} operating system.
   `;

src/prompt.ts

@@ -5,6 +5,7 @@ import {
   getExplanation,
   getRevision,
   getScriptAndInfo,
+  extractCommand,
 } from './helpers/completion';
 import { getConfig } from './helpers/config';
 import { projectName } from './helpers/constants';
@@ -115,35 +116,33 @@ export async function prompt({
   const thePrompt = usePrompt || (await getPrompt());
   const spin = p.spinner();
   spin.start(i18n.t(`Loading...`));
-  const { readInfo, readScript } = await getScriptAndInfo({
+  const { readScript } = await getScriptAndInfo({
     prompt: thePrompt,
     key,
     model,
     apiEndpoint,
   });
   spin.stop(`${i18n.t('Your script')}:`);
   console.log('');
-  const script = await readScript(process.stdout.write.bind(process.stdout));
+  const rawScript = await readScript(process.stdout.write.bind(process.stdout));
+  const script = extractCommand(rawScript);
   console.log('');
   console.log('');
   console.log(dim('•'));
   if (!skipCommandExplanation) {
     spin.start(i18n.t(`Getting explanation...`));
-    const info = await readInfo(process.stdout.write.bind(process.stdout));
-    if (!info) {
-      const { readExplanation } = await getExplanation({
-        script,
-        key,
-        model,
-        apiEndpoint,
-      });
-      spin.stop(`${i18n.t('Explanation')}:`);
-      console.log('');
-      await readExplanation(process.stdout.write.bind(process.stdout));
-      console.log('');
-      console.log('');
-      console.log(dim('•'));
-    }
+    const { readExplanation } = await getExplanation({
+      script,
+      key,
+      model,
+      apiEndpoint,
+    });
+    spin.stop(`${i18n.t('Explanation')}:`);
+    console.log('');
+    await readExplanation(process.stdout.write.bind(process.stdout));
+    console.log('');
+    console.log('');
+    console.log(dim('•'));
   }
 
   await runOrReviseFlow(script, key, model, apiEndpoint, silentMode);
@@ -238,7 +237,8 @@ async function revisionFlow(
   spin.stop(`${i18n.t(`Your new script`)}:`);
 
   console.log('');
-  const script = await readScript(process.stdout.write.bind(process.stdout));
+  const rawScript = await readScript(process.stdout.write.bind(process.stdout));
+  const script = extractCommand(rawScript);
   console.log('');
   console.log('');
   console.log(dim('•'));