@mastra/rag 1.2.2 → 1.2.3-alpha.1

package/src/tools/vector-query.ts

@@ -1,169 +0,0 @@
-import { createTool } from '@mastra/core/tools';
-import type { MastraVector } from '@mastra/core/vector';
-import type { EmbeddingModel } from 'ai';
-import { z } from 'zod';
-
-import { rerank, rerankWithScorer } from '../rerank';
-import type { RerankConfig, RerankResult } from '../rerank';
-import { vectorQuerySearch, defaultVectorQueryDescription, filterSchema, outputSchema, baseSchema } from '../utils';
-import type { RagTool } from '../utils';
-import { convertToSources } from '../utils/convert-sources';
-import type { VectorQueryToolOptions } from './types';
-
-export const createVectorQueryTool = (options: VectorQueryToolOptions) => {
-  const { id, description } = options;
-  const storeName = options['vectorStoreName'] ? options.vectorStoreName : 'DirectVectorStore';
-
-  const toolId = id || `VectorQuery ${storeName} ${options.indexName} Tool`;
-  const toolDescription = description || defaultVectorQueryDescription();
-  const inputSchema = options.enableFilter ? filterSchema : z.object(baseSchema).passthrough();
-
-  return createTool({
-    id: toolId,
-    description: toolDescription,
-    inputSchema,
-    outputSchema,
-    execute: async ({ context, mastra, runtimeContext }) => {
-      const indexName: string = runtimeContext.get('indexName') ?? options.indexName;
-      const vectorStoreName: string =
-        'vectorStore' in options ? storeName : (runtimeContext.get('vectorStoreName') ?? storeName);
-      const includeVectors: boolean = runtimeContext.get('includeVectors') ?? options.includeVectors ?? false;
-      const includeSources: boolean = runtimeContext.get('includeSources') ?? options.includeSources ?? true;
-      const reranker: RerankConfig = runtimeContext.get('reranker') ?? options.reranker;
-      const databaseConfig = runtimeContext.get('databaseConfig') ?? options.databaseConfig;
-      const model: EmbeddingModel<string> = runtimeContext.get('model') ?? options.model;
-
-      if (!indexName) throw new Error(`indexName is required, got: ${indexName}`);
-      if (!vectorStoreName) throw new Error(`vectorStoreName is required, got: ${vectorStoreName}`); // won't fire
-
-      const topK: number = runtimeContext.get('topK') ?? context.topK ?? 10;
-      const filter: Record<string, any> = runtimeContext.get('filter') ?? context.filter;
-      const queryText = context.queryText;
-      const enableFilter = !!runtimeContext.get('filter') || (options.enableFilter ?? false);
-
-      const logger = mastra?.getLogger();
-      if (!logger) {
-        console.warn(
-          '[VectorQueryTool] Logger not initialized: no debug or error logs will be recorded for this tool execution.',
-        );
-      }
-      if (logger) {
-        logger.debug('[VectorQueryTool] execute called with:', { queryText, topK, filter, databaseConfig });
-      }
-      try {
-        const topKValue =
-          typeof topK === 'number' && !isNaN(topK)
-            ? topK
-            : typeof topK === 'string' && !isNaN(Number(topK))
-              ? Number(topK)
-              : 10;
-
-        let vectorStore: MastraVector | undefined = undefined;
-        if ('vectorStore' in options) {
-          vectorStore = options.vectorStore;
-        } else if (mastra) {
-          vectorStore = mastra.getVector(vectorStoreName);
-        }
-        if (!vectorStore) {
-          if (logger) {
-            logger.error('Vector store not found', { vectorStoreName });
-          }
-          return { relevantContext: [], sources: [] };
-        }
-        // Get relevant chunks from the vector database
-        let queryFilter = {};
-        if (enableFilter && filter) {
-          queryFilter = (() => {
-            try {
-              return typeof filter === 'string' ? JSON.parse(filter) : filter;
-            } catch (error) {
-              // Log the error and use empty object
-              if (logger) {
-                logger.warn('Failed to parse filter as JSON, using empty filter', { filter, error });
-              }
-              return {};
-            }
-          })();
-        }
-        if (logger) {
-          logger.debug('Prepared vector query parameters', { queryText, topK: topKValue, queryFilter, databaseConfig });
-        }
-
-        const { results } = await vectorQuerySearch({
-          indexName,
-          vectorStore,
-          queryText,
-          model,
-          queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
-          topK: topKValue,
-          includeVectors,
-          databaseConfig,
-        });
-        if (logger) {
-          logger.debug('vectorQuerySearch returned results', { count: results.length });
-        }
-
-        if (reranker) {
-          if (logger) {
-            logger.debug('Reranking results', { rerankerModel: reranker.model, rerankerOptions: reranker.options });
-          }
-
-          let rerankedResults: RerankResult[] = [];
-
-          if (typeof reranker?.model === 'object' && 'getRelevanceScore' in reranker?.model) {
-            rerankedResults = await rerankWithScorer({
-              results,
-              query: queryText,
-              scorer: reranker.model,
-              options: {
-                ...reranker.options,
-                topK: reranker.options?.topK || topKValue,
-              },
-            });
-          } else {
-            rerankedResults = await rerank(results, queryText, reranker.model, {
-              ...reranker.options,
-              topK: reranker.options?.topK || topKValue,
-            });
-          }
-
-          if (logger) {
-            logger.debug('Reranking complete', { rerankedCount: rerankedResults.length });
-          }
-
-          const relevantChunks = rerankedResults.map(({ result }) => result?.metadata);
-
-          if (logger) {
-            logger.debug('Returning reranked relevant context chunks', { count: relevantChunks.length });
-          }
-
-          const sources = includeSources ? convertToSources(rerankedResults) : [];
-
-          return { relevantContext: relevantChunks, sources };
-        }
-
-        const relevantChunks = results.map(result => result?.metadata);
-
-        if (logger) {
-          logger.debug('Returning relevant context chunks', { count: relevantChunks.length });
-        }
-        // `sources` exposes the full retrieval objects
-        const sources = includeSources ? convertToSources(results) : [];
-        return {
-          relevantContext: relevantChunks,
-          sources,
-        };
-      } catch (err) {
-        if (logger) {
-          logger.error('Unexpected error in VectorQueryTool execute', {
-            error: err,
-            errorMessage: err instanceof Error ? err.message : String(err),
-            errorStack: err instanceof Error ? err.stack : undefined,
-          });
-        }
-        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/convert-sources.ts

@@ -1,43 +0,0 @@
-import type { QueryResult } from '@mastra/core/vector';
-import type { RankedNode } from '../graph-rag';
-import type { RerankResult } from '../rerank';
-
-type SourceInput = QueryResult | RankedNode | RerankResult;
-
-/**
- * Convert an array of source inputs (QueryResult, RankedNode, or RerankResult) to an array of sources.
- * @param results Array of source inputs to convert.
- * @returns Array of sources.
- */
-export const convertToSources = (results: SourceInput[]) => {
-  return results.map(result => {
-    // RankedNode
-    if ('content' in result) {
-      return {
-        id: result.id,
-        vector: result.embedding || [],
-        score: result.score,
-        metadata: result.metadata,
-        document: result.content || '',
-      };
-    }
-    // RerankResult
-    if ('result' in result) {
-      return {
-        id: result.result.id,
-        vector: result.result.vector || [],
-        score: result.score,
-        metadata: result.result.metadata,
-        document: result.result.document || '',
-      };
-    }
-    // QueryResult
-    return {
-      id: result.id,
-      vector: result.vector || [],
-      score: result.score,
-      metadata: result.metadata,
-      document: result.document || '',
-    };
-  });
-};

package/src/utils/default-settings.ts

@@ -1,38 +0,0 @@
-export const defaultVectorQueryDescription = () =>
-  `Access the knowledge base to find information needed to answer user questions.`;
-
-export const defaultGraphRagDescription = () =>
-  `Access and analyze relationships between information in the knowledge base to answer complex questions about connections and patterns.`;
-
-export const queryTextDescription = `The text query to search for in the vector database.
-- ALWAYS provide a non-empty query string
-- Must contain the user's question or search terms
-- Example: "market data" or "financial reports"
-- If the user's query is about a specific topic, use that topic as the queryText
-- Cannot be an empty string
-- Do not include quotes, just the text itself
-- Required for all searches`;
-
-export const topKDescription = `Controls how many matching documents to return.
-- ALWAYS provide a value
-- If no value is provided, use the default (10)
-- Must be a valid and positive number
-- Cannot be NaN
-- Uses provided value if specified
-- Default: 10 results (use this if unsure)
-- Higher values (like 20) provide more context
-- Lower values (like 3) focus on best matches
-- Based on query requirements`;
-
-export const filterDescription = `JSON-formatted criteria to refine search results.
-- ALWAYS provide a filter value
-- If no filter is provided, use the default ("{}")
-- MUST be a valid, complete JSON object with proper quotes and brackets
-- Uses provided filter if specified
-- Default: "{}" (no filtering)
-- Example for no filtering: "filter": "{}"
-- Example: '{"category": "health"}'
-- Based on query intent
-- Do NOT use single quotes or unquoted properties
-- IMPORTANT: Always ensure JSON is properly closed with matching brackets
-- Multiple filters can be combined`;

package/src/utils/index.ts

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

package/src/utils/tool-schemas.ts

@@ -1,38 +0,0 @@
-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']>;
-};