@llumiverse/drivers 0.9.2 → 0.11.0

package/src/mistral/index.ts

@@ -1,5 +1,6 @@
-import { AIModel, AbstractDriver, Completion, DriverOptions, ExecutionOptions, PromptFormats, PromptSegment } from "@llumiverse/core";
+import { AIModel, AbstractDriver, Completion, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, PromptSegment } from "@llumiverse/core";
 import { transformSSEStream } from "@llumiverse/core/async";
+import { OpenAITextMessage, formatOpenAILikePrompt, getJSONSafetyNotice } from "@llumiverse/core/formatters";
 import { FetchClient } from "api-fetch-client";
 import { CompletionRequestParams, ListModelsResponse, ResponseFormat } from "./types.js";
 
@@ -13,10 +14,9 @@ interface MistralAIDriverOptions extends DriverOptions {
     endpoint_url?: string;
 }
 
-export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, LLMMessage[]> {
+export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, OpenAITextMessage[]> {
     provider: string;
     apiKey: string;
-    defaultFormat: PromptFormats;
     //client: MistralClient;
     client: FetchClient;
     endpointUrl?: string;
@@ -24,7 +24,6 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, LLMM
     constructor(options: MistralAIDriverOptions) {
         super(options);
         this.provider = "MistralAI";
-        this.defaultFormat = PromptFormats.genericTextLLM;
         this.apiKey = options.apiKey;
         //this.client = new MistralClient(options.apiKey, options.endpointUrl);
         this.client = new FetchClient(options.endpoint_url || ENDPOINT).withHeaders({
@@ -34,47 +33,41 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, LLMM
 
     getResponseFormat = (_options: ExecutionOptions): ResponseFormat | undefined => {
 
+        // const responseFormatJson: ResponseFormat = {
+        //     type: "json_object",
+        // } as ResponseFormat
 
-        /*const responseFormatJson: ResponseFormat = {
-            type: "json_object",
-        } as ResponseFormat
+        // const responseFormatText: ResponseFormat = {
+        //     type: "text",
+        // } as ResponseFormat;
 
-        const responseFormatText: ResponseFormat = {
-            type: "text",
-        } as ResponseFormat;
-        */
 
-        //return _options.resultSchema ? responseFormatJson : responseFormatText;
+        // return _options.resultSchema ? responseFormatJson : responseFormatText;
 
         //TODO remove this when Mistral properly supports the parameters - it makes an error for now
+        // some models like mixtral mistrall tiny or medium are throwing an error when using the response_format parameter
         return undefined
     }
 
-    createPrompt(segments: PromptSegment[], opts: ExecutionOptions): LLMMessage[] {
-        // use same format as OpenAI as that's what MistralAI uses
-        const prompts = super.createPrompt(segments, { ...opts, format: PromptFormats.openai })
-
+    protected formatPrompt(segments: PromptSegment[], opts: ExecutionOptions): OpenAITextMessage[] {
+        const messages = formatOpenAILikePrompt(segments);
         //Add JSON instruction is schema is provided
         if (opts.resultSchema) {
-            const content = "The user is explicitely instructing that the result should be a JSON object.\nThe schema is as follows: \n" + JSON.stringify(opts.resultSchema);
-            prompts.push({
+            messages.push({
                 role: "user",
-                content: content
+                content: "IMPORTANT: " + getJSONSafetyNotice(opts.resultSchema)
             });
         }
-
-        return prompts;
-
+        return messages;
     }
 
-    async requestCompletion(messages: LLMMessage[], options: ExecutionOptions): Promise<Completion<any>> {
-
+    async requestCompletion(messages: OpenAITextMessage[], options: ExecutionOptions): Promise<Completion<any>> {
         const res = await this.client.post('/v1/chat/completions', {
             payload: _makeChatCompletionRequest({
                 model: options.model,
                 messages: messages,
-                maxTokens: options.max_tokens ?? 1024,
-                temperature: options.temperature ?? 0.7,
+                maxTokens: options.max_tokens,
+                temperature: options.temperature,
                 responseFormat: this.getResponseFormat(options),
             })
         })
@@ -91,14 +84,13 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, LLMM
         };
     }
 
-    async requestCompletionStream(messages: LLMMessage[], options: ExecutionOptions): Promise<AsyncIterable<string>> {
-
+    async requestCompletionStream(messages: OpenAITextMessage[], options: ExecutionOptions): Promise<AsyncIterable<string>> {
         const stream = await this.client.post('/v1/chat/completions', {
             payload: _makeChatCompletionRequest({
                 model: options.model,
                 messages: messages,
-                maxTokens: options.max_tokens ?? 1024,
-                temperature: options.temperature ?? 0.7,
+                maxTokens: options.max_tokens,
+                temperature: options.temperature,
                 responseFormat: this.getResponseFormat(options),
                 stream: true
             }),
@@ -121,32 +113,33 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, LLMM
                 name: m.id,
                 description: undefined,
                 provider: m.owned_by,
-                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.");
-    }
 
-}
+    async generateEmbeddings({ content, model = "mistral-embed" }: EmbeddingsOptions): Promise<EmbeddingsResult> {
+        const r = await this.client.post('/v1/embeddings', {
+            payload: {
+                model,
+                input: [content],
+                encoding_format: "float"
+            },
+        });
+        return {
+            values: r.data[0].embedding,
+            model,
+            token_count: r.usage.total_tokens
+        }
+    }
 
-interface LLMMessage {
-    role: string;
-    content: string;
 }
 
-
 /**
  * Creates a chat completion request
  * @param {*} model

package/src/openai.ts

@@ -5,16 +5,17 @@ import {
     Completion,
     DataSource,
     DriverOptions,
+    EmbeddingsOptions,
+    EmbeddingsResult,
     ExecutionOptions,
     ExecutionTokenUsage,
     ModelType,
-    PromptFormats,
-    PromptSegment,
     TrainingJob,
     TrainingJobStatus,
     TrainingOptions,
-    TrainingPromptOptions
+    TrainingPromptOptions,
 } from "@llumiverse/core";
+import { formatOpenAILikePrompt } from "@llumiverse/core/formatters";
 import { asyncMap } from "@llumiverse/core/async";
 import OpenAI from "openai";
 import { Stream } from "openai/streaming";
@@ -39,18 +40,13 @@ export class OpenAIDriver extends AbstractDriver<
     generatedContentTypes: string[] = ["text/plain"];
     service: OpenAI;
     provider = BuiltinProviders.openai;
-    defaultFormat = PromptFormats.openai;
 
     constructor(opts: OpenAIDriverOptions) {
         super(opts);
         this.service = new OpenAI({
             apiKey: opts.apiKey,
         });
-    }
-
-    createPrompt(segments: PromptSegment[], opts: ExecutionOptions): OpenAI.Chat.Completions.ChatCompletionMessageParam[] {
-        // openai only supports opanai format -  force the format
-        return super.createPrompt(segments, { ...opts, format: PromptFormats.openai })
+        this.formatPrompt = formatOpenAILikePrompt;
     }
 
     extractDataFromResponse(
@@ -211,7 +207,7 @@ export class OpenAIDriver extends AbstractDriver<
     }
 
 
-    async generateEmbeddings(content: string, model: string = "text-embedding-ada-002"): Promise<{ embeddings: number[], model: string; }> {
+    async generateEmbeddings({ content, model = "text-embedding-ada-002" }: EmbeddingsOptions): Promise<EmbeddingsResult> {
         const res = await this.service.embeddings.create({
             input: content,
             model: model,
@@ -223,7 +219,7 @@ export class OpenAIDriver extends AbstractDriver<
             throw new Error("No embedding found");
         }
 
-        return { embeddings, model };
+        return { values: embeddings, model } as EmbeddingsResult;
     }
 
 }

package/src/replicate.ts

@@ -5,9 +5,9 @@ import {
     Completion,
     DataSource,
     DriverOptions,
+    EmbeddingsResult,
     ExecutionOptions,
     ModelSearchPayload,
-    PromptFormats,
     TrainingJob,
     TrainingJobStatus,
     TrainingOptions
@@ -36,7 +36,6 @@ export interface ReplicateDriverOptions extends DriverOptions {
 export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
     provider = BuiltinProviders.replicate;
     service: Replicate;
-    defaultFormat = PromptFormats.genericTextLLM;
 
     static parseModelId(modelId: string) {
         const [owner, modelPart] = modelId.split("/");
@@ -73,7 +72,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         const predictionData = {
             input: {
                 prompt: prompt,
-                max_new_tokens: options.max_tokens || 1024,
+                max_new_tokens: options.max_tokens,
                 temperature: options.temperature,
             },
             version: model.version,
@@ -113,7 +112,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         const predictionData = {
             input: {
                 prompt: prompt,
-                max_new_tokens: options.max_tokens || 1024,
+                max_new_tokens: options.max_tokens,
                 temperature: options.temperature,
             },
             version: model.version,
@@ -282,8 +281,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         return models;
     }
 
-    generateEmbeddings(content: string, model?: string): Promise<{ embeddings: number[], model: string; }> {
-        this.logger?.debug(`[Replicate] Generating embeddings for ${content} on ${model}`);
+    async generateEmbeddings(): Promise<EmbeddingsResult> {
         throw new Error("Method not implemented.");
     }
 

package/src/test/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, AIModelStatus, CompletionStream, Driver, ExecutionOptions, ExecutionResponse, ModelType, PromptOptions, PromptSegment, TrainingJob } from "@llumiverse/core";
+import { AIModel, AIModelStatus, CompletionStream, Driver, EmbeddingsResult, 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";
@@ -83,7 +83,7 @@ export class TestDriver implements Driver<PromptSegment[]> {
     validateConnection(): Promise<boolean> {
         throw new Error("Method not implemented.");
     }
-    generateEmbeddings(): Promise<{ embeddings: number[]; model: string; }> {
+    generateEmbeddings(): Promise<EmbeddingsResult> {
         throw new Error("Method not implemented.");
     }
 

package/src/togetherai/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, AbstractDriver, Completion, DriverOptions, ExecutionOptions, PromptFormats } from "@llumiverse/core";
+import { AIModel, AbstractDriver, Completion, DriverOptions, EmbeddingsResult, ExecutionOptions } from "@llumiverse/core";
 import { transformSSEStream } from "@llumiverse/core/async";
 import { FetchClient } from "api-fetch-client";
 import { TogetherModelInfo } from "./interfaces.js";
@@ -10,13 +10,11 @@ interface TogetherAIDriverOptions extends DriverOptions {
 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}`
@@ -37,8 +35,8 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
                 model: options.model,
                 prompt: prompt,
                 response_format: this.getResponseFormat(options),
-                max_tokens: options.max_tokens ?? 1024,
-                temperature: options.temperature ?? 0.7,
+                max_tokens: options.max_tokens,
+                temperature: options.temperature,
                 stop: [
                     "</s>",
                     "[/INST]"
@@ -64,8 +62,8 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
             payload: {
                 model: options.model,
                 prompt: prompt,
-                max_tokens: options.max_tokens ?? 1024,
-                temperature: options.temperature ?? 0.7,
+                max_tokens: options.max_tokens,
+                temperature: options.temperature,
                 response_format: this.getResponseFormat(options),
                 stream: true,
                 stop: [
@@ -93,7 +91,6 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
                 name: m.display_name,
                 description: m.description,
                 provider: this.provider,
-                formats: [PromptFormats.genericTextLLM],
             }
         });
 
@@ -101,14 +98,10 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
 
     }
 
-    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; }> {
+    generateEmbeddings(): Promise<EmbeddingsResult> {
         throw new Error("Method not implemented.");
     }
 

package/src/vertexai/embeddings/embeddings-text.ts

@@ -0,0 +1,52 @@
+
+import { EmbeddingsResult } from '@llumiverse/core';
+import { VertexAIDriver } from '../index.js';
+
+export interface TextEmbeddingsOptions {
+    model?: string;
+    task_type?: "RETRIEVAL_QUERY" | "RETRIEVAL_DOCUMENT" | "SEMANTIC_SIMILARITY" | "CLASSIFICATION" | "CLUSTERING",
+    title?: string, // the title for the embedding
+    content: string // the text to generate embeddings for
+}
+
+interface EmbedingsForTextPrompt {
+    instances: TextEmbeddingsOptions[]
+}
+
+interface TextEmbeddingsResult {
+    predictions: [
+        {
+            embeddings: TextEmbeddings
+        }
+    ]
+}
+
+interface TextEmbeddings {
+    statistics: {
+        truncated: boolean,
+        token_count: number
+    },
+    values: [number]
+}
+
+export async function getEmbeddingsForText(driver: VertexAIDriver, options: TextEmbeddingsOptions): Promise<EmbeddingsResult> {
+    const prompt = {
+        instances: [{
+            task_type: options.task_type,
+            title: options.title,
+            content: options.content
+        }]
+    } as EmbedingsForTextPrompt;
+
+    const model = options.model || "textembedding-gecko@latest";
+
+    const result = await driver.fetchClient.post(`/publishers/google/models/${model}:predict`, {
+        payload: prompt
+    }) as TextEmbeddingsResult;
+
+    return {
+        ...result.predictions[0].embeddings,
+        model,
+        token_count: result.predictions[0].embeddings.statistics?.token_count
+    };
+}

package/src/vertexai/index.ts

@@ -1,9 +1,9 @@
-//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 { AIModel, AbstractDriver, BuiltinProviders, Completion, DriverOptions, EmbeddingsResult, ExecutionOptions, ModelSearchPayload, PromptOptions, PromptSegment } from "@llumiverse/core";
 import { FetchClient } from "api-fetch-client";
+import { TextEmbeddingsOptions, getEmbeddingsForText } from "./embeddings/embeddings-text.js";
 import { BuiltinModels, getModelDefinition } from "./models.js";
-//import { GoogleAuth } from "google-auth-library";
+
 
 export interface VertexAIDriverOptions extends DriverOptions {
     project: string;
@@ -12,7 +12,6 @@ export interface VertexAIDriverOptions extends DriverOptions {
 
 export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, GenerateContentRequest> {
     provider = BuiltinProviders.vertexai;
-    defaultFormat = PromptFormats.genericTextLLM;
 
     //aiplatform: v1.ModelServiceClient;
     vertexai: VertexAI;
@@ -31,7 +30,8 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Genera
             region: this.options.region,
             project: this.options.project,
         }).withAuthCallback(async () => {
-            const token = await this.vertexai.preview.token;
+            //@ts-ignore
+            const token = await this.vertexai.preview.googleAuth.getAccessToken();
             return `Bearer ${token}`;
         });
         // this.aiplatform = new v1.ModelServiceClient({
@@ -72,14 +72,13 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Genera
 
         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.");
+
+    async generateEmbeddings(options: TextEmbeddingsOptions): Promise<EmbeddingsResult> {
+        return getEmbeddingsForText(this, options);
     }
 
 }

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

@@ -1,6 +1,6 @@
 import { AIModel, ModelType, PromptOptions, PromptRole, PromptSegment } from "@llumiverse/core";
+import { getJSONSafetyNotice } from "@llumiverse/core/formatters";
 import { VertexAIDriver } from "../index.js";
-import { getJSONSafetyNotice } from "../utils/prompts.js";
 import { AbstractPalmModelDefinition, NonStreamingPromptBase, PalmResponseMetadata, StreamingPromptBase } from "./palm-model-base.js";
 
 export interface CodeyChatMessage {

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

@@ -1,6 +1,6 @@
 import { AIModel, ModelType, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { formatTextPrompt } from "@llumiverse/core/formatters";
 import { VertexAIDriver } from "../index.js";
-import { getPromptAsText } from "../utils/prompts.js";
 import { AbstractPalmModelDefinition, NonStreamingPromptBase, PalmResponseMetadata, StreamingPromptBase } from "./palm-model-base.js";
 
 
@@ -50,7 +50,7 @@ export class CodeyTextDefinition extends AbstractPalmModelDefinition<CodeyTextPr
     createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): CodeyTextPrompt {
         return {
             instances: [{
-                prefix: getPromptAsText(segments, opts)
+                prefix: formatTextPrompt(segments, opts.resultSchema)
             }],
             parameters: {
                 // put defauilts here

package/src/vertexai/models/gemini.ts

@@ -1,10 +1,10 @@
-import { Content, GenerateContentRequest, GenerativeModel, HarmBlockThreshold, HarmCategory, TextPart } from "@google-cloud/vertexai";
+import { Content, GenerateContentRequest, 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 {
+function getGenerativeModel(driver: VertexAIDriver, options: ExecutionOptions) {
     return driver.vertexai.preview.getGenerativeModel({
         model: options.model,
         //TODO pass in the options
@@ -106,8 +106,8 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentReq
         const response = await r.response;
         const usage = response.usageMetadata;
         const token_usage: ExecutionTokenUsage = {
-            prompt: usage?.prompt_token_count,
-            result: usage?.candidates_token_count,
+            prompt: usage?.promptTokenCount,
+            result: usage?.candidatesTokenCount,
             total: usage?.totalTokenCount,
         }
 

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

@@ -2,9 +2,25 @@ import { AIModel, Completion, ExecutionOptions, PromptOptions, PromptSegment } f
 import { transformSSEStream } from "@llumiverse/core/async";
 import { VertexAIDriver } from "../index.js";
 import { ModelDefinition } from "../models.js";
-import { PromptParamatersBase } from "../utils/prompts.js";
 import { generateStreamingPrompt } from "../utils/tensor.js";
 
+
+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 interface NonStreamingPromptBase<InstanceType = any> {
     instances: InstanceType[];
     parameters: PromptParamatersBase;

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

@@ -1,6 +1,6 @@
 import { AIModel, ModelType, PromptOptions, PromptRole, PromptSegment } from "@llumiverse/core";
+import { getJSONSafetyNotice } from "@llumiverse/core/formatters";
 import { VertexAIDriver } from "../index.js";
-import { getJSONSafetyNotice } from "../utils/prompts.js";
 import { AbstractPalmModelDefinition, NonStreamingPromptBase, PalmResponseMetadata, StreamingPromptBase } from "./palm-model-base.js";
 
 export interface Palm2ChatMessage {

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

@@ -1,6 +1,6 @@
 import { AIModel, ModelType, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { formatTextPrompt } from "@llumiverse/core/formatters";
 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<{
@@ -50,7 +50,7 @@ export class Palm2TextDefinition extends AbstractPalmModelDefinition<Palm2TextPr
     createNonStreamingPrompt(_driver: VertexAIDriver, segments: PromptSegment[], opts: PromptOptions): Palm2TextPrompt {
         return {
             instances: [{
-                prompt: getPromptAsText(segments, opts)
+                prompt: formatTextPrompt(segments, opts.resultSchema)
             }],
             parameters: {
                 // put defauilts here