@llumiverse/drivers 0.11.0 → 0.12.1

package/src/bedrock/index.ts

@@ -91,29 +91,33 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         const body = decoder.decode(response.body);
         const result = JSON.parse(body);
 
-        const getText = () => {
-            if (result.completion) {
-                return result.completion;
-            } else if (result.generation) {
-                return result.generation;
+        const getTextAnsStopReason = (): string[] => {
+            if (result.generation) {
+                // LLAMA2
+                return [result.generation, result.stop_reason]; // comes in coirrect format (stop, length)
             } else if (result.generations) {
-                return result.generations[0].text;
+                // COHERE                
+                return [result.generations[0].text, cohereFinishReason(result.generations[0].finish_reason)];
             } else if (result.completions) {
                 //A21
-                return result.completions[0].data?.text;
-            } else if (result.content) { // calude
-                return result.content[0]?.text || '';
-                //result.stop_reason --> the stop reason
+                return [result.completions[0].data?.text, a21FinishReason(result.completions[0].finishReason?.reason)];
+            } else if (result.content) {
+                // anthropic claude                
+                return [result.content[0]?.text || '', claudeFinishReason(result.stop_reason)];
             } else if (result.outputs) {
                 // mistral
-                return result.outputs[0]?.text;
-                //result.outputs[0]?.stop_reason --> the stop reason
+                return [result.outputs[0]?.text, result.outputs[0]?.stop_reason]; // the stop reason is in the expected format ("stop" and "length")
+            } else if (result.results) {
+                // Amazon Titan
+                return [result.results[0]?.outputText ?? '', titanFinishReason(result.results[0]?.completionReason)];
+            } else if (result.completion) { // TODO: who uses this?
+                return [result.completion];
             } else {
-                return result.toString();
+                return [result.toString()];
             }
         };
 
-        const text = getText();
+        const [text, finish_reason] = getTextAnsStopReason();
 
         const promptLength = typeof prompt === 'string' ? prompt.length :
             (prompt.system || '').length + prompt.messages.reduce((acc, m) => acc + m.content.length, 0);
@@ -123,7 +127,8 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                 result: text?.length,
                 prompt: promptLength,
                 total: text?.length + promptLength,
-            }
+            },
+            finish_reason
         }
     }
 
@@ -136,7 +141,11 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             contentType: "application/json",
             body: JSON.stringify(payload),
         });
-        return this.extractDataFromResponse(prompt, res);
+        const completion = this.extractDataFromResponse(prompt, res);
+        if (options.include_original_response) {
+            completion.original_response = res;
+        }
+        return completion;
     }
 
     protected async canStream(options: ExecutionOptions): Promise<boolean> {
@@ -167,9 +176,9 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
 
             return transformAsyncIterator(res.body, (stream: ResponseStream) => {
                 const segment = JSON.parse(decoder.decode(stream.chunk?.bytes));
-                if (segment.delta) {
+                if (segment.delta) { // who is this?
                     return segment.delta.text || '';
-                } else if (segment.completion) {
+                } else if (segment.completion) { // who is this?
                     return segment.completion;
                 } else if (segment.completions) {
                     return segment.completions[0].data?.text;
@@ -181,6 +190,11 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                     // mistral.mixtral-8x7b-instruct-v0:1
                     return segment.outputs[0].text;
                     //segment.outputs[0].stop_reason;
+                } else if (segment.outputText) {
+                    // Amazon Titan
+                    return segment.outputText;
+                    //completionReason
+                    // token count too
                 } else {
                     segment.toString();
                 }
@@ -230,12 +244,12 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             } as CohereRequestPayload;
         } else if (contains(options.model, "amazon")) {
             return {
-                inputText: prompt,
+                inputText: "User: " + (prompt as string) + "\nBot:", // see https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html#model-parameters-titan-request-response
                 textGenerationConfig: {
                     temperature: options.temperature,
-                    topP: 0.9,
+                    topP: options.top_p,
                     maxTokenCount: options.max_tokens,
-                    stopSequences: ["\n"],
+                    //stopSequences: ["\n"],
                 },
             } as AmazonRequestPayload;
         } else if (contains(options.model, "mistral")) {
@@ -502,4 +516,38 @@ function jobInfo(job: GetModelCustomizationJobCommandOutput, jobId: string): Tra
 }
 
 
+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 cohereFinishReason(reason: string | undefined) {
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'COMPLETE': return "stop";
+        case 'MAX_TOKENS': return "length";
+        default: return reason;
+    }
+}
+
+function a21FinishReason(reason: string | undefined) {
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'endoftext': return "stop";
+        case 'length': return "length";
+        default: return reason;
+    }
+}
+
+function titanFinishReason(reason: string | undefined) {
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'FINISH': return "stop";
+        case 'LENGTH': return "length";
+        default: return reason;
+    }
+}

package/src/groq/index.ts

@@ -0,0 +1,134 @@
+import { AIModel, AbstractDriver, BuiltinProviders, Completion, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, PromptSegment } from "@llumiverse/core";
+import { transformAsyncIterator } from "@llumiverse/core/async";
+import { OpenAITextMessage, formatOpenAILikePrompt, getJSONSafetyNotice } from "@llumiverse/core/formatters";
+import Groq from "groq-sdk";
+
+
+interface GroqDriverOptions extends DriverOptions {
+    apiKey: string;
+    endpoint_url?: string;
+}
+
+
+export class GroqDriver extends AbstractDriver<GroqDriverOptions, OpenAITextMessage[]> {
+    provider: string;
+    apiKey: string;
+    client: Groq;
+    endpointUrl?: string;
+
+    constructor(options: GroqDriverOptions) {
+        super(options);
+        this.provider = BuiltinProviders.groq;
+        this.apiKey = options.apiKey;
+        this.client = new Groq({
+            apiKey: options.apiKey,
+            baseURL: options.endpoint_url
+        });
+    }
+
+    // protected canStream(options: ExecutionOptions): Promise<boolean> {
+    //     if (options.resultSchema) {
+    //         // not yet streamign json responses
+    //         return Promise.resolve(false);
+    //     } else {
+    //         return Promise.resolve(true);
+    //     }
+    // }
+
+    getResponseFormat(_options: ExecutionOptions): Groq.Chat.Completions.CompletionCreateParams.ResponseFormat | undefined {
+        //TODO: when forcing json_object type the streaming is not supported.
+        // either implement canStream as above or comment the code below:
+        // const responseFormatJson: Groq.Chat.Completions.CompletionCreateParams.ResponseFormat = {
+        //     type: "json_object",
+        // }
+
+        // return _options.resultSchema ? responseFormatJson : undefined;
+        return undefined;
+    }
+
+    protected formatPrompt(segments: PromptSegment[], opts: ExecutionOptions): OpenAITextMessage[] {
+        const messages = formatOpenAILikePrompt(segments);
+        //Add JSON instruction is schema is provided
+        if (opts.resultSchema) {
+            messages.push({
+                role: "user",
+                content: "IMPORTANT: " + getJSONSafetyNotice(opts.resultSchema)
+            });
+        }
+        return messages;
+    }
+
+    async requestCompletion(messages: OpenAITextMessage[], options: ExecutionOptions): Promise<Completion<any>> {
+
+
+        const res = await this.client.chat.completions.create({
+            model: options.model,
+            messages: messages,
+            max_tokens: options.max_tokens,
+            temperature: options.temperature,
+            response_format: this.getResponseFormat(options),
+        });
+
+
+        const choice = res.choices[0];
+        const result = choice.message.content;
+
+        return {
+            result: result,
+            token_usage: {
+                prompt: res.usage?.prompt_tokens,
+                result: res.usage?.completion_tokens,
+                total: res.usage?.total_tokens,
+            },
+            finish_reason: choice.finish_reason,
+            original_response: options.include_original_response ? res : undefined,
+        };
+    }
+
+    async requestCompletionStream(messages: OpenAITextMessage[], options: ExecutionOptions): Promise<AsyncIterable<string>> {
+
+        const res = await this.client.chat.completions.create({
+            model: options.model,
+            messages: messages,
+            max_tokens: options.max_tokens,
+            temperature: options.temperature,
+            response_format: this.getResponseFormat(options),
+            stream: true
+        });
+
+        return transformAsyncIterator(res, (res) => res.choices[0].delta.content || '');
+
+    }
+
+    async listModels(): Promise<AIModel<string>[]> {
+        const models = await this.client.models.list();
+
+        if (!models.data) {
+            throw new Error("No models found");
+        }
+
+        const aimodels = models.data?.map(m => {
+            if (!m.id) {
+                throw new Error("Model id is missing");
+            }
+            return {
+                id: m.id,
+                name: m.id,
+                description: undefined,
+                provider: this.provider,
+                owner: m.owned_by || '',
+            }
+        });
+
+        return aimodels;
+    }
+
+    validateConnection(): Promise<boolean> {
+        throw new Error("Method not implemented.");
+    }
+
+    async generateEmbeddings({ }: EmbeddingsOptions): Promise<EmbeddingsResult> {
+        throw new Error("Method not implemented.");
+    }
+
+}

package/src/huggingface_ie.ts

@@ -92,6 +92,10 @@ export class HuggingFaceIEDriver extends AbstractDriver<HuggingFaceIEDriverOptio
             },
         });
 
+        let finish_reason = res.details?.finish_reason;
+        if (finish_reason === "eos_token") {
+            finish_reason = "stop";
+        }
         return {
             result: res.generated_text,
             token_usage: {
@@ -99,6 +103,8 @@ export class HuggingFaceIEDriver extends AbstractDriver<HuggingFaceIEDriverOptio
                 prompt: prompt.length,
                 total: res.generated_text.length + prompt.length,
             },
+            finish_reason,
+            original_response: options.include_original_response ? res : undefined,
         };
 
     }

package/src/index.ts

@@ -6,4 +6,4 @@ export * from "./replicate.js";
 export * from "./test/index.js";
 export * from "./togetherai/index.js";
 export * from "./vertexai/index.js";
-
+export * from "./groq/index.js";

package/src/mistral/index.ts

@@ -1,8 +1,8 @@
-import { AIModel, AbstractDriver, Completion, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, PromptSegment } from "@llumiverse/core";
+import { AIModel, AbstractDriver, BuiltinProviders, 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";
+import { ChatCompletionResponse, CompletionRequestParams, ListModelsResponse, ResponseFormat } from "./types.js";
 
 //TODO retry on 429
 //const RETRY_STATUS_CODES = [429, 500, 502, 503, 504];
@@ -23,7 +23,7 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
 
     constructor(options: MistralAIDriverOptions) {
         super(options);
-        this.provider = "MistralAI";
+        this.provider = BuiltinProviders.mistralai;
         this.apiKey = options.apiKey;
         //this.client = new MistralClient(options.apiKey, options.endpointUrl);
         this.client = new FetchClient(options.endpoint_url || ENDPOINT).withHeaders({
@@ -70,9 +70,10 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
                 temperature: options.temperature,
                 responseFormat: this.getResponseFormat(options),
             })
-        })
+        }) as ChatCompletionResponse;
 
-        const result = res.choices[0]?.message.content;
+        const choice = res.choices[0];
+        const result = choice.message.content;
 
         return {
             result: result,
@@ -80,7 +81,9 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
                 prompt: res.usage.prompt_tokens,
                 result: res.usage.completion_tokens,
                 total: res.usage.total_tokens,
-            }
+            },
+            finish_reason: choice.finish_reason,
+            original_response: options.include_original_response ? res : undefined,
         };
     }
 
@@ -112,7 +115,8 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
                 id: m.id,
                 name: m.id,
                 description: undefined,
-                provider: m.owned_by,
+                provider: this.provider,
+                owner: m.owned_by,
             }
         });
 

package/src/mistral/types.ts

@@ -79,7 +79,7 @@ export interface ChatCompletionResponseChoice {
         role: string;
         content: string;
     };
-    finish_reason: string;
+    finish_reason: string; // "stop" "length" "model_length" "error" "tool_calls"
 }
 
 export interface ChatCompletionResponseChunkChoice {
@@ -89,7 +89,7 @@ export interface ChatCompletionResponseChunkChoice {
         content?: string;
         tool_calls?: ToolCalls[];
     };
-    finish_reason: string;
+    finish_reason: string; // "stop" "length" "model_length" "error" "tool_calls"
 }
 
 export interface ChatCompletionResponse {

package/src/openai.ts

@@ -15,8 +15,8 @@ import {
     TrainingOptions,
     TrainingPromptOptions,
 } from "@llumiverse/core";
-import { formatOpenAILikePrompt } from "@llumiverse/core/formatters";
 import { asyncMap } from "@llumiverse/core/async";
+import { formatOpenAILikePrompt } from "@llumiverse/core/formatters";
 import OpenAI from "openai";
 import { Stream } from "openai/streaming";
 
@@ -59,16 +59,20 @@ export class OpenAIDriver extends AbstractDriver<
             total: result.usage?.total_tokens,
         };
 
+        const choice = result.choices[0];
+        const finish_reason = choice.finish_reason;
+
         //if no schema, return content
         if (!options.resultSchema) {
             return {
-                result: result.choices[0]?.message.content as string,
+                result: choice.message.content as string,
                 token_usage: tokenInfo,
+                finish_reason
             }
         }
 
         //we have a schema: get the content and return after validation
-        const data = result.choices[0]?.message.function_call?.arguments as any;
+        const data = choice?.message.tool_calls?.[0].function.arguments;
         if (!data) {
             this.logger?.error("[OpenAI] Response is not valid", result);
             throw new Error("Response is not valid: no data");
@@ -76,7 +80,8 @@ export class OpenAIDriver extends AbstractDriver<
 
         return {
             result: data,
-            token_usage: tokenInfo
+            token_usage: tokenInfo,
+            finish_reason
         };
     }
 
@@ -84,7 +89,7 @@ export class OpenAIDriver extends AbstractDriver<
         const mapFn = options.resultSchema
             ? (chunk: OpenAI.Chat.Completions.ChatCompletionChunk) => {
                 return (
-                    chunk.choices[0]?.delta?.function_call?.arguments ?? ""
+                    chunk.choices[0]?.delta?.tool_calls?.[0].function?.arguments ?? ""
                 );
             }
             : (chunk: OpenAI.Chat.Completions.ChatCompletionChunk) => {
@@ -98,17 +103,22 @@ export class OpenAIDriver extends AbstractDriver<
             temperature: options.temperature,
             n: 1,
             max_tokens: options.max_tokens,
-            functions: options.resultSchema
+            tools: options.resultSchema
                 ? [
                     {
-                        name: "format_output",
-                        parameters: options.resultSchema as any,
-                    },
+                        function: {
+                            name: "format_output",
+                            parameters: options.resultSchema as any,
+                        },
+                        type: "function"
+                    } as OpenAI.Chat.ChatCompletionTool,
                 ]
                 : undefined,
-            function_call: options.resultSchema
-                ? { name: "format_output" }
-                : undefined,
+            tool_choice: options.resultSchema
+                ? {
+                    type: 'function',
+                    function: { name: "format_output" }
+                } : undefined,
         })) as Stream<OpenAI.Chat.Completions.ChatCompletionChunk>;
 
         return asyncMap(stream, mapFn);
@@ -118,9 +128,12 @@ export class OpenAIDriver extends AbstractDriver<
         const functions = options.resultSchema
             ? [
                 {
-                    name: "format_output",
-                    parameters: options.resultSchema as any,
-                },
+                    function: {
+                        name: "format_output",
+                        parameters: options.resultSchema as any,
+                    },
+                    type: 'function'
+                } as OpenAI.Chat.ChatCompletionTool,
             ]
             : undefined;
 
@@ -131,13 +144,23 @@ export class OpenAIDriver extends AbstractDriver<
             temperature: options.temperature,
             n: 1,
             max_tokens: options.max_tokens,
-            functions: functions,
-            function_call: options.resultSchema
-                ? { name: "format_output" }
-                : undefined,
+            tools: functions,
+            tool_choice: options.resultSchema
+                ? {
+                    type: 'function',
+                    function: { name: "format_output" }
+                } : undefined,
+            // functions: functions,
+            // function_call: options.resultSchema
+            //     ? { name: "format_output" }
+            //     : undefined,
         });
 
-        return this.extractDataFromResponse(options, res);
+        const completion = this.extractDataFromResponse(options, res);
+        if (options.include_original_response) {
+            completion.original_response = res;
+        }
+        return completion;
     }
 
     createTrainingPrompt(options: TrainingPromptOptions): string {

package/src/replicate.ts

@@ -138,6 +138,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
                 prompt: prompt.length,
                 total: res.output.length + prompt.length,
             },
+            original_response: options.include_original_response ? res : undefined,
         };
     }
 

package/src/togetherai/index.ts

@@ -1,7 +1,7 @@
 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";
+import { TextCompletion, TogetherModelInfo } from "./interfaces.js";
 
 interface TogetherAIDriverOptions extends DriverOptions {
     apiKey: string;
@@ -42,9 +42,9 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
                     "[/INST]"
                 ],
             }
-        })
-
-        const text = res.choices[0]?.text ?? '';
+        }) as TextCompletion;
+        const choice = res.choices[0];
+        const text = choice.text ?? '';
         const usage = res.usage || {};
         return {
             result: text,
@@ -53,6 +53,8 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
                 result: usage.completion_tokens,
                 total: usage.total_tokens,
             },
+            finish_reason: choice.finish_reason,
+            original_response: options.include_original_response ? res : undefined,
         }
     }
 

package/src/togetherai/interfaces.ts

@@ -85,4 +85,20 @@ export interface TogetherModelInfo {
     link: string;
     descriptionLink: string;
     depth: Depth;
+}
+
+export interface TextCompletion {
+    id: string;
+    choices: {
+        text: string,
+        finish_reason: string, // stop | length ?
+    }[];
+    usage: {
+        prompt_tokens: number;
+        completion_tokens: number;
+        total_tokens: number;
+    }
+    created: number;
+    model: string;
+    object: string;
 }

package/src/vertexai/models/gemini.ts

@@ -1,4 +1,4 @@
-import { Content, GenerateContentRequest, HarmBlockThreshold, HarmCategory, TextPart } from "@google-cloud/vertexai";
+import { Content, FinishReason, 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";
@@ -7,12 +7,13 @@ import { ModelDefinition } from "../models.js";
 function getGenerativeModel(driver: VertexAIDriver, options: ExecutionOptions) {
     return driver.vertexai.preview.getGenerativeModel({
         model: options.model,
-        //TODO pass in the options
+        //TODO pass in the options        
         safety_settings: [{
             category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
             threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE
         }],
         generation_config: {
+            candidate_count: 1,
             temperature: options.temperature,
             max_output_tokens: options.max_tokens
         },
@@ -111,9 +112,14 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentReq
             total: usage?.totalTokenCount,
         }
 
-        let result: any;
+        let finish_reason: string | undefined, result: any;
         const candidate = response.candidates[0];
         if (candidate) {
+            switch (candidate.finishReason) {
+                case FinishReason.MAX_TOKENS: finish_reason = "length"; break;
+                case FinishReason.STOP: finish_reason = "stop"; break;
+                default: finish_reason = candidate.finishReason;
+            }
             const content = candidate.content;
             if (content) {
                 result = collectTextParts(content);
@@ -126,8 +132,10 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentReq
 
         return {
             result: result ?? '',
-            token_usage
-        };
+            token_usage,
+            finish_reason,
+            original_response: options.include_original_response ? response : undefined,
+        } as Completion;
     }
 
     async requestCompletionStream(driver: VertexAIDriver, prompt: GenerateContentRequest, options: ExecutionOptions): Promise<AsyncIterable<string>> {

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

@@ -89,7 +89,9 @@ export abstract class AbstractPalmModelDefinition<NonStreamingPromptT extends No
                 prompt: inputTokens,
                 result: outputTokens,
                 total: inputTokens && outputTokens ? inputTokens + outputTokens : undefined,
-            }
+            },
+            //finish_reason not available
+            original_response: options.include_original_response ? response : undefined,
         } as Completion;
 
     }