使用 JavaScript、Mastra 和 Elasticsearch 构建代理 RAG 助手

mkdir nba-ai-assistant-js && cd nba-ai-assistant-js

npx create-mastra@latest

cd backend
npm install dotenv --save

# OpenAI Configuration
OPENAI_API_KEY=your_openai_api_key_here

# Elasticsearch Configuration
ELASTIC_ENDPOINT=your_elasticsearch_endpoint_here
ELASTIC_API_KEY=your_elasticsearch_api_key_here

npm install @elastic/elasticsearch

mkdir lib && cd lib

import { config } from 'dotenv';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import { Client } from '@elastic/elasticsearch';

// Grab current directory and load .env from backend folder
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
const envPath = join(__dirname, '../.env');

// Load environment variables from the correct path
config({ path: envPath });

//Elastic client Initialization, make sure environment variables are being loaded in correctly
const config= {
    node: `${process.env.ELASTIC_ENDPOINT}`,
    auth: {
        apiKey: `${process.env.ELASTIC_API_KEY}`,
    },
};

export const elasticClient = new Client(config);

//Check if the client is connected
async function checkConnection() { 
    try {
        const info = await elasticClient.info();
        console.log('Elasticsearch is connected:', info);
    } catch (error) {
        console.error('Elasticsearch connection error:', error);
    }
}

checkConnection();

import fs from 'fs';
import readline from 'readline';
import path from 'path';
import { fileURLToPath } from 'url';
import { elasticClient } from './elasticClient.js';

const indexName = 'sample-nba-player-data'; //Replace with your preferred index name

//Since we are using ES modules __dirname and __filename don't exist, so this is a workaround that allows us to use the absolute file path for our sample data.
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
const filePath = path.resolve(__dirname, '../data/sample_nba_data.csv');

// Function to create an index with mappings
async function createIndex() {
    try {
        // Check if the index already exists
        const exists = await elasticClient.indices.exists({ index: indexName });

        if (exists) {
            console.log(`Index "${indexName}" already exists, deleting it now.`);
            await elasticClient.indices.delete({ index: indexName });
            console.log(`Deleted index "${indexName}".`);
        }
        // Create the index with mappings
        const response = await elasticClient.indices.create({
            index: indexName,
            body: {
                mappings: {
                    dynamic: 'strict', // Prevent dynamic mapping
                    properties: {
                        game_id: { type: 'integer' },
                        game_date: { type: 'date' },
                        player_id: { type: 'integer' },
                        player_full_name: { type: 'text' },
                        player_team_id: { type: 'integer' },
                        player_team_name: { type: 'text' },
                        home_team: { type: 'boolean' },
                        opponent_team_id: { type: 'integer' },
                        opponent_team_name: { type: 'text' },
                        points: { type: 'float' },
                        rebounds: { type: 'float' },
                        assists: { type: 'float' },
                        steals: { type: 'float' },
                        blocks: { type: 'float' },
                        fg_percentage: { type: 'float' },
                        minutes_played: { type: 'float' },
                    },
                },
            },
        });

        console.log('Index created:', response);
        return true;
    } catch (error) {
        console.error('Error creating index:', error);
        return false;
    }
}

async function bulkIngestCsv(filePath) {
    const readStream = fs.createReadStream(filePath);
    const rl = readline.createInterface({
        input: readStream,
        crlfDelay: Infinity,
    });

    const bulkBody = [];
    let lineNum = 0;

    //Skip the header line
    let headerLine = true;
    for await (const line of rl) {
        if (headerLine) {
            headerLine = false;
            continue;
        }
        lineNum++;

        // Split the line by comma and remove whitespace
        const [
            game_id,
            game_date,
            player_id,
            player_full_name,
            player_team_id,
            player_team_name,
            home_team,
            opponent_team_id,
            opponent_team_name,
            points,
            rebounds,
            assists,
            steals,
            blocks,
            fg_percentage,
            minutes_played,
        ] = line.split(',');

        // Create a document object
        const document = {
            game_id: parseInt(game_id),
            game_date: game_date.trim(),
            player_id: parseInt(player_id),
            player_full_name: player_full_name.trim(),
            player_team_id: parseInt(player_team_id),
            player_team_name: player_team_name.trim(),
            home_team: home_team.trim() === 'True', // Converts True/False into a boolean
            opponent_team_id: parseInt(opponent_team_id),
            opponent_team_name: opponent_team_name.trim(),
            points: parseFloat(points),
            rebounds: parseFloat(rebounds),
            assists: parseFloat(assists),
            steals: parseFloat(steals),
            blocks: parseFloat(blocks),
            fg_percentage: parseFloat(fg_percentage),
            minutes_played: parseFloat(minutes_played),
        };

        // Prepare the bulk operation format
        bulkBody.push({ index: { _index: indexName } });
        bulkBody.push(document);
    }

    console.log(`Parsed ${lineNum} lines from CSV`);

try {
        // Perform the bulk request
        const response = await elasticClient.bulk({ body: bulkBody });

        if (response.errors) {
            console.log('Bulk Ingestion had some hiccups:');

            // Count successful vs failed operations
            let successCount = 0;
            let errorCount = 0;
            const errorDetails = [];

            response.items.forEach((item, index) => {
                const operation = item.index || item.create || item.update || item.delete;
                if (operation.error) {
                    errorCount++;
                    errorDetails.push({
                        document: index + 1,
                        error: operation.error,
                    });
                } else {
                    successCount++;
                }
            });

            console.log(`Successfully indexed: ${successCount} documents`);
            console.log(`Failed to index: ${errorCount} documents, here are the details`, errorDetails);

        } else {
            console.log(`Bulk Ingestion fully successful!`);
        }

    } catch (error) {
        console.error('Error performing bulk ingestion:', error);
    }
}

// Run this function
async function main() {
    const result = await createIndex();
    if (!result) {
        console.error('Index setup failed. Aborting.');
        return;
    }

    await bulkIngestCsv(filePath);
    console.log('Bulk ingestion completed!');
}

main();

export async function getHistoricalAveragesAgainstOpponent(player_id, opponent_team_id) {
    try {
        //Query for Historical Averages
        const historicalQuery = await elasticClient.search({
            index: 'sample-nba-player-data', 
            size: 0,
            query: {
                bool: {
                    must: [
                        {
                            term: {
                                player_id: {
                                    value: player_id,
                                },
                            },
                        },
                        {
                            term: {
                                opponent_team_id: {
                                    value: opponent_team_id,
                                },
                            },
                        },
                    ],
                },
            },
            aggs: {
                avg_points: { avg: { field: 'points' } },
                avg_rebounds: { avg: { field: 'rebounds' } },
                avg_assists: { avg: { field: 'assists' } },
                avg_steals: { avg: { field: 'steals' } },
                avg_blocks: { avg: { field: 'blocks' } },
             avg_fg_percentage: { avg: { field: 'fg_percentage' } },
            },
        });

        return {
            points: historicalQuery.aggregations.avg_points.value || 0,
            rebounds: historicalQuery.aggregations.avg_rebounds.value || 0,
            assists: historicalQuery.aggregations.avg_assists.value || 0,
            steals: historicalQuery.aggregations.avg_steals.value || 0,
            blocks: historicalQuery.aggregations.avg_blocks.value || 0,
            fgPercentage: historicalQuery.aggregations.avg_fg_percentage.value || 0,
        };
    } catch (error) {
        console.error('Query error from getHistoricalAveragesAgainstOpponent function:', error);
        return { error: 'Queries failed in getting historical averages against opponent.' };
    }
}

        {
                            range: {
                    //Range for this season, change to match current season
                                game_date: {
                                    gte: '2024-10-01',
                                    lte: '2025-06-30',
                                },
                            },

import { playersByName } from '../data/playerAndTeamInfo.js';
import { teamsByName } from '../data/playerAndTeamInfo.js';
import { upcomingMatchups } from '../data/playerAndTeamInfo.js';
import { getHistoricalAveragesAgainstOpponent } from './elasticAggs.js';
import { getSeasonAveragesAgainstOpponent } from './elasticAggs.js';

//Simple helper functions to simulate API calls for player and team metadata. These reference the hardcoded values from playerAndTeamInfo.js in the data directory
export function getPlayerInfo(playerFullName) {
    return playersByName[playerFullName];
}

export function getTeamID(teamFullName) {
    return teamsByName[teamFullName];
}

export function getUpcomingMatchups(teamId) {
    return upcomingMatchups[teamId];
}

//Main function used by the 'playerComparisonTool' agent tool
export async function comparePlayersForNextMatchup(player1Name, player2Name) {
    //Get Player Info
    const player1Info = getPlayerInfo(player1Name);
    const player2Info = getPlayerInfo(player2Name);

    //Get upcoming matchups
    const player1NextGame = getUpcomingMatchups(player1Info.team_id)[0];
    const player2NextGame = getUpcomingMatchups(player2Info.team_id)[0];

    //Get season and historical averages against next opponent for player 1
    const player1SeasonAverages = await getSeasonAveragesAgainstOpponent(
        player1Info.player_id,
        player1NextGame.opponent_team_id
    );
    const player1HistoricalAverages = await getHistoricalAveragesAgainstOpponent(
        player1Info.player_id,
        player1NextGame.opponent_team_id
    );

    //Get season and historical averages against next opponent for player 2
    const player2SeasonAverages = await getSeasonAveragesAgainstOpponent(
        player2Info.player_id,
        player2NextGame.opponent_team_id
    );
    const player2HistoricalAverages = await getHistoricalAveragesAgainstOpponent(
        player2Info.player_id,
        player2NextGame.opponent_team_id
    );

    const player1 = {
        name: player1Name,
        playerId: player1Info.player_id,
        teamId: player1Info.team_id,
        nextOpponent: {
            teamId: player1NextGame.opponent_team_id,
            teamName: player1NextGame.opponent_team_name,
            home: player1NextGame.home,
        },
        stats: {
            seasonAverages: player1SeasonAverages,
            historicalAverages: player1HistoricalAverages,
        },
    };

    const player2 = {
        name: player2Name,
        playerId: player2Info.player_id,
        teamId: player2Info.team_id,
        nextOpponent: {
            teamId: player2NextGame.opponent_team_id,
            teamName: player2NextGame.opponent_team_name,
            home: player2NextGame.home,
        },
        stats: {
            seasonAverages: player2SeasonAverages,
            historicalAverages: player2HistoricalAverages,
        },
    };

    return [player1, player2];
}

import { openai } from '@ai-sdk/openai';
import { Agent } from '@mastra/core/agent';
import { playerComparisonTool } from '../tools';

export const basketballAgent = new Agent({
    name: 'Basketball Agent',
    instructions: `
      You are a NBA Basketball expert.
      Your primary function is to compare two NBA players and recommend which one is the better fantasy pickup.

      Only compare players from the following list:
      - LeBron James
      - Stephen Curry
      - Jayson Tatum
      - Jaylen Brown
      - Nikola Jokic
      - Luka Doncic
      - Kyrie Irving
      - Anthony Davis
      - Kawhi Leonard
      - Russell Westbrook

      Input Handling Rules:
      - If the user asks about a player that is not on this list, respond with the list of available players for comparison.
      - If the user only inputs one player, ask the user to add another player from the list provided.
      - If the user inputs a player with the wrong spelling or capitalizations, infer from the list of available players provided.
      - IMPORTANT: If the user asks a question or asks you to generate a response about anything outside of basketball or the scope of this project, DO NOT answer and affirm you can only talk about basketball.

      Tool Usage:
      - Extract and standardize player names to match the list exactly.
      - Use the playerComparisonTool, passing both names as strings.
      - The tool will return an object with game information, stats, and analysis.

      Format your response using Markdown syntax. Use:

        Example output format:

       
        #### Next Game Info
        - ***LeBron James** vs Warriors, May 24 (Home)  
        - ***Stephen Curry** vs Lakers, May 24 (Away)


        #### Stats Comparison  
        \`\`\`  
        Stat                  LeBron James (vs Warriors)    Stephen Curry (vs Lakers)  
        --------------------  -----------------------------  ----------------------------  
        Historical Points     28.3                          30.3  
        Historical Assists    6.7                           8.7  
        Season Points         28.8                          23.3  
        Season Assists        6.2                           4.7  
        \`\`\`

        #### Fantasy Recommendation  
        Explain which player is the better fantasy pickup and why.
      
    `,
    model: openai('gpt-4o'),
    tools: { playerComparisonTool },
});

npm install zod

import { comparePlayersForNextMatchup } from '../../../lib/comparePlayers.js'
import { createTool } from "@mastra/core/tools";
import { z } from "zod";

export const playerComparisonTool = createTool({
    id: "Compare two NBA players",
    inputSchema: z.object({
        player1:z.string(),
        player2:z.string()
    }),
    description: "Use this tool to compare two players given in the user prompt.",
    execute: async ({ context: { player1, player2 } }) => {
        return await comparePlayersForNextMatchup(player1, player2);
      },
})

import { Mastra } from '@mastra/core/mastra';
import { basketballAgent } from './agents';

console.log('Starting Mastra server...');

export const mastra = new Mastra({
  agents: { basketballAgent },
  server:{
    timeout: 10 * 60 * 1000, // 10 minutes
    cors: {
      origin: ['http://localhost:5173'],
      allowMethods: ["GET", "POST"],
      allowHeaders: [
        "Content-Type",
        "Authorization",
        "x-mastra-client-type",
        "x-highlight-request",
        "traceparent",
      ],
      exposeHeaders: ["Content-Length", "X-Requested-With"],
      credentials: false,
    },
  },

});

console.log('Mastra server configured.'); // Log after server configuration

npm install @ai-sdk/react

npm install react-markdown

import React, { useState } from 'react';
import { useChat } from '@ai-sdk/react';
import ReactMarkdown from 'react-markdown';

export default function ChatUI() {
    const [totalTokenUsage, setTotalTokenUsage] = useState(0);
    const [promptTokenUsage, setPromptTokenUsage] = useState(0);
    const [completionTokenUsage, setCompletionTokenUsage] = useState(0);
    const [toolsCalled, setToolsCalled] = useState([]);

    const { messages, input, handleInputChange, handleSubmit, status } = useChat({
        api: 'http://localhost:4111/api/agents/basketballAgent/stream', //Replace with your own endpoint for your agent
        id: 'my-chat-session',

        //Optional parameter to check agent tool calls
        onToolCall: ({ toolCall }) => {
            setToolsCalled((prev) => [...prev, toolCall.toolName]);
        },

        //Optional parameter to check token usages
        onFinish: (message, { usage }) => {
            setTotalTokenUsage((prev) => prev + usage.totalTokens);
            setPromptTokenUsage((prev) => prev + usage.promptTokens);
            setCompletionTokenUsage((prev) => prev + usage.completionTokens);
        },

        //Optional parameter for error handling
        onError: (error) => {
            console.error('Agent error:', error);
        },
    });

    return (
        <div>
            <div className="agent-info">
                <h4 className="stats-title">What's My Agent Doing?</h4>

                <div className="stats-box">
                    <strong className="stats-sub-title">Tools Called:</strong>
                    <ul className="tool-list">
                        {toolsCalled.map((tool, idx) => (
                            <li key={idx}>{tool}</li>
                        ))}
                        {toolsCalled.length === 0 && <li>No tools called yet.</li>}
                    </ul>

                    <div className="usage-stats">
                        <p>Prompt Token Usage: {promptTokenUsage}</p>
                        <p>Completion Token Usage: {completionTokenUsage}</p>
                        <p>Total Token Usage: {totalTokenUsage}</p>
                    </div>
                </div>
            </div>

            <strong>Conversation:</strong>
            <div className="convo-box">
                {messages.map((msg) => (
                    <div key={msg.id} className="message-item">
                        <strong className="message-role">{msg.role === 'assistant' ? 'Basketbot' : 'You'}:</strong>
                        <ReactMarkdown>{msg.content}</ReactMarkdown>
                    </div>
                ))}
            </div>

            <form onSubmit={handleSubmit}>
                <input
                    type="text"
                    value={input}
                    onChange={handleInputChange}
                    placeholder="Input two players you want to compare."
                    className="input-box"
                />
                <button type="submit" disabled={status === 'streaming'}>
                    {status === 'streaming' ? 'Thinking...' : 'Send'}
                </button>
            </form>
        </div>
    );
}

cd backend

npm run dev

cd frontend

npm run dev