@llumiverse/core 0.10.0 → 0.11.0

package/src/CompletionStream.ts

@@ -60,7 +60,6 @@ export class DefaultCompletionStream<PromptT = any> implements CompletionStream<
 export class FallbackCompletionStream<PromptT = any> implements CompletionStream<PromptT> {
 
     prompt: PromptT;
-    chunks: string[];
     completion: ExecutionResponse<PromptT> | undefined;
 
     constructor(public driver: AbstractDriver<DriverOptions, PromptT>,
@@ -68,22 +67,16 @@ export class FallbackCompletionStream<PromptT = any> implements CompletionStream
         public options: ExecutionOptions) {
         this.driver = driver;
         this.prompt = this.driver.createPrompt(segments, options);
-        this.chunks = [];
     }
 
     async *[Symbol.asyncIterator]() {
         // reset state
         this.completion = undefined;
-        if (this.chunks.length > 0) {
-            this.chunks = [];
-        }
         this.driver.logger.debug(
             `[${this.driver.provider}] Streaming is not supported, falling back to blocking execution`
         );
         const completion = await this.driver._execute(this.prompt, this.options);
-
-        const content = completion.result === 'string' ? completion.result : JSON.stringify(completion.result);
-        this.chunks.push(content);
+        const content = typeof completion.result === 'string' ? completion.result : JSON.stringify(completion.result);
         yield content;
 
         this.completion = completion;

package/src/Driver.ts

@@ -5,7 +5,7 @@
  */
 
 import { DefaultCompletionStream, FallbackCompletionStream } from "./CompletionStream.js";
-import { PromptFormatters } from "./formatters/index.js";
+import { formatLlama2Prompt, formatTextPrompt } from "./formatters/index.js";
 import {
     AIModel,
     Completion,
@@ -18,7 +18,6 @@ import {
     ExecutionResponse,
     Logger,
     ModelSearchPayload,
-    PromptFormats,
     PromptOptions,
     PromptSegment,
     TrainingJob,
@@ -54,7 +53,7 @@ export function createLogger(logger: Logger | "console" | undefined) {
 
 function applyExecutionDefaults(options: ExecutionOptions): ExecutionOptions {
     return {
-        max_tokens: 1024,
+        max_tokens: 4096,
         temperature: 0.7,
         ...options
     }
@@ -106,7 +105,6 @@ export abstract class AbstractDriver<OptionsT extends DriverOptions = DriverOpti
     logger: Logger;
 
     abstract provider: string; // the provider name
-    abstract defaultFormat: PromptFormats;
 
     constructor(opts: OptionsT) {
         this.options = opts;
@@ -181,11 +179,22 @@ export abstract class AbstractDriver<OptionsT extends DriverOptions = DriverOpti
         }
     }
 
+    /**
+     * Override this method to provide a custom prompt formatter
+     * @param segments 
+     * @param options 
+     * @returns 
+     */
+    protected formatPrompt(segments: PromptSegment[], opts: PromptOptions): PromptT {
+        if (/\bllama2?\b/i.test(opts.model)) {
+            return formatLlama2Prompt(segments, opts.resultSchema) as PromptT;
+        } else {
+            return formatTextPrompt(segments, opts.resultSchema) as PromptT;
+        }
+    }
+
     public createPrompt(segments: PromptSegment[], opts: PromptOptions): PromptT {
-        return PromptFormatters[opts.format || this.defaultFormat](
-            segments,
-            opts.resultSchema
-        );
+        return opts.format ? opts.format(segments, opts.resultSchema) : this.formatPrompt(segments, opts);
     }
 
     /**

package/src/formatters/claude.ts

@@ -1,5 +1,6 @@
 import { JSONSchema4 } from "json-schema";
 import { PromptRole, PromptSegment } from "../index.js";
+import { getJSONSafetyNotice } from "./commons.js";
 
 export interface ClaudeMessage {
     role: 'user' | 'assistant',
@@ -19,7 +20,7 @@ export interface ClaudeMessagesPrompt {
  * A formatter user by Bedrock to format prompts for claude related models
  */
 
-export function claudeMessages(segments: PromptSegment[], schema?: JSONSchema4): ClaudeMessagesPrompt {
+export function formatClaudePrompt(segments: PromptSegment[], schema?: JSONSchema4): ClaudeMessagesPrompt {
     const system: string[] = [];
     const safety: string[] = [];
     const messages: ClaudeMessage[] = [];
@@ -35,13 +36,9 @@ export function claudeMessages(segments: PromptSegment[], schema?: JSONSchema4):
     }
 
     if (schema) {
-        safety.push(`You must answer using the following JSONSchema:
----
-${JSON.stringify(schema)}
----`);
+        safety.push(getJSONSafetyNotice(schema));
     }
 
-
     // messages must contains at least 1 item. If the prompt doesn;t contains a user message (but only system messages)
     // we need to put the system messages in the messages array
 

package/src/formatters/commons.ts

@@ -0,0 +1,5 @@
+import { JSONSchema4 } from "json-schema";
+
+export function getJSONSafetyNotice(schema: JSONSchema4) {
+    return "The answer must be a JSON object using the following JSON Schema:\n" + JSON.stringify(schema);
+}

package/src/formatters/generic.ts

@@ -1,34 +1,66 @@
 import { JSONSchema4 } from "json-schema";
-import { PromptRole, PromptSegment } from "../index.js";
+import { PromptRole, PromptSegment } from "../types.js";
+import { getJSONSafetyNotice } from "./commons.js";
 
-export function genericColonSeparator(
-    messages: PromptSegment[],
-    schema?: JSONSchema4,
-    labels: {
-        user: string;
-        assistant: string;
-        system: string;
-    } = { user: "User", assistant: "Assistant", system: "System" }
-) {
-    const promptMessages = [];
-    for (const m of messages) {
-        if (m.role === PromptRole.user) {
-            promptMessages.push(`${labels?.user}: ${m.content.trim()}`);
-        }
-        if (m.role === PromptRole.assistant) {
-            promptMessages.push(`${labels.assistant}: ${m.content.trim()}`);
+interface Labels {
+    user: string,
+    system: string,
+    assistant: string,
+    safety: string,
+    instruction: string
+}
+
+export function createTextPromptFormatter(labels: Labels = {
+    user: "USER",
+    system: "CONTEXT",
+    assistant: "ASSISTANT",
+    safety: "IMPORTANT",
+    instruction: "INSTRUCTION"
+}) {
+    return function genericTextPrompt(segments: PromptSegment[], schema?: JSONSchema4): 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(`${labels.user}: ${segment.content}`);
+                    } else {
+                        content.push(segment.content);
+                    }
+                    break;
+                case PromptRole.assistant:
+                    content.push(`${labels.assistant}: ${segment.content}`);
+                    break;
+                case PromptRole.system:
+                    context.push(segment.content);
+                    break;
+                case PromptRole.safety:
+                    safety.push(segment.content);
+                    break;
+            }
         }
-        if (m.role === PromptRole.system) {
-            promptMessages.push(`${labels.system}: ${m.content.trim()}`);
+
+        if (schema) {
+            safety.push(getJSONSafetyNotice(schema));
         }
-    }
 
-    if (schema) {
-        promptMessages.push(`${labels.system}: You must answer using the following JSONSchema:
-        ---
-        ${JSON.stringify(schema)}
-        ---`);
+        const out = [];
+        if (context.length > 0) {
+            out.push(`${labels.system}: ${context.join('\n')}`);
+        }
+        if (content.length > 0) {
+            const prefix = context.length > 0 && !isChat ? `${labels.instruction}: ` : '';
+            out.push(prefix + content.join('\n'));
+        }
+        if (safety.length > 0) {
+            out.push(`${labels.safety}: ${safety.join('\n')}`);
+        }
+        return out.join('\n');
     }
-
-    return promptMessages.join("\n\n");
 }
+
+const formatTextPrompt = createTextPromptFormatter();
+export { formatTextPrompt };

package/src/formatters/index.ts

@@ -1,33 +1,10 @@
 import { JSONSchema4 } from "json-schema";
-import { genericColonSeparator } from "./generic.js";
-import { llama2 } from "./llama2.js";
-import { openAI } from "./openai.js";
-import {
-    PromptFormats,
-    PromptSegment
-} from "../types.js";
-import { claudeMessages } from "./claude.js";
+import { PromptSegment } from "../types.js";
 
-export function inferFormatterFromModelName(modelName: string): PromptFormats {
-    const name = modelName.toLowerCase();
-    if (name.includes("llama")) {
-        return PromptFormats.llama2;
-    } else if (name.includes("gpt")) {
-        return PromptFormats.openai;
-    } else if (name.includes("claude")) {
-        return PromptFormats.claude;
-    } else {
-        return PromptFormats.genericTextLLM;
-    }
-}
-
-export const PromptFormatters: Record<
-    PromptFormats,
-    (messages: PromptSegment[], schema?: JSONSchema4) => any
-> = {
-    openai: openAI,
-    llama2: llama2,
-    claude: claudeMessages,
-    genericTextLLM: genericColonSeparator,
-};
+export type PromptFormatter<T = any> = (messages: PromptSegment[], schema?: JSONSchema4) => T;
 
+export * from "./commons.js"
+export * from "./generic.js";
+export * from "./llama2.js";
+export * from "./claude.js";
+export * from "./openai.js";

package/src/formatters/llama2.ts

@@ -1,7 +1,7 @@
 import { JSONSchema4 } from "json-schema";
 import { PromptRole, PromptSegment } from "../index.js";
 
-export function llama2(messages: PromptSegment[], schema?: JSONSchema4) {
+export function formatLlama2Prompt(messages: PromptSegment[], schema?: JSONSchema4) {
     const BOS = "<s>";
     const EOS = "</s>";
     const INST = "[INST]";

package/src/formatters/openai.ts

@@ -1,17 +1,25 @@
 import { PromptRole } from "../index.js";
 import { PromptSegment } from "../types.js";
-import OpenAI from "openai";
 
-export function openAI(segments: PromptSegment[]) {
-    const system: OpenAI.Chat.ChatCompletionMessageParam[] = [];
-    const others: OpenAI.Chat.ChatCompletionMessageParam[] = [];
-    const safety: OpenAI.Chat.ChatCompletionMessageParam[] = [];
+export interface OpenAITextMessage {
+    content: string;
+    role: "system" | "user" | "assistant";
+}
+/**
+ * OpenAI text only prompts
+ * @param segments 
+ * @returns 
+ */
+export function formatOpenAILikePrompt(segments: PromptSegment[]) {
+    const system: OpenAITextMessage[] = [];
+    const others: OpenAITextMessage[] = [];
+    const safety: OpenAITextMessage[] = [];
 
     for (const msg of segments) {
         if (msg.role === PromptRole.system) {
             system.push({ content: msg.content, role: "system" });
         } else if (msg.role === PromptRole.safety) {
-            safety.push({ content: msg.content, role: "system" });
+            safety.push({ content: "IMPORTANT: " + msg.content, role: "system" });
         } else {
             others.push({ content: msg.content, role: "user" });
         }

package/src/index.ts

@@ -1,4 +1,3 @@
 export * from "./Driver.js";
-export * from "./formatters/index.js";
 export * from "./json.js";
 export * from "./types.js";

package/src/json.ts

@@ -6,15 +6,15 @@ function extractJsonFromText(text: string): string {
     return text.replace(/\\n/g, "");
 }
 
-export function extractAndParseJSON(text: string): Json {
+export function extractAndParseJSON(text: string): JSONValue {
     return parseJSON(extractJsonFromText(text));
 }
 
-export type JsonPrimative = string | number | boolean | null;
-export type JsonArray = Json[];
-export type JsonObject = { [key: string]: Json };
-export type JsonComposite = JsonArray | JsonObject;
-export type Json = JsonPrimative | JsonComposite;
+export type JSONPrimitive = string | number | boolean | null;
+export type JSONArray = JSONValue[];
+export type JSONObject = { [key: string]: JSONValue };
+export type JSONComposite = JSONArray | JSONObject;
+export type JSONValue = JSONPrimitive | JSONComposite;
 
 
 
@@ -184,7 +184,7 @@ export class JsonParser {
 }
 
 
-export function parseJSON(text: string): Json {
+export function parseJSON(text: string): JSONValue {
     text = text.trim();
     try {
         return JSON.parse(text);

package/src/types.ts

@@ -1,6 +1,6 @@
 import { JSONSchema4 } from "json-schema";
-import { Readable } from "stream";
-import { JsonObject } from "./json.js";
+import { JSONObject } from "./json.js";
+import { PromptFormatter } from "./formatters/index.js";
 
 export interface EmbeddingsOptions {
     /**
@@ -81,7 +81,11 @@ export interface DriverOptions {
 
 export interface PromptOptions {
     model: string;
-    format?: PromptFormats;
+    /**
+     * A custom formatter to use for format the final model prompt from the input prompt segments. 
+     * If no one is specified the driver will choose a formatter compatible with the target model
+     */
+    format?: PromptFormatter;
     resultSchema?: JSONSchema4;
 }
 export interface ExecutionOptions extends PromptOptions {
@@ -171,13 +175,6 @@ export enum ModelType {
 // ============== Built-in formats and drivers =====================
 //TODO
 
-export enum PromptFormats {
-    openai = "openai",
-    llama2 = "llama2",
-    claude = "claude",
-    genericTextLLM = "genericTextLLM",
-}
-
 export enum BuiltinProviders {
     openai = 'openai',
     huggingface_ie = 'huggingface_ie',
@@ -193,21 +190,22 @@ export enum BuiltinProviders {
 // ============== training =====================
 
 
+
 export interface DataSource {
     name: string;
-    getStream(): Readable;
+    getStream(): ReadableStream<Uint8Array | string>;
     getURL(): Promise<string>;
 }
 
 export interface TrainingOptions {
     name: string; // the new model name
     model: string; // the model to train 
-    params?: JsonObject; // the training parameters
+    params?: JSONObject; // the training parameters
 }
 
 export interface TrainingPromptOptions {
     segments: PromptSegment[];
-    completion: string | JsonObject;
+    completion: string | JSONObject;
     model: string; // the model to train 
     schema?: JSONSchema4; // the resuilt schema f any
 }