1 files changed, 54 insertions, 0 deletions
diff --git a/src/server/ApiManagers/AssistantManager.ts b/src/server/ApiManagers/AssistantManager.ts
index 36468157a..f69ca1383 100644
--- a/src/server/ApiManagers/AssistantManager.ts
+++ b/src/server/ApiManagers/AssistantManager.ts
@@ -9,6 +9,9 @@ import { Method } from '../RouteManager';
 import ApiManager, { Registration } from './ApiManager';
 import axios from 'axios';
 import { Chunk } from '../../client/views/nodes/ChatBox/types';
+import { UnstructuredClient } from 'unstructured-client';
+import { PartitionResponse } from 'unstructured-client/sdk/models/operations';
+import { ChunkingStrategy, Strategy } from 'unstructured-client/sdk/models/shared';
 
 export enum Directory {
     parsed_files = 'parsed_files',
@@ -42,6 +45,11 @@ export default class AssistantManager extends ApiManager {
             apiKey: process.env._CLIENT_OPENAI_KEY, // Use client key so don't have to set key seperately for client and server.
             dangerouslyAllowBrowser: true,
         });
+        const unstructuredClient = new UnstructuredClient({
+            security: {
+                apiKeyAuth: process.env._CLIENT_UNSTRUCTURED_API_KEY!,
+            },
+        });
 
         register({
             method: Method.POST,
@@ -187,5 +195,51 @@ export default class AssistantManager extends ApiManager {
                 res.send({ formattedChunks: content });
             },
         });
+
+        register({
+            method: Method.POST,
+            subscription: '/chunkDocument',
+            secureHandler: async ({ req, res }) => {
+                const { file_path } = req.body;
+                const public_path = path.join(publicDirectory, file_path);
+                const file_name = path.basename(file_path);
+
+                try {
+                    // Read file data and convert to base64
+                    const file_data = await fs.promises.readFile(public_path);
+
+                    try {
+                        const result = await unstructuredClient.general.partition({
+                            partitionParameters: {
+                                files: {
+                                    content: file_data,
+                                    fileName: file_name,
+                                },
+                                strategy: Strategy.Auto,
+                                chunkingStrategy: ChunkingStrategy.ByTitle,
+                                extractImageBlockTypes: ['Image', 'Table'],
+                            },
+                        });
+
+                        if (result.statusCode === 200) {
+                            console.log(result.elements);
+                            const jsonElements = JSON.stringify(result.elements, null, 2);
+                            // Print the processed data.
+                            console.log(jsonElements);
+                            res.send({ document_json: jsonElements });
+                        } else {
+                            console.error(`Unexpected status code: ${result.statusCode}`);
+                            res.status(result.statusCode).send({ error: 'Failed to process the document', details: result });
+                        }
+                    } catch (e: any) {
+                        console.error('Error during partitioning:', e);
+                        res.status(500).send({ error: 'Failed to partition the document', details: e.message });
+                    }
+                } catch (error: any) {
+                    console.error('Error reading file:', error);
+                    res.status(500).send({ error: 'Failed to read the file', details: error.message });
+                }
+            },
+        });
     }
 }