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/BaseTool.ts   |  4 +--
 .../views/nodes/ChatBox/tools/CalculateTool.ts     |  2 +-
 src/client/views/nodes/ChatBox/tools/RAGTool.ts    | 29 +++++++++++++++-------
 .../views/nodes/ChatBox/tools/WikipediaTool.ts     |  2 +-
 4 files changed, 24 insertions(+), 13 deletions(-)

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

diff --git a/src/client/views/nodes/ChatBox/tools/BaseTool.ts b/src/client/views/nodes/ChatBox/tools/BaseTool.ts
index 3511d9528..903161bd5 100644
--- a/src/client/views/nodes/ChatBox/tools/BaseTool.ts
+++ b/src/client/views/nodes/ChatBox/tools/BaseTool.ts
@@ -1,6 +1,6 @@
 import { Tool } from '../types';
 
-export abstract class BaseTool implements Tool {
+export abstract class BaseTool<T extends Record<string, any> = Record<string, any>> implements Tool<T> {
     constructor(
         public name: string,
         public description: string,
@@ -9,7 +9,7 @@ export abstract class BaseTool implements Tool {
         public briefSummary: string
     ) {}
 
-    abstract execute(args: Record<string, any>): Promise<any>;
+    abstract execute(args: T): Promise<any>;
 
     getActionRule(): Record<string, any> {
         return {
diff --git a/src/client/views/nodes/ChatBox/tools/CalculateTool.ts b/src/client/views/nodes/ChatBox/tools/CalculateTool.ts
index b881d90fa..818332c44 100644
--- a/src/client/views/nodes/ChatBox/tools/CalculateTool.ts
+++ b/src/client/views/nodes/ChatBox/tools/CalculateTool.ts
@@ -1,6 +1,6 @@
 import { BaseTool } from './BaseTool';
 
-export class CalculateTool extends BaseTool {
+export class CalculateTool extends BaseTool<{ expression: string }> {
     constructor() {
         super(
             'calculate',
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` });
diff --git a/src/client/views/nodes/ChatBox/tools/WikipediaTool.ts b/src/client/views/nodes/ChatBox/tools/WikipediaTool.ts
index 0aef58f61..8ef2830d4 100644
--- a/src/client/views/nodes/ChatBox/tools/WikipediaTool.ts
+++ b/src/client/views/nodes/ChatBox/tools/WikipediaTool.ts
@@ -1,7 +1,7 @@
 import { BaseTool } from './BaseTool';
 import axios from 'axios';
 
-export class WikipediaTool extends BaseTool {
+export class WikipediaTool extends BaseTool<{ title: string }> {
     constructor() {
         super(
             'wikipedia',
-- 
cgit v1.2.3-70-g09d2