@mastra/rag 0.0.0-commonjs-20250227130920

package/src/utils/vector-search.ts

@@ -0,0 +1,41 @@
+import type { MastraVector, QueryResult } from '@mastra/core/vector';
+import { embed } from 'ai';
+import type { EmbeddingModel } from 'ai';
+
+interface VectorQuerySearchParams {
+  indexName: string;
+  vectorStore: MastraVector;
+  queryText: string;
+  model: EmbeddingModel<string>;
+  queryFilter?: any;
+  topK: number;
+  includeVectors?: boolean;
+  maxRetries?: number;
+}
+
+interface VectorQuerySearchResult {
+  results: QueryResult[];
+  queryEmbedding: number[];
+}
+
+// Helper function to handle vector query search
+export const vectorQuerySearch = async ({
+  indexName,
+  vectorStore,
+  queryText,
+  model,
+  queryFilter = {},
+  topK,
+  includeVectors = false,
+  maxRetries = 2,
+}: VectorQuerySearchParams): Promise<VectorQuerySearchResult> => {
+  const { embedding } = await embed({
+    value: queryText,
+    model,
+    maxRetries,
+  });
+  // Get relevant chunks from the vector database
+  const results = await vectorStore.query(indexName, embedding, topK, queryFilter, includeVectors);
+
+  return { results, queryEmbedding: embedding };
+};

package/tsconfig.json

@@ -0,0 +1,5 @@
+{
+  "extends": "../../tsconfig.node.json",
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "**/*.test.ts"]
+}

package/vitest.config.ts

@@ -0,0 +1,11 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['src/**/*.test.ts'],
+    exclude: ['**/node_modules/**', '**/dist/**'],
+    globals: true,
+    setupFiles: ['dotenv/config'],
+  },
+});