@struktoai/mirage-core/dist/core/wandb/client.js

A Unified Virtual Filesystem For AI Agents

import { apiRequest, stallTimeout } from "../api/client.js";
import { eacces, enoent } from "../../errors/fs.js";
import { WandbAPIError } from "./errors.js";
import { PROJECTS, RUNS, RUN, FILE, FILES, HISTORY, HISTORY_KEYS } from "./queries.js";
function responseError(response) {
    if ([401, 403].includes(response.status))
        return eacces('W&B authentication or authorization failed');
    return new WandbAPIError(`W&B HTTP ${String(response.status)}`, response.status);
}
export class WandbClient {
    config;
    constructor(config) {
        this.config = config;
    }
    headers() {
        if (!this.config.apiKey)
            return {};
        const bytes = new TextEncoder().encode(`api:${this.config.apiKey}`);
        return {
            Authorization: `Basic ${btoa(Array.from(bytes, (b) => String.fromCharCode(b)).join(''))}`,
        };
    }
    async request(query, variables) {
        const result = (await apiRequest('POST', this.config.baseUrl.replace(/\/$/, '') + '/graphql', {
            timeout: this.config.timeout,
            errorOf: responseError,
            headers: { ...this.headers(), 'Content-Type': 'application/json' },
            json: { query, variables },
        }));
        if (result.errors?.length)
            throw new WandbAPIError('W&B GraphQL request failed');
        if (result.data === undefined || result.data === null)
            throw new WandbAPIError('W&B response has no data');
        return result.data;
    }
    async pages(query, variables, keys, limit) {
        let cursor = null;
        const seen = new Set();
        const rows = [];
        for (let page = 0; page < this.config.maxPages; page++) {
            let data = await this.request(query, {
                ...variables,
                cursor,
                perPage: this.config.pageSize,
            });
            for (const key of keys) {
                data = data[key];
                if (data === null)
                    throw enoent('W&B object not found');
            }
            const connection = data;
            for (const item of connection.edges.map((edge) => edge.node))
                rows.push(item);
            if (!connection.pageInfo.hasNextPage)
                return rows;
            if (limit !== undefined && rows.length >= limit)
                return rows;
            cursor = connection.pageInfo.endCursor;
            if (!cursor || seen.has(cursor))
                throw new WandbAPIError('W&B pagination did not advance');
            seen.add(cursor);
        }
        throw new WandbAPIError('W&B pagination limit exceeded');
    }
    projects(entity, limit) {
        return this.pages(PROJECTS, { entity }, ['models'], limit);
    }
    runs(entity, project, limit) {
        return this.pages(RUNS, { entity, project }, ['project', 'runs'], limit);
    }
    async run(variables, query = RUN) {
        const data = await this.request(query, {
            ...variables,
        });
        if (!data.project?.run)
            throw enoent('W&B run not found');
        return data.project.run;
    }
    files(variables) {
        return this.pages(FILES, { ...variables }, ['project', 'run', 'files']);
    }
    async file(variables, name) {
        const data = await this.request(FILE, { ...variables, names: [name] });
        if (!data.project?.run)
            throw enoent('W&B run not found');
        return data.project.run.files.edges[0]?.node ?? null;
    }
    async *history(variables) {
        const run = await this.run(variables, HISTORY_KEYS);
        const last = run.historyKeys?.lastStep ?? -1;
        if (!Number.isSafeInteger(last) || last < -1)
            throw new WandbAPIError('W&B invalid last history step');
        const size = this.config.pageSize;
        if (Math.floor((last + size) / size) > this.config.maxPages)
            throw new WandbAPIError('W&B history pagination limit exceeded');
        for (let start = 0; start <= last; start += size) {
            const stop = Math.min(start + size, last + 1);
            const queryStart = stop - start === 1 ? Math.max(0, start - 1) : start;
            const queryStop = Math.max(stop, queryStart + 2);
            const data = await this.request(HISTORY, {
                ...variables,
                minStep: queryStart,
                maxStep: queryStop,
                pageSize: Math.max(size, queryStop - queryStart),
            });
            if (!data.project?.run)
                throw enoent('W&B run not found');
            for (const raw of data.project.run.history) {
                const row = JSON.parse(raw);
                if ((queryStart === start && queryStop === stop) ||
                    (typeof row._step === 'number' && start <= row._step && row._step < stop))
                    yield row;
            }
        }
    }
    async *download(url) {
        const target = new URL(url, this.config.baseUrl + '/');
        if (!['http:', 'https:'].includes(target.protocol) || target.username || target.password)
            throw new WandbAPIError('W&B invalid download URL');
        const headers = target.origin === new URL(this.config.baseUrl).origin ? this.headers() : {};
        const response = await stallTimeout(this.config.timeout)(target, { headers });
        if (!response.ok)
            throw responseError(response);
        if (!response.body)
            throw new WandbAPIError('W&B download has no body');
        const reader = response.body.getReader();
        try {
            let next = await reader.read();
            while (!next.done) {
                yield next.value;
                next = await reader.read();
            }
        }
        finally {
            await reader.cancel();
            reader.releaseLock();
        }
    }
}
//# sourceMappingURL=client.js.map