Merge pull request #54 from PinataCloud/chore/update-vectors

chore: Refactored vectors to use uploadUrl and include returnFile for queries
PinataCloud · Dec 4, 2024 · 4a1f407 · 4a1f407
2 parents a5b6447 + 106bac1
commit 4a1f407
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -390,6 +390,7 @@ export type VectorizeFileResponse = {
 export type VectorizeQuery = {
 	groupId: string;
 	query: string;
+	returnFile?: boolean;
 };
 
 export type VectorQueryMatch = {