@llumiverse/drivers 0.8.0

package/src/vertexai/models/palm-model-base.ts

@@ -0,0 +1,122 @@
+import { AIModel, Completion, ExecutionOptions, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { VertexAIDriver } from "../index.js";
+import { ModelDefinition } from "../models.js";
+import { PromptParamatersBase } from "../utils/prompts.js";
+import { generateStreamingPrompt } from "../utils/tensor.js";
+import { ServerSentEvent } from "api-fetch-client";
+
+export interface NonStreamingPromptBase<InstanceType = any> {
+    instances: InstanceType[];
+    parameters: PromptParamatersBase;
+}
+
+export interface StreamingPromptBase<InputType = any> {
+    inputs: { structVal: InputType }[];
+    parameters: {
+        structVal: {
+            temperature?: { floatval: number },
+            maxOutputTokens?: { intVal: number },
+            //TODO more params
+            [key: string]: Record<string, any> | undefined
+        }
+    }
+}
+
+export interface PalmResponseMetadata {
+    tokenMetadata: {
+        outputTokenCount: {
+            totalBillableCharacters: number,
+            totalTokens: number
+        },
+        inputTokenCount: {
+            totalBillableCharacters: number,
+            totalTokens: number
+        }
+    }
+}
+
+
+export abstract class AbstractPalmModelDefinition<NonStreamingPromptT extends NonStreamingPromptBase, StreamingPromptT extends StreamingPromptBase> implements ModelDefinition<NonStreamingPromptT | StreamingPromptT> {
+
+    abstract model: AIModel;
+
+    abstract versions?: string[];
+
+    abstract createNonStreamingPrompt(driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions): NonStreamingPromptT;
+
+    abstract extractContentFromResponse(response: any): string;
+
+    abstract extractContentFromResponseChunk(chunk: any): string;
+
+    createPrompt(driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions) {
+        return this.createNonStreamingPrompt(driver, segments, options);
+    }
+
+    async requestCompletion(driver: VertexAIDriver, prompt: NonStreamingPromptT | StreamingPromptT, options: ExecutionOptions): Promise<Completion<any>> {
+        const nonStreamingPrompt = prompt as NonStreamingPromptT;
+        Object.assign((nonStreamingPrompt).parameters, {
+            temperature: options.temperature,
+            maxOutputTokens: options.max_tokens,
+        });
+
+        const response = await driver.fetchClient.post(`/publishers/google/models/${this.model.id}:predict`, {
+            payload: nonStreamingPrompt
+        });
+
+        const metadata = response.metadata as PalmResponseMetadata;
+        const inputTokens = metadata.tokenMetadata.inputTokenCount.totalTokens;
+        const outputTokens = metadata.tokenMetadata.outputTokenCount.totalTokens;
+        const result = this.extractContentFromResponse(response);
+        return {
+            result,
+            token_usage: {
+                prompt: inputTokens,
+                result: outputTokens,
+                total: inputTokens && outputTokens ? inputTokens + outputTokens : undefined,
+            }
+        } as Completion;
+
+    }
+
+    async requestCompletionStream(driver: VertexAIDriver, prompt: NonStreamingPromptT | StreamingPromptT, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+        const inPrompt = prompt as NonStreamingPromptT;
+        Object.assign(inPrompt.parameters, {
+            temperature: options.temperature,
+            maxOutputTokens: options.max_tokens,
+        });
+
+        const path = `/publishers/google/models/${this.model.id}:serverStreamingPredict?alt=sse`;
+
+        const newPrompt = generateStreamingPrompt(inPrompt);
+
+        // we need to modify the existing prompt since it is not the final one
+        const outPrompt = prompt as StreamingPromptT;
+        delete (outPrompt as any).instances;
+        outPrompt.inputs = newPrompt.inputs;
+        outPrompt.parameters = newPrompt.parameters;
+
+        const eventStrean = await driver.fetchClient.post(path, {
+            payload: newPrompt,
+            reader: 'sse'
+        });
+        return eventStrean.pipeThrough(new ChunkTransformStream(this));
+
+    }
+
+}
+
+class ChunkTransformStream<NonStreamingPromptT extends NonStreamingPromptBase, StreamingPromptT extends StreamingPromptBase> extends TransformStream {
+    constructor(def: AbstractPalmModelDefinition<NonStreamingPromptT, StreamingPromptT>) {
+        super({
+            transform(event: ServerSentEvent, controller: TransformStreamDefaultController) {
+                if (event.type === 'event' && event.data) {
+                    const data = JSON.parse(event.data);
+                    const stringChunk = def.extractContentFromResponseChunk(data);
+                    controller.enqueue(Array.isArray(stringChunk) ? stringChunk.join('') : stringChunk);
+                }
+            }
+        })
+    }
+}
+
+

package/src/vertexai/models/palm2-chat.ts

@@ -0,0 +1,119 @@
+import { AIModel, ModelType, PromptOptions, PromptRole, PromptSegment } from "@llumiverse/core";
+import { VertexAIDriver } from "../index.js";
+import { getJSONSafetyNotice } from "../utils/prompts.js";
+import { AbstractPalmModelDefinition, NonStreamingPromptBase, PalmResponseMetadata, StreamingPromptBase } from "./palm-model-base.js";
+
+export interface Palm2ChatMessage {
+    author: string,
+    content: string
+}
+
+export type Palm2ChatPrompt = NonStreamingPromptBase<{
+    context?: string;
+    examples?: [
+        {
+            input: { content: string },
+            output: { content: string }
+        }
+    ],
+    messages: Palm2ChatMessage[];
+}>
+
+
+export type Palm2ChatStreamingPrompt = StreamingPromptBase<{
+    messages: {
+        listVal: {
+            structVal: {
+                author: {
+                    stringVal: string
+                },
+                content: {
+                    stringVal: string
+                }
+            }
+        }[]
+    }
+}>
+
+interface Palm2ChatResponsePrediction {
+    candidates: Palm2ChatMessage[],
+    safetyAttributes: {
+        scores: number[],
+        blocked: boolean,
+        categories: string[],
+        errors: number[],
+    },
+    citationMetadata: {
+        citations: any[]
+    },
+    logprobs: any,
+}
+
+export interface Palm2ChatResponse {
+    predictions: Palm2ChatResponsePrediction[],
+    metadata: PalmResponseMetadata
+}
+
+export class Palm2ChatDefinition extends AbstractPalmModelDefinition<Palm2ChatPrompt, Palm2ChatStreamingPrompt> {
+    versions: string[] = [];
+    model: AIModel = {
+        id: "chat-bison",
+        name: "PaLM 2 for Chat",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Chat,
+        canStream: true,
+    }
+
+    createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): Palm2ChatPrompt {
+        const system: string[] = [];
+        const messages: Palm2ChatMessage[] = [];
+        const safety: string[] = [];
+        for (const segment of segments) {
+            switch (segment.role) {
+                case PromptRole.user:
+                    messages.push({ author: 'user', content: segment.content });
+                    break;
+                case PromptRole.assistant:
+                    messages.push({ author: 'assistant', content: segment.content });
+                    break;
+                case PromptRole.system:
+                    system.push(segment.content);
+                    break;
+                case PromptRole.safety:
+                    safety.push(segment.content);
+                    break;
+            }
+        }
+
+        if (opts.resultSchema) {
+            safety.push(getJSONSafetyNotice(opts.resultSchema));
+        }
+
+        const context = []
+        if (system.length > 0) {
+            context.push(system.join('\n'));
+        }
+        if (safety.length > 0) {
+            context.push('IMPORTANT: ' + safety.join('\n'));
+        }
+
+        return {
+            instances: [{
+                context: context.length > 0 ? context.join('\n') : undefined,
+                messages
+            }],
+            parameters: {
+                // put defauilts here
+            }
+        } as Palm2ChatPrompt;
+    }
+
+    extractContentFromResponse(response: Palm2ChatResponse): string {
+        return response.predictions[0].candidates[0].content ?? '';
+    }
+
+    extractContentFromResponseChunk(chunk: any): string {
+        return chunk.outputs[0]?.structVal.candidates.listVal[0].structVal.content.stringVal || ''
+    }
+}

package/src/vertexai/models/palm2-text.ts

@@ -0,0 +1,69 @@
+import { AIModel, ModelType, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { VertexAIDriver } from "../index.js";
+import { getPromptAsText } from "../utils/prompts.js";
+import { AbstractPalmModelDefinition, NonStreamingPromptBase, PalmResponseMetadata, StreamingPromptBase } from "./palm-model-base.js";
+
+export type Palm2TextPrompt = NonStreamingPromptBase<{
+    prompt: string
+}>
+
+
+export type Palm2TextStreamingPrompt = StreamingPromptBase<{
+    prompt: {
+        stringVal: string
+    }
+}>
+
+
+interface Palm2TextResponsePrediction {
+    content: string,
+    safetyAttributes: {
+        scores: number[],
+        safetyRatings: {
+            severity: string,
+            probabilityScore: number,
+            severityScore: number,
+            category: string
+        }[]
+    },
+    citationMetadata: {
+        citations: any[]
+    }
+}
+
+export interface Palm2TextResponse {
+    predictions: Palm2TextResponsePrediction[],
+    metadata: PalmResponseMetadata
+}
+
+export class Palm2TextDefinition extends AbstractPalmModelDefinition<Palm2TextPrompt, Palm2TextStreamingPrompt> {
+    versions: string[] = [];
+    model: AIModel = {
+        id: "text-bison",
+        name: "PaLM 2 for Text",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Text,
+        canStream: true,
+    }
+
+    createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): Palm2TextPrompt {
+        return {
+            instances: [{
+                prompt: getPromptAsText(segments, opts)
+            }],
+            parameters: {
+                // put defauilts here
+            }
+        } as Palm2TextPrompt;
+    }
+
+    extractContentFromResponse(response: Palm2TextResponse): string {
+        return response.predictions[0].content ?? '';
+    }
+
+    extractContentFromResponseChunk(chunk: any): string {
+        return chunk.outputs[0]?.structVal.content.stringVal || '';
+    }
+
+}

package/src/vertexai/models.ts

@@ -0,0 +1,104 @@
+import { AIModel, Completion, ExecutionOptions, ModelType, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { GeminiModelDefinition } from "./models/gemini.js";
+import { VertexAIDriver } from "./index.js";
+import { Palm2TextDefinition } from "./models/palm2-text.js";
+import { Palm2ChatDefinition } from "./models/palm2-chat.js";
+import { CodeyChatDefinition } from "./models/codey-chat.js";
+import { CodeyTextDefinition } from "./models/codey-text.js";
+
+
+const Models: Record<string, ModelDefinition> = {
+    "gemini-pro": new GeminiModelDefinition(),
+    "text-bison": new Palm2TextDefinition(),
+    "chat-bison": new Palm2ChatDefinition(),
+    "code-bison": new CodeyTextDefinition(),
+    "codechat-bison": new CodeyChatDefinition(),
+}
+
+export interface ModelDefinition<PromptT = any> {
+    model: AIModel;
+    versions?: string[]; // the versions of the model that are available. ex: ['001', '002']
+    createPrompt: (driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions) => PromptT;
+    requestCompletion: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<Completion>;
+    requestCompletionStream: (driver: VertexAIDriver, promp: PromptT, options: ExecutionOptions) => Promise<AsyncIterable<string>>;
+}
+
+export function getModelName(model: string) {
+    const i = model.lastIndexOf('@');
+    return i > -1 ? model.substring(0, i) : model;
+}
+
+export function getModelDefinition(model: string): ModelDefinition {
+    const def = Models[getModelName(model)];
+    if (!def) {
+        throw new Error(`Unknown model ${model}`);
+    }
+    return def;
+}
+
+export function getAIModels() {
+    return Object.values(Models).map(m => m.model);
+}
+
+// Builtin models. VertexAI doesn't provide an API to list models. so we have to hardcode them here.
+export const BuiltinModels: AIModel<string>[] = [
+    {
+        id: "gemini-pro",
+        name: "Gemini Pro",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Text,
+    },
+    { // ChatModel.from_pretrained("chat-bison@002")
+        id: "chat-bison",
+        name: "PaLM 2 Chat Bison",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Chat,
+    },
+    {
+        id: "text-bison", // versions 001, 002
+        name: "PaLM 2 Text Bison",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Text,
+    },
+    {
+        id: "code-gecko",
+        name: "Codey for Code Completion",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Code,
+    },
+    {
+        id: "code-bison",
+        name: "Codey for Code Generation",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Code,
+    },
+    {
+        id: "codechat-bison",
+        name: "Codey for Code Chat",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Code,
+    },
+    {
+        id: "tablextembedding-gecko",
+        name: "Gecko Text Embeddings",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Embedding,
+    },
+    {
+        id: "textembedding-gecko-multilingual",
+        name: "Gecko Multilingual Text Embeddings",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Embedding,
+    },
+
+
+
+]

package/src/vertexai/utils/prompts.ts

@@ -0,0 +1,66 @@
+import { PromptOptions, PromptRole, PromptSegment } from "@llumiverse/core";
+import { JSONSchema4 } from "json-schema";
+
+export interface PromptParamatersBase {
+    temperature?: number,
+    maxOutputTokens?: number,
+    topK?: number,
+    topP?: number,
+    groundingConfig?: string,
+    stopSequences?: string[],
+    candidateCount?: number,
+    logprobs?: number,
+    presencePenalty?: number,
+    frequencyPenalty?: number,
+    logitBias?: Record<string, number>,
+    seed?: number,
+    echo?: boolean
+}
+
+export function getJSONSafetyNotice(schema: JSONSchema4) {
+    return "The answer must be a JSON object using the following JSON Schema:\n" + JSON.stringify(schema)
+}
+
+export function getPromptAsText(segments: PromptSegment[], options: PromptOptions): string {
+    const isChat = segments.find(m => m.role === PromptRole.assistant);
+    const context: string[] = [];
+    const content: string[] = [];
+    const safety: string[] = [];
+    for (const segment of segments) {
+        switch (segment.role) {
+            case PromptRole.user:
+                if (isChat) {
+                    content.push('USER: ' + segment.content);
+                } else {
+                    content.push(segment.content);
+                }
+                break;
+            case PromptRole.assistant:
+                content.push('ASSISTANT: ' + segment.content);
+                break;
+            case PromptRole.system:
+                context.push(segment.content);
+                break;
+            case PromptRole.safety:
+                safety.push(segment.content);
+                break;
+        }
+    }
+
+    if (options.resultSchema) {
+        safety.push(getJSONSafetyNotice(options.resultSchema));
+    }
+
+    const out = [];
+    if (context.length > 0) {
+        out.push('CONTEXT: ' + context.join('\n'));
+    }
+    if (content.length > 0) {
+        const prefix = context.length > 0 && !isChat ? 'INSTRUCTION: ' : '';
+        out.push(prefix + content.join('\n'));
+    }
+    if (safety.length > 0) {
+        out.push('IMPORTANT: ' + safety.join('\n'));
+    }
+    return out.join('\n');
+}

package/src/vertexai/utils/tensor.ts

@@ -0,0 +1,82 @@
+
+
+function getTensorType(val: any) {
+    if (val == null) {
+        return null;
+    }
+    if (Array.isArray(val)) {
+        if (val.length > 0) {
+            val = val[0]
+        } else {
+            return 'listVal';
+        }
+    }
+    const type = typeof val;
+    if (type === 'string') {
+        return 'stringVal';
+    }
+    if (type === 'number') {
+        if (val % 1 === 0) { // is interger
+            return 'intVal';
+        } else {
+            return 'floatVal';
+        }
+    }
+    if (type === 'boolean') {
+        return 'boolVal';
+    }
+    if (type === 'object') {
+        return 'structVal';
+    }
+    return undefined;
+}
+
+export function formatArrayAsTensor(arr: any[]): any {
+    return arr.map(item => {
+        const type = typeof item;
+        if (type === 'string' || type === 'number' || type === 'boolean') {
+            return item; // primitve values 
+        } else {
+            return formatObjectAsTensor(item)
+        }
+    });
+}
+
+/**
+ * Do not support nested array or nested object arrays
+ * @param obj 
+ * @returns 
+ */
+function formatObjectAsTensor(obj: any, isField = false): any {
+    const struct: any = {};
+    const keys = Object.keys(obj);
+    for (const key of keys) {
+        //console.log('###', key);
+        const val = obj[key];
+        const type = getTensorType(val);
+        if (type === 'structVal') {
+            if (Array.isArray(val)) {
+                struct[key] = { listVal: formatArrayAsTensor(val) };
+            } else {
+                struct[key] = {
+                    [type]: formatObjectAsTensor(val, true)
+                }
+            }
+        } else if (type) {
+            struct[key] = {
+                [type]: val
+            }
+        }
+    }
+    return isField ? struct : {
+        structVal: struct
+    };
+}
+
+
+export function generateStreamingPrompt(prompt: { instances: any, parameters: any }): any {
+    return {
+        inputs: prompt.instances.map((inst: any) => formatObjectAsTensor(inst)),
+        parameters: formatObjectAsTensor(prompt.parameters)
+    }
+}