@mastra/upstash 0.1.6-alpha.1 → 0.1.6-alpha.4

package/src/vector/index.ts

@@ -1,6 +1,12 @@
-import type { Filter } from '@mastra/core/filter';
 import { MastraVector } from '@mastra/core/vector';
-import type { QueryResult } from '@mastra/core/vector';
+import type {
+  CreateIndexParams,
+  ParamsToArgs,
+  QueryResult,
+  QueryVectorParams,
+  UpsertVectorParams,
+} from '@mastra/core/vector';
+import type { VectorFilter } from '@mastra/core/vector/filter';
 import { Index } from '@upstash/vector';
 
 import { UpstashFilterTranslator } from './filter';
@@ -16,12 +22,11 @@ export class UpstashVector extends MastraVector {
     });
   }
 
-  async upsert(
-    indexName: string,
-    vectors: number[][],
-    metadata?: Record<string, any>[],
-    ids?: string[],
-  ): Promise<string[]> {
+  async upsert(...args: ParamsToArgs<UpsertVectorParams>): Promise<string[]> {
+    const params = this.normalizeArgs<UpsertVectorParams>('upsert', args);
+
+    const { indexName, vectors, metadata, ids } = params;
+
     const generatedIds = ids || vectors.map(() => crypto.randomUUID());
 
     const points = vectors.map((vector, index) => ({
@@ -36,26 +41,20 @@ export class UpstashVector extends MastraVector {
     return generatedIds;
   }
 
-  transformFilter(filter?: Filter) {
+  transformFilter(filter?: VectorFilter) {
     const translator = new UpstashFilterTranslator();
     return translator.translate(filter);
   }
 
-  async createIndex(
-    _indexName: string,
-    _dimension: number,
-    _metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
-  ): Promise<void> {
+  async createIndex(..._args: ParamsToArgs<CreateIndexParams>): Promise<void> {
     console.log('No need to call createIndex for Upstash');
   }
 
-  async query(
-    indexName: string,
-    queryVector: number[],
-    topK: number = 10,
-    filter?: Filter,
-    includeVector: boolean = false,
-  ): Promise<QueryResult[]> {
+  async query(...args: ParamsToArgs<QueryVectorParams>): Promise<QueryResult[]> {
+    const params = this.normalizeArgs<QueryVectorParams>('query', args);
+
+    const { indexName, queryVector, topK = 10, filter, includeVector = false } = params;
+
     const ns = this.client.namespace(indexName);
 
     const filterString = this.transformFilter(filter);