@llumiverse/drivers 0.14.0 → 0.16.0

package/src/vertexai/index.ts

@@ -1,9 +1,15 @@
 import { GenerateContentRequest, VertexAI } from "@google-cloud/vertexai";
-import { AIModel, AbstractDriver, Completion, DriverOptions, EmbeddingsResult, ExecutionOptions, ModelSearchPayload, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { AIModel, AbstractDriver, Completion, CompletionChunkObject, DriverOptions, EmbeddingsResult, ExecutionOptions, ImageGeneration, Modalities, ModelSearchPayload, PromptSegment } from "@llumiverse/core";
 import { FetchClient } from "api-fetch-client";
-import { GoogleAuthOptions } from "google-auth-library";
+import { GoogleAuth, GoogleAuthOptions } from "google-auth-library";
+import { JSONClient } from "google-auth-library/build/src/auth/googleauth.js";
 import { TextEmbeddingsOptions, getEmbeddingsForText } from "./embeddings/embeddings-text.js";
-import { BuiltinModels, getModelDefinition } from "./models.js";
+import { getModelDefinition } from "./models.js";
+import { EmbeddingsOptions } from "@llumiverse/core";
+import { getEmbeddingsForImages } from "./embeddings/embeddings-image.js";
+import { v1beta1 } from '@google-cloud/aiplatform';
+import { AnthropicVertex } from '@anthropic-ai/vertex-sdk';
+import { ImagenModelDefinition, ImagenPrompt } from "./models/imagen.js";
 
 
 export interface VertexAIDriverOptions extends DriverOptions {
@@ -12,19 +18,31 @@ export interface VertexAIDriverOptions extends DriverOptions {
     googleAuthOptions?: GoogleAuthOptions;
 }
 
-export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, GenerateContentRequest> {
+//General Prompt type for VertexAI
+export type VertexAIPrompt = GenerateContentRequest | ImagenPrompt;
+
+export function trimModelName(model: string) {
+    const i = model.lastIndexOf('@');
+    return i > -1 ? model.substring(0, i) : model;
+}
+
+export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, VertexAIPrompt> {
     static PROVIDER = "vertexai";
     provider = VertexAIDriver.PROVIDER;
 
-    //aiplatform: v1.ModelServiceClient;
+    aiplatform: v1beta1.ModelServiceClient;
     vertexai: VertexAI;
     fetchClient: FetchClient;
-
-    constructor(
-        options: VertexAIDriverOptions
-    ) {
+    authClient: JSONClient | GoogleAuth<JSONClient>;
+    anthropicClient: AnthropicVertex | undefined;
+    
+    constructor( options: VertexAIDriverOptions) {
         super(options);
-        //this.aiplatform = new v1.ModelServiceClient();
+
+        this.anthropicClient = undefined;
+
+        this.authClient = options.googleAuthOptions?.authClient ?? new GoogleAuth(options.googleAuthOptions);
+
         this.vertexai = new VertexAI({
             project: this.options.project,
             location: this.options.region,
@@ -35,54 +53,113 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Genera
             project: this.options.project,
         }).withAuthCallback(async () => {
             //@ts-ignore
-            const token = await this.options.googleAuthOptions?.authClient?.getAccessToken();
+            const token = await this.authClient.getAccessToken();
             return `Bearer ${token}`;
         });
-        // this.aiplatform = new v1.ModelServiceClient({
-        //     projectId: this.options.project,
-        //     apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
-        // });
+        this.aiplatform = new v1beta1.ModelServiceClient({
+            projectId: this.options.project,
+            apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
+        });
+    }
+
+    public getAnthropicClient() : AnthropicVertex {
+        //Lazy initialisation
+        if (!this.anthropicClient) {
+            this.anthropicClient = new AnthropicVertex({region: "us-east5", projectId: process.env.GOOGLE_PROJECT_ID});
+        }
+        return this.anthropicClient;
     }
 
     protected canStream(options: ExecutionOptions): Promise<boolean> {
+        if (options.output_modality == Modalities.image) {
+            return Promise.resolve(false);
+        }
         return Promise.resolve(getModelDefinition(options.model).model.can_stream === true);
     }
 
-    public createPrompt(segments: PromptSegment[], options: PromptOptions): Promise<GenerateContentRequest> {
+    public createPrompt(segments: PromptSegment[], options: ExecutionOptions): Promise<VertexAIPrompt> {
+        if (options.model.includes("imagen")) {
+            return new ImagenModelDefinition(options.model).createPrompt(this, segments, options);
+        }
         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 requestTextCompletion(prompt: VertexAIPrompt, options: ExecutionOptions): Promise<Completion<any>> {
+        return getModelDefinition(options.model).requestTextCompletion(this, prompt, options);
     }
-    async requestCompletionStream(prompt: GenerateContentRequest, options: ExecutionOptions): Promise<AsyncIterable<string>> {
-        return getModelDefinition(options.model).requestCompletionStream(this, prompt, options);
+    async requestTextCompletionStream(prompt: VertexAIPrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
+        return getModelDefinition(options.model).requestTextCompletionStream(this, prompt, options);
+    }
+
+    async requestImageGeneration(_prompt: ImagenPrompt, _options: ExecutionOptions): Promise <Completion<ImageGeneration>> {
+        const splits = _options.model.split("/");
+        const modelName = trimModelName(splits[splits.length - 1]);
+        return new ImagenModelDefinition(modelName).requestImageGeneration(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
+        let models: AIModel<string>[] = [];
+        const modelGarden = new v1beta1.ModelGardenServiceClient({
+            projectId: this.options.project,
+            apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
+        });
+
+        //Project specific deployed models
+        const [response] = await this.aiplatform.listModels({
+            parent: `projects/${this.options.project}/locations/${this.options.region}`,
+        });
+        models = models.concat(response.map(model => ({
+            id: model.name?.split('/').pop() ?? '',
+            name: model.displayName ?? '',
+            provider: 'vertexai',
+        })));
+
+        //Model Garden Publisher models - Pretrained models
+        const publishers = ['google', 'anthropic']
+        const supportedModels = {google: ['gemini','imagen'], anthropic: ['claude']}
+
+        for (const publisher of publishers) {
+            const [response] = await modelGarden.listPublisherModels({
+                parent: `publishers/${publisher}`,
+                orderBy: 'name',
+                //filter: `name eq name`,
+                listAllVersions: true,
+            });
+
+            models = models.concat(response.map(model => ({
+                id: model.name ?? '',
+                name: model.name?.split('/').pop() ?? '',
+                provider: 'vertexai',
+                owner: publisher,
+            })).filter(model => {
+                const modelFamily = supportedModels[publisher as keyof typeof supportedModels];
+                for (const family of modelFamily) {
+                    if (model.name.includes(family)) {
+                        return true;
+                    }
+                }
+            }));
+        }
+
+        return models;
     }
 
     validateConnection(): Promise<boolean> {
         throw new Error("Method not implemented.");
     }
 
-    async generateEmbeddings(options: TextEmbeddingsOptions): Promise<EmbeddingsResult> {
-        return getEmbeddingsForText(this, options);
+    async generateEmbeddings(options: EmbeddingsOptions): Promise<EmbeddingsResult> {
+        if (options.image || options.model?.includes("multimodal")) {
+            if (options.text && options.image) {
+                throw new Error("Text and Image simultaneous embedding not implemented. Submit seperately");
+            }
+            return getEmbeddingsForImages(this, options);
+        }
+        const text_options: TextEmbeddingsOptions = {
+            content: options.text ?? '',
+            model: options.model,
+        }
+        return getEmbeddingsForText(this, text_options);
     }
 
 }

package/src/vertexai/models/claude.ts

@@ -0,0 +1,281 @@
+import * as AnthropicAPI from '@anthropic-ai/sdk';
+import { ContentBlock, Message, TextBlockParam } from "@anthropic-ai/sdk/resources/index.js";
+import { AIModel, Completion, CompletionChunkObject, ExecutionOptions, JSONObject, ModelType, PromptOptions, PromptRole, PromptSegment, ToolUse } from "@llumiverse/core";
+import { asyncMap } from "@llumiverse/core/async";
+import { VertexAIClaudeOptions } from "../../../../core/src/options/vertexai.js";
+import { VertexAIDriver } from "../index.js";
+import { ModelDefinition } from "../models.js";
+
+type MessageParam = AnthropicAPI.Anthropic.MessageParam;
+
+interface ClaudePrompt {
+    messages: MessageParam[];
+    system: TextBlockParam[];
+}
+
+function getFullModelName(model: string): string {
+    if (model.includes("claude-3-5-sonnet-v2")) {
+        return "claude-3-5-sonnet-v2@20241022"
+    } else if (model.includes("claude-3-5-sonnet")) {
+        return "claude-3-5-sonnet@20240620"
+    } else if (model.includes("claude-3-5-haiku")) {
+        return "claude-3-5-haiku@20241022"
+    } else if (model.includes("claude-3-opus")) {
+        return "claude-3-opus@20240229"
+    } else if (model.includes("claude-3-sonnet")) {
+        return "claude-3-sonnet@20240229"
+    } else if (model.includes("claude-3-haike")) {
+        return "claude-3-haiku@20240307"
+    }
+    return model;
+}
+
+function claudeFinishReason(reason: string | undefined) {
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'end_turn': return "stop";
+        case 'max_tokens': return "length";
+        default: return reason; //stop_sequence
+    }
+}
+
+function collectTextParts(content: any) {
+    const out = [];
+
+    for (const block of content) {
+        if (block?.text) {
+            out.push(block.text);
+        }
+    }
+    return out.join('\n');
+}
+
+function maxToken(max_tokens: number | undefined, model: string): number {
+    const contains = (str: string, substr: string) => str.indexOf(substr) !== -1;
+    if (max_tokens) {
+        return max_tokens;
+    } else if (contains(model, "claude-3-5")) {
+        return 8192;
+    } else {
+        return 4096
+    }
+}
+
+export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
+
+    model: AIModel
+
+    constructor(modelId: string) {
+        this.model = {
+            id: modelId,
+            name: modelId,
+            provider: 'vertexai',
+            type: ModelType.Text,
+            can_stream: true,
+        } as AIModel;
+    }
+
+    async createPrompt(_driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions): Promise<ClaudePrompt> {
+        // Convert the prompt to the format expected by the Claude API
+        const systemSegments: TextBlockParam[] = segments
+            .filter(segment => segment.role === PromptRole.system)
+            .map(segment => ({
+                text: segment.content,
+                type: 'text'
+            }));
+
+        const safetySegments: TextBlockParam[] = segments
+            .filter(segment => segment.role === PromptRole.safety)
+            .map(segment => ({
+                text: segment.content,
+                type: 'text'
+            }));
+
+        if (options.result_schema) {
+            const schemaSegments: TextBlockParam = {
+                text: "The answer must be a JSON object using the following JSON Schema:\n" + JSON.stringify(options.result_schema),
+                type: 'text'
+            }
+            safetySegments.push(schemaSegments);
+        }
+
+        const messages: MessageParam[] = segments.filter(segment =>
+            segment.role == PromptRole.user
+            || segment.role == PromptRole.assistant
+            || segment.role === PromptRole.tool)
+            .map(segment => {
+                if (segment.role === PromptRole.tool) {
+                    if (!segment.tool_use_id) {
+                        throw new Error("Tool prompt segment must have a tool_use_id");
+                    }
+                    return {
+                        role: 'user',
+                        content: [
+                            {
+                                type: 'tool_result',
+                                tool_use_id: segment.tool_use_id,
+                                content: segment.content || undefined
+                            }
+                        ]
+                    }
+                } else {
+                    return {
+                        role: segment.role !== PromptRole.user ? 'assistant' : 'user',
+                        content: segment.content
+                    }
+                }
+            });
+
+        const system = systemSegments.concat(safetySegments);
+
+        return {
+            messages: messages,
+            system: system
+        }
+    }
+
+    async requestTextCompletion(driver: VertexAIDriver, prompt: ClaudePrompt, options: ExecutionOptions): Promise<Completion> {
+        const client = driver.getAnthropicClient();
+        const splits = options.model.split("/");
+        const modelName = splits[splits.length - 1];
+        options = { ...options, model: modelName };
+        options.model_options = options.model_options as VertexAIClaudeOptions;
+
+        if (options.model_options?._option_id !== "vertexai-claude") {
+            driver.logger.warn("Invalid model options", { options: options.model_options });
+        }
+
+        let conversation = updateConversation(options.conversation as ClaudePrompt, prompt);
+
+        const result = await client.messages.create({
+            ...conversation, // messages, system,
+            tools: options.tools, // we are using the same shape as claude for tools
+            temperature: options.model_options?.temperature,
+            model: modelName,
+            max_tokens: maxToken(options.model_options?.max_tokens, modelName),
+            top_p: options.model_options?.top_p,
+            top_k: options.model_options?.top_k,
+            stop_sequences: options.model_options?.stop_sequence,
+            thinking: options.model_options?.thinking_mode ?
+                {
+                    budget_tokens: options.model_options?.thinking_budget_tokens ?? 1024,
+                    type: "enabled"
+                } : {
+                    type: "disabled"
+                }
+        }) as Message;
+
+        const text = collectTextParts(result.content);
+        const tool_use = collectTools(result.content);
+
+        conversation = updateConversation(options.conversation as ClaudePrompt, createPromptFromResponse(result));
+
+        return {
+            chat: [prompt, { role: result.role, content: result.content }],
+            result: text ?? '',
+            tool_use,
+            token_usage: {
+                prompt: result?.usage.input_tokens,
+                result: result?.usage.output_tokens,
+                total: result?.usage.input_tokens + result?.usage.output_tokens
+            },
+            // make sure we set finish_reason to the correct value (claude is normally setting this by itself)
+            finish_reason: tool_use ? "tool_use" : claudeFinishReason(result?.stop_reason ?? ''),
+            conversation
+        } as Completion;
+    }
+
+    async requestTextCompletionStream(driver: VertexAIDriver, prompt: ClaudePrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
+        const client = driver.getAnthropicClient();
+        const splits = options.model.split("/");
+        const modelName = splits[splits.length - 1];
+        options = { ...options, model: modelName };
+        options.model_options = options.model_options as VertexAIClaudeOptions;
+
+        if (options.model_options?._option_id !== "vertexai-claude") {
+            driver.logger.warn("Invalid model options", { options: options.model_options });
+        }
+
+        const response_stream = await client.messages.stream({
+            ...prompt, // messages, system,
+            tools: options.tools, // we are using the same shape as claude for tools
+            temperature: options.model_options?.temperature,
+            model: modelName,
+            max_tokens: maxToken(options.model_options?.max_tokens, modelName),
+            top_p: options.model_options?.top_p,
+            top_k: options.model_options?.top_k,
+            stop_sequences: options.model_options?.stop_sequence,
+            thinking: options.model_options?.thinking_mode ?
+                {
+                    budget_tokens: options.model_options?.thinking_budget_tokens ?? 1024,
+                    type: "enabled"
+                } : {
+                    type: "disabled"
+                }
+        });
+
+        //Streaming does not give information on the input tokens,
+        //So we use a seperate call to get the input tokens.
+        //Non-critical and model name sensitive so we put it in a try catch block
+        let count_tokens = { input_tokens: 0 };
+        try {
+            count_tokens = await client.messages.countTokens({
+                ...prompt,  // messages, system
+                model: getFullModelName(modelName),
+            });
+        } catch (e) {
+            driver.logger.warn("Failed to get token count for model " + modelName);
+        }
+
+        const stream = asyncMap(response_stream, async (item: any) => {
+            return {
+                result: item?.delta?.text ?? '',
+                token_usage: { prompt: count_tokens.input_tokens, result: item?.usage?.output_tokens },
+                finish_reason: claudeFinishReason(item?.delta?.stop_reason ?? ''),
+            }
+        });
+
+        return stream;
+    }
+}
+
+export function collectTools(content: ContentBlock[]): ToolUse[] | undefined {
+    const out: ToolUse[] = [];
+
+    for (const block of content) {
+        if (block?.type === "tool_use") {
+            out.push({
+                id: block.id,
+                name: block.name,
+                input: block.input as JSONObject,
+            });
+        }
+    }
+
+    return out.length > 0 ? out : undefined;
+}
+
+function createPromptFromResponse(response: Message): ClaudePrompt {
+    return {
+        messages: [{
+            role: PromptRole.assistant,
+            content: response.content,
+        }],
+        system: []
+    }
+}
+
+/**
+ * Update the converatation messages
+ * @param prompt
+ * @param response
+ * @returns
+ */
+function updateConversation(conversation: ClaudePrompt | undefined | null, prompt: ClaudePrompt): ClaudePrompt {
+    const baseSystemMessages = conversation ? conversation.system : [];
+    const baseMessages = conversation ? conversation.messages : []
+    return {
+        messages: baseMessages.concat(prompt.messages || []),
+        system: baseSystemMessages.concat(prompt.system || [])
+    };
+}