使用 Mastra 和 Elasticsearch 构建具有语义召回功能的知识代理

npm create mastra@latest

npm install ai @ai-sdk/openai @elastic/elasticsearch dotenv

ELASTICSEARCH_ENDPOINT="your-endpoint-here"
ELASTICSEARCH_API_KEY="your-key-here"
OPENAI_API_KEY="your-key-here"

export interface ElasticVectorConfig extends ClientOptions {
    /**
     * Explicitly specify if connecting to Elasticsearch Serverless.
     * If not provided, will be auto-detected on first use.
     */
    isServerless?: boolean;
    
    /**
     * Maximum documents to count accurately when describing indices.
     * Higher values provide accurate counts but may impact performance on large indices.
     * 
     * @default 10000
     */
    maxCountAccuracy?: number;
}

export class ElasticVector extends MastraVector {
    private client: Client;
    private isServerless: boolean | undefined;
    private deploymentChecked: boolean = false;
    private readonly maxCountAccuracy: number;

    constructor(config: ElasticVectorConfig) {
        super();
        this.client = new Client(config);
        this.isServerless = config.isServerless;
        this.maxCountAccuracy = config.maxCountAccuracy ?? 10000;
    }
}

private async detectServerless(): Promise<boolean> {
    // Return cached result if already detected
    if (this.deploymentChecked) {
        return this.isServerless ?? false;
    }

    // Use explicit configuration if provided
    if (this.isServerless !== undefined) {
        this.deploymentChecked = true;
        this.logger?.info(
            `Using explicit deployment type: ${this.isServerless ? 'Serverless' : 'Standard'}`
        );
        return this.isServerless;
    }

    try {
        const info = await this.client.info();
        
        // Primary detection: build flavor (most reliable)
        const isBuildFlavorServerless = info.version?.build_flavor === 'serverless';
        
        // Secondary detection: tagline (fallback)
        const isTaglineServerless = info.tagline?.toLowerCase().includes('serverless') ?? false;
        
        this.isServerless = isBuildFlavorServerless || isTaglineServerless;
        this.deploymentChecked = true;
        
        this.logger?.info(
            `Auto-detected ${this.isServerless ? 'Serverless' : 'Standard'} Elasticsearch deployment`,
            { 
                buildFlavor: info.version?.build_flavor, 
                version: info.version?.number,
                detectionMethod: isBuildFlavorServerless ? 'build_flavor' : 'tagline'
            }
        );
        
        return this.isServerless;
    } catch (error) {
        this.logger?.warn(
            'Could not auto-detect deployment type, assuming Standard Elasticsearch. ' +
            'Set isServerless: true explicitly in config if using Serverless.',
            { error: error instanceof Error ? error.message : String(error) }
        );
        this.isServerless = false;
        this.deploymentChecked = true;
        return false;
    }
}

// Option 1: Auto-detect (recommended)
const vector = new ElasticVector({
    node: 'https://your-cluster.es.cloud',
    auth: { apiKey: 'your-api-key' }
});
// Detection happens automatically on first index operation

// Option 2: Explicit configuration (faster startup)
const vector = new ElasticVector({
    node: 'https://your-serverless.es.cloud',
    auth: { apiKey: 'your-api-key' },
    isServerless: true  // Skips auto-detection
});

async createIndex(params: CreateIndexParams): Promise<void> {
    const { indexName, dimension = 1536, metric = 'cosine' } = params;

    try {
        const exists = await this.client.indices.exists({ index: indexName });

        if (exists) {
            try {
                await this.validateExistingIndex(indexName, dimension, metric);
                this.logger?.info(`Index "${indexName}" already exists and is valid`);
                return;
            } catch (validationError) {
                throw new Error(
                    `Index "${indexName}" exists but does not match the required configuration: ${
                        validationError instanceof Error ? validationError.message : String(validationError)
                    }`
                );
            }
        }

        const isServerless = await this.detectServerless();
        const similarity = this.mapMetricToSimilarity(metric);

        const indexConfig: any = {
            index: indexName,
            mappings: {
                properties: {
                    vector: {
                        type: 'dense_vector',
                        dims: dimension,
                        index: true,
                        similarity: similarity,
                    },
                    metadata: {
                        type: 'object',
                        enabled: true,
                        dynamic: true, // Allows flexible metadata structures
                    },
                },
            },
        };

        // Only configure shards/replicas for non-serverless deployments
        // Serverless manages infrastructure automatically
        if (!isServerless) {
            indexConfig.settings = {
                number_of_shards: 1,
                number_of_replicas: 0, // Increase for production HA deployments
            };
        }

        await this.client.indices.create(indexConfig);

        this.logger?.info(
            `Created ${isServerless ? 'Serverless' : 'Standard'} Elasticsearch index "${indexName}"`,
            { dimension, metric, similarity }
        );
    } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        this.logger?.error(`Failed to create index "${indexName}": ${errorMessage}`);
        throw new Error(`Failed to create index "${indexName}": ${errorMessage}`);
    }
}

async upsert(params: UpsertVectorParams): Promise<string[]> {
    const { indexName, vectors, metadata = [], ids } = params;

    try {
        // Generate unique IDs if not provided
        const vectorIds = ids || vectors.map((_, i) => 
            `vec_${Date.now()}_${i}_${Math.random().toString(36).substr(2, 9)}`
        );

        const operations = vectors.flatMap((vec, index) => [
            { index: { _index: indexName, _id: vectorIds[index] } },
            {
                vector: vec,
                metadata: metadata[index] || {},
            },
        ]);

        const response = await this.client.bulk({
            refresh: true,
            operations,
        });

        if (response.errors) {
            const erroredItems = response.items.filter((item: any) => item.index?.error);
            const erroredIds = erroredItems.map((item: any) => item.index?._id);
            const errorDetails = erroredItems.slice(0, 3).map((item: any) => ({
                id: item.index?._id,
                error: item.index?.error?.reason || item.index?.error,
                type: item.index?.error?.type
            }));
            
            const errorMessage = `Failed to upsert ${erroredIds.length}/${vectors.length} vectors`;
            console.error(`${errorMessage}. Sample errors:`, JSON.stringify(errorDetails, null, 2));
            this.logger?.error(errorMessage, { 
                failedCount: erroredIds.length, 
                totalCount: vectors.length,
                sampleErrors: errorDetails 
            });
            
            // Still return successfully inserted IDs
            const successfulIds = vectorIds.filter((id, idx) => 
                !erroredIds.includes(id)
            );
            
            if (successfulIds.length === 0) {
                throw new Error(`${errorMessage}. All operations failed. See logs for details.`);
            }
            
            return successfulIds;
        }

        this.logger?.info(`Successfully upserted ${vectors.length} vectors to "${indexName}"`);
        return vectorIds;
    } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        this.logger?.error(`Failed to upsert vectors to "${indexName}": ${errorMessage}`);
        throw new Error(`Failed to upsert vectors to "${indexName}": ${errorMessage}`);
    }
}

async query(params: QueryVectorParams<any>): Promise<QueryResult[]> {
    const { indexName, queryVector, topK = 10, filter, includeVector = false } = params;

    try {
        const knnQuery: any = {
            field: 'vector',
            query_vector: queryVector,
            k: topK,
            num_candidates: Math.max(topK * 10, 100), // Search more candidates for better recall
        };

        // Apply metadata filters if provided
        if (filter) {
            knnQuery.filter = this.buildElasticFilter(filter);
        }

        const sourceFields = ['metadata'];
        if (includeVector) {
            sourceFields.push('vector');
        }

        const response = await this.client.search({
            index: indexName,
            knn: knnQuery,
            size: topK,
            _source: sourceFields,
        });

        const results = response.hits.hits.map((hit: any) => ({
            id: hit._id,
            score: hit._score || 0,
            metadata: hit._source?.metadata || {},
            vector: includeVector ? hit._source?.vector : undefined,
        }));

        this.logger?.debug(`Query returned ${results.length} results from "${indexName}"`);
        return results;
    } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        this.logger?.error(`Failed to query vectors from "${indexName}": ${errorMessage}`);
        throw new Error(`Failed to query vectors from "${indexName}": ${errorMessage}`);
    }
}

import { Agent } from '@mastra/core/agent';
import { Memory } from '@mastra/memory';
import { openai } from '@ai-sdk/openai';
import { Client } from '@elastic/elasticsearch';
import { ElasticVector } from '../stores/elastic-store';
import dotenv from "dotenv";

dotenv.config();

const ELASTICSEARCH_ENDPOINT = process.env.ELASTICSEARCH_ENDPOINT;
const ELASTICSEARCH_API_KEY = process.env.ELASTICSEARCH_API_KEY;

//Error check for undefined credentials
if (!ELASTICSEARCH_ENDPOINT || !ELASTICSEARCH_API_KEY) {
  throw new Error('Missing Elasticsearch credentials');
}

//Check to see if a connection can be established
const testClient = new Client({
  node: ELASTICSEARCH_ENDPOINT,
  auth: { 
    apiKey: ELASTICSEARCH_API_KEY 
  },
});

try {
  await testClient.ping();
  console.log('Connected to Elasticsearch successfully');
} catch (error: unknown) {
  if (error instanceof Error) {
    console.error('Failed to connect to Elasticsearch:', error.message);
  } else {
    console.error('Failed to connect to Elasticsearch:', error);
  }
  process.exit(1);
}
//Initialize the Elasticsearch vector store
const vectorStore = new ElasticVector({
  node: ELASTICSEARCH_ENDPOINT,
  auth: {
    apiKey: ELASTICSEARCH_API_KEY,
  },
//Optional: Explicitly set to true if using Elasticsearch Serverless to skip auto-detection and improve startup time
//isServerless: true,
});

export const knowledgeAgent = new Agent({
    name: 'KnowledgeAgent',
    instructions: 'You are a helpful knowledge assistant.',
    model: openai('gpt-4o'),
    memory: new Memory({

        vector: vectorStore,

        //embedder used to create embeddings for each message
        embedder: 'openai/text-embedding-3-small',

        //set semantic recall options
        options: {
            semanticRecall: {
                topK: 3, // retrieve 3 similar messages
                messageRange: 2, // include 2 messages before/after each match
                scope: 'resource',
            },
        },
    }),
});

export const mastra = new Mastra({
  agents: { knowledgeAgent },
  storage: new LibSQLStore({
    // stores observability, scores, ... into memory storage, if it needs to persist, change to file:../mastra.db
    url: ":memory:",
  }),
  logger: new PinoLogger({
    name: 'Mastra',
    level: 'info',
  }),
  telemetry: {
    // Telemetry is deprecated and will be removed in the Nov 4th release
    enabled: false, 
  },
  observability: {
    // Enables DefaultExporter and CloudExporter for AI tracing
    default: { enabled: true }, 
  },
});

npm run dev