@mastra/upstash 0.12.3 → 0.12.4-alpha.1

package/src/vector/index.test.ts

@@ -37,7 +37,7 @@ function waitUntilVectorsIndexed(vector: UpstashVector, indexName: string, expec
  */
 describe.skipIf(!process.env.UPSTASH_VECTOR_URL || !process.env.UPSTASH_VECTOR_TOKEN)('UpstashVector', () => {
   let vectorStore: UpstashVector;
-  const VECTOR_DIMENSION = 1536;
+  const VECTOR_DIMENSION = 1024;
   const testIndexName = 'default';
   const filterIndexName = 'filter-index';
 
@@ -250,7 +250,7 @@ describe.skipIf(!process.env.UPSTASH_VECTOR_URL || !process.env.UPSTASH_VECTOR_T
     it('should describe an index correctly', async () => {
       const stats = await vectorStore.describeIndex({ indexName: 'mastra_default' });
       expect(stats).toEqual({
-        dimension: 1536,
+        dimension: 1024,
         metric: 'cosine',
         count: 0,
       });

package/src/vector/index.ts

@@ -7,16 +7,12 @@ import type {
   DescribeIndexParams,
   IndexStats,
   QueryResult,
-  QueryVectorParams,
-  UpdateVectorParams,
-  UpsertVectorParams,
 } from '@mastra/core/vector';
 import { Index } from '@upstash/vector';
 
 import { UpstashFilterTranslator } from './filter';
 import type { UpstashVectorFilter } from './filter';
-
-type UpstashQueryVectorParams = QueryVectorParams<UpstashVectorFilter>;
+import type { UpstashUpsertVectorParams, UpstashQueryVectorParams, UpstashUpdateVectorParams } from './types';
 
 export class UpstashVector extends MastraVector<UpstashVectorFilter> {
   private client: Index;
@@ -40,12 +36,19 @@ export class UpstashVector extends MastraVector<UpstashVectorFilter> {
    * @param {UpsertVectorParams} params - The parameters for the upsert operation.
    * @returns {Promise<string[]>} A promise that resolves to the IDs of the upserted vectors.
    */
-  async upsert({ indexName: namespace, vectors, metadata, ids }: UpsertVectorParams): Promise<string[]> {
+  async upsert({
+    indexName: namespace,
+    vectors,
+    metadata,
+    ids,
+    sparseVectors,
+  }: UpstashUpsertVectorParams): Promise<string[]> {
     const generatedIds = ids || vectors.map(() => crypto.randomUUID());
 
     const points = vectors.map((vector, index) => ({
       id: generatedIds[index]!,
       vector,
+      ...(sparseVectors?.[index] && { sparseVector: sparseVectors[index] }),
       metadata: metadata?.[index],
     }));
 
@@ -97,6 +100,9 @@ export class UpstashVector extends MastraVector<UpstashVectorFilter> {
     topK = 10,
     filter,
     includeVector = false,
+    sparseVector,
+    fusionAlgorithm,
+    queryMode,
   }: UpstashQueryVectorParams): Promise<QueryResult[]> {
     try {
       const ns = this.client.namespace(namespace);
@@ -105,9 +111,12 @@ export class UpstashVector extends MastraVector<UpstashVectorFilter> {
       const results = await ns.query({
         topK,
         vector: queryVector,
+        ...(sparseVector && { sparseVector }),
         includeVectors: includeVector,
         includeMetadata: true,
         ...(filterString ? { filter: filterString } : {}),
+        ...(fusionAlgorithm && { fusionAlgorithm }),
+        ...(queryMode && { queryMode }),
       });
 
       // Map the results to our expected format
@@ -209,47 +218,37 @@ export class UpstashVector extends MastraVector<UpstashVectorFilter> {
    * @returns A promise that resolves when the update is complete.
    * @throws Will throw an error if no updates are provided or if the update operation fails.
    */
-  async updateVector({ indexName: namespace, id, update }: UpdateVectorParams): Promise<void> {
-    try {
-      if (!update.vector && !update.metadata) {
-        throw new Error('No update data provided');
-      }
+  async updateVector({ indexName: namespace, id, update }: UpstashUpdateVectorParams): Promise<void> {
+    if (!update.vector && !update.metadata && !update.sparseVector) {
+      throw new MastraError({
+        id: 'STORAGE_UPSTASH_VECTOR_UPDATE_VECTOR_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.THIRD_PARTY,
+        details: { namespace, id },
+        text: 'No update data provided',
+      });
+    }
 
-      // The upstash client throws an exception as: 'This index requires dense vectors' when
-      // only metadata is present in the update object.
-      if (!update.vector && update.metadata) {
-        throw new Error('Both vector and metadata must be provided for an update');
-      }
-    } catch (error) {
-      throw new MastraError(
-        {
-          id: 'STORAGE_UPSTASH_VECTOR_UPDATE_VECTOR_FAILED',
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.THIRD_PARTY,
-          details: { namespace, id },
-        },
-        error,
-      );
+    // The upstash client throws an exception as: 'This index requires dense/sparse vectors' when
+    // only metadata is present in the update object.
+    if (!update.vector && !update.sparseVector && update.metadata) {
+      throw new MastraError({
+        id: 'STORAGE_UPSTASH_VECTOR_UPDATE_VECTOR_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.THIRD_PARTY,
+        details: { namespace, id },
+        text: 'Both vector and metadata must be provided for an update',
+      });
     }
 
     try {
-      const updatePayload: any = { id: id };
-      if (update.vector) {
-        updatePayload.vector = update.vector;
-      }
-      if (update.metadata) {
-        updatePayload.metadata = update.metadata;
-      }
+      const points: any = { id };
 
-      const points = {
-        id: updatePayload.id,
-        vector: updatePayload.vector,
-        metadata: updatePayload.metadata,
-      };
+      if (update.vector) points.vector = update.vector;
+      if (update.metadata) points.metadata = update.metadata;
+      if (update.sparseVector) points.sparseVector = update.sparseVector;
 
-      await this.client.upsert(points, {
-        namespace,
-      });
+      await this.client.upsert(points, { namespace });
     } catch (error) {
       throw new MastraError(
         {

package/src/vector/types.ts

@@ -0,0 +1,26 @@
+import type { UpsertVectorParams, QueryVectorParams, UpdateVectorParams } from '@mastra/core/vector';
+import type { FusionAlgorithm, QueryMode } from '@upstash/vector';
+import type { UpstashVectorFilter } from './filter';
+
+export interface UpstashSparseVector {
+  indices: number[];
+  values: number[];
+}
+
+export interface UpstashUpsertVectorParams extends UpsertVectorParams {
+  sparseVectors?: UpstashSparseVector[];
+}
+
+export interface UpstashQueryVectorParams extends QueryVectorParams<UpstashVectorFilter> {
+  sparseVector?: UpstashSparseVector;
+  fusionAlgorithm?: FusionAlgorithm;
+  queryMode?: QueryMode;
+}
+
+export interface UpstashUpdateVectorParams extends UpdateVectorParams {
+  update: {
+    vector?: number[];
+    metadata?: Record<string, any>;
+    sparseVector?: UpstashSparseVector;
+  };
+}