@mastra/rag 1.2.2 → 1.2.3-alpha.1

package/src/rerank/relevance/cohere/index.ts

@@ -1,56 +0,0 @@
-import type { RelevanceScoreProvider } from '@mastra/core/relevance';
-
-interface CohereRerankingResponse {
-  results: Array<{
-    index: number;
-    relevance_score: number;
-  }>;
-  id: string;
-  meta: {
-    api_version: {
-      version: string;
-      is_experimental: boolean;
-    };
-    billed_units: {
-      search_units: number;
-    };
-  };
-}
-
-export class CohereRelevanceScorer implements RelevanceScoreProvider {
-  private model: string;
-  private apiKey?: string;
-  constructor(model: string, apiKey?: string) {
-    this.apiKey = apiKey;
-    this.model = model;
-  }
-
-  async getRelevanceScore(query: string, text: string): Promise<number> {
-    const response = await fetch(`https://api.cohere.com/v2/rerank`, {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: `Bearer ${this.apiKey}`,
-      },
-      body: JSON.stringify({
-        query,
-        documents: [text],
-        model: this.model,
-        top_n: 1,
-      }),
-    });
-
-    if (!response.ok) {
-      throw new Error(`Cohere API error: ${response.status} ${await response.text()}`);
-    }
-
-    const data = (await response.json()) as CohereRerankingResponse;
-    const relevanceScore = data.results[0]?.relevance_score;
-
-    if (!relevanceScore) {
-      throw new Error('No relevance score found on Cohere response');
-    }
-
-    return relevanceScore;
-  }
-}

package/src/rerank/relevance/index.ts

@@ -1,3 +0,0 @@
-export * from './cohere';
-export * from './mastra-agent';
-export * from './zeroentropy';

package/src/rerank/relevance/mastra-agent/index.ts

@@ -1,32 +0,0 @@
-import { Agent } from '@mastra/core/agent';
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import { createSimilarityPrompt } from '@mastra/core/relevance';
-import type { RelevanceScoreProvider } from '@mastra/core/relevance';
-
-// Mastra Agent implementation
-export class MastraAgentRelevanceScorer implements RelevanceScoreProvider {
-  private agent: Agent;
-
-  constructor(name: string, model: MastraLanguageModel) {
-    this.agent = new Agent({
-      name: `Relevance Scorer ${name}`,
-      instructions: `You are a specialized agent for evaluating the relevance of text to queries.
-Your task is to rate how well a text passage answers a given query.
-Output only a number between 0 and 1, where:
-1.0 = Perfectly relevant, directly answers the query
-0.0 = Completely irrelevant
-Consider:
-- Direct relevance to the question
-- Completeness of information
-- Quality and specificity
-Always return just the number, no explanation.`,
-      model,
-    });
-  }
-
-  async getRelevanceScore(query: string, text: string): Promise<number> {
-    const prompt = createSimilarityPrompt(query, text);
-    const response = await this.agent.generate(prompt);
-    return parseFloat(response.text);
-  }
-}

package/src/rerank/relevance/zeroentropy/index.ts

@@ -1,26 +0,0 @@
-import type { RelevanceScoreProvider } from '@mastra/core/relevance';
-import ZeroEntropy from 'zeroentropy';
-
-// ZeroEntropy implementation
-export class ZeroEntropyRelevanceScorer implements RelevanceScoreProvider {
-  private client: ZeroEntropy;
-  private model: string;
-
-  constructor(model?: string, apiKey?: string) {
-    this.client = new ZeroEntropy({
-      apiKey: apiKey || process.env.ZEROENTROPY_API_KEY || '',
-    });
-    this.model = model || 'zerank-1';
-  }
-
-  async getRelevanceScore(query: string, text: string): Promise<number> {
-    const response = await this.client.models.rerank({
-      query,
-      documents: [text],
-      model: this.model,
-      top_n: 1,
-    });
-
-    return response.results[0]?.relevance_score ?? 0;
-  }
-}

package/src/tools/README.md

@@ -1,153 +0,0 @@
-# Vector Query Tool with Database-Specific Configurations
-
-The `createVectorQueryTool` function now supports database-specific configurations to handle unique properties and optimizations for different vector databases.
-
-## Database Configuration Types
-
-### Pinecone Configuration
-
-```typescript
-import { createVectorQueryTool } from '@mastra/rag/tools';
-
-const pineconeVectorTool = createVectorQueryTool({
-  id: 'pinecone-search',
-  indexName: 'my-index',
-  vectorStoreName: 'pinecone',
-  model: embedModel,
-  databaseConfig: {
-    pinecone: {
-      namespace: 'my-namespace', // Pinecone namespace
-      sparseVector: {
-        // For hybrid search
-        indices: [0, 1, 2],
-        values: [0.1, 0.2, 0.3],
-      },
-    },
-  },
-});
-```
-
-### pgVector Configuration
-
-```typescript
-const pgVectorTool = createVectorQueryTool({
-  id: 'pgvector-search',
-  indexName: 'my-index',
-  vectorStoreName: 'postgres',
-  model: embedModel,
-  databaseConfig: {
-    pgvector: {
-      minScore: 0.7, // Minimum similarity score
-      ef: 200, // HNSW search parameter
-      probes: 10, // IVFFlat probe parameter
-    },
-  },
-});
-```
-
-### Chroma Configuration
-
-```typescript
-const chromaTool = createVectorQueryTool({
-  id: 'chroma-search',
-  indexName: 'my-index',
-  vectorStoreName: 'chroma',
-  model: embedModel,
-  databaseConfig: {
-    chroma: {
-      where: {
-        // Metadata filtering
-        category: 'documents',
-      },
-      whereDocument: {
-        // Document content filtering
-        $contains: 'important',
-      },
-    },
-  },
-});
-```
-
-## Runtime Configuration Override
-
-You can also override database configurations at runtime using the runtime context:
-
-```typescript
-import { RuntimeContext } from '@mastra/core/runtime-context';
-
-const runtimeContext = new RuntimeContext();
-
-// Override Pinecone namespace at runtime
-runtimeContext.set('databaseConfig', {
-  pinecone: {
-    namespace: 'runtime-namespace',
-  },
-});
-
-await vectorTool.execute({
-  context: { queryText: 'search query' },
-  mastra,
-  runtimeContext,
-});
-```
-
-## Extensibility for New Databases
-
-The system is designed to be extensible. For new vector databases, you can:
-
-1. Add configuration types:
-
-```typescript
-export interface NewDatabaseConfig {
-  customParam1?: string;
-  customParam2?: number;
-}
-
-export type DatabaseConfig = {
-  pinecone?: PineconeConfig;
-  pgvector?: PgVectorConfig;
-  chroma?: ChromaConfig;
-  newdatabase?: NewDatabaseConfig; // Add your config here
-  [key: string]: any;
-};
-```
-
-2. The configuration will be automatically passed through to the vector store's query method.
-
-## Type Safety
-
-All database configurations are fully typed, providing IntelliSense and compile-time checking:
-
-```typescript
-const config: DatabaseConfig = {
-  pinecone: {
-    namespace: 'valid-namespace',
-    sparseVector: {
-      indices: [1, 2, 3],
-      values: [0.1, 0.2, 0.3],
-    },
-  },
-  pgvector: {
-    minScore: 0.8,
-    ef: 100,
-    probes: 5,
-  },
-};
-```
-
-## Migration Guide
-
-Existing code will continue to work without changes. To add database-specific configurations:
-
-```diff
-const vectorTool = createVectorQueryTool({
-  indexName: 'my-index',
-  vectorStoreName: 'pinecone',
-  model: embedModel,
-+ databaseConfig: {
-+   pinecone: {
-+     namespace: 'my-namespace'
-+   }
-+ }
-});
-```

package/src/tools/document-chunker.ts

@@ -1,34 +0,0 @@
-import { createTool } from '@mastra/core/tools';
-import { z } from 'zod';
-
-import type { MDocument, ChunkParams } from '../document';
-
-const DEFAULT_CHUNK_PARAMS = {
-  strategy: 'recursive' as const,
-  maxSize: 512,
-  overlap: 50,
-  separators: ['\n'],
-} satisfies ChunkParams;
-
-export const createDocumentChunkerTool = ({
-  doc,
-  params = DEFAULT_CHUNK_PARAMS,
-}: {
-  doc: MDocument;
-  params?: ChunkParams;
-}): ReturnType<typeof createTool> => {
-  return createTool({
-    id: `Document Chunker ${params.strategy} ${params.maxSize}`,
-    inputSchema: z.object({}),
-    description: `Chunks document using ${params.strategy} strategy with maxSize ${params.maxSize} and ${
-      params.overlap || 0
-    } overlap`,
-    execute: async () => {
-      const chunks = await doc.chunk(params);
-
-      return {
-        chunks,
-      };
-    },
-  });
-};

package/src/tools/graph-rag.test.ts

@@ -1,115 +0,0 @@
-import { RuntimeContext } from '@mastra/core/runtime-context';
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { GraphRAG } from '../graph-rag';
-import { vectorQuerySearch } from '../utils';
-import { createGraphRAGTool } from './graph-rag';
-
-vi.mock('../utils', async importOriginal => {
-  const actual: any = await importOriginal();
-  return {
-    ...actual,
-    vectorQuerySearch: vi.fn().mockResolvedValue({
-      results: [
-        { metadata: { text: 'foo' }, vector: [1, 2, 3] },
-        { metadata: { text: 'bar' }, vector: [4, 5, 6] },
-      ],
-      queryEmbedding: [1, 2, 3],
-    }),
-  };
-});
-
-vi.mock('../graph-rag', async importOriginal => {
-  const actual: any = await importOriginal();
-  return {
-    ...actual,
-    GraphRAG: vi.fn().mockImplementation(() => {
-      return {
-        createGraph: vi.fn(),
-        query: vi.fn(() => [
-          { content: 'foo', metadata: { text: 'foo' } },
-          { content: 'bar', metadata: { text: 'bar' } },
-        ]),
-      };
-    }),
-  };
-});
-
-const mockModel = { name: 'test-model' } as any;
-const mockMastra = {
-  getVector: vi.fn(storeName => ({
-    [storeName]: {},
-  })),
-  getLogger: vi.fn(() => ({
-    debug: vi.fn(),
-    warn: vi.fn(),
-    info: vi.fn(),
-    error: vi.fn(),
-  })),
-};
-
-describe('createGraphRAGTool', () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
-  it('validates input schema', () => {
-    const tool = createGraphRAGTool({
-      id: 'test',
-      model: mockModel,
-      vectorStoreName: 'testStore',
-      indexName: 'testIndex',
-    });
-    expect(() => tool.inputSchema?.parse({ queryText: 'foo', topK: 10 })).not.toThrow();
-    expect(() => tool.inputSchema?.parse({})).toThrow();
-  });
-
-  describe('runtimeContext', () => {
-    it('calls vectorQuerySearch and GraphRAG with runtimeContext params', async () => {
-      const tool = createGraphRAGTool({
-        id: 'test',
-        model: mockModel,
-        indexName: 'testIndex',
-        vectorStoreName: 'testStore',
-      });
-      const runtimeContext = new RuntimeContext();
-      runtimeContext.set('indexName', 'anotherIndex');
-      runtimeContext.set('vectorStoreName', 'anotherStore');
-      runtimeContext.set('topK', 5);
-      runtimeContext.set('filter', { foo: 'bar' });
-      runtimeContext.set('randomWalkSteps', 99);
-      runtimeContext.set('restartProb', 0.42);
-      const result = await tool.execute({
-        context: { queryText: 'foo', topK: 2 },
-        mastra: mockMastra as any,
-        runtimeContext,
-      });
-      expect(result.relevantContext).toEqual(['foo', 'bar']);
-      expect(result.sources.length).toBe(2);
-      expect(vectorQuerySearch).toHaveBeenCalledWith(
-        expect.objectContaining({
-          indexName: 'anotherIndex',
-          vectorStore: {
-            anotherStore: {},
-          },
-          queryText: 'foo',
-          model: mockModel,
-          queryFilter: { foo: 'bar' },
-          topK: 5,
-          includeVectors: true,
-        }),
-      );
-      // GraphRAG createGraph and query should be called
-      expect(GraphRAG).toHaveBeenCalled();
-      const instance = (GraphRAG as any).mock.results[0].value;
-      expect(instance.createGraph).toHaveBeenCalled();
-      expect(instance.query).toHaveBeenCalledWith(
-        expect.objectContaining({
-          query: [1, 2, 3],
-          topK: 5,
-          randomWalkSteps: 99,
-          restartProb: 0.42,
-        }),
-      );
-    });
-  });
-});

package/src/tools/graph-rag.ts

@@ -1,154 +0,0 @@
-import { createTool } from '@mastra/core/tools';
-import { z } from 'zod';
-
-import { GraphRAG } from '../graph-rag';
-import { vectorQuerySearch, defaultGraphRagDescription, filterSchema, outputSchema, baseSchema } from '../utils';
-import type { RagTool } from '../utils';
-import { convertToSources } from '../utils/convert-sources';
-import type { GraphRagToolOptions } from './types';
-import { defaultGraphOptions } from './types';
-
-export const createGraphRAGTool = (options: GraphRagToolOptions) => {
-  const { model, id, description } = options;
-
-  const toolId = id || `GraphRAG ${options.vectorStoreName} ${options.indexName} Tool`;
-  const toolDescription = description || defaultGraphRagDescription();
-  const graphOptions = {
-    ...defaultGraphOptions,
-    ...(options.graphOptions || {}),
-  };
-  // Initialize GraphRAG
-  const graphRag = new GraphRAG(graphOptions.dimension, graphOptions.threshold);
-  let isInitialized = false;
-
-  const inputSchema = options.enableFilter ? filterSchema : z.object(baseSchema).passthrough();
-
-  return createTool({
-    id: toolId,
-    inputSchema,
-    outputSchema,
-    description: toolDescription,
-    execute: async ({ context, mastra, runtimeContext }) => {
-      const indexName: string = runtimeContext.get('indexName') ?? options.indexName;
-      const vectorStoreName: string = runtimeContext.get('vectorStoreName') ?? options.vectorStoreName;
-      if (!indexName) throw new Error(`indexName is required, got: ${indexName}`);
-      if (!vectorStoreName) throw new Error(`vectorStoreName is required, got: ${vectorStoreName}`);
-      const includeSources: boolean = runtimeContext.get('includeSources') ?? options.includeSources ?? true;
-      const randomWalkSteps: number | undefined = runtimeContext.get('randomWalkSteps') ?? graphOptions.randomWalkSteps;
-      const restartProb: number | undefined = runtimeContext.get('restartProb') ?? graphOptions.restartProb;
-      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(
-          '[GraphRAGTool] Logger not initialized: no debug or error logs will be recorded for this tool execution.',
-        );
-      }
-      if (logger) {
-        logger.debug('[GraphRAGTool] execute called with:', { queryText, topK, filter });
-      }
-      try {
-        const topKValue =
-          typeof topK === 'number' && !isNaN(topK)
-            ? topK
-            : typeof topK === 'string' && !isNaN(Number(topK))
-              ? Number(topK)
-              : 10;
-        const vectorStore = mastra?.getVector(vectorStoreName);
-
-        if (!vectorStore) {
-          if (logger) {
-            logger.error('Vector store not found', { vectorStoreName });
-          }
-          return { relevantContext: [], sources: [] };
-        }
-
-        let queryFilter = {};
-        if (enableFilter) {
-          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:', { queryFilter, topK: topKValue });
-        }
-        const { results, queryEmbedding } = await vectorQuerySearch({
-          indexName,
-          vectorStore,
-          queryText,
-          model,
-          queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : undefined,
-          topK: topKValue,
-          includeVectors: true,
-        });
-        if (logger) {
-          logger.debug('vectorQuerySearch returned results', { count: results.length });
-        }
-
-        // Initialize graph if not done yet
-        if (!isInitialized) {
-          // Get all chunks and embeddings for graph construction
-          const chunks = results.map(result => ({
-            text: result?.metadata?.text,
-            metadata: result.metadata ?? {},
-          }));
-          const embeddings = results.map(result => ({
-            vector: result.vector || [],
-          }));
-
-          if (logger) {
-            logger.debug('Initializing graph', { chunkCount: chunks.length, embeddingCount: embeddings.length });
-          }
-          graphRag.createGraph(chunks, embeddings);
-          isInitialized = true;
-        } else if (logger) {
-          logger.debug('Graph already initialized, skipping graph construction');
-        }
-
-        // Get reranked results using GraphRAG
-        const rerankedResults = graphRag.query({
-          query: queryEmbedding,
-          topK: topKValue,
-          randomWalkSteps,
-          restartProb,
-        });
-        if (logger) {
-          logger.debug('GraphRAG query returned results', { count: rerankedResults.length });
-        }
-        // Extract and combine relevant chunks
-        const relevantChunks = rerankedResults.map(result => result.content);
-        if (logger) {
-          logger.debug('Returning relevant context chunks', { count: relevantChunks.length });
-        }
-        // `sources` exposes the full retrieval objects
-        const sources = includeSources ? convertToSources(rerankedResults) : [];
-        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/tools/index.ts

@@ -1,3 +0,0 @@
-export * from './document-chunker';
-export * from './graph-rag';
-export * from './vector-query';

package/src/tools/types.ts

@@ -1,110 +0,0 @@
-import type { MastraVector } from '@mastra/core/vector';
-import type { EmbeddingModel } from 'ai';
-import type { RerankConfig } from '../rerank';
-
-export interface PineconeConfig {
-  namespace?: string;
-  sparseVector?: {
-    indices: number[];
-    values: number[];
-  };
-}
-
-export interface PgVectorConfig {
-  minScore?: number;
-  ef?: number; // HNSW search parameter
-  probes?: number; // IVFFlat probe parameter
-}
-
-// Chroma types
-type LiteralValue = string | number | boolean;
-type ListLiteralValue = LiteralValue[];
-type LiteralNumber = number;
-type LogicalOperator = '$and' | '$or';
-type InclusionOperator = '$in' | '$nin';
-type WhereOperator = '$gt' | '$gte' | '$lt' | '$lte' | '$ne' | '$eq';
-type OperatorExpression = {
-  [key in WhereOperator | InclusionOperator | LogicalOperator]?: LiteralValue | ListLiteralValue;
-};
-type BaseWhere = {
-  [key: string]: LiteralValue | OperatorExpression;
-};
-type LogicalWhere = {
-  [key in LogicalOperator]?: Where[];
-};
-type Where = BaseWhere | LogicalWhere;
-type WhereDocumentOperator = '$contains' | '$not_contains' | LogicalOperator;
-type WhereDocument = {
-  [key in WhereDocumentOperator]?: LiteralValue | LiteralNumber | WhereDocument[];
-};
-
-export interface ChromaConfig {
-  // Add Chroma-specific configs here if needed
-  where?: Where;
-  whereDocument?: WhereDocument;
-}
-
-// Union type for all database-specific configs
-export type DatabaseConfig = {
-  pinecone?: PineconeConfig;
-  pgvector?: PgVectorConfig;
-  chroma?: ChromaConfig;
-  // Add other database configs as needed
-  [key: string]: any; // Allow for future database extensions
-};
-
-export type VectorQueryToolOptions = {
-  id?: string;
-  description?: string;
-  indexName: string;
-  model: EmbeddingModel<string>;
-  enableFilter?: boolean;
-  includeVectors?: boolean;
-  includeSources?: boolean;
-  reranker?: RerankConfig;
-  /** Database-specific configuration options */
-  databaseConfig?: DatabaseConfig;
-} & (
-  | {
-      vectorStoreName: string;
-    }
-  | {
-      vectorStoreName?: string;
-      vectorStore: MastraVector;
-    }
-);
-
-export type GraphRagToolOptions = {
-  id?: string;
-  description?: string;
-  indexName: string;
-  vectorStoreName: string;
-  model: EmbeddingModel<string>;
-  enableFilter?: boolean;
-  includeSources?: boolean;
-  graphOptions?: {
-    dimension?: number;
-    randomWalkSteps?: number;
-    restartProb?: number;
-    threshold?: number;
-  };
-};
-
-/**
- * Default options for GraphRAG
- * @default
- * ```json
- * {
- *   "dimension": 1536,
- *   "randomWalkSteps": 100,
- *   "restartProb": 0.15,
- *   "threshold": 0.7
- * }
- * ```
- */
-export const defaultGraphOptions = {
-  dimension: 1536,
-  randomWalkSteps: 100,
-  restartProb: 0.15,
-  threshold: 0.7,
-};