@llumiverse/drivers 0.8.0

package/src/replicate.ts

@@ -0,0 +1,333 @@
+import {
+    AIModel,
+    AbstractDriver,
+    BuiltinProviders,
+    Completion,
+    DataSource,
+    DriverOptions,
+    ExecutionOptions,
+    ModelSearchPayload,
+    PromptFormats,
+    TrainingJob,
+    TrainingJobStatus,
+    TrainingOptions
+} from "@llumiverse/core";
+import { EventStream } from "@llumiverse/core/async";
+import EventSource from "eventsource";
+import Replicate, { Prediction } from "replicate";
+
+let cachedTrainableModels: AIModel[] | undefined;
+let cachedTrainableModelsTimestamp: number = 0;
+
+const supportFineTunning = new Set([
+    "meta/llama-2-70b-chat",
+    "meta/llama-2-13b-chat",
+    "meta/llama-2-7b-chat",
+    "meta/llama-2-7b",
+    "meta/llama-2-70b",
+    "meta/llama-2-13b",
+    "mistralai/mistral-7b-v0.1"
+]);
+
+export interface ReplicateDriverOptions extends DriverOptions {
+    apiKey: string;
+}
+
+export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
+    provider = BuiltinProviders.replicate;
+    service: Replicate;
+    defaultFormat = PromptFormats.genericTextLLM;
+
+    static parseModelId(modelId: string) {
+        const [owner, modelPart] = modelId.split("/");
+        const i = modelPart.indexOf(':');
+        if (i === -1) {
+            throw new Error("Invalid model id. Expected format: owner/model:version");
+        }
+        return {
+            owner, model: modelPart.slice(0, i), version: modelPart.slice(i + 1)
+        }
+    }
+
+    constructor(options: ReplicateDriverOptions) {
+        super(options);
+        this.service = new Replicate({
+            auth: options.apiKey,
+        });
+    }
+
+    extractDataFromResponse(prompt: string, response: Prediction): Completion {
+        const text = response.output.join("");
+        return {
+            result: text,
+            token_usage: {
+                result: response.output.length,
+                prompt: prompt.length,
+                total: response.output.length + prompt.length,
+            },
+        };
+    }
+
+    async requestCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+        const model = ReplicateDriver.parseModelId(options.model);
+        const predictionData = {
+            input: {
+                prompt: prompt,
+                max_new_tokens: options.max_tokens || 1024,
+                temperature: options.temperature,
+            },
+            version: model.version,
+            stream: true, //streaming described here https://replicate.com/blog/streaming
+        };
+
+        const prediction =
+            await this.service.predictions.create(predictionData);
+
+        const stream = new EventStream<string>();
+
+        const source = new EventSource(prediction.urls.stream!);
+        source.addEventListener("output", (e: any) => {
+            stream.push(e.data);
+        });
+        source.addEventListener("error", (e: any) => {
+            let error: any;
+            try {
+                error = JSON.parse(e.data);
+            } catch (error) {
+                error = JSON.stringify(e);
+            }
+            this.logger?.error(e, error, "Error in SSE stream");
+        });
+        source.addEventListener("done", () => {
+            try {
+                stream.close(""); // not using e.data which is {}
+            } finally {
+                source.close();
+            }
+        });
+        return stream;
+    }
+
+    async requestCompletion(prompt: string, options: ExecutionOptions) {
+        const model = ReplicateDriver.parseModelId(options.model);
+        const predictionData = {
+            input: {
+                prompt: prompt,
+                max_new_tokens: options.max_tokens || 1024,
+                temperature: options.temperature,
+            },
+            version: model.version,
+            //TODO stream
+            //stream: stream,     //streaming described here https://replicate.com/blog/streaming
+        };
+
+        const prediction =
+            await this.service.predictions.create(predictionData);
+
+        //TODO stream
+        //if we're streaming, return right away for the stream handler to handle
+        //        if (stream) return prediction;
+
+        //not streaming, wait for the result
+        const res = await this.service.wait(prediction, {});
+
+        const text = res.output.join("");
+        return {
+            result: text,
+            token_usage: {
+                result: res.output.length,
+                prompt: prompt.length,
+                total: res.output.length + prompt.length,
+            },
+        };
+    }
+
+    async startTraining(dataset: DataSource, options: TrainingOptions): Promise<TrainingJob> {
+        if (options.name.indexOf("/") === -1) {
+            throw new Error("Invalid target model name. Expected format: owner/model");
+        }
+        const { owner, model, version } = ReplicateDriver.parseModelId(options.model);
+        const job = await this.service.trainings.create(owner, model, version, {
+            destination: options.name as any,
+            input: {
+                train_data: dataset.getURL(),
+            },
+        })
+        return jobInfo(job, options.name);
+    }
+
+    /**
+     * This method is not returning a consistent TrainingJob like the one returned by startTraining
+     * Instead of returning the full model name `owner/model:version` it returns only the version `version
+     * @param jobId 
+     * @returns 
+     */
+    async cancelTraining(jobId: string): Promise<TrainingJob> {
+        const job = await this.service.trainings.cancel(jobId);
+        return jobInfo(job);
+    }
+
+    /**
+     * This method is not returning a consistent TrainingJob like the one returned by startTraining
+     * Instead of returning the full model name `owner/model:version` it returns only the version `version
+     * @param jobId 
+     * @returns 
+     */
+    async getTrainingJob(jobId: string): Promise<TrainingJob> {
+        const job = await this.service.trainings.get(jobId);
+        return jobInfo(job);
+    }
+
+    // ========= management API =============
+
+    async validateConnection(): Promise<boolean> {
+        try {
+            await this.service.predictions.list();
+            return true;
+        } catch (error) {
+            return false;
+        }
+    }
+
+    async _listTrainableModels(): Promise<AIModel[]> {
+        const promises = Array.from(supportFineTunning).map(id => {
+            const [owner, model] = id.split('/');
+            return this.service.models.get(owner, model)
+        });
+        const results = await Promise.all(promises);
+        return results.filter(m => !!m.latest_version).map(m => {
+            const fullName = m.owner + '/' + m.name;
+            const v = m.latest_version!;
+            return {
+                id: fullName + ':' + v.id,
+                name:
+                    fullName + "@" + v.cog_version + ":" + v.id.slice(0, 6),
+                provider: this.provider,
+                owner: m.owner,
+                description: m.description,
+            } as AIModel;
+        });
+    }
+
+    async listTrainableModels(): Promise<AIModel[]> {
+        if (!cachedTrainableModels || Date.now() > cachedTrainableModelsTimestamp + 12 * 3600 * 1000) { // 12 hours
+            cachedTrainableModels = await this._listTrainableModels();
+            cachedTrainableModelsTimestamp = Date.now();
+        }
+        return cachedTrainableModels;
+    }
+
+    async listModels(params: ModelSearchPayload): Promise<AIModel[]> {
+        if (!params.text) {
+            return this.listTrainableModels();
+        }
+        const [owner, model] = params.text.split("/");
+        if (!owner || !model) {
+            throw new Error("Invalid model name. Expected format: owner/model");
+        }
+
+        return this.listModelVersions(owner, model);
+    }
+
+    async listModelVersions(owner: string, model: string): Promise<AIModel[]> {
+        const [rModel, versions] = await Promise.all([
+            this.service.models.get(owner, model),
+            this.service.models.versions.list(owner, model),
+        ]);
+
+        if (!rModel || !versions || versions.length === 0) {
+            throw new Error("Model not found or no versions avaialble");
+        }
+
+        const models: AIModel[] = (versions as any).results.map((v: any) => {
+            const fullName = rModel.owner + '/' + rModel.name;
+            return {
+                id: fullName + ':' + v.id,
+                name:
+                    fullName + "@" + v.cog_version + ":" + v.id.slice(0, 6),
+                provider: this.provider,
+                owner: rModel.owner,
+                description: rModel.description,
+                canTrain: supportFineTunning.has(fullName),
+            } as AIModel;
+        });
+
+        //set latest version
+        //const idx = models.findIndex(m => m.id === rModel.latest_version?.id);
+        //models[idx].name = rModel.name + "@latest"
+
+        return models;
+    }
+
+    async searchModels(params: ModelSearchPayload): Promise<AIModel[]> {
+        const res = await this.service.request("models/search", {
+            params: {
+                query: params.text,
+            },
+        });
+
+        const rModels = ((await res.json()) as any).models;
+
+        const models: AIModel[] = rModels.map((v: any) => {
+            return {
+                id: v.name,
+                name: v.name,
+                provider: this.provider,
+                owner: v.username,
+                description: v.description,
+                has_versions: true,
+            };
+        });
+
+        return models;
+    }
+
+    generateEmbeddings(content: string, model?: string): Promise<{ embeddings: number[], model: string; }> {
+        this.logger?.debug(`[Replicate] Generating embeddings for ${content} on ${model}`);
+        throw new Error("Method not implemented.");
+    }
+
+}
+
+function jobInfo(job: Prediction, modelName?: string): TrainingJob {
+    // 'starting' | 'processing' | 'succeeded' | 'failed' | 'canceled'
+    const jobStatus = job.status;
+    let details: string | undefined;
+    let status = TrainingJobStatus.running;
+    if (jobStatus === 'succeeded') {
+        status = TrainingJobStatus.succeeded;
+    } else if (jobStatus === 'failed') {
+        status = TrainingJobStatus.failed;
+        const error = job.error;
+        if (typeof error === 'string') {
+            details = error;
+        } else {
+            const parts = [];
+            if (error.code) {
+                parts.push(error.code + ' - ');
+            }
+            if (error.message) {
+                parts.push(error.message);
+            }
+            if (parts.length) {
+                details = parts.join(' ');
+            } else {
+                details = JSON.stringify(error);
+            }
+        }
+        details = job.error ? `${job.error.code} - ${job.error.message} ${job.error.param ? " [" + job.error.param + "]" : ""}` : "error";
+    } else if (jobStatus === 'canceled') {
+        status = TrainingJobStatus.cancelled;
+    } else {
+        status = TrainingJobStatus.running;
+        details = job.status;
+    }
+
+    return {
+        id: job.id,
+        status,
+        details,
+        model: modelName ? modelName + ':' + job.version : job.version
+    } as TrainingJob;
+
+}

package/src/test/TestErrorCompletionStream.ts

@@ -0,0 +1,17 @@
+import { CompletionStream, ExecutionOptions, ExecutionResponse, PromptSegment } from "@llumiverse/core";
+import { sleep, throwError } from "./utils.js";
+
+export class TestErrorCompletionStream implements CompletionStream<PromptSegment[]> {
+
+    completion: ExecutionResponse<PromptSegment[]> | undefined;
+
+    constructor(public segments: PromptSegment[],
+        public options: ExecutionOptions) {
+    }
+    async *[Symbol.asyncIterator]() {
+        yield "Started TestError. Next we will thrown an error.\n";
+        sleep(1000);
+        throwError("Testing stream completion error.", this.segments);
+    }
+}
+

package/src/test/TestValidationErrorCompletionStream.ts

@@ -0,0 +1,21 @@
+import { CompletionStream, ExecutionOptions, ExecutionResponse, PromptSegment } from "@llumiverse/core";
+import { createValidationErrorCompletion, sleep } from "./utils.js";
+
+
+export class TestValidationErrorCompletionStream implements CompletionStream<PromptSegment[]> {
+
+    completion: ExecutionResponse<PromptSegment[]> | undefined;
+
+    constructor(public segments: PromptSegment[],
+        public options: ExecutionOptions) {
+    }
+    async *[Symbol.asyncIterator]() {
+        yield "Started TestValidationError.\n";
+        await sleep(1000);
+        yield "chunk1\n"
+        await sleep(1000);
+        yield "chunk2\n"
+        await sleep(1000);
+        this.completion = createValidationErrorCompletion(this.segments);
+    }
+}

package/src/test/index.ts

@@ -0,0 +1,102 @@
+import { AIModel, AIModelStatus, CompletionStream, Driver, ExecutionOptions, ExecutionResponse, ModelType, PromptOptions, PromptSegment, TrainingJob } from "@llumiverse/core";
+import { TestErrorCompletionStream } from "./TestErrorCompletionStream.js";
+import { TestValidationErrorCompletionStream } from "./TestValidationErrorCompletionStream.js";
+import { createValidationErrorCompletion, sleep, throwError } from "./utils.js";
+
+export * from "./TestErrorCompletionStream.js";
+export * from "./TestValidationErrorCompletionStream.js";
+
+export enum TestDriverModels {
+    executionError = "execution-error",
+    validationError = "validation-error",
+}
+
+export class TestDriver implements Driver<PromptSegment[]> {
+    provider = "test";
+
+    createTrainingPrompt(): string {
+        throw new Error("Method not implemented.");
+    }
+
+    startTraining(): Promise<TrainingJob> {
+        throw new Error("Method not implemented.");
+    }
+
+    cancelTraining(): Promise<TrainingJob> {
+        throw new Error("Method not implemented.");
+    }
+
+    getTrainingJob(_jobId: string): Promise<TrainingJob> {
+        throw new Error("Method not implemented.");
+    }
+
+    createPrompt(segments: PromptSegment[], _opts: PromptOptions): PromptSegment[] {
+        return segments;
+    }
+    execute(segments: PromptSegment[], options: ExecutionOptions): Promise<ExecutionResponse<PromptSegment[]>> {
+        switch (options.model) {
+            case TestDriverModels.executionError:
+                return this.executeError(segments, options);
+            case TestDriverModels.validationError:
+                return this.executeValidationError(segments, options);
+            default:
+                throwError("[test driver] Unknown model: " + options.model, segments)
+        }
+    }
+    async stream(segments: PromptSegment[], options: ExecutionOptions): Promise<CompletionStream<PromptSegment[]>> {
+        switch (options.model) {
+            case TestDriverModels.executionError:
+                return new TestErrorCompletionStream(segments, options);
+            case TestDriverModels.validationError:
+                return new TestValidationErrorCompletionStream(segments, options);
+            default:
+                throwError("[test driver] Unknown model: " + options.model, segments)
+        }
+    }
+
+    async listTrainableModels(): Promise<AIModel<string>[]> {
+        return [];
+    }
+
+    async listModels(): Promise<AIModel<string>[]> {
+        return [
+            {
+                id: TestDriverModels.executionError,
+                name: "Execution Error",
+                type: ModelType.Test,
+                provider: this.provider,
+                status: AIModelStatus.Available,
+                description: "Test execution errors",
+                tags: [],
+            },
+            {
+                id: TestDriverModels.validationError,
+                name: "Validation Error",
+                type: ModelType.Test,
+                provider: this.provider,
+                status: AIModelStatus.Available,
+                description: "Test validation errors",
+                tags: [],
+            },
+        ]
+    }
+    validateConnection(): Promise<boolean> {
+        throw new Error("Method not implemented.");
+    }
+    generateEmbeddings(): Promise<{ embeddings: number[]; model: string; }> {
+        throw new Error("Method not implemented.");
+    }
+
+    // ============== execution error ==================
+    async executeError(segments: PromptSegment[], _options: ExecutionOptions): Promise<ExecutionResponse<PromptSegment[]>> {
+        await sleep(1000);
+        throwError("Testing stream completion error.", segments);
+    }
+    // ============== validation error ==================
+    async executeValidationError(segments: PromptSegment[], _options: ExecutionOptions): Promise<ExecutionResponse<PromptSegment[]>> {
+        await sleep(3000);
+        return createValidationErrorCompletion(segments);
+    }
+
+}
+

package/src/test/utils.ts

@@ -0,0 +1,28 @@
+import { ExecutionResponse, PromptSegment } from "@llumiverse/core";
+
+export function throwError(message: string, prompt: PromptSegment[]): never {
+    const err = new Error(message);
+    (err as any).prompt = prompt;
+    throw err;
+}
+
+export function createValidationErrorCompletion(segments: PromptSegment[]) {
+    return {
+        result: "An invalid result",
+        prompt: segments,
+        execution_time: 3000,
+        error: {
+            code: "validation_error",
+            message: "Result cannot be validated!",
+        },
+        token_usage: {
+            result: 10,
+            prompt: 10,
+            total: 20,
+        },
+    } as ExecutionResponse<PromptSegment[]>;
+}
+
+export function sleep(ms: number) {
+    return new Promise(resolve => setTimeout(resolve, ms));
+}

package/src/togetherai/index.ts

@@ -0,0 +1,105 @@
+import { AIModel, AbstractDriver, Completion, DriverOptions, ExecutionOptions, PromptFormats } from "@llumiverse/core";
+import { FetchClient, ServerSentEvent } from "api-fetch-client";
+import { TogetherModelInfo } from "./interfaces.js";
+
+interface TogetherAIDriverOptions extends DriverOptions {
+    apiKey: string;
+}
+
+export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, string> {
+    provider: string;
+    apiKey: string;
+    defaultFormat: PromptFormats;
+    fetchClient: FetchClient;
+
+    constructor(options: TogetherAIDriverOptions) {
+        super(options);
+        this.provider = "togetherai";
+        this.defaultFormat = PromptFormats.genericTextLLM;
+        this.apiKey = options.apiKey;
+        this.fetchClient = new FetchClient('https://api.together.xyz').withHeaders({
+            authorization: `Bearer ${this.apiKey}`
+        });
+    }
+
+    async requestCompletion(prompt: string, options: ExecutionOptions): Promise<Completion<any>> {
+        const res = await this.fetchClient.post('/v1/completions', {
+            payload: {
+                model: options.model,
+                prompt: prompt,
+                max_tokens: options.max_tokens ?? 1024,
+                temperature: options.temperature ?? 0.7,
+            }
+        })
+
+        const text = res.choices[0]?.text ?? '';
+        const usage = res.usage || {};
+        return {
+            result: text,
+            token_usage: {
+                prompt: usage.prompt_tokens,
+                result: usage.completion_tokens,
+                total: usage.total_tokens,
+            },
+        }
+    }
+
+    async requestCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+
+        const stream = await this.fetchClient.post('/v1/completions', {
+            payload: {
+                model: options.model,
+                prompt: prompt,
+                max_tokens: options.max_tokens ?? 1024,
+                temperature: options.temperature ?? 0.7,
+                stream: true,
+            },
+            reader: 'sse'
+        })
+
+        return stream.pipeThrough(new TransformStream<ServerSentEvent, string>({
+            transform(event: ServerSentEvent, controller) {
+                if (event.type === 'event' && event.data && event.data !== '[DONE]') {
+                    try {
+                        const data = JSON.parse(event.data);
+                        controller.enqueue(data.choices[0]?.text ?? '');
+                    } catch (err) {
+                        // double check for the last event whicb is not a JSON - at this time togetherai returrns the string [DONE]
+                        // do nothing - happens if data is not a JSON - the last event data is the [DONE] string
+                    }
+                }
+            }
+        }));
+
+    }
+
+    async listModels(): Promise<AIModel<string>[]> {
+        const models: TogetherModelInfo[] = await this.fetchClient.get("/models/info");
+        //        logObject('#### LIST MODELS RESULT IS', models[0]);
+
+        const aimodels = models.map(m => {
+            return {
+                id: m.name,
+                name: m.display_name,
+                description: m.description,
+                provider: this.provider,
+                formats: [PromptFormats.genericTextLLM],
+            }
+        });
+
+        return aimodels;
+
+    }
+
+    listTrainableModels(): Promise<AIModel<string>[]> {
+        throw new Error("Method not implemented.");
+    }
+    validateConnection(): Promise<boolean> {
+        throw new Error("Method not implemented.");
+    }
+    //@ts-ignore
+    generateEmbeddings(content: string, model?: string | undefined): Promise<{ embeddings: number[]; model: string; }> {
+        throw new Error("Method not implemented.");
+    }
+
+}

package/src/togetherai/interfaces.ts

@@ -0,0 +1,88 @@
+interface ModelInstanceConfig {
+    appearsIn: any[];
+    order: number;
+}
+
+interface Config {
+    stop: string[];
+    prompt_format: string;
+    chat_template: string;
+}
+
+interface Pricing {
+    input: number;
+    output: number;
+    hourly: number;
+}
+
+interface Instance {
+    avzone: string;
+    cluster: string;
+}
+
+interface Ask {
+    [key: string]: number;
+}
+
+interface Gpu {
+    [key: string]: number;
+}
+
+interface Price {
+    base: number;
+    finetune: number;
+    hourly: number;
+    input: number;
+    output: number;
+}
+
+interface Stat {
+    avzone: string;
+    cluster: string;
+    capacity: number;
+    qps: number;
+    throughput_in: number;
+    throughput_out: number;
+    error_rate: number;
+    retry_rate: number;
+}
+
+interface Depth {
+    num_asks: number;
+    num_bids: number;
+    num_running: number;
+    asks: Ask;
+    asks_updated: string;
+    gpus: Gpu;
+    qps: number;
+    permit_required: boolean;
+    price: Price;
+    throughput_in: number;
+    throughput_out: number;
+    stats: Stat[];
+}
+
+export interface TogetherModelInfo {
+    modelInstanceConfig: ModelInstanceConfig;
+    _id: string;
+    name: string;
+    display_name: string;
+    display_type: string;
+    description: string;
+    license: string;
+    creator_organization: string;
+    hardware_label: string;
+    num_parameters: number;
+    show_in_playground: boolean;
+    isFeaturedModel: boolean;
+    context_length: number;
+    config: Config;
+    pricing: Pricing;
+    created_at: string;
+    update_at: string;
+    instances: Instance[];
+    access: string;
+    link: string;
+    descriptionLink: string;
+    depth: Depth;
+}