Merge pull request #458 from kaovilai/thiknagain

di-sukharev · web-flow · commit b3700ae685f2 · 2025-04-14T11:19:20.000+03:00
Extends #445 to other providers which many provides deepseek
diff --git a/src/engine/anthropic.ts b/src/engine/anthropic.ts
@@ -8,6 +8,7 @@ import axios from 'axios';
 import chalk from 'chalk';
 import { OpenAI } from 'openai';
 import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
+import { removeContentTags } from '../utils/removeContentTags';
 import { tokenCount } from '../utils/tokenCount';
 import { AiEngine, AiEngineConfig } from './Engine';
 
@@ -54,8 +55,8 @@ export class AnthropicEngine implements AiEngine {
       const data = await this.client.messages.create(params);
 
       const message = data?.content[0].text;
-
-      return message;
+      let content = message;
+      return removeContentTags(content, 'think');
     } catch (error) {
       const err = error as Error;
       outro(`${chalk.red('✖')} ${err?.message || err}`);
diff --git a/src/engine/azure.ts b/src/engine/azure.ts
@@ -7,6 +7,7 @@ import axios from 'axios';
 import chalk from 'chalk';
 import { OpenAI } from 'openai';
 import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
+import { removeContentTags } from '../utils/removeContentTags';
 import { tokenCount } from '../utils/tokenCount';
 import { AiEngine, AiEngineConfig } from './Engine';
 
@@ -52,7 +53,9 @@ export class AzureEngine implements AiEngine {
       if (message?.content === null) {
         return undefined;
       }
-      return message?.content;
+      
+      let content = message?.content;
+      return removeContentTags(content, 'think');
     } catch (error) {
       outro(`${chalk.red('✖')} ${this.config.model}`);
 
diff --git a/src/engine/deepseek.ts b/src/engine/deepseek.ts
@@ -1,6 +1,7 @@
 import axios from 'axios';
 import { OpenAI } from 'openai';
 import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
+import { removeContentTags } from '../utils/removeContentTags';
 import { tokenCount } from '../utils/tokenCount';
 import { OpenAiEngine, OpenAiConfig } from './openAi';
 
@@ -41,8 +42,8 @@ export class DeepseekEngine extends OpenAiEngine {
       const completion = await this.client.chat.completions.create(params);
 
       const message = completion.choices[0].message;
-
-      return message?.content;
+      let content = message?.content;
+      return removeContentTags(content, 'think');
     } catch (error) {
       const err = error as Error;
       if (
diff --git a/src/engine/flowise.ts b/src/engine/flowise.ts
@@ -1,5 +1,6 @@
 import axios, { AxiosInstance } from 'axios';
 import { OpenAI } from 'openai';
+import { removeContentTags } from '../utils/removeContentTags';
 import { AiEngine, AiEngineConfig } from './Engine';
 
 interface FlowiseAiConfig extends AiEngineConfig {}
@@ -36,7 +37,8 @@ export class FlowiseEngine implements AiEngine {
     try {
       const response = await this.client.post('', payload);
       const message = response.data;
-      return message?.text;
+      let content = message?.text;
+      return removeContentTags(content, 'think');
     } catch (err: any) {
       const message = err.response?.data?.error ?? err.message;
       throw new Error('local model issues. details: ' + message);
diff --git a/src/engine/gemini.ts b/src/engine/gemini.ts
@@ -7,6 +7,7 @@ import {
 } from '@google/generative-ai';
 import axios from 'axios';
 import { OpenAI } from 'openai';
+import { removeContentTags } from '../utils/removeContentTags';
 import { AiEngine, AiEngineConfig } from './Engine';
 
 interface GeminiConfig extends AiEngineConfig {}
@@ -71,7 +72,8 @@ export class GeminiEngine implements AiEngine {
         }
       });
 
-      return result.response.text();
+      const content = result.response.text();
+      return removeContentTags(content, 'think');
     } catch (error) {
       const err = error as Error;
       if (
diff --git a/src/engine/mistral.ts b/src/engine/mistral.ts
@@ -1,27 +1,21 @@
 import axios from 'axios';
-import { Mistral } from '@mistralai/mistralai';
 import { OpenAI } from 'openai';
 import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
+import { removeContentTags } from '../utils/removeContentTags';
 import { tokenCount } from '../utils/tokenCount';
 import { AiEngine, AiEngineConfig } from './Engine';
-import {
-  AssistantMessage as MistralAssistantMessage,
-  SystemMessage as MistralSystemMessage,
-  ToolMessage as MistralToolMessage,
-  UserMessage as MistralUserMessage
-} from '@mistralai/mistralai/models/components';
 
+// Using any for Mistral types to avoid type declaration issues
 export interface MistralAiConfig extends AiEngineConfig {}
-export type MistralCompletionMessageParam = Array<
-| (MistralSystemMessage & { role: "system" })
-| (MistralUserMessage & { role: "user" })
-| (MistralAssistantMessage & { role: "assistant" })
-| (MistralToolMessage & { role: "tool" })
->
+export type MistralCompletionMessageParam = Array<any>;
+
+// Import Mistral dynamically to avoid TS errors
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const Mistral = require('@mistralai/mistralai').Mistral;
 
 export class MistralAiEngine implements AiEngine {
   config: MistralAiConfig;
-  client: Mistral;
+  client: any; // Using any type for Mistral client to avoid TS errors
 
   constructor(config: MistralAiConfig) {
     this.config = config;
@@ -64,7 +58,8 @@ export class MistralAiEngine implements AiEngine {
       if (!message || !message.content)
         throw Error('No completion choice available.')
 
-      return message.content as string;
+      let content = message.content as string;
+      return removeContentTags(content, 'think');
     } catch (error) {
       const err = error as Error;
       if (
diff --git a/src/engine/mlx.ts b/src/engine/mlx.ts
@@ -1,7 +1,7 @@
 import axios, { AxiosInstance } from 'axios';
 import { OpenAI } from 'openai';
+import { removeContentTags } from '../utils/removeContentTags';
 import { AiEngine, AiEngineConfig } from './Engine';
-import { chown } from 'fs';
 
 interface MLXConfig extends AiEngineConfig {}
 
@@ -37,11 +37,11 @@ export class MLXEngine implements AiEngine {
                 
                 const choices = response.data.choices;
                 const message = choices[0].message;
-                
-                return message?.content;
+                let content = message?.content;
+                return removeContentTags(content, 'think');
             } catch (err: any) {
                 const message = err.response?.data?.error ?? err.message;
                 throw new Error(`MLX provider error: ${message}`);
             }
         }
-}
+}
diff --git a/src/engine/ollama.ts b/src/engine/ollama.ts
@@ -1,5 +1,6 @@
 import axios, { AxiosInstance } from 'axios';
 import { OpenAI } from 'openai';
+import { removeContentTags } from '../utils/removeContentTags';
 import { AiEngine, AiEngineConfig } from './Engine';
 
 interface OllamaConfig extends AiEngineConfig {}
@@ -35,12 +36,7 @@ export class OllamaEngine implements AiEngine {
 
       const { message } = response.data;
       let content = message?.content;
-
-      if (content && content.includes('<think>')) {
-        return content.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
-      }
-
-      return content;
+      return removeContentTags(content, 'think');
     } catch (err: any) {
       const message = err.response?.data?.error ?? err.message;
       throw new Error(`Ollama provider error: ${message}`);
diff --git a/src/engine/openAi.ts b/src/engine/openAi.ts
@@ -1,6 +1,7 @@
 import axios from 'axios';
 import { OpenAI } from 'openai';
 import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
+import { removeContentTags } from '../utils/removeContentTags';
 import { tokenCount } from '../utils/tokenCount';
 import { AiEngine, AiEngineConfig } from './Engine';
 
@@ -45,8 +46,8 @@ export class OpenAiEngine implements AiEngine {
       const completion = await this.client.chat.completions.create(params);
 
       const message = completion.choices[0].message;
-
-      return message?.content;
+      let content = message?.content;
+      return removeContentTags(content, 'think');
     } catch (error) {
       const err = error as Error;
       if (
diff --git a/src/utils/removeContentTags.ts b/src/utils/removeContentTags.ts
@@ -0,0 +1,51 @@
+/**
+ * Removes content wrapped in specified tags from a string
+ * @param content The content string to process
+ * @param tag The tag name without angle brackets (e.g., 'think' for '<think></think>')
+ * @returns The content with the specified tags and their contents removed, and trimmed
+ */
+export function removeContentTags<T extends string | null | undefined>(content: T, tag: string): T {
+  if (!content || typeof content !== 'string') {
+    return content;
+  }
+  
+  // Dynamic implementation for other cases
+  const openTag = `<${tag}>`;
+  const closeTag = `</${tag}>`;
+  
+  // Parse the content and remove tags
+  let result = '';
+  let skipUntil: number | null = null;
+  let depth = 0;
+  
+  for (let i = 0; i < content.length; i++) {
+    // Check for opening tag
+    if (content.substring(i, i + openTag.length) === openTag) {
+      depth++;
+      if (depth === 1) {
+        skipUntil = content.indexOf(closeTag, i + openTag.length);
+        i = i + openTag.length - 1; // Skip the opening tag
+        continue;
+      }
+    }
+    // Check for closing tag
+    else if (content.substring(i, i + closeTag.length) === closeTag && depth > 0) {
+      depth--;
+      if (depth === 0) {
+        i = i + closeTag.length - 1; // Skip the closing tag
+        skipUntil = null;
+        continue;
+      }
+    }
+    
+    // Only add character if not inside a tag
+    if (skipUntil === null) {
+      result += content[i];
+    }
+  }
+  
+  // Normalize spaces (replace multiple spaces with a single space)
+  result = result.replace(/\s+/g, ' ').trim();
+  
+  return result as unknown as T;
+}
diff --git a/test/unit/removeContentTags.test.ts b/test/unit/removeContentTags.test.ts
@@ -0,0 +1,57 @@
+import { removeContentTags } from '../../src/utils/removeContentTags';
+
+describe('removeContentTags', () => {
+  it('should remove content wrapped in specified tags', () => {
+    const content = 'This is <think>something to hide</think> visible content';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('This is visible content');
+  });
+
+  it('should handle multiple tag occurrences', () => {
+    const content = '<think>hidden</think> visible <think>also hidden</think> text';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('visible text');
+  });
+
+  it('should handle multiline content within tags', () => {
+    const content = 'Start <think>hidden\nover multiple\nlines</think> End';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('Start End');
+  });
+
+  it('should return content as is when tag is not found', () => {
+    const content = 'Content without any tags';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('Content without any tags');
+  });
+
+  it('should work with different tag names', () => {
+    const content = 'This is <custom>something to hide</custom> visible content';
+    const result = removeContentTags(content, 'custom');
+    expect(result).toBe('This is visible content');
+  });
+
+  it('should handle null content', () => {
+    const content = null;
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe(null);
+  });
+
+  it('should handle undefined content', () => {
+    const content = undefined;
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe(undefined);
+  });
+
+  it('should trim the result', () => {
+    const content = '  <think>hidden</think> visible  ';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('visible');
+  });
+
+  it('should handle nested tags correctly', () => {
+    const content = 'Outside <think>Inside <think>Nested</think></think> End';
+    const result = removeContentTags(content, 'think');
+    expect(result).toBe('Outside End');
+  });
+});
