n8n/packages/@n8n/nodes-langchain/nodes/vendors/GoogleGemini/helpers/utils.ts

import axios from 'axios';
import type { IDataObject, IExecuteFunctions } from 'n8n-workflow';
import { NodeOperationError } from 'n8n-workflow';
import { Readable } from 'node:stream';
import type Stream from 'node:stream';

import type { FileSearchOperation } from './interfaces';
import { apiRequest } from '../transport';

const OPERATION_CHECK_INTERVAL = 1000;

interface File {
	name: string;
	uri: string;
	mimeType: string;
	state: string;
	error?: { message: string };
}

interface FileStreamData {
	stream: Stream;
	mimeType: string;
}

interface FileBufferData {
	buffer: Buffer;
	mimeType: string;
}

interface UploadStreamConfig {
	endpoint: string;
	mimeType: string;
	body?: IDataObject;
}

const CHUNK_SIZE = 256 * 1024;

export async function downloadFile(
	this: IExecuteFunctions,
	url: string,
	fallbackMimeType?: string,
	qs?: IDataObject,
) {
	const downloadResponse = (await this.helpers.httpRequest({
		method: 'GET',
		url,
		qs,
		returnFullResponse: true,
		encoding: 'arraybuffer',
	})) as { body: ArrayBuffer; headers: IDataObject };

	const mimeType =
		(downloadResponse.headers?.['content-type'] as string)?.split(';')?.[0] ?? fallbackMimeType;
	const fileContent = Buffer.from(downloadResponse.body);
	return {
		fileContent,
		mimeType,
	};
}

export async function uploadFile(this: IExecuteFunctions, fileContent: Buffer, mimeType: string) {
	const numBytes = fileContent.length.toString();
	const uploadInitResponse = (await apiRequest.call(this, 'POST', '/upload/v1beta/files', {
		headers: {
			'X-Goog-Upload-Protocol': 'resumable',
			'X-Goog-Upload-Command': 'start',
			'X-Goog-Upload-Header-Content-Length': numBytes,
			'X-Goog-Upload-Header-Content-Type': mimeType,
			'Content-Type': 'application/json',
		},
		option: {
			returnFullResponse: true,
		},
	})) as { headers: IDataObject };
	const uploadUrl = uploadInitResponse.headers['x-goog-upload-url'] as string;

	const uploadResponse = (await this.helpers.httpRequest({
		method: 'POST',
		url: uploadUrl,
		headers: {
			'Content-Length': numBytes,
			'X-Goog-Upload-Offset': '0',
			'X-Goog-Upload-Command': 'upload, finalize',
		},
		body: fileContent,
	})) as { file: File };

	while (uploadResponse.file.state !== 'ACTIVE' && uploadResponse.file.state !== 'FAILED') {
		await new Promise((resolve) => setTimeout(resolve, OPERATION_CHECK_INTERVAL));
		uploadResponse.file = (await apiRequest.call(
			this,
			'GET',
			`/v1beta/${uploadResponse.file.name}`,
		)) as File;
	}

	if (uploadResponse.file.state === 'FAILED') {
		throw new NodeOperationError(
			this.getNode(),
			uploadResponse.file.error?.message ?? 'Unknown error',
			{
				description: 'Error uploading file',
			},
		);
	}

	return { fileUri: uploadResponse.file.uri, mimeType: uploadResponse.file.mimeType };
}

async function getFileStreamFromUrlOrBinary(
	this: IExecuteFunctions,
	i: number,
	downloadUrl?: string,
	fallbackMimeType?: string,
	qs?: IDataObject,
): Promise<FileStreamData | FileBufferData> {
	if (downloadUrl) {
		const downloadResponse = await axios.get(downloadUrl, {
			params: qs,
			responseType: 'stream',
		});

		const contentType = downloadResponse.headers['content-type'] as string | undefined;
		const mimeType = contentType?.split(';')?.[0] ?? fallbackMimeType ?? 'application/octet-stream';

		return {
			stream: downloadResponse.data as Stream,
			mimeType,
		};
	}

	const binaryPropertyName = this.getNodeParameter('binaryPropertyName', i, 'data');
	if (!binaryPropertyName) {
		throw new NodeOperationError(
			this.getNode(),
			'Binary property name or download URL is required',
			{
				description: 'Error uploading file',
			},
		);
	}

	const binaryData = this.helpers.assertBinaryData(i, binaryPropertyName);
	if (!binaryData.id) {
		const buffer = await this.helpers.getBinaryDataBuffer(i, binaryPropertyName);
		return {
			buffer,
			mimeType: binaryData.mimeType,
		};
	}

	return {
		stream: await this.helpers.getBinaryStream(binaryData.id, CHUNK_SIZE),
		mimeType: binaryData.mimeType,
	};
}

async function uploadStream(
	this: IExecuteFunctions,
	stream: Stream,
	config: UploadStreamConfig,
): Promise<{ body: IDataObject }> {
	const { endpoint, mimeType, body } = config;

	const uploadInitResponse = (await apiRequest.call(this, 'POST', endpoint, {
		headers: {
			'X-Goog-Upload-Protocol': 'resumable',
			'X-Goog-Upload-Command': 'start',
			'X-Goog-Upload-Header-Content-Type': mimeType,
			'Content-Type': 'application/json',
		},
		body,
		option: { returnFullResponse: true },
	})) as { headers: IDataObject };

	const uploadUrl = uploadInitResponse.headers['x-goog-upload-url'] as string;
	if (!uploadUrl) {
		throw new NodeOperationError(this.getNode(), 'Failed to get upload URL');
	}

	return (await this.helpers.httpRequest({
		method: 'POST',
		url: uploadUrl,
		headers: {
			'X-Goog-Upload-Offset': '0',
			'X-Goog-Upload-Command': 'upload, finalize',
			'Content-Type': mimeType,
		},
		body: stream,
		returnFullResponse: true,
	})) as { body: IDataObject };
}

export async function transferFile(
	this: IExecuteFunctions,
	i: number,
	downloadUrl?: string,
	fallbackMimeType?: string,
	qs?: IDataObject,
) {
	const fileData = await getFileStreamFromUrlOrBinary.call(
		this,
		i,
		downloadUrl,
		fallbackMimeType,
		qs,
	);

	if ('buffer' in fileData) {
		return await uploadFile.call(this, fileData.buffer, fileData.mimeType);
	}

	const { stream, mimeType } = fileData;
	const uploadResponse = (await uploadStream.call(this, stream, {
		endpoint: '/upload/v1beta/files',
		mimeType,
	})) as { body: { file: File } };

	let file = uploadResponse.body.file;

	while (file.state !== 'ACTIVE' && file.state !== 'FAILED') {
		await new Promise((resolve) => setTimeout(resolve, OPERATION_CHECK_INTERVAL));
		file = (await apiRequest.call(this, 'GET', `/v1beta/${file.name}`)) as File;
	}

	if (file.state === 'FAILED') {
		throw new NodeOperationError(this.getNode(), file.error?.message ?? 'Unknown error', {
			description: 'Error uploading file',
		});
	}

	return { fileUri: file.uri, mimeType: file.mimeType };
}

export async function createFileSearchStore(this: IExecuteFunctions, displayName: string) {
	return (await apiRequest.call(this, 'POST', '/v1beta/fileSearchStores', {
		body: { displayName },
	})) as IDataObject;
}

export async function uploadToFileSearchStore(
	this: IExecuteFunctions,
	i: number,
	fileSearchStoreName: string,
	displayName: string,
	downloadUrl?: string,
	fallbackMimeType?: string,
	qs?: IDataObject,
) {
	const fileData = await getFileStreamFromUrlOrBinary.call(
		this,
		i,
		downloadUrl,
		fallbackMimeType,
		qs,
	);

	let stream: Stream;
	let mimeType: string;

	if ('buffer' in fileData) {
		stream = Readable.from(fileData.buffer);
		mimeType = fileData.mimeType;
	} else {
		stream = fileData.stream;
		mimeType = fileData.mimeType;
	}

	const uploadResponse = (await uploadStream.call(this, stream, {
		endpoint: `/upload/v1beta/${fileSearchStoreName}:uploadToFileSearchStore`,
		mimeType,
		body: { displayName, mimeType },
	})) as { body: { name: string } };

	const operationName = uploadResponse.body.name;
	let operation = (await apiRequest.call(
		this,
		'GET',
		`/v1beta/${operationName}`,
	)) as FileSearchOperation;

	while (!operation.done) {
		await new Promise((resolve) => setTimeout(resolve, OPERATION_CHECK_INTERVAL));
		operation = (await apiRequest.call(
			this,
			'GET',
			`/v1beta/${operationName}`,
		)) as FileSearchOperation;
	}

	if (operation.error) {
		throw new NodeOperationError(this.getNode(), operation.error.message ?? 'Unknown error', {
			description: 'Error uploading file to File Search store',
		});
	}

	return operation.response;
}

export async function listFileSearchStores(
	this: IExecuteFunctions,
	pageSize?: number,
	pageToken?: string,
) {
	const qs: IDataObject = {};
	if (pageSize !== undefined) {
		qs.pageSize = pageSize;
	}
	if (pageToken) {
		qs.pageToken = pageToken;
	}

	return (await apiRequest.call(this, 'GET', '/v1beta/fileSearchStores', { qs })) as IDataObject;
}

export async function deleteFileSearchStore(
	this: IExecuteFunctions,
	name: string,
	force?: boolean,
) {
	const qs: IDataObject = {};
	if (force !== undefined) {
		qs.force = force;
	}

	return (await apiRequest.call(this, 'DELETE', `/v1beta/${name}`, { qs })) as IDataObject;
}