hub/testing/suite.js

const WebSocket = require('ws');

const USE_LLM = true;

const hostname = 'localhost:5000';
const ws = new WebSocket(`ws://${hostname}/ws`);
var nonceCounter = 1;
var nonceCallbacks = {};
var messageQueue = [];
var messageResolverQueue = [];

// Setup

const credentials = { username: "test-user", password: "123456" };

const tokenPromise = new Promise(async resolve => {
    try {
        const response = await fetch(`http://${hostname}/user/login`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(credentials),
        });
        const data = await response.json();
        resolve(data.token);
    } catch (error) {
        console.error("Login Failed", error);
        process.exit(1);
    }
});

ws.on('open', async function open() {
    const token = await tokenPromise;
    console.log("Token:", token);
    const packet = JSON.stringify({token});
    ws.send(packet);

    runTest();
});

ws.on('message', function incoming(rawMessage) {
    const parsed = JSON.parse(rawMessage);
    const callback = nonceCallbacks[parsed.nonce];
    if (callback) {
        callback(parsed);
    } else {
        if (messageResolverQueue.length > 0) {
            const resolver = messageResolverQueue.shift();
            resolver(parsed);
        } else {
            messageQueue.push(parsed);
        }
    }
});

async function send (action, message) {
    const nonce = nonceCounter++;
    let promise = new Promise((resolve, reject) => {
        nonceCallbacks[nonce] = res => {
            if (res.error) reject(Error(res.error));
            else resolve(res);
        };;
    });

    const packet = JSON.stringify({ nonce, action, ...message });
    ws.send(packet);
    return promise;
}

function waitMessage() {
    return new Promise((resolve) => {
        if (messageQueue.length > 0) {
            resolve(messageQueue.shift());
        } else {
            messageResolverQueue.push(resolve);
        }
    });
}


// Testing

async function runTest() {
    console.log("Initial Chats",
        await send('get-chats')
    );

    console.log("Agent Kinds",
        await send('get-agent-kinds')
    );

    let createAgentResponse = await send('create-agent', {body: {
        name: 'First Agent',
        description: 'First agent does nothing',
        kind: 'basic',
        data: {}
    }});
    console.log("Create Agent", createAgentResponse);

    let agents = (await send('get-agents')).agents;
    console.log("Get Agents", agents);

    let mainAgentId = createAgentResponse.agentId;
    let mainAgent = agents.find(agent => agent.id == mainAgentId)

    if (USE_LLM) {
        let promptResponse = await send('send-agent-prompt', {
            body: {
                agentId: mainAgentId,
                content: "Hello World!",
            }
        });
        console.log("Prompt Response", promptResponse);
    }

    // Update agent
    console.log("Updating agent...");
    await send('update-agent', {
        body: {
            agentId: mainAgentId,
            name: mainAgent.name + " (UPDATED)",
            description: "This is a placeholder description for the agent.",
            data: {
                system_prompt: "first system prompt",
                // These are not config fields for 'basic' agent, they cannot be seen by the user
                field1: "some data",
                field2: "more data",
            },
        }
    });
    await send('update-agent', {
        body: {
            agentId: mainAgentId,
            data: {
                system_prompt: "second system prompt",
                field2: "replaced data",
            },
        }
    });
    console.log("Updated Agent:", (
        (await send('get-agents')).agents.find(agent => agent.id == mainAgentId)
    ));
    console.log("Agent Data:", (
        (await send('get-agent-config-fields', {body: {agentId: mainAgentId}}))
    ));

    await testChats();

    if (USE_LLM) {
        await testGral(mainAgentId);
    }

    process.exit();
}

async function testChats () {
    let response = await send('get-chats')
    let chat = response.chats[0];

    if (USE_LLM) {
        let msgRes = await send('send-chat-message', {
            body: {
                chatId: chat.id,
                content: "Whats up!",
            }
        });
        console.log('message sent', msgRes);
        console.log('Incoming:', await waitMessage());

        // Multi turn conversation
        await send('send-chat-message', {
            body: {
                chatId: chat.id,
                content: "Spell the initials of my first message.",
            }
        });
        console.log('Incoming:', await waitMessage());

        console.log('Messages in Chat', chat.id, await send('get-chat-messages', {
            body: {chatId: chat.id}
        }));
    }
}

async function testGral (agentId) {
    let res = await send('call-agent-action', {
        body: {
            agentId,
            agentAction: 'gral_message',
            arguments: {
                message_history: [{
                    'role': 'user',
                    'content': 'Add a name argument to the sayHello function in main.py',
                }],
                file_map: {
                    'main.py': "def sayHello():\n    print('Hello World')\n"
                },
            },
        }
    });

    console.log('Gral Result:', res);
}