@mastra/rag 0.10.1-alpha.0 → 0.10.1-alpha.1

package/src/tools/graph-rag.ts

@@ -3,13 +3,8 @@ import type { EmbeddingModel } from 'ai';
 import { z } from 'zod';
 
 import { GraphRAG } from '../graph-rag';
-import {
-  vectorQuerySearch,
-  defaultGraphRagDescription,
-  filterDescription,
-  topKDescription,
-  queryTextDescription,
-} from '../utils';
+import { vectorQuerySearch, defaultGraphRagDescription, filterSchema, outputSchema, baseSchema } from '../utils';
+import type { RagTool } from '../utils';
 import { convertToSources } from '../utils/convert-sources';
 
 export const createGraphRAGTool = ({
@@ -47,38 +42,12 @@ export const createGraphRAGTool = ({
   const graphRag = new GraphRAG(graphOptions.dimension, graphOptions.threshold);
   let isInitialized = false;
 
-  const baseSchema = {
-    queryText: z.string().describe(queryTextDescription),
-    topK: z.coerce.number().describe(topKDescription),
-  };
-  const inputSchema = enableFilter
-    ? z
-        .object({
-          ...baseSchema,
-          filter: z.coerce.string().describe(filterDescription),
-        })
-        .passthrough()
-    : z.object(baseSchema).passthrough();
+  const inputSchema = enableFilter ? filterSchema : z.object(baseSchema).passthrough();
+
   return createTool({
     id: toolId,
     inputSchema,
-    // Output schema includes `sources`, which exposes the full set of retrieved chunks (QueryResult objects)
-    // Each source contains all information needed to reference
-    // the original document, chunk, and similarity score.
-    outputSchema: z.object({
-      // Array of metadata or content for compatibility with prior usage
-      relevantContext: z.any(),
-      // Array of full retrieval result objects
-      sources: z.array(
-        z.object({
-          id: z.string(), // Unique chunk/document identifier
-          metadata: z.any(), // All metadata fields (document ID, etc.)
-          vector: z.array(z.number()), // Embedding vector (if available)
-          score: z.number(), // Similarity score for this retrieval
-          document: z.string(), // Full chunk/document text (if available)
-        }),
-      ),
-    }),
+    outputSchema,
     description: toolDescription,
     execute: async ({ context: { queryText, topK, filter }, mastra }) => {
       const logger = mastra?.getLogger();
@@ -188,5 +157,6 @@ export const createGraphRAGTool = ({
         return { relevantContext: [], sources: [] };
       }
     },
-  });
+    // Use any for output schema as the structure of the output causes type inference issues
+  }) as RagTool<typeof inputSchema, any>;
 };

package/src/tools/vector-query.test.ts

@@ -2,23 +2,13 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { vectorQuerySearch } from '../utils';
 import { createVectorQueryTool } from './vector-query';
 
-// Mock dependencies
-vi.mock('@mastra/core/tools', () => ({
-  createTool: vi.fn(({ inputSchema, execute }) => ({
-    inputSchema,
-    execute,
-    // Return a simplified version of the tool for testing
-    __inputSchema: inputSchema,
-  })),
-}));
-
-vi.mock('../utils', () => ({
-  vectorQuerySearch: vi.fn().mockResolvedValue({ results: [] }),
-  defaultVectorQueryDescription: () => 'Default vector query description',
-  queryTextDescription: 'Query text description',
-  filterDescription: 'Filter description',
-  topKDescription: 'Top K description',
-}));
+vi.mock('../utils', async importOriginal => {
+  const actual: any = await importOriginal();
+  return {
+    ...actual,
+    vectorQuerySearch: vi.fn().mockResolvedValue({ results: [] }),
+  };
+});
 
 describe('createVectorQueryTool', () => {
   const mockModel = { name: 'test-model' } as any;
@@ -60,21 +50,21 @@ describe('createVectorQueryTool', () => {
       });
 
       // Get the Zod schema
-      const schema = tool.__inputSchema;
+      const schema = tool.inputSchema;
 
       // Test with no filter (should be valid)
       const validInput = {
         queryText: 'test query',
         topK: 5,
       };
-      expect(() => schema.parse(validInput)).not.toThrow();
+      expect(() => schema?.parse(validInput)).not.toThrow();
 
       // Test with filter (should throw - unexpected property)
       const inputWithFilter = {
         ...validInput,
         filter: '{"field": "value"}',
       };
-      expect(() => schema.parse(inputWithFilter)).not.toThrow();
+      expect(() => schema?.parse(inputWithFilter)).not.toThrow();
     });
 
     it('should handle filter when enableFilter is true', () => {
@@ -86,7 +76,7 @@ describe('createVectorQueryTool', () => {
       });
 
       // Get the Zod schema
-      const schema = tool.__inputSchema;
+      const schema = tool.inputSchema;
 
       // Test various filter inputs that should coerce to string
       const testCases = [
@@ -105,7 +95,7 @@ describe('createVectorQueryTool', () => {
 
       testCases.forEach(({ filter }) => {
         expect(() =>
-          schema.parse({
+          schema?.parse({
             queryText: 'test query',
             topK: 5,
             filter,
@@ -115,12 +105,12 @@ describe('createVectorQueryTool', () => {
 
       // Verify that all parsed values are strings
       testCases.forEach(({ filter }) => {
-        const result = schema.parse({
+        const result = schema?.parse({
           queryText: 'test query',
           topK: 5,
           filter,
         });
-        expect(typeof result.filter).toBe('string');
+        expect(typeof result?.filter).toBe('string');
       });
     });
 
@@ -135,7 +125,7 @@ describe('createVectorQueryTool', () => {
 
       // Should reject unexpected property
       expect(() =>
-        toolWithoutFilter.__inputSchema.parse({
+        toolWithoutFilter.inputSchema?.parse({
           queryText: 'test query',
           topK: 5,
           unexpectedProp: 'value',
@@ -152,7 +142,7 @@ describe('createVectorQueryTool', () => {
 
       // Should reject unexpected property even with valid filter
       expect(() =>
-        toolWithFilter.__inputSchema.parse({
+        toolWithFilter.inputSchema?.parse({
           queryText: 'test query',
           topK: 5,
           filter: '{}',
@@ -179,6 +169,7 @@ describe('createVectorQueryTool', () => {
           topK: 5,
         },
         mastra: mockMastra as any,
+        runtimeContext: {} as any,
       });
 
       // Check that vectorQuerySearch was called with undefined queryFilter
@@ -208,6 +199,7 @@ describe('createVectorQueryTool', () => {
           filter: filterJson,
         },
         mastra: mockMastra as any,
+        runtimeContext: {} as any,
       });
 
       // Check that vectorQuerySearch was called with the parsed filter
@@ -237,6 +229,7 @@ describe('createVectorQueryTool', () => {
           filter: stringFilter,
         },
         mastra: mockMastra as any,
+        runtimeContext: {} as any,
       });
 
       // Since this is not a valid filter, it should be ignored

package/src/tools/vector-query.ts

@@ -4,13 +4,8 @@ import { z } from 'zod';
 
 import { rerank } from '../rerank';
 import type { RerankConfig } from '../rerank';
-import {
-  vectorQuerySearch,
-  defaultVectorQueryDescription,
-  filterDescription,
-  topKDescription,
-  queryTextDescription,
-} from '../utils';
+import { vectorQuerySearch, defaultVectorQueryDescription, filterSchema, outputSchema, baseSchema } from '../utils';
+import type { RagTool } from '../utils';
 import { convertToSources } from '../utils/convert-sources';
 
 export const createVectorQueryTool = ({
@@ -36,39 +31,11 @@ export const createVectorQueryTool = ({
 }) => {
   const toolId = id || `VectorQuery ${vectorStoreName} ${indexName} Tool`;
   const toolDescription = description || defaultVectorQueryDescription();
-  // Create base schema with required fields
-  const baseSchema = {
-    queryText: z.string().describe(queryTextDescription),
-    topK: z.coerce.number().describe(topKDescription),
-  };
-  const inputSchema = enableFilter
-    ? z
-        .object({
-          ...baseSchema,
-          filter: z.coerce.string().describe(filterDescription),
-        })
-        .passthrough()
-    : z.object(baseSchema).passthrough();
+  const inputSchema = enableFilter ? filterSchema : z.object(baseSchema).passthrough();
   return createTool({
     id: toolId,
     inputSchema,
-    // Output schema includes `sources`, which exposes the full set of retrieved chunks (QueryResult objects)
-    // Each source contains all information needed to reference
-    // the original document, chunk, and similarity score.
-    outputSchema: z.object({
-      // Array of metadata or content for compatibility with prior usage
-      relevantContext: z.any(),
-      // Array of full retrieval result objects
-      sources: z.array(
-        z.object({
-          id: z.string(), // Unique chunk/document identifier
-          metadata: z.any(), // All metadata fields (document ID, etc.)
-          vector: z.array(z.number()), // Embedding vector (if available)
-          score: z.number(), // Similarity score for this retrieval
-          document: z.string(), // Full chunk/document text (if available)
-        }),
-      ),
-    }),
+    outputSchema,
     description: toolDescription,
     execute: async ({ context: { queryText, topK, filter }, mastra }) => {
       const logger = mastra?.getLogger();
@@ -167,5 +134,6 @@ export const createVectorQueryTool = ({
         return { relevantContext: [], sources: [] };
       }
     },
-  });
+    // Use any for output schema as the structure of the output causes type inference issues
+  }) as RagTool<typeof inputSchema, any>;
 };

package/src/utils/index.ts

@@ -1,2 +1,3 @@
 export * from './vector-search';
 export * from './default-settings';
+export * from './tool-schemas';

package/src/utils/tool-schemas.ts

@@ -0,0 +1,38 @@
+import type { Tool } from '@mastra/core/tools';
+import { z } from 'zod';
+import { queryTextDescription, topKDescription, filterDescription } from './default-settings';
+
+export const baseSchema = {
+  queryText: z.string().describe(queryTextDescription),
+  topK: z.coerce.number().describe(topKDescription),
+};
+
+// Output schema includes `sources`, which exposes the full set of retrieved chunks (QueryResult objects)
+// Each source contains all information needed to reference
+// the original document, chunk, and similarity score.
+export const outputSchema = z.object({
+  // Array of metadata or content for compatibility with prior usage
+  relevantContext: z.any(),
+  // Array of full retrieval result objects
+  sources: z.array(
+    z.object({
+      id: z.string(), // Unique chunk/document identifier
+      metadata: z.any(), // All metadata fields (document ID, etc.)
+      vector: z.array(z.number()), // Embedding vector (if available)
+      score: z.number(), // Similarity score for this retrieval
+      document: z.string(), // Full chunk/document text (if available)
+    }),
+  ),
+});
+
+export const filterSchema = z.object({
+  ...baseSchema,
+  filter: z.coerce.string().describe(filterDescription),
+});
+
+export type RagTool<
+  TInput extends z.ZodType<any, z.ZodTypeDef, any> | undefined,
+  TOutput extends z.ZodType<any, z.ZodTypeDef, any> | undefined,
+> = Tool<TInput, TOutput> & {
+  execute: NonNullable<Tool<TInput, TOutput>['execute']>;
+};