From 106bac185f24231b611c4d388ab782910c5fc1ee Mon Sep 17 00:00:00 2001
From: Steve <stevedsimkins@gmail.com>
Date: Wed, 4 Dec 2024 11:07:28 -0500
Subject: [PATCH] chore: Refactored vectors to use uploadUrl and include
 returnFile for queries

---
 src/core/files/deleteFileVectors.ts |  4 ++--
 src/core/files/vectorizeFile.ts     |  4 ++--
 src/core/files/vectorizeQuery.ts    | 20 +++++++++++++++++---
 src/core/pinataSDK.ts               |  4 +++-
 src/core/types.ts                   |  1 +
 5 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/core/files/deleteFileVectors.ts b/src/core/files/deleteFileVectors.ts
index 14fb28e..1d3db15 100644
--- a/src/core/files/deleteFileVectors.ts
+++ b/src/core/files/deleteFileVectors.ts
@@ -30,8 +30,8 @@ export const deleteFileVectors = async (
 
 	let endpoint: string = "https://uploads.pinata.cloud/v3";
 
-	if (config.endpointUrl) {
-		endpoint = config.endpointUrl;
+	if (config.uploadUrl) {
+		endpoint = config.uploadUrl;
 	}
 
 	try {
diff --git a/src/core/files/vectorizeFile.ts b/src/core/files/vectorizeFile.ts
index 72a8c3c..a73cd39 100644
--- a/src/core/files/vectorizeFile.ts
+++ b/src/core/files/vectorizeFile.ts
@@ -30,8 +30,8 @@ export const vectorizeFile = async (
 
 	let endpoint: string = "https://uploads.pinata.cloud/v3";
 
-	if (config.endpointUrl) {
-		endpoint = config.endpointUrl;
+	if (config.uploadUrl) {
+		endpoint = config.uploadUrl;
 	}
 
 	try {
diff --git a/src/core/files/vectorizeQuery.ts b/src/core/files/vectorizeQuery.ts
index c6f7820..0d1f9f4 100644
--- a/src/core/files/vectorizeQuery.ts
+++ b/src/core/files/vectorizeQuery.ts
@@ -1,4 +1,5 @@
 import type {
+	GetCIDResponse,
 	PinataConfig,
 	VectorizeQuery,
 	VectorizeQueryResponse,
@@ -10,10 +11,12 @@ import {
 	ValidationError,
 } from "../../utils/custom-errors";
 
+import { getCid } from "../gateway/getCid";
+
 export const vectorizeQuery = async (
 	config: PinataConfig | undefined,
 	options: VectorizeQuery,
-): Promise<VectorizeQueryResponse> => {
+): Promise<VectorizeQueryResponse | GetCIDResponse> => {
 	if (!config) {
 		throw new ValidationError("Pinata configuration is missing");
 	}
@@ -34,8 +37,8 @@ export const vectorizeQuery = async (
 
 	let endpoint: string = "https://uploads.pinata.cloud/v3";
 
-	if (config.endpointUrl) {
-		endpoint = config.endpointUrl;
+	if (config.uploadUrl) {
+		endpoint = config.uploadUrl;
 	}
 
 	const body = JSON.stringify({
@@ -67,8 +70,19 @@ export const vectorizeQuery = async (
 				errorData,
 			);
 		}
+
 		const res = await request.json();
 		const resData: VectorizeQueryResponse = res.data;
+
+		if (options.returnFile) {
+			if (resData.matches.length === 0) {
+				throw new PinataError(`No files returned in query to fetch`);
+			}
+			const cid = resData.matches[0].cid;
+			const fileRes: GetCIDResponse = await getCid(config, cid, undefined);
+			return fileRes;
+		}
+
 		return resData;
 	} catch (error) {
 		if (error instanceof PinataError) {
diff --git a/src/core/pinataSDK.ts b/src/core/pinataSDK.ts
index e9d9887..12dc057 100644
--- a/src/core/pinataSDK.ts
+++ b/src/core/pinataSDK.ts
@@ -191,7 +191,9 @@ class Files {
 		return vectorizeFile(this.config, fileId);
 	}
 
-	queryVectors(options: VectorizeQuery): Promise<VectorizeQueryResponse> {
+	queryVectors(
+		options: VectorizeQuery,
+	): Promise<VectorizeQueryResponse | GetCIDResponse> {
 		return vectorizeQuery(this.config, options);
 	}
 
diff --git a/src/core/types.ts b/src/core/types.ts
index e73c31c..be7732d 100644
--- a/src/core/types.ts
+++ b/src/core/types.ts
@@ -390,6 +390,7 @@ export type VectorizeFileResponse = {
 export type VectorizeQuery = {
 	groupId: string;
 	query: string;
+	returnFile?: boolean;
 };
 
 export type VectorQueryMatch = {