feat: Skip extract file content with txt or markdown file ✨

Author: lethemanh

packages/cozy-viewer/src/Panel/AI/AIAssistantPanel.jsx

@@ -23,7 +23,7 @@ import { useAlert } from 'cozy-ui/transpiled/react/providers/Alert'
 
 import { SUMMARY_SYSTEM_PROMPT, getSummaryUserPrompt } from './prompts'
 import styles from './styles.styl'
-import { roughTokensEstimation } from '../../helpers'
+import { roughTokensEstimation, isTextFile } from '../../helpers'
 import { useViewer } from '../../providers/ViewerProvider'
 
 const AIAssistantPanel = ({ className }) => {
@@ -49,25 +49,36 @@ const AIAssistantPanel = ({ className }) => {
     }
   }
 
+  const extractFileContent = async (client, fileBlob, file) => {
+    if (isTextFile(file.mime)) {
+      return await fileBlob.text()
+    }
+
+    const rawTextContent = await extractText(client, fileBlob, {
+      name: file.name,
+      mime: file.mime
+    })
+    return rawTextContent ? JSON.stringify(rawTextContent) : ''
+  }
+
+  const validateContentSize = textContent => {
+    const summaryConfig = flag('drive.summary')
+    if (
+      summaryConfig?.maxTokens &&
+      roughTokensEstimation(textContent) > summaryConfig.maxTokens
+    ) {
+      const error = new Error('DOCUMENT_TOO_LARGE')
+      error.code = 'DOCUMENT_TOO_LARGE'
+      throw error
+    }
+  }
+
   const summarizeFile = async ({ client, file, stream = false, model }) => {
     try {
       const fileBlob = await fetchBlobFileById(client, file?._id)
+      const textContent = await extractFileContent(client, fileBlob, file)
 
-      const rawTextContent = await extractText(client, fileBlob, {
-        name: file.name,
-        mime: file.mime
-      })
-      const textContent = rawTextContent ? JSON.stringify(rawTextContent) : ''
-
-      const summaryConfig = flag('drive.summary')
-      if (
-        summaryConfig?.maxTokens &&
-        roughTokensEstimation(textContent) > summaryConfig.maxTokens
-      ) {
-        const error = new Error('DOCUMENT_TOO_LARGE')
-        error.code = 'DOCUMENT_TOO_LARGE'
-        throw error
-      }
+      validateContentSize(textContent)
 
       const messages = [
         { role: 'system', content: SUMMARY_SYSTEM_PROMPT },

packages/cozy-viewer/src/helpers.js

@@ -84,6 +84,20 @@ export const roughTokensEstimation = text => {
   return Math.ceil(text.length / 4)
 }
 
+/**
+ * Check if a file is a text-based file type
+ * @param {string} mime - MIME type of the file
+ * @returns {boolean} Whether the file is a text file
+ */
+export const isTextFile = mime => {
+  const normalizedMime = mime?.toLowerCase() || ''
+  return (
+    normalizedMime === 'text/plain' ||
+    normalizedMime === 'text/markdown' ||
+    normalizedMime.startsWith('text/')
+  )
+}
+
 /**
  * Check if a file is compatible with AI summary feature
  * Compatible file types are defined in the drive.summary flag