cucumber-ai/dist/data-agent/index.js

Write automated tests using natural language

"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DataAgent = void 0;
const node_fs_1 = __importDefault(require("node:fs"));
const node_path_1 = require("node:path");
const index_js_1 = require("@modelcontextprotocol/sdk/client/index.js");
const sse_js_1 = require("@modelcontextprotocol/sdk/client/sse.js");
const stdio_js_1 = require("@modelcontextprotocol/sdk/client/stdio.js");
const types_js_1 = require("@modelcontextprotocol/sdk/types.js");
const app_root_path_1 = require("app-root-path");
class DataAgent {
    constructor(context) {
        this.context = context;
        this.clients = [];
        this.tools = [];
        this.toolMap = {};
        this.started = false;
        const configPath = (0, node_path_1.join)(app_root_path_1.path, "config.json");
        if (node_fs_1.default.existsSync(configPath)) {
            this.config = require(configPath);
        }
        else {
            this.config = { mcpServer: {} };
        }
    }
    async start() {
        if (this.started) {
            await this.stop();
        }
        this.systemPrompt = node_fs_1.default.readFileSync((0, node_path_1.join)(__dirname, "system.prompt.md"), "utf-8");
        await this.startServers();
        await this.collectTools();
        this.registerActions(this.context.getActions());
        this.started = true;
    }
    async collectTools() {
        for (const client of this.clients) {
            const mcpToolList = await client.listTools();
            for (const tool of mcpToolList.tools) {
                this.toolMap[tool.name] = client;
                this.tools.push({
                    type: "function",
                    function: {
                        name: tool.name,
                        description: tool.description,
                        parameters: tool.inputSchema,
                    },
                });
            }
        }
    }
    async startServers() {
        for (const mcpName in this.config.mcpServer) {
            const client = new index_js_1.Client({
                name: `mcp-client-${mcpName}`,
                version: "0.1.0",
            });
            const mcpServer = this.config.mcpServer[mcpName];
            const clientTransport = mcpServer.type === "sse"
                ? new sse_js_1.SSEClientTransport(new URL(mcpServer.url))
                : new stdio_js_1.StdioClientTransport({
                    command: mcpServer.command,
                    args: mcpServer.args,
                    env: {
                        ...process.env,
                        ...mcpServer.env,
                    },
                });
            await client.connect(clientTransport);
            this.clients.push(client);
        }
    }
    async stop() {
        for (const client of this.clients) {
            await client.close();
        }
        this.clients = [];
        this.tools = [];
        this.toolMap = {};
        this.unregisterActions(this.context.getActions());
        this.started = false;
    }
    async ask(prompt, opts = {}) {
        const callTool = async (toolCall) => {
            const result = types_js_1.CallToolResultSchema.parse(await this.toolMap[toolCall.function.name].callTool({
                name: toolCall.function.name,
                arguments: JSON.parse(toolCall.function.arguments),
            }));
            return result.content;
        };
        return await this.context.getToolExecutor().execute(prompt, {
            callTool,
            useCache: opts.useCache ?? this.context.isCacheEnabled(),
            systemPrompt: this.systemPrompt,
            cacheKey: "data-agent",
            tools: this.tools,
        });
    }
    registerActions(actions) {
        actions.register("data", async (text) => await this.ask(text));
    }
    unregisterActions(actions) {
        actions.unregister("data");
    }
}
exports.DataAgent = DataAgent;
//# sourceMappingURL=index.js.map