@mastra/rag 0.1.0-alpha.86 → 0.1.0-alpha.87

package/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @mastra/rag
 
+## 0.1.0-alpha.87
+
+### Minor Changes
+
+- 8b416d9: Breaking changes
+
+### Patch Changes
+
+- 9c10484: update all packages
+- Updated dependencies [9c10484]
+- Updated dependencies [8b416d9]
+  - @mastra/core@0.2.0-alpha.94
+
 ## 0.1.0-alpha.86
 
 ### Patch Changes

package/dist/index.d.ts

@@ -1,11 +1,8 @@
 import { LLM, TitleExtractorPrompt, TitleCombinePrompt, SummaryPrompt, QuestionExtractPrompt, Document } from 'llamaindex';
 import { TiktokenEncoding, TiktokenModel } from 'js-tiktoken';
-import * as _mastra_core from '@mastra/core';
-import { EmbeddingOptions, EmbedResult, EmbedManyResult, ModelConfig } from '@mastra/core';
 import { QueryResult } from '@mastra/core/vector';
-import * as _mastra_core_tools from '@mastra/core/tools';
-import { z } from 'zod';
-import { EmbeddingOptions as EmbeddingOptions$1 } from '@mastra/core/embeddings';
+import { LanguageModelV1, EmbeddingModel } from 'ai';
+import { createTool } from '@mastra/core/tools';
 
 declare enum Language {
     CPP = "cpp",
@@ -118,9 +115,6 @@ declare class MDocument {
     getMetadata(): Record<string, any>[];
 }
 
-declare function embed(chunk: Document | string, options: EmbeddingOptions): Promise<EmbedResult<string>>;
-declare function embedMany(chunks: (Document | string)[], options: EmbeddingOptions): Promise<EmbedManyResult<string>>;
-
 type WeightConfig = {
     semantic?: number;
     vector?: number;
@@ -151,9 +145,9 @@ interface RerankerFunctionOptions {
 }
 interface RerankConfig {
     options?: RerankerOptions;
-    model: ModelConfig;
+    model: LanguageModelV1;
 }
-declare function rerank(results: QueryResult[], query: string, modelConfig: ModelConfig, options: RerankerFunctionOptions): Promise<RerankResult[]>;
+declare function rerank(results: QueryResult[], query: string, model: LanguageModelV1, options: RerankerFunctionOptions): Promise<RerankResult[]>;
 
 /**
  * TODO: GraphRAG Enhancements
@@ -214,12 +208,12 @@ declare class GraphRAG {
 declare const createDocumentChunkerTool: ({ doc, params, }: {
     doc: MDocument;
     params?: ChunkParams;
-}) => _mastra_core_tools.Tool<"Document Chunker undefined undefined" | `Document Chunker undefined ${number}` | "Document Chunker markdown undefined" | `Document Chunker markdown ${number}` | "Document Chunker latex undefined" | `Document Chunker latex ${number}` | "Document Chunker html undefined" | `Document Chunker html ${number}` | "Document Chunker recursive undefined" | `Document Chunker recursive ${number}` | "Document Chunker character undefined" | `Document Chunker character ${number}` | "Document Chunker token undefined" | `Document Chunker token ${number}` | "Document Chunker json undefined" | `Document Chunker json ${number}`, z.ZodObject<{}, "strip", z.ZodTypeAny, {}, {}>, undefined, _mastra_core.ToolExecutionContext<z.ZodObject<{}, "strip", z.ZodTypeAny, {}, {}>, _mastra_core.WorkflowContext<any>>>;
+}) => ReturnType<typeof createTool>;
 
-declare const createGraphRAGTool: ({ vectorStoreName, indexName, options, enableFilter, graphOptions, id, description, }: {
+declare const createGraphRAGTool: ({ vectorStoreName, indexName, model, enableFilter, graphOptions, id, description, }: {
     vectorStoreName: string;
     indexName: string;
-    options: EmbeddingOptions$1;
+    model: EmbeddingModel<string>;
     enableFilter?: boolean;
     graphOptions?: {
         dimension?: number;
@@ -229,77 +223,17 @@ declare const createGraphRAGTool: ({ vectorStoreName, indexName, options, enable
     };
     id?: string;
     description?: string;
-}) => _mastra_core_tools.Tool<string, z.ZodObject<{
-    queryText: z.ZodString;
-    topK: z.ZodNumber;
-    filter: z.ZodString;
-}, "strip", z.ZodTypeAny, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}>, z.ZodObject<{
-    relevantContext: z.ZodAny;
-}, "strip", z.ZodTypeAny, {
-    relevantContext?: any;
-}, {
-    relevantContext?: any;
-}>, _mastra_core.ToolExecutionContext<z.ZodObject<{
-    queryText: z.ZodString;
-    topK: z.ZodNumber;
-    filter: z.ZodString;
-}, "strip", z.ZodTypeAny, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}>, _mastra_core.WorkflowContext<any>>>;
+}) => ReturnType<typeof createTool>;
 
-declare const createVectorQueryTool: ({ vectorStoreName, indexName, options, enableFilter, reranker, id, description, }: {
+declare const createVectorQueryTool: ({ vectorStoreName, indexName, model, enableFilter, reranker, id, description, }: {
     vectorStoreName: string;
     indexName: string;
-    options: EmbeddingOptions$1;
+    model: EmbeddingModel<string>;
     enableFilter?: boolean;
     reranker?: RerankConfig;
     id?: string;
     description?: string;
-}) => _mastra_core_tools.Tool<string, z.ZodObject<{
-    queryText: z.ZodString;
-    topK: z.ZodNumber;
-    filter: z.ZodString;
-}, "strip", z.ZodTypeAny, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}>, z.ZodObject<{
-    relevantContext: z.ZodAny;
-}, "strip", z.ZodTypeAny, {
-    relevantContext?: any;
-}, {
-    relevantContext?: any;
-}>, _mastra_core.ToolExecutionContext<z.ZodObject<{
-    queryText: z.ZodString;
-    topK: z.ZodNumber;
-    filter: z.ZodString;
-}, "strip", z.ZodTypeAny, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}, {
-    filter: string;
-    topK: number;
-    queryText: string;
-}>, _mastra_core.WorkflowContext<any>>>;
+}) => ReturnType<typeof createTool>;
 
 /**
  * Vector store specific prompts that detail supported operators and examples.
@@ -319,4 +253,4 @@ declare const defaultFilter = "You MUST generate for each query:\n    filter: qu
 declare const defaultVectorQueryDescription: (vectorStoreName: string, indexName: string) => string;
 declare const defaultGraphRagDescription: (vectorStoreName: string, indexName: string) => string;
 
-export { ASTRA_PROMPT, CHROMA_PROMPT, GraphRAG, LIBSQL_PROMPT, MDocument, PGVECTOR_PROMPT, PINECONE_PROMPT, QDRANT_PROMPT, type RerankConfig, type RerankResult, type RerankerFunctionOptions, type RerankerOptions, UPSTASH_PROMPT, VECTORIZE_PROMPT, createDocumentChunkerTool, createGraphRAGTool, createVectorQueryTool, defaultFilter, defaultGraphRagDescription, defaultTopK, defaultVectorQueryDescription, embed, embedMany, rerank };
+export { ASTRA_PROMPT, CHROMA_PROMPT, GraphRAG, LIBSQL_PROMPT, MDocument, PGVECTOR_PROMPT, PINECONE_PROMPT, QDRANT_PROMPT, type RerankConfig, type RerankResult, type RerankerFunctionOptions, type RerankerOptions, UPSTASH_PROMPT, VECTORIZE_PROMPT, createDocumentChunkerTool, createGraphRAGTool, createVectorQueryTool, defaultFilter, defaultGraphRagDescription, defaultTopK, defaultVectorQueryDescription, rerank };

package/dist/index.js

@@ -1,10 +1,10 @@
 import { Document, SummaryExtractor, QuestionsAnsweredExtractor, KeywordExtractor, TitleExtractor, IngestionPipeline } from 'llamaindex';
 import { parse } from 'node-html-better-parser';
 import { encodingForModel, getEncoding } from 'js-tiktoken';
-import { embed as embed$1, embedMany as embedMany$1 } from '@mastra/core/embeddings';
-import { MastraAgentRelevanceScorer, CohereRelevanceScorer } from '@mastra/core/relevance';
+import { CohereRelevanceScorer, MastraAgentRelevanceScorer } from '@mastra/core/relevance';
 import { createTool } from '@mastra/core/tools';
 import { z } from 'zod';
+import { embed } from 'ai';
 
 // src/document/document.ts
 
@@ -1272,15 +1272,6 @@ var MDocument = class _MDocument {
     return this.chunks.map((doc) => doc.metadata);
   }
 };
-function getText(input) {
-  return input instanceof Document ? input.getText() : input;
-}
-function embed(chunk, options) {
-  return embed$1(getText(chunk), options);
-}
-function embedMany(chunks, options) {
-  return embedMany$1(chunks.map(getText), options);
-}
 var DEFAULT_WEIGHTS = {
   semantic: 0.4,
   vector: 0.4,
@@ -1299,17 +1290,12 @@ function adjustScores(score, queryAnalysis) {
   const featureStrengthAdjustment = queryAnalysis.magnitude > 5 ? 1.05 : 1;
   return score * magnitudeAdjustment * featureStrengthAdjustment;
 }
-async function rerank(results, query, modelConfig, options) {
-  const { provider } = modelConfig;
+async function rerank(results, query, model, options) {
   let semanticProvider;
-  if ("model" in modelConfig) {
-    semanticProvider = new MastraAgentRelevanceScorer(provider, "CUSTOM_MODEL", modelConfig.model);
-  } else if (provider === "COHERE" && "name" in modelConfig && modelConfig.name === "rerank-v3.5") {
-    semanticProvider = new CohereRelevanceScorer(modelConfig.name, modelConfig.apiKey);
-  } else if ("name" in modelConfig) {
-    semanticProvider = new MastraAgentRelevanceScorer(provider, modelConfig.name);
+  if (model.modelId === "rerank-v3.5") {
+    semanticProvider = new CohereRelevanceScorer(model.modelId);
   } else {
-    throw new Error("Invalid model configuration");
+    semanticProvider = new MastraAgentRelevanceScorer(model.provider, model);
   }
   const { queryEmbedding, topK = 3 } = options;
   const weights = {
@@ -1571,18 +1557,21 @@ var createDocumentChunkerTool = ({
     }
   });
 };
-
-// src/utils/vector-search.ts
 var vectorQuerySearch = async ({
   indexName,
   vectorStore,
   queryText,
-  options,
+  model,
   queryFilter = {},
   topK,
-  includeVectors = false
+  includeVectors = false,
+  maxRetries = 2
 }) => {
-  const { embedding } = await embed(queryText, options);
+  const { embedding } = await embed({
+    value: queryText,
+    model,
+    maxRetries
+  });
   const results = await vectorStore.query(indexName, embedding, topK, queryFilter, includeVectors);
   return { results, queryEmbedding: embedding };
 };
@@ -1621,7 +1610,7 @@ var defaultGraphRagDescription = (vectorStoreName, indexName) => `Fetches and re
 var createGraphRAGTool = ({
   vectorStoreName,
   indexName,
-  options,
+  model,
   enableFilter = false,
   graphOptions = {
     dimension: 1536,
@@ -1667,7 +1656,7 @@ var createGraphRAGTool = ({
           indexName,
           vectorStore,
           queryText,
-          options,
+          model,
           queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
           topK,
           includeVectors: true
@@ -1703,7 +1692,7 @@ var createGraphRAGTool = ({
 var createVectorQueryTool = ({
   vectorStoreName,
   indexName,
-  options,
+  model,
   enableFilter = false,
   reranker,
   id,
@@ -1742,7 +1731,7 @@ var createVectorQueryTool = ({
           indexName,
           vectorStore,
           queryText,
-          options,
+          model,
           queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
           topK
         });
@@ -2485,4 +2474,4 @@ Example Complex Query:
   "inStock": true
 }`;
 
-export { ASTRA_PROMPT, CHROMA_PROMPT, GraphRAG, LIBSQL_PROMPT, MDocument, PGVECTOR_PROMPT, PINECONE_PROMPT, QDRANT_PROMPT, UPSTASH_PROMPT, VECTORIZE_PROMPT, createDocumentChunkerTool, createGraphRAGTool, createVectorQueryTool, defaultFilter, defaultGraphRagDescription, defaultTopK, defaultVectorQueryDescription, embed, embedMany, rerank };
+export { ASTRA_PROMPT, CHROMA_PROMPT, GraphRAG, LIBSQL_PROMPT, MDocument, PGVECTOR_PROMPT, PINECONE_PROMPT, QDRANT_PROMPT, UPSTASH_PROMPT, VECTORIZE_PROMPT, createDocumentChunkerTool, createGraphRAGTool, createVectorQueryTool, defaultFilter, defaultGraphRagDescription, defaultTopK, defaultVectorQueryDescription, rerank };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/rag",
-  "version": "0.1.0-alpha.86",
+  "version": "0.1.0-alpha.87",
   "description": "",
   "type": "module",
   "main": "dist/index.js",
@@ -25,13 +25,19 @@
     "node-html-better-parser": "^1.4.7",
     "pathe": "^2.0.2",
     "zod": "^3.24.1",
-    "@mastra/core": "^0.2.0-alpha.93"
+    "@mastra/core": "^0.2.0-alpha.94"
+  },
+  "peerDependencies": {
+    "ai": "^4.0.0"
   },
   "devDependencies": {
+    "@ai-sdk/openai": "latest",
+    "@ai-sdk/cohere": "latest",
     "@babel/preset-env": "^7.26.0",
     "@babel/preset-typescript": "^7.26.0",
     "@tsconfig/recommended": "^1.0.7",
     "@types/node": "^22.9.0",
+    "ai": "^4.0.34",
     "tsup": "^8.0.1",
     "vitest": "^3.0.4"
   },

package/src/document/document.test.ts

@@ -1,4 +1,6 @@
-import { embedMany } from '../embeddings';
+import { createOpenAI } from '@ai-sdk/openai';
+import { embedMany } from 'ai';
+import { describe, it, expect } from 'vitest';
 
 import { MDocument } from './document';
 import { Language } from './types';
@@ -14,6 +16,10 @@ Welcome to our comprehensive guide on modern web development. This resource cove
 - Senior developers seeking a refresher on current best practices
 `;
 
+const openai = createOpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
 describe('MDocument', () => {
   describe('basics', () => {
     let chunks: MDocument['chunks'];
@@ -48,10 +54,9 @@ describe('MDocument', () => {
     }, 15000);
 
     it('embed - create embedding from chunk', async () => {
-      const embeddings = await embedMany(chunks, {
-        provider: 'OPEN_AI',
-        model: 'text-embedding-3-small',
-        maxRetries: 3,
+      const embeddings = await embedMany({
+        values: chunks.map(chunk => chunk.text),
+        model: openai.embedding('text-embedding-3-small'),
       });
 
       console.log(embeddings);

package/src/index.ts

@@ -1,5 +1,4 @@
 export * from './document/document';
-export * from './embeddings';
 export * from './rerank';
 export { GraphRAG } from './graph-rag';
 export * from './tools';

package/src/rerank/index.test.ts

@@ -1,3 +1,4 @@
+import { cohere } from '@ai-sdk/cohere';
 import { CohereRelevanceScorer } from '@mastra/core/relevance';
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 
@@ -25,12 +26,7 @@ describe('rerank', () => {
       { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
     ];
     await expect(
-      rerank(
-        results,
-        'test query',
-        { provider: 'COHERE', name: 'rerank-v3.5' },
-        { weights: { semantic: 0.5, vector: 0.3, position: 0.5 } },
-      ),
+      rerank(results, 'test query', cohere('rerank-v3.5'), { weights: { semantic: 0.5, vector: 0.3, position: 0.5 } }),
     ).rejects.toThrow('Weights must add up to 1');
   });
 
@@ -41,15 +37,7 @@ describe('rerank', () => {
       { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
     ];
 
-    const rerankedResults = await rerank(
-      results,
-      'test query',
-      {
-        provider: 'COHERE',
-        name: 'rerank-v3.5',
-      },
-      { topK: 2 },
-    );
+    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), { topK: 2 });
 
     expect(rerankedResults).toHaveLength(2);
     expect(rerankedResults[0]).toStrictEqual({
@@ -83,22 +71,14 @@ describe('rerank', () => {
       { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
     ];
 
-    const rerankedResults = await rerank(
-      results,
-      'test query',
-      {
-        provider: 'COHERE',
-        name: 'rerank-v3.5',
-      },
-      {
-        weights: {
-          semantic: 0.5,
-          vector: 0.4,
-          position: 0.1,
-        },
-        topK: 2,
+    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), {
+      weights: {
+        semantic: 0.5,
+        vector: 0.4,
+        position: 0.1,
       },
-    );
+      topK: 2,
+    });
 
     expect(rerankedResults).toHaveLength(2);
     expect(rerankedResults[0]).toStrictEqual({
@@ -130,18 +110,10 @@ describe('rerank', () => {
       { id: '2', metadata: { text: 'Test result 2' }, score: 0.6 },
     ];
 
-    const rerankedResults = await rerank(
-      results,
-      'test query',
-      {
-        provider: 'COHERE',
-        name: 'rerank-v3.5',
-      },
-      {
-        queryEmbedding: [0.5, 0.3, -0.2, 0.4],
-        topK: 2,
-      },
-    );
+    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), {
+      queryEmbedding: [0.5, 0.3, -0.2, 0.4],
+      topK: 2,
+    });
 
     // Ensure query embedding analysis is being applied (we don't know exact score without knowing internals)
     expect(rerankedResults).toHaveLength(2);

package/src/rerank/index.ts

@@ -1,6 +1,6 @@
-import { type ModelConfig } from '@mastra/core';
 import { MastraAgentRelevanceScorer, CohereRelevanceScorer, RelevanceScoreProvider } from '@mastra/core/relevance';
 import { QueryResult } from '@mastra/core/vector';
+import { LanguageModelV1 } from 'ai';
 
 // Default weights for different scoring components (must add up to 1)
 const DEFAULT_WEIGHTS = {
@@ -46,7 +46,7 @@ export interface RerankerFunctionOptions {
 
 export interface RerankConfig {
   options?: RerankerOptions;
-  model: ModelConfig;
+  model: LanguageModelV1;
 }
 
 // Calculate position score based on position in original list
@@ -85,19 +85,14 @@ function adjustScores(score: number, queryAnalysis: { magnitude: number; dominan
 export async function rerank(
   results: QueryResult[],
   query: string,
-  modelConfig: ModelConfig,
+  model: LanguageModelV1,
   options: RerankerFunctionOptions,
 ): Promise<RerankResult[]> {
-  const { provider } = modelConfig;
   let semanticProvider: RelevanceScoreProvider;
-  if ('model' in modelConfig) {
-    semanticProvider = new MastraAgentRelevanceScorer(provider, 'CUSTOM_MODEL', modelConfig.model);
-  } else if (provider === 'COHERE' && 'name' in modelConfig && modelConfig.name === 'rerank-v3.5') {
-    semanticProvider = new CohereRelevanceScorer(modelConfig.name, modelConfig.apiKey);
-  } else if ('name' in modelConfig) {
-    semanticProvider = new MastraAgentRelevanceScorer(provider, modelConfig.name);
+  if (model.modelId === 'rerank-v3.5') {
+    semanticProvider = new CohereRelevanceScorer(model.modelId);
   } else {
-    throw new Error('Invalid model configuration');
+    semanticProvider = new MastraAgentRelevanceScorer(model.provider, model);
   }
   const { queryEmbedding, topK = 3 } = options;
   const weights = {

package/src/tools/document-chunker.ts

@@ -14,7 +14,7 @@ export const createDocumentChunkerTool = ({
 }: {
   doc: MDocument;
   params?: ChunkParams;
-}) => {
+}): ReturnType<typeof createTool> => {
   return createTool({
     id: `Document Chunker ${params.strategy} ${params.size}`,
     inputSchema: z.object({}),

package/src/tools/graph-rag.ts

@@ -1,5 +1,5 @@
-import { type EmbeddingOptions } from '@mastra/core/embeddings';
 import { createTool } from '@mastra/core/tools';
+import { EmbeddingModel } from 'ai';
 import { z } from 'zod';
 
 import { GraphRAG } from '../graph-rag';
@@ -8,7 +8,7 @@ import { vectorQuerySearch, defaultGraphRagDescription } from '../utils';
 export const createGraphRAGTool = ({
   vectorStoreName,
   indexName,
-  options,
+  model,
   enableFilter = false,
   graphOptions = {
     dimension: 1536,
@@ -21,7 +21,7 @@ export const createGraphRAGTool = ({
 }: {
   vectorStoreName: string;
   indexName: string;
-  options: EmbeddingOptions;
+  model: EmbeddingModel<string>;
   enableFilter?: boolean;
   graphOptions?: {
     dimension?: number;
@@ -31,7 +31,7 @@ export const createGraphRAGTool = ({
   };
   id?: string;
   description?: string;
-}) => {
+}): ReturnType<typeof createTool> => {
   const toolId = id || `GraphRAG ${vectorStoreName} ${indexName} Tool`;
   const toolDescription = description || defaultGraphRagDescription(vectorStoreName, indexName);
 
@@ -73,7 +73,7 @@ export const createGraphRAGTool = ({
           indexName,
           vectorStore,
           queryText,
-          options,
+          model,
           queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
           topK,
           includeVectors: true,

package/src/tools/vector-query.ts

@@ -1,5 +1,5 @@
-import { type EmbeddingOptions } from '@mastra/core/embeddings';
 import { createTool } from '@mastra/core/tools';
+import { EmbeddingModel } from 'ai';
 import { z } from 'zod';
 
 import { rerank, RerankConfig } from '../rerank';
@@ -8,7 +8,7 @@ import { vectorQuerySearch, defaultVectorQueryDescription } from '../utils';
 export const createVectorQueryTool = ({
   vectorStoreName,
   indexName,
-  options,
+  model,
   enableFilter = false,
   reranker,
   id,
@@ -16,12 +16,12 @@ export const createVectorQueryTool = ({
 }: {
   vectorStoreName: string;
   indexName: string;
-  options: EmbeddingOptions;
+  model: EmbeddingModel<string>;
   enableFilter?: boolean;
   reranker?: RerankConfig;
   id?: string;
   description?: string;
-}) => {
+}): ReturnType<typeof createTool> => {
   const toolId = id || `VectorQuery ${vectorStoreName} ${indexName} Tool`;
   const toolDescription = description || defaultVectorQueryDescription(vectorStoreName, indexName);
 
@@ -61,7 +61,7 @@ export const createVectorQueryTool = ({
           indexName,
           vectorStore,
           queryText,
-          options,
+          model,
           queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
           topK,
         });

package/src/utils/vector-search.ts

@@ -1,16 +1,15 @@
-import { type EmbeddingOptions } from '@mastra/core/embeddings';
 import { type MastraVector, type QueryResult } from '@mastra/core/vector';
-
-import { embed } from '../embeddings';
+import { embed, EmbeddingModel } from 'ai';
 
 interface VectorQuerySearchParams {
   indexName: string;
   vectorStore: MastraVector;
   queryText: string;
-  options: EmbeddingOptions;
+  model: EmbeddingModel<string>;
   queryFilter?: any;
   topK: number;
   includeVectors?: boolean;
+  maxRetries?: number;
 }
 
 interface VectorQuerySearchResult {
@@ -23,12 +22,17 @@ export const vectorQuerySearch = async ({
   indexName,
   vectorStore,
   queryText,
-  options,
+  model,
   queryFilter = {},
   topK,
   includeVectors = false,
+  maxRetries = 2,
 }: VectorQuerySearchParams): Promise<VectorQuerySearchResult> => {
-  const { embedding } = await embed(queryText, options);
+  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);
 

package/src/embeddings/index.ts

@@ -1,17 +0,0 @@
-import { type EmbedManyResult, type EmbedResult, type EmbeddingOptions } from '@mastra/core';
-import { embed as embedCore, embedMany as embedManyCore } from '@mastra/core/embeddings';
-import { Document as Chunk } from 'llamaindex';
-
-function getText(input: Chunk | string): string {
-  return input instanceof Chunk ? input.getText() : input;
-}
-
-// Added explicit return type as it was not being inferred correctly
-export function embed(chunk: Chunk | string, options: EmbeddingOptions): Promise<EmbedResult<string>> {
-  return embedCore(getText(chunk), options);
-}
-
-// Added explicit return type as it was not being inferred correctly
-export function embedMany(chunks: (Chunk | string)[], options: EmbeddingOptions): Promise<EmbedManyResult<string>> {
-  return embedManyCore(chunks.map(getText), options);
-}