From aa8b1248408846d6a158f8df1c76fa3015ce3aac Mon Sep 17 00:00:00 2001
From: "A.J. Shulman" <Shulman.aj@gmail.com>
Date: Wed, 10 Jul 2024 16:35:11 -0400
Subject: Fixing bugs and attempting to get it to work

---
 src/client/views/nodes/ChatBox/tools/RAGTool.ts | 29 +++++++++++++++++--------
 1 file changed, 20 insertions(+), 9 deletions(-)

(limited to 'src/client/views/nodes/ChatBox/tools/RAGTool.ts')

diff --git a/src/client/views/nodes/ChatBox/tools/RAGTool.ts b/src/client/views/nodes/ChatBox/tools/RAGTool.ts
index 84d5430e7..185efa0ba 100644
--- a/src/client/views/nodes/ChatBox/tools/RAGTool.ts
+++ b/src/client/views/nodes/ChatBox/tools/RAGTool.ts
@@ -1,8 +1,9 @@
 import { BaseTool } from './BaseTool';
 import { Vectorstore } from '../vectorstore/VectorstoreUpload';
 import { Chunk } from '../types';
+import * as fs from 'fs';
 
-export class RAGTool extends BaseTool {
+export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
     constructor(
         private vectorstore: Vectorstore,
         summaries: string
@@ -59,16 +60,26 @@ export class RAGTool extends BaseTool {
         for (const chunk of relevantChunks) {
             content.push({
                 type: 'text',
-                text: `<chunk chunk_id=${chunk.id} chunk_type=${chunk.metadata.type === 'image' ? 'image' : 'text'}>`,
+                text: `<chunk chunk_id=${chunk.id} chunk_type=${chunk.metadata.type === 'image' || chunk.metadata.type === 'table' ? 'image' : 'text'}>`,
             });
 
-            if (chunk.metadata.type === 'image') {
-                // Implement image loading and base64 encoding here
-                // For now, we'll just add a placeholder
-                content.push({
-                    type: 'image_url',
-                    image_url: { url: chunk.metadata.file_path },
-                });
+            if (chunk.metadata.type === 'image' || chunk.metadata.type === 'table') {
+                try {
+                    const imageBuffer = fs.readFileSync(chunk.metadata.file_path);
+                    const base64Image = imageBuffer.toString('base64');
+                    if (base64Image) {
+                        content.push({
+                            type: 'image_url',
+                            image_url: {
+                                url: `data:image/jpeg;base64,${base64Image}`,
+                            },
+                        });
+                    } else {
+                        console.log(`Failed to encode image for chunk ${chunk.id}`);
+                    }
+                } catch (error) {
+                    console.error(`Error reading image file for chunk ${chunk.id}:`, error);
+                }
             }
 
             content.push({ type: 'text', text: `${chunk.metadata.text}\n</chunk>\n` });
-- 
cgit v1.2.3-70-g09d2