✨ feat: Add configurable PDF processing method with Unstructured

src/config/knowledge.ts

@@ -1,19 +1,18 @@
 import { createEnv } from '@t3-oss/env-nextjs';
 import { z } from 'zod';
 
-export const getKnowledgeConfig = () => {
-  return createEnv({
-    runtimeEnv: {
-      DEFAULT_FILES_CONFIG: process.env.DEFAULT_FILES_CONFIG,
-      UNSTRUCTURED_API_KEY: process.env.UNSTRUCTURED_API_KEY,
-      UNSTRUCTURED_SERVER_URL: process.env.UNSTRUCTURED_SERVER_URL,
-    },
-    server: {
-      DEFAULT_FILES_CONFIG: z.string().optional(),
-      UNSTRUCTURED_API_KEY: z.string().optional(),
-      UNSTRUCTURED_SERVER_URL: z.string().optional(),
-    },
-  });
-};
-
-export const knowledgeEnv = getKnowledgeConfig();
+export const knowledgeEnv = createEnv({
+  runtimeEnv: {
+    DEFAULT_FILES_CONFIG: process.env.DEFAULT_FILES_CONFIG,
+    UNSTRUCTURED_API_KEY: process.env.UNSTRUCTURED_API_KEY,
+    UNSTRUCTURED_SERVER_URL: process.env.UNSTRUCTURED_SERVER_URL,
+    USE_UNSTRUCTURED_FOR_PDF: process.env.USE_UNSTRUCTURED_FOR_PDF,
+    FILE_TYPE_CHUNKING_RULES: process.env.FILE_TYPE_CHUNKING_RULES || 'pdf=unstructured',
+  },
+  server: {
+    DEFAULT_FILES_CONFIG: z.string().optional(),
+    UNSTRUCTURED_API_KEY: z.string().optional(),
+    UNSTRUCTURED_SERVER_URL: z.string().optional(),
+    USE_UNSTRUCTURED_FOR_PDF: z.string().optional(),
+  },
+});

src/server/modules/ContentChunk/index.ts

@@ -1,8 +1,10 @@
 import { ChunkingLoader } from 'src/libs/langchain';
 import { Strategy } from 'unstructured-client/sdk/models/shared';
+import { ChunkingRuleParser } from './rules';
 
 import { NewChunkItem, NewUnstructuredChunkItem } from '@/database/schemas';
 import { ChunkingStrategy, Unstructured } from '@/libs/unstructured';
+import { knowledgeEnv } from '@/config/knowledge';
 
 export interface ChunkContentParams {
   content: Uint8Array;
@@ -19,23 +21,67 @@ interface ChunkResult {
 export class ContentChunk {
   private unstructuredClient: Unstructured;
   private langchainClient: ChunkingLoader;
+  private chunkingRules: Record<string, ChunkingService[]>;
 
   constructor() {
     this.unstructuredClient = new Unstructured();
     this.langchainClient = new ChunkingLoader();
+    this.chunkingRules = ChunkingRuleParser.parse(knowledgeEnv.FILE_TYPE_CHUNKING_RULES);
   }
 
-  isUsingUnstructured(params: ChunkContentParams) {
-    return params.fileType === 'application/pdf' && params.mode === 'hi-res';
+  private getChunkingServices(fileType: string): ChunkingService[] {
+    const ext = fileType.split('/').pop()?.toLowerCase() || '';
+    return this.chunkingRules[ext] || ['default'];
   }
 
   async chunkContent(params: ChunkContentParams): Promise<ChunkResult> {
-    if (this.isUsingUnstructured(params))
-      return await this.chunkByUnstructured(params.filename, params.content);
-
+    const services = this.getChunkingServices(params.fileType);
+
+    for (const service of services) {
+      try {
+        switch (service) {
+          case 'unstructured':
+            if (this.canUseUnstructured()) {
+              return await this.chunkByUnstructured(params.filename, params.content);
+            }
+            break;
+
+          case 'doc2x':
+            // Future implementation
+            break;
+
+          case 'default':
+          default:
+            return await this.chunkByLangChain(params.filename, params.content);
+        }
+      } catch (error) {
+        // If this is the last service, throw the error
+        if (service === services[services.length - 1]) throw error;
+        // Otherwise continue to next service
+        console.error(`Chunking failed with service ${service}:`, error);
+        continue;
+      }
+    }
+
+    // Fallback to langchain if no service succeeded
     return await this.chunkByLangChain(params.filename, params.content);
   }
 
+  private canUseUnstructured(): boolean {
+    return !!(
+      knowledgeEnv.USE_UNSTRUCTURED_FOR_PDF &&
+      knowledgeEnv.UNSTRUCTURED_API_KEY &&
+      knowledgeEnv.UNSTRUCTURED_SERVER_URL
+    );
+  }
+
+  isUsingUnstructured(params: ChunkContentParams) {
+    return params.fileType === 'application/pdf' &&
+           !!knowledgeEnv.USE_UNSTRUCTURED_FOR_PDF &&
+           !!knowledgeEnv.UNSTRUCTURED_API_KEY &&
+           !!knowledgeEnv.UNSTRUCTURED_SERVER_URL;
+  }
+
   private chunkByUnstructured = async (
     filename: string,
     content: Uint8Array,

src/server/modules/ContentChunk/rules.ts

@@ -0,0 +1,29 @@
+type ChunkingService = 'unstructured' | 'doc2x' | 'default';
+
+interface FileTypeRule {
+  fileType: string;
+  services: ChunkingService[];
+}
+
+export class ChunkingRuleParser {
+  static parse(rulesStr: string): Record<string, ChunkingService[]> {
+    const rules: Record<string, ChunkingService[]> = {};
+
+    // Split by semicolon for different file types
+    const fileTypeRules = rulesStr.split(';');
+
+    for (const rule of fileTypeRules) {
+      const [fileType, services] = rule.split('=');
+      if (!fileType || !services) continue;
+
+      // Split services by comma and validate each service
+      rules[fileType.toLowerCase()] = services
+        .split(',')
+        .map(s => s.trim().toLowerCase())
+        .filter((s): s is ChunkingService =>
+          ['unstructured', 'doc2x', 'default'].includes(s));
+    }
+
+    return rules;
+  }
+}