@llumiverse/drivers 0.8.0

package/src/vertexai/README.md

@@ -0,0 +1,257 @@
+## Palm2 for text
+See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text
+
+### Request
+`POST https://us-central1-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/us-central1/publishers/google/models/text-bison:predict`
+
+```json
+{
+  "instances": [
+    {
+      "prompt": string
+    }
+  ],
+  "parameters": {
+    "temperature": number,
+    "maxOutputTokens": integer,
+    "topK": integer,
+    "topP": number,
+    "groundingConfig": string,
+    "stopSequences": [ string ],
+    "candidateCount": integer,
+    "logprobs": integer,
+    "presencePenalty": float,
+    "frequencyPenalty": float,
+    "logitBias": map<string, float>,
+    "echo": boolean,
+    "seed": integer
+  }
+}
+```
+
+### Properties
+
+#### prompt
+Text input to generate model response. Prompts can include preamble, questions, suggestions, instructions, or examples.
+
+### Response
+```json
+{
+  "predictions":[
+    {
+      "content": string,
+      "citationMetadata": {
+        "citations": [
+          {
+            "startIndex": integer,
+            "endIndex": integer,
+            "url": string,
+            "title": string,
+            "license": string,
+            "publicationDate": string
+          }
+        ]
+      },
+      "logprobs": {
+        "tokenLogProbs": [ float ],
+        "tokens": [ string ],
+        "topLogProbs": [ { map<string, float> } ]
+      },
+      "safetyAttributes": {
+        "categories": [ string ],
+        "blocked": boolean,
+        "scores": [ float ],
+        "errors": [ int ]
+      }
+    }
+  ],
+  "metadata": {
+    "tokenMetadata": {
+      "input_token_count": {
+        "total_tokens": integer,
+        "total_billable_characters": integer
+      },
+      "output_token_count": {
+        "total_tokens": integer,
+        "total_billable_characters": integer
+      }
+    }
+  }
+}
+```
+
+## Code completion
+
+See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/code-generation
+
+### Request:
+`POST https://us-central1-googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/publishers/google/models/code-gecko:predict`
+
+```json
+{
+  "instances":[
+    {
+      "prefix": string,
+      "suffix": string
+    }
+  ],
+  "parameters": {
+    "temperature": number,
+    "maxOutputTokens": integer,
+    "candidateCount": integer,
+    "stopSequences": [ string ],
+    "logprobs": integer,
+    "presencePenalty": float,
+    "frequencyPenalty": float,
+    "logitBias": map<string, float>,
+    "echo": boolean,
+    "seed": integer
+  }
+}
+```
+### Properties
+
+#### prefix (required)
+
+For code models, prefix represents the beginning of a piece of meaningful programming code or a natural language prompt that describes code to be generated. The model attempts to fill in the code in between the prefix and suffix.
+
+#### suffix (optional)
+For code completion, suffix represents the end of a piece of meaningful programming code. The model attempts to fill in the code in between the prefix and suffix.
+
+
+### Response
+
+```json
+{
+  "predictions": [
+    {
+      "content": string,
+      "citationMetadata": {
+        "citations": [
+          {
+            "startIndex": integer,
+            "endIndex": integer,
+            "url": string,
+            "title": string,
+            "license": string,
+            "publicationDate": string
+          }
+        ]
+      },
+      "logprobs": {
+        "tokenLogProbs": [ float ],
+        "tokens": [ string ],
+        "topLogProbs": [ { map<string, float> } ]
+      },
+      "safetyAttributes":{
+        "categories": [ string ],
+        "blocked": boolean,
+        "scores": [ float ],
+        "errors": [ int ]
+      },
+      "score": float
+    }
+  ]
+}
+```
+
+
+## Code generation
+
+See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/code-generation
+
+### Request
+`POST https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/publishers/google/models/code-bison:predict`
+
+
+```json
+{
+  "instances": [
+    { "prefix": string }
+  ],
+  "parameters": {
+    "temperature": number,
+    "maxOutputTokens": integer,
+    "candidateCount": integer,
+    "stopSequences": [ string ]
+    "logprobs": integer,
+    "presencePenalty": float,
+    "frequencyPenalty": float,
+    "logitBias": map<string, float>,
+    "echo": boolean,
+    "seed": integer
+  }
+}
+```
+### Properties
+
+#### prefix (required)
+
+For code models, prefix represents the beginning of a piece of meaningful programming code or a natural language prompt that describes code to be generated. The model attempts to fill in the code in between the prefix and suffix.
+
+### Response
+```json
+{
+  "predictions": [
+    {
+      "content": string,
+      "citationMetadata": {
+        "citations": [
+          {
+            "startIndex": integer,
+            "endIndex": integer,
+            "url": string,
+            "title": string,
+            "license": string,
+            "publicationDate": string
+          }
+        ]
+      },
+      "logprobs": {
+        "tokenLogProbs": [ float ],
+        "tokens": [ string ],
+        "topLogProbs": [ { map<string, float> } ]
+      },
+      "safetyAttributes": {
+        "categories": [ string ],
+        "blocked": false,
+        "scores": [ float ],
+        "errors": [ int ]
+      },
+      "score": float
+    }
+  ]
+}
+```
+
+## Embeddings
+See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings
+
+### Request
+`POST https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/publishers/google/models/textembedding-gecko:predict`
+
+```json
+{
+  "instances": [
+    { "content": "TEXT"}
+  ],
+}
+```
+
+### Response
+
+```json
+{
+  "predictions": [
+    {
+      "embeddings": {
+        "statistics": {
+          "truncated": boolean,
+          "token_count": integer
+        },
+        "values": [ number ]
+      }
+    }
+  ]
+}
+```

package/src/vertexai/debug.ts

@@ -0,0 +1,6 @@
+import util from 'util';
+
+export function logObject(prefix: string, obj: any) {
+    const fullObj = util.inspect(obj, { showHidden: false, depth: null, colors: true });
+    console.log(prefix, fullObj)
+}

package/src/vertexai/index.ts

@@ -0,0 +1,99 @@
+//import { v1 } from "@google-cloud/aiplatform";
+import { GenerateContentRequest, VertexAI } from "@google-cloud/vertexai";
+import { AIModel, AbstractDriver, BuiltinProviders, Completion, DriverOptions, ExecutionOptions, ModelSearchPayload, PromptFormats, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { FetchClient } from "api-fetch-client";
+import { BuiltinModels, getModelDefinition } from "./models.js";
+//import { GoogleAuth } from "google-auth-library";
+
+export interface VertexAIDriverOptions extends DriverOptions {
+    project: string;
+    region: string;
+}
+
+export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, GenerateContentRequest> {
+    provider = BuiltinProviders.vertexai;
+    defaultFormat = PromptFormats.genericTextLLM;
+
+    //aiplatform: v1.ModelServiceClient;
+    vertexai: VertexAI;
+    fetchClient: FetchClient;
+
+    constructor(
+        options: VertexAIDriverOptions
+    ) {
+        super(options);
+        //this.aiplatform = new v1.ModelServiceClient();
+        this.vertexai = new VertexAI({
+            project: this.options.project,
+            location: this.options.region,
+        });
+        this.fetchClient = createFetchClient({
+            region: this.options.region,
+            project: this.options.project,
+        }).withAuthCallback(async () => {
+            const token = await this.vertexai.preview.token;
+            return `Bearer ${token}`;
+        });
+        // this.aiplatform = new v1.ModelServiceClient({
+        //     projectId: this.options.project,
+        //     apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
+        // });
+    }
+
+    protected canStream(options: ExecutionOptions): Promise<boolean> {
+        return Promise.resolve(getModelDefinition(options.model).model.canStream === true);
+    }
+
+    public createPrompt(segments: PromptSegment[], options: PromptOptions): GenerateContentRequest {
+        return getModelDefinition(options.model).createPrompt(this, segments, options);
+    }
+
+    async requestCompletion(prompt: GenerateContentRequest, options: ExecutionOptions): Promise<Completion<any>> {
+        return getModelDefinition(options.model).requestCompletion(this, prompt, options);
+    }
+    async requestCompletionStream(prompt: GenerateContentRequest, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+        return getModelDefinition(options.model).requestCompletionStream(this, prompt, options);
+    }
+
+    async listModels(_params?: ModelSearchPayload): Promise<AIModel<string>[]> {
+        return BuiltinModels;
+        // try {
+        //     const response = await this.fetchClient.get('/publishers/google/models/gemini-pro');
+        //     console.log('>>>>>>>>', response);
+        // } catch (err: any) {
+        //     console.error('+++++VETREX ERROR++++++', err);
+        //     throw err;
+        // }
+
+        // TODO uncomment this to use apiplatform instead of the fetch client
+        // const response = await this.aiplatform.listModels({
+        //     parent: `projects/${this.options.project}/locations/${this.options.region}`,
+        // });
+
+        return []; //TODO
+    }
+    listTrainableModels(): Promise<AIModel<string>[]> {
+        throw new Error("Method not implemented.");
+    }
+    validateConnection(): Promise<boolean> {
+        throw new Error("Method not implemented.");
+    }
+    generateEmbeddings(_content: string, _model?: string | undefined): Promise<{ embeddings: number[]; model: string; }> {
+        throw new Error("Method not implemented.");
+    }
+
+}
+
+//'us-central1-aiplatform.googleapis.com',
+const API_BASE_PATH = 'aiplatform.googleapis.com';
+function createFetchClient({ region, project, apiEndpoint, apiVersion = 'v1' }: {
+    region: string;
+    project: string;
+    apiEndpoint?: string;
+    apiVersion?: string;
+}) {
+    const vertexBaseEndpoint = apiEndpoint ?? `${region}-${API_BASE_PATH}`;
+    return new FetchClient(`https://${vertexBaseEndpoint}/${apiVersion}/projects/${project}/locations/${region}`).withHeaders({
+        'Content-Type': 'application/json',
+    });
+}

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

@@ -0,0 +1,115 @@
+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 CodeyChatMessage {
+    author: string,
+    content: string
+}
+
+export type CodeyChatPrompt = NonStreamingPromptBase<{
+    context?: string;
+    messages: CodeyChatMessage[];
+}>
+
+
+export type CodeyChatStreamingPrompt = StreamingPromptBase<{
+    messages: {
+        listVal: {
+            structVal: {
+                author: {
+                    stringVal: string
+                },
+                content: {
+                    stringVal: string
+                }
+            }
+        }[]
+    }
+}>
+
+
+interface CodeyChatResponsePrediction {
+    candidates: CodeyChatMessage[],
+    safetyAttributes: {
+        scores: number[],
+        blocked: boolean,
+        categories: string[],
+        errors: number[],
+    },
+    citationMetadata: {
+        citations: any[]
+    },
+    logprobs: any,
+}
+
+export interface CodeyChatResponse {
+    predictions: CodeyChatResponsePrediction[],
+    metadata: PalmResponseMetadata
+}
+
+export class CodeyChatDefinition extends AbstractPalmModelDefinition<CodeyChatPrompt, CodeyChatStreamingPrompt> {
+    versions: string[] = [];
+    model: AIModel = {
+        id: "codechat-bison",
+        name: "Codey for Code Chat",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Chat,
+        canStream: true,
+    }
+
+    createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): CodeyChatPrompt {
+        const system: string[] = [];
+        const messages: CodeyChatMessage[] = [];
+        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 CodeyChatPrompt;
+    }
+
+    extractContentFromResponse(response: CodeyChatResponse): 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/codey-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 CodeyTextPrompt = NonStreamingPromptBase<{
+    prefix: string
+}>
+
+
+export type CodeyTextStreamingPrompt = StreamingPromptBase<{
+    prefix: {
+        stringVal: string
+    }
+}>
+
+interface CodeyTextResponsePrediction {
+    content: string,
+    safetyAttributes: {
+        scores: number[],
+        safetyRatings: {
+            severity: string,
+            probabilityScore: number,
+            severityScore: number,
+            category: string
+        }[]
+    },
+    citationMetadata: {
+        citations: any[]
+    }
+}
+
+export interface CodeyTextResponse {
+    predictions: CodeyTextResponsePrediction[],
+    metadata: PalmResponseMetadata
+}
+
+export class CodeyTextDefinition extends AbstractPalmModelDefinition<CodeyTextPrompt, CodeyTextStreamingPrompt> {
+    versions: string[] = [];
+    model: AIModel = {
+        id: "code-bison",
+        name: "Codey for Code Generation",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Text,
+        canStream: true,
+    }
+
+    createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): CodeyTextPrompt {
+        return {
+            instances: [{
+                prefix: getPromptAsText(segments, opts)
+            }],
+            parameters: {
+                // put defauilts here
+            }
+        } as CodeyTextPrompt;
+    }
+
+    extractContentFromResponse(response: CodeyTextResponse): string {
+        return response.predictions[0].content ?? '';
+    }
+
+    extractContentFromResponseChunk(chunk: any): string {
+        return chunk.outputs[0]?.structVal.content.stringVal || '';
+    }
+
+}

package/src/vertexai/models/gemini.ts

@@ -0,0 +1,152 @@
+import { Content, GenerateContentRequest, GenerativeModel, HarmBlockThreshold, HarmCategory, TextPart } from "@google-cloud/vertexai";
+import { AIModel, Completion, ExecutionOptions, ExecutionTokenUsage, ModelType, PromptOptions, PromptRole, PromptSegment } from "@llumiverse/core";
+import { asyncMap } from "@llumiverse/core/async";
+import { VertexAIDriver } from "../index.js";
+import { ModelDefinition } from "../models.js";
+
+function getGenerativeModel(driver: VertexAIDriver, options: ExecutionOptions): GenerativeModel {
+    return driver.vertexai.preview.getGenerativeModel({
+        model: options.model,
+        //TODO pass in the options
+        safety_settings: [{
+            category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+            threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE
+        }],
+        generation_config: {
+            temperature: options.temperature,
+            max_output_tokens: options.max_tokens
+        },
+    });
+}
+
+
+function collectTextParts(content: Content) {
+    const out = [];
+    const parts = content.parts;
+    if (parts) {
+        for (const part of parts) {
+            if (part.text) {
+                out.push(part.text);
+            }
+        }
+    }
+    return out.join('\n');
+}
+
+export class GeminiModelDefinition implements ModelDefinition<GenerateContentRequest> {
+
+    model: AIModel = {
+        id: "gemini-pro",
+        name: "Gemini Pro",
+        provider: "vertexai",
+        owner: "google",
+        type: ModelType.Text,
+        canStream: true,
+    }
+
+    createPrompt(_driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions): GenerateContentRequest {
+        const schema = options.resultSchema;
+        const contents: Content[] = [];
+        const safety: string[] = [];
+
+        let lastUserContent: Content | undefined = undefined;
+        for (const msg of segments) {
+            if (msg.role === PromptRole.safety) {
+                safety.push(msg.content);
+            } else {
+                const role = msg.role === PromptRole.assistant ? "model" : "user";
+                if (lastUserContent && lastUserContent.role === role) {
+                    lastUserContent.parts.push({ text: msg.content } as TextPart);
+                } else {
+                    const content: Content = {
+                        role,
+                        parts: [{ text: msg.content } as TextPart],
+                    }
+                    if (role === 'user') {
+                        lastUserContent = content;
+                    }
+                    contents.push(content);
+                }
+            }
+        }
+
+        let tools: any = undefined;
+        if (schema) {
+
+            // tools = [{
+            //     function_declarations: [{
+            //         name: "validate_json_response",
+            //         description: "Validate the given JSON response",
+            //         parameters: schema as any,
+            //     }]
+            // } as Tool];
+
+            safety.push("The answer must be a JSON object using the following JSON Schema:\n" + JSON.stringify(schema));
+        }
+
+        if (safety.length > 0) {
+            const content = safety.join('\n');
+            if (lastUserContent) {
+                lastUserContent.parts.push({ text: content } as TextPart);
+            } else {
+                contents.push({
+                    role: 'user',
+                    parts: [{ text: content } as TextPart],
+                })
+            }
+        }
+
+        // put system mesages first and safety last
+        return { contents, tools } as GenerateContentRequest;
+    }
+
+    async requestCompletion(driver: VertexAIDriver, prompt: GenerateContentRequest, options: ExecutionOptions): Promise<Completion> {
+        const model = getGenerativeModel(driver, options);
+        const r = await model.generateContent(prompt);
+        const response = await r.response;
+        const usage = response.usageMetadata;
+        const token_usage: ExecutionTokenUsage = {
+            prompt: usage?.prompt_token_count,
+            result: usage?.candidates_token_count,
+            total: usage?.totalTokenCount,
+        }
+
+        let result: any;
+        const candidate = response.candidates[0];
+        if (candidate) {
+            const content = candidate.content;
+            if (content) {
+                result = collectTextParts(content);
+                // if (options.resultSchema) {
+                //     result = candidate.;
+                // } else {
+                // }
+            }
+        }
+
+        return {
+            result: result ?? '',
+            token_usage
+        };
+    }
+
+    async requestCompletionStream(driver: VertexAIDriver, prompt: GenerateContentRequest, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+        const model = getGenerativeModel(driver, options);
+        const streamingResp = await model.generateContentStream(prompt);
+
+        const stream = asyncMap(streamingResp.stream, async (item) => {
+            if (item.candidates.length > 0) {
+                for (const candidate of item.candidates) {
+                    if (candidate.content?.role === 'model') {
+                        const text = collectTextParts(candidate.content);
+                        if (text) return text;
+                    }
+                }
+            }
+            return '';
+        });
+
+        return stream;
+    }
+
+}