ai-desktop-demo/electron.cjs

301 lines
9.2 KiB
JavaScript
Raw Normal View History

2024-07-28 20:12:20 -04:00
const { app, BrowserWindow, ipcMain, protocol, dialog } = require('electron');
const { BedrockRuntimeClient, InvokeModelCommand, InvokeModelWithResponseStreamCommand } = require("@aws-sdk/client-bedrock-runtime");
const { fromIni } = require("@aws-sdk/credential-provider-ini");
const fs = require('fs').promises;
const path = require('path');
const url = require('url');
const axios = require('axios');
const cheerio = require('cheerio');
let bedrockClient;
async function initializeBedrockClient() {
bedrockClient = new BedrockRuntimeClient({
region: "us-east-1",
credentials: fromIni({ profile: 'default' })
});
}
function createWindow() {
const win = new BrowserWindow({
width: 800,
height: 600,
webPreferences: {
nodeIntegration: true,
contextIsolation: false,
},
});
if (process.env.NODE_ENV === 'development') {
win.loadURL('http://localhost:5173');
} else {
win.loadFile(path.join(__dirname, 'build', 'index.html'));
}
if (process.env.NODE_ENV === 'development') {
win.webContents.openDevTools();
}
}
app.whenReady().then(async () => {
await initializeBedrockClient();
// Register protocol
protocol.registerFileProtocol('safe-file', (request, callback) => {
const filePath = url.fileURLToPath('file://' + request.url.slice('safe-file://'.length));
callback(filePath);
});
createWindow();
});
ipcMain.handle('fetch-url-content', async (event, url) => {
try {
const response = await axios.get(url);
const $ = cheerio.load(response.data);
// Extract text content from the page
const content = $('body').text().trim();
// Limit content to a reasonable length (e.g., 1000 characters)
const limitedContent = content.substring(0, 1000);
event.sender.send('url-content-fetched', limitedContent);
} catch (error) {
console.error('Error fetching URL:', error);
event.sender.send('url-fetch-error', error.message);
}
});
async function readDirectory(dirPath) {
const fileTree = [];
const files = await fs.readdir(dirPath, { withFileTypes: true });
for (const file of files) {
const filePath = path.join(dirPath, file.name);
if (file.isDirectory()) {
const subDir = await readDirectory(filePath);
fileTree.push({ name: file.name, type: 'directory', children: subDir });
} else {
let content = '';
if (file.name.match(/\.(txt|md|js|py|html|css|json)$/i)) {
content = await fs.readFile(filePath, 'utf-8');
}
fileTree.push({ name: file.name, type: 'file', content });
}
}
return fileTree;
}
function formatFileTree(fileTree, indent = '') {
let result = '';
for (const item of fileTree) {
result += `${indent}${item.type === 'directory' ? '📁' : '📄'} ${item.name}\n`;
if (item.type === 'directory' && item.children) {
result += formatFileTree(item.children, indent + ' ');
}
}
return result;
}
ipcMain.handle('select-directory', async (event) => {
const result = await dialog.showOpenDialog({ properties: ['openDirectory'] });
if (result.canceled) {
return null;
}
const dirPath = result.filePaths[0];
const fileTree = await readDirectory(dirPath);
const formattedTree = formatFileTree(fileTree);
const fileContents = JSON.stringify(fileTree, null, 2);
return { tree: formattedTree, contents: fileContents };
});
ipcMain.handle('generate-response-stream', async (event, { prompt, imagePath, model, temperature, maxTokens, context, fileTree, fileContents }) => {
if (!bedrockClient) {
throw new Error('Bedrock client not initialized');
}
console.log("DEBUG: Generating response for prompt:", prompt);
try {
const requestBody = {
anthropic_version: "bedrock-2023-05-31",
max_tokens: maxTokens,
temperature: temperature,
messages: [
{
role: "user",
content: [
{ type: "text", text: `Context: ${context || ''}\n\nFile Tree:\n${fileTree || ''}\n\nFile Contents:\n${fileContents || ''}\n\nUser: ${prompt}` }
]
}
]
};
if (imagePath) {
const imageBuffer = await fs.readFile(imagePath);
const base64Image = imageBuffer.toString('base64');
requestBody.messages[0].content.unshift({
type: "image",
source: {
type: "base64",
media_type: "image/jpeg",
data: base64Image
}
});
}
const params = {
modelId: "anthropic.claude-3-sonnet-20240229-v1:0", // or your preferred Claude 3 model
contentType: "application/json",
accept: "application/json",
body: JSON.stringify(requestBody),
};
const command = new InvokeModelWithResponseStreamCommand(params);
const response = await bedrockClient.send(command);
for await (const chunk of response.body) {
if (chunk.chunk && chunk.chunk.bytes) {
const decodedChunk = JSON.parse(new TextDecoder().decode(chunk.chunk.bytes));
if (decodedChunk.type === 'content_block_delta' && decodedChunk.delta && decodedChunk.delta.text) {
event.sender.send('response-chunk', decodedChunk.delta.text);
}
}
}
event.sender.send('response-end');
} catch (error) {
console.error("ERROR: Failed to generate text with Bedrock:", error);
event.sender.send('response-error', error.message);
}
});
ipcMain.handle('generate-response', async (event, { prompt, imagePath }) => {
if (!bedrockClient) {
throw new Error('Bedrock client not initialized');
}
console.log("DEBUG: Generating response for prompt:", prompt);
try {
const requestBody = {
anthropic_version: "bedrock-2023-05-31",
max_tokens: 20000,
temperature: 0.7,
top_p: 0.9,
messages: [
{
role: "user",
content: []
}
]
};
if (imagePath) {
const imageBuffer = await fs.readFile(imagePath);
const base64Image = imageBuffer.toString('base64');
requestBody.messages[0].content.push({
type: "image",
source: {
type: "base64",
media_type: "image/jpeg",
data: base64Image
}
});
}
requestBody.messages[0].content.push({
type: "text",
text: prompt
});
const params = {
modelId: "anthropic.claude-3-sonnet-20240229-v1:0",
contentType: "application/json",
accept: "application/json",
body: JSON.stringify(requestBody),
};
const command = new InvokeModelCommand(params);
const response = await bedrockClient.send(command);
const responseBody = JSON.parse(new TextDecoder().decode(response.body));
console.log("DEBUG: Raw response from Bedrock:", JSON.stringify(responseBody));
if (responseBody.content && Array.isArray(responseBody.content) && responseBody.content.length > 0) {
const generatedText = responseBody.content[0].text;
if (generatedText) {
console.log("DEBUG: Generated text:", generatedText);
return generatedText;
} else {
console.log("WARNING: Generated text is empty. Full response:", JSON.stringify(responseBody));
return "No response generated";
}
} else {
console.log("WARNING: Unexpected response format. Full response:", JSON.stringify(responseBody));
return "Unexpected response format";
}
} catch (error) {
console.error("ERROR: Failed to generate text with Bedrock:", error);
return "Error: " + error.message;
}
});
ipcMain.handle('generate-image', async (event, { prompt }) => {
if (!bedrockClient) {
throw new Error('Bedrock client not initialized');
}
console.log("DEBUG: Generating image for prompt:", prompt);
try {
const requestBody = {
text_prompts: [{ text: prompt }],
cfg_scale: 10,
steps: 50,
seed: Math.floor(Math.random() * 1000000),
};
const params = {
modelId: "stability.stable-diffusion-xl-v1",
contentType: "application/json",
accept: "application/json",
body: JSON.stringify(requestBody),
};
const command = new InvokeModelCommand(params);
const response = await bedrockClient.send(command);
const responseBody = JSON.parse(new TextDecoder().decode(response.body));
if (responseBody.artifacts && responseBody.artifacts.length > 0) {
const imageBase64 = responseBody.artifacts[0].base64;
const imagePath = path.join(app.getPath('userData'), `generated_image_${Date.now()}.png`);
await fs.writeFile(imagePath, imageBase64, 'base64');
return `safe-file://${imagePath}`; // Return the safe-file URL instead of the direct file path
} else {
throw new Error('No image generated');
}
} catch (error) {
console.error("ERROR: Failed to generate image with Bedrock:", error);
throw error;
}
});
app.on('window-all-closed', () => {
if (process.platform !== 'darwin') {
app.quit();
}
});
app.on('activate', () => {
if (BrowserWindow.getAllWindows().length === 0) {
createWindow();
}
});