From d845f603cdb994098ee7cd6eabc44f8b9e25145e Mon Sep 17 00:00:00 2001 From: siyuanf Date: Fri, 5 Jun 2026 14:13:19 -0700 Subject: [PATCH] add ENABLE_THINKING setting --- .gitignore | 8 ++++++++ README.md | 11 ++++++++++ src/commands/chat.ts | 5 +++++ src/helpers/completion.ts | 42 +++++++++++++++++++++++++++++---------- src/helpers/config.ts | 23 +++++++++++++++++++++ src/prompt.ts | 38 ++++++++++++++++++++++++++++++----- 6 files changed, 111 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index a1ca6fed..9f71f659 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,11 @@ *.stTheme.cache *.sublime-workspace *.sublime-project +logs +*.log +npm-debug.log* +yarn-debug.log* +yarn-error.log* +lerna-debug.log* +node_modules/ +dist diff --git a/README.md b/README.md index c630871a..a3db1759 100644 --- a/README.md +++ b/README.md @@ -127,6 +127,16 @@ or save the option as a preference using this command: ai config set SILENT_MODE=true ``` +### Thinking mode + +For OpenAI-compatible endpoints that support `chat_template_kwargs.enable_thinking`, you can explicitly enable or disable thinking: + +```bash +ai config set ENABLE_THINKING=false +``` + +When unset, AI Shell does not send `chat_template_kwargs`. + ### Custom API endpoint You can custom OpenAI API endpoint to set OPENAI_API_ENDPOINT（default: `https://api.openai.com/v1`） @@ -181,6 +191,7 @@ To get an interactive UI like below: │ ○ OpenAI Key │ ○ OpenAI API Endpoint │ ○ Silent Mode +│ ○ Enable Thinking │ ● Model (gpt-4o-mini) │ ○ Language │ ○ Cancel diff --git a/src/commands/chat.ts b/src/commands/chat.ts index 898010f9..574c49ab 100644 --- a/src/commands/chat.ts +++ b/src/commands/chat.ts @@ -20,6 +20,7 @@ export default command( OPENAI_KEY: key, OPENAI_API_ENDPOINT: apiEndpoint, MODEL: model, + ENABLE_THINKING: enableThinking, } = await getConfig(); const chatHistory: ChatCompletionRequestMessage[] = []; @@ -51,6 +52,7 @@ export default command( key, model, apiEndpoint, + enableThinking, }); infoSpin.stop(`${green('AI Shell:')}`); @@ -77,12 +79,14 @@ async function getResponse({ key, model, apiEndpoint, + enableThinking, }: { prompt: string | ChatCompletionRequestMessage[]; number?: number; model?: string; key: string; apiEndpoint: string; + enableThinking?: boolean; }) { const stream = await generateCompletion({ prompt, @@ -90,6 +94,7 @@ async function getResponse({ model, number, apiEndpoint, + enableThinking, }); const iterableStream = streamToIterable(stream); diff --git a/src/helpers/completion.ts b/src/helpers/completion.ts index 33171961..78c446ae 100644 --- a/src/helpers/completion.ts +++ b/src/helpers/completion.ts @@ -34,11 +34,13 @@ export async function getScriptAndInfo({ key, model, apiEndpoint, + enableThinking, }: { prompt: string; key: string; model?: string; apiEndpoint: string; + enableThinking?: boolean; }) { const fullPrompt = getFullPrompt(prompt); const stream = await generateCompletion({ @@ -47,6 +49,7 @@ export async function getScriptAndInfo({ key, model, apiEndpoint, + enableThinking, }); const iterableStream = streamToIterable(stream); return { @@ -61,26 +64,37 @@ export async function generateCompletion({ key, model, apiEndpoint, + enableThinking, }: { prompt: string | ChatCompletionRequestMessage[]; number?: number; model?: string; key: string; apiEndpoint: string; + enableThinking?: boolean; }) { const openAi = getOpenAi(key, apiEndpoint); try { - const completion = await openAi.createChatCompletion( - { - model: model || 'gpt-4o-mini', - messages: Array.isArray(prompt) - ? prompt - : [{ role: 'user', content: prompt }], - n: Math.min(number, 10), - stream: true, - }, - { responseType: 'stream' } - ); + const request = { + model: model || 'gpt-4o-mini', + messages: Array.isArray(prompt) + ? prompt + : [{ role: 'user', content: prompt }], + n: Math.min(number, 10), + stream: true, + }; + + if (enableThinking !== undefined) { + Object.assign(request, { + chat_template_kwargs: { + enable_thinking: enableThinking, + }, + }); + } + + const completion = await openAi.createChatCompletion(request as any, { + responseType: 'stream', + }); return completion.data as unknown as IncomingMessage; } catch (err) { @@ -141,11 +155,13 @@ export async function getExplanation({ key, model, apiEndpoint, + enableThinking, }: { script: string; key: string; model?: string; apiEndpoint: string; + enableThinking?: boolean; }) { const prompt = getExplanationPrompt(script); const stream = await generateCompletion({ @@ -154,6 +170,7 @@ export async function getExplanation({ number: 1, model, apiEndpoint, + enableThinking, }); const iterableStream = streamToIterable(stream); return { readExplanation: readData(iterableStream) }; @@ -165,12 +182,14 @@ export async function getRevision({ key, model, apiEndpoint, + enableThinking, }: { prompt: string; code: string; key: string; model?: string; apiEndpoint: string; + enableThinking?: boolean; }) { const fullPrompt = getRevisionPrompt(prompt, code); const stream = await generateCompletion({ @@ -179,6 +198,7 @@ export async function getRevision({ number: 1, model, apiEndpoint, + enableThinking, }); const iterableStream = streamToIterable(stream); return { diff --git a/src/helpers/config.ts b/src/helpers/config.ts index cf3549af..9139ab72 100644 --- a/src/helpers/config.ts +++ b/src/helpers/config.ts @@ -48,6 +48,13 @@ const configParsers = { SILENT_MODE(mode?: string) { return String(mode).toLowerCase() === 'true'; }, + ENABLE_THINKING(mode?: string) { + if (mode === undefined || mode === '') { + return undefined; + } + + return String(mode).toLowerCase() === 'true'; + }, OPENAI_API_ENDPOINT(apiEndpoint?: string) { return apiEndpoint || 'https://api.openai.com/v1'; }, @@ -142,6 +149,14 @@ export const showConfigUI = async () => { ? config.SILENT_MODE.toString() : i18n.t('(not set)'), }, + { + label: i18n.t('Enable Thinking'), + value: 'ENABLE_THINKING', + hint: + config.ENABLE_THINKING === undefined + ? i18n.t('(not set)') + : config.ENABLE_THINKING.toString(), + }, { label: i18n.t('Model'), value: 'MODEL', @@ -187,6 +202,14 @@ export const showConfigUI = async () => { }); if (p.isCancel(silentMode)) return; await setConfigs([['SILENT_MODE', silentMode ? 'true' : 'false']]); + } else if (choice === 'ENABLE_THINKING') { + const enableThinking = await p.confirm({ + message: i18n.t('Enable thinking?'), + }); + if (p.isCancel(enableThinking)) return; + await setConfigs([ + ['ENABLE_THINKING', enableThinking ? 'true' : 'false'], + ]); } else if (choice === 'MODEL') { const { OPENAI_KEY: key, OPENAI_API_ENDPOINT: apiEndpoint } = await getConfig(); diff --git a/src/prompt.ts b/src/prompt.ts index 145c07b3..3a4b90e7 100644 --- a/src/prompt.ts +++ b/src/prompt.ts @@ -106,6 +106,7 @@ export async function prompt({ SILENT_MODE, OPENAI_API_ENDPOINT: apiEndpoint, MODEL: model, + ENABLE_THINKING: enableThinking, } = await getConfig(); const skipCommandExplanation = silentMode || SILENT_MODE; @@ -120,6 +121,7 @@ export async function prompt({ key, model, apiEndpoint, + enableThinking, }); spin.stop(`${i18n.t('Your script')}:`); console.log(''); @@ -136,6 +138,7 @@ export async function prompt({ key, model, apiEndpoint, + enableThinking, }); spin.stop(`${i18n.t('Explanation')}:`); console.log(''); @@ -146,7 +149,14 @@ export async function prompt({ } } - await runOrReviseFlow(script, key, model, apiEndpoint, silentMode); + await runOrReviseFlow( + script, + key, + model, + apiEndpoint, + silentMode, + enableThinking + ); } async function runOrReviseFlow( @@ -154,7 +164,8 @@ async function runOrReviseFlow( key: string, model: string, apiEndpoint: string, - silentMode?: boolean + silentMode?: boolean, + enableThinking?: boolean ) { const emptyScript = script.trim() === ''; @@ -191,7 +202,14 @@ async function runOrReviseFlow( label: '🔁 ' + i18n.t('Revise'), hint: i18n.t('Give feedback via prompt and get a new result'), value: async () => { - await revisionFlow(script, key, model, apiEndpoint, silentMode); + await revisionFlow( + script, + key, + model, + apiEndpoint, + silentMode, + enableThinking + ); }, }, { @@ -223,7 +241,8 @@ async function revisionFlow( key: string, model: string, apiEndpoint: string, - silentMode?: boolean + silentMode?: boolean, + enableThinking?: boolean ) { const revision = await promptForRevision(); const spin = p.spinner(); @@ -234,6 +253,7 @@ async function revisionFlow( key, model, apiEndpoint, + enableThinking, }); spin.stop(`${i18n.t(`Your new script`)}:`); @@ -251,6 +271,7 @@ async function revisionFlow( key, model, apiEndpoint, + enableThinking, }); infoSpin.stop(`${i18n.t('Explanation')}:`); @@ -261,7 +282,14 @@ async function revisionFlow( console.log(dim('•')); } - await runOrReviseFlow(script, key, model, apiEndpoint, silentMode); + await runOrReviseFlow( + script, + key, + model, + apiEndpoint, + silentMode, + enableThinking + ); } export const parseAssert = (name: string, condition: any, message: string) => {