Merge ced37e6a456bac238064d58ab0334fdd9e856e24 into 7e86f19aac93768e10bcd3a117b04c510c839708

chainforge/react-server/package.json

@@ -7,20 +7,20 @@
   },
   "dependencies": {
     "@anthropic-ai/sdk": "^0.4.4",
+    "@aws-sdk/client-bedrock-runtime": "^3.602.0",
     "@azure/openai": "^1.0.0-beta.2",
     "@codemirror/lang-javascript": "^6.1.5",
     "@codemirror/lang-python": "^6.1.2",
     "@emoji-mart/data": "^1.1.2",
     "@emoji-mart/react": "^1.1.1",
     "@fontsource/geist-mono": "^5.0.1",
-    "@google-ai/generativelanguage": "^0.2.0",
+    "@google-ai/generativelanguage": "^2.6.0",
     "@google/generative-ai": "^0.1.3",
-    "@mantine/core": "^6.0.9",
-    "@mantine/dates": "^6.0.13",
-    "@mantine/dropzone": "^6.0.19",
+    "@mantine/core": "^6.0.21",
+    "@mantine/dates": "^6.0.21",
+    "@mantine/dropzone": "^6.0.21",
     "@mantine/form": "^6.0.11",
-    "@mantine/prism": "^6.0.15",
-    "@mirai73/bedrock-fm": "^0.4.10",
+    "@mantine/prism": "^6.0.21",
     "@reactflow/background": "^11.2.0",
     "@reactflow/controls": "^11.1.11",
     "@reactflow/core": "^11.7.0",
@@ -86,8 +86,8 @@
     "react-dom": "^18.2.0",
     "react-edit-text": "^5.1.0",
     "react-plotly.js": "^2.6.0",
-    "react-scripts": "5.0.1",
-    "reactflow": "^11.0",
+    "react-scripts": "^5.0.1",
+    "reactflow": "^11.0.0",
     "request": "^2.88.2",
     "socket.io-client": "^4.6.1",
     "stream-browserify": "^3.0.0",
@@ -133,6 +133,7 @@
   "devDependencies": {
     "@craco/craco": "^7.1.0",
     "@types/lodash": "^4.17.0",
+    "@types/markdown-it": "^14.1.1",
     "@types/papaparse": "^5.3.14",
     "@types/react-beautiful-dnd": "^13.1.8",
     "@types/react-edit-text": "^5.0.4",

chainforge/react-server/src/ModelSettingSchemas.tsx

@@ -1476,9 +1476,9 @@ const BedrockClaudeSettings: ModelSettingsDict = {
           "anthropic.claude-3-haiku-20240307-v1:0": "claude-3-haiku",
         },
       },
-      system_msg: {
+      system: {
         type: "string",
-        title: "system_msg",
+        title: "system",
         description: "A system message to use with the model",
         default: "",
       },
@@ -1492,24 +1492,14 @@ const BedrockClaudeSettings: ModelSettingsDict = {
         maximum: 1,
         multipleOf: 0.01,
       },
-      max_tokens_to_sample: {
+      max_tokens: {
         type: "integer",
-        title: "max_tokens_to_sample",
+        title: "max_tokens",
         description:
           "A maximum number of tokens to generate before stopping. Lower this if you want shorter responses. By default, ChainForge uses the value 1024, although the Anthropic API does not specify a default value.",
         default: 1024,
         minimum: 1,
       },
-      custom_prompt_wrapper: {
-        type: "string",
-        title: "Prompt Wrapper (ChainForge)",
-        description:
-          // eslint-disable-next-line no-template-curly-in-string
-          'Anthropic models expect prompts in the form "\\n\\nHuman: ${prompt}\\n\\nAssistant:". ChainForge wraps all prompts in this template by default. If you wish to' +
-          // eslint-disable-next-line no-template-curly-in-string
-          "explore custom prompt wrappers that deviate, write a Python template here with a single variable, ${prompt}, where the actual prompt text should go. Otherwise, leave this field blank. (Note that you should enter newlines as newlines, not escape codes like \\n.)",
-        default: "",
-      },
       stop_sequences: {
         type: "string",
         title: "stop_sequences",
@@ -1551,7 +1541,7 @@ const BedrockClaudeSettings: ModelSettingsDict = {
       "ui:help": "Defaults to 1.0.",
       "ui:widget": "range",
     },
-    max_tokens_to_sample: {
+    max_tokens: {
       "ui:help": "Defaults to 1024.",
     },
     top_k: {
@@ -1564,11 +1554,6 @@ const BedrockClaudeSettings: ModelSettingsDict = {
       "ui:widget": "textarea",
       "ui:help": 'Defaults to one stop sequence, "\\n\\nHuman: "',
     },
-    custom_prompt_wrapper: {
-      "ui:widget": "textarea",
-      "ui:help":
-        'Defaults to Anthropic\'s internal wrapper "\\n\\nHuman: {prompt}\\n\\nAssistant".',
-    },
   },
 
   postprocessors: {
@@ -1704,6 +1689,123 @@ const BedrockJurassic2Settings: ModelSettingsDict = {
   },
 };
 
+const BedrockJambaInstruct: ModelSettingsDict = {
+  fullName: "Jamba Instruct (Ai21) via Amazon Bedrock",
+  schema: {
+    type: "object",
+    required: ["shortname"],
+    properties: {
+      shortname: {
+        type: "string",
+        title: "Nickname",
+        description:
+          "Unique identifier to appear in ChainForge. Keep it short.",
+        default: "JambaInstruct",
+      },
+      model: {
+        type: "string",
+        title: "Model Version",
+        description:
+          "Select a version of Jamba Instruct to query. For more details on the differences, see the AI21 API documentation.",
+        enum: [NativeLLM.Bedrock_Jamba_Instruct],
+        default: NativeLLM.Bedrock_Jamba_Instruct,
+      },
+      temperature: {
+        type: "number",
+        title: "temperature",
+        description:
+          "Amount of randomness injected into the response. Ranges from 0 to 1. Use temp closer to 0 for analytical / multiple choice, and temp closer to 1 for creative and generative tasks.",
+        default: 1,
+        minimum: 0,
+        maximum: 1,
+        multipleOf: 0.01,
+      },
+      max_tokens: {
+        type: "integer",
+        title: "max_tokens",
+        description:
+          "The maximum number of tokens to generate for each response.",
+        default: 1024,
+        minimum: 1,
+      },
+      top_p: {
+        type: "number",
+        title: "top_p",
+        description:
+          "Does nucleus sampling, in which we compute the cumulative distribution over all the options for each subsequent token in decreasing probability order and cut it off once it reaches a particular probability specified by top_p. Defaults to -1, which disables it. Note that you should either alter temperature or top_p, but not both.",
+        default: 1,
+        minimum: 0.01,
+        maximum: 1,
+        multipleOf: 0.001,
+      },
+      stop_sequences: {
+        type: "string",
+        title: "stop",
+        description:
+          'Enclose stop sequences in double-quotes "" and use whitespace to separate them.',
+        default: "",
+      },
+      frequency_penalty: {
+        type: "number",
+        title: "frequency_penalty",
+        description:
+          "How much to penalize new tokens based on their existing frequency in the text. Must be between -2.0 and 2.0.",
+        minimum: -2.0,
+        maximum: 2.0,
+        multipleOf: 0.01,
+        default: 0,
+      },
+      presence_penalty: {
+        type: "number",
+        title: "presence_penalty",
+        description:
+          "How much to penalize new tokens based on whether they appear in the text so far. Must be between -2.0 and 2.0.",
+        minimum: -2.0,
+        maximum: 2.0,
+        multipleOf: 0.01,
+        default: 0,
+      },
+    },
+  },
+  postprocessors: {
+    stop_sequences: (str) => {
+      if (typeof str !== "string" || str.trim().length === 0) return [];
+      return str
+        .match(/"((?:[^"\\]|\\.)*)"/g)
+        ?.map((s) => s.substring(1, s.length - 1)); // split on double-quotes but exclude escaped double-quotes inside the group
+    },
+  },
+  uiSchema: {
+    "ui:submitButtonOptions": UI_SUBMIT_BUTTON_SPEC,
+    shortname: {
+      "ui:autofocus": true,
+    },
+    model: {
+      "ui:help": "Defaults to Jurassic 2 Ultra. ",
+    },
+    temperature: {
+      "ui:help": "Defaults to 1.0.",
+      "ui:widget": "range",
+    },
+    max_tokens: {
+      "ui:help": "Defaults to 1024.",
+    },
+    top_p: {
+      "ui:help": "Defaults to 1.",
+    },
+    frequency_penalty: {
+      "ui:help": "Defaults to 0.",
+    },
+    presence_penalty: {
+      "ui:help": "Defaults to 0.",
+    },
+    stop_sequences: {
+      "ui:widget": "textarea",
+      "ui:help": "Defaults to no sequence",
+    },
+  },
+};
+
 const BedrockTitanSettings: ModelSettingsDict = {
   fullName: "Titan (Amazon) via Amazon Bedrock",
   schema: {
@@ -1964,14 +2066,6 @@ const MistralSettings: ModelSettingsDict = {
           'Enclose stop sequences in double-quotes "" and use whitespace to separate them.',
         default: "",
       },
-      top_k: {
-        type: "integer",
-        title: "top_k",
-        description:
-          "The number of top-scoring tokens to consider for each generation step.",
-        minimum: 0,
-        default: 0,
-      },
       top_p: {
         type: "number",
         title: "top_p",
@@ -2335,6 +2429,7 @@ export const ModelSettings: Dict<ModelSettingsDict> = {
   ollama: OllamaSettings,
   "br.anthropic.claude": BedrockClaudeSettings,
   "br.ai21.j2": BedrockJurassic2Settings,
+  "br.ai21.jamba": BedrockJambaInstruct,
   "br.amazon.titan": BedrockTitanSettings,
   "br.cohere.command": BedrockCommandTextSettings,
   "br.mistral.mistral": MistralSettings,

chainforge/react-server/src/backend/models.ts

@@ -108,6 +108,7 @@ export enum NativeLLM {
   Bedrock_Claude_Instant_1 = "anthropic.claude-instant-v1",
   Bedrock_Jurassic_Ultra = "ai21.j2-ultra",
   Bedrock_Jurassic_Mid = "ai21.j2-mid",
+  Bedrock_Jamba_Instruct = "ai21.jamba-instruct-v1:0",
   Bedrock_Titan_Light = "amazon.titan-text-lite-v1",
   Bedrock_Titan_Large = "amazon.titan-tg1-large",
   Bedrock_Titan_Express = "amazon.titan-text-express-v1",
@@ -261,6 +262,7 @@ export const RATE_LIMIT_BY_MODEL: { [key in LLM]?: number } = {
   [NativeLLM.GEMINI_PRO]: 60,
   [NativeLLM.Bedrock_Jurassic_Mid]: 400,
   [NativeLLM.Bedrock_Jurassic_Ultra]: 25,
+  [NativeLLM.Bedrock_Jamba_Instruct]: 100,
   [NativeLLM.Bedrock_Titan_Light]: 800,
   [NativeLLM.Bedrock_Titan_Express]: 400, // 400 RPM
   [NativeLLM.Bedrock_Claude_2]: 500, // 500 RPM

chainforge/react-server/src/backend/utils.ts

@@ -43,13 +43,14 @@ import {
 } from "@azure/openai";
 import { GoogleGenerativeAI } from "@google/generative-ai";
 import { UserForcedPrematureExit } from "./errors";
-import {
-  fromModelId,
-  ChatMessage as BedrockChatMessage,
-} from "@mirai73/bedrock-fm";
 import StorageCache from "./cache";
 import Compressor from "compressorjs";
-import { Models } from "@mirai73/bedrock-fm/lib/bedrock";
+import {
+  BedrockRuntimeClient,
+  ConverseCommand,
+  Message,
+} from "@aws-sdk/client-bedrock-runtime";
+import { captureRejectionSymbol } from "events";
 
 const ANTHROPIC_HUMAN_PROMPT = "\n\nHuman:";
 const ANTHROPIC_AI_PROMPT = "\n\nAssistant:";
@@ -1293,26 +1294,14 @@ export async function call_ollama_provider(
 }
 
 /** Convert OpenAI chat history to Bedrock format */
-function to_bedrock_chat_history(
-  chat_history: ChatHistory,
-): BedrockChatMessage[] {
-  const role_map: Dict<string> = {
-    assistant: "ai",
-    user: "human",
-  };
-
+function to_bedrock_chat_history(chat_history: ChatHistory): Message[] {
   // Transform the ChatMessage format in the chat_history array to what is expected by Bedrock
-  return chat_history.map((msg) =>
-    transformDict(
-      msg,
-      undefined,
-      (key) => (key === "content" ? "message" : key),
-      (key: string, val: string): string => {
-        if (key === "role") return val in role_map ? role_map[val] : val;
-        return val;
-      },
-    ),
-  ) as BedrockChatMessage[];
+  return chat_history
+    .filter((msg) => msg.role !== "system")
+    .map((msg) => ({
+      role: msg.role,
+      content: [{ text: msg.content }],
+    })) as Message[];
 }
 
 /**
@@ -1338,6 +1327,13 @@ export async function call_bedrock(
     );
   }
 
+  const query: Dict = {
+    model: model,
+    stream: false,
+    temperature,
+    ...params, // 'the rest' of the settings, passed from the front-end settings
+  };
+
   const modelName: string = model.toString();
   let stopWords = [];
   if (
@@ -1357,15 +1353,9 @@ export async function call_bedrock(
 
   delete params?.stop_sequences;
 
-  const query: Dict = {
-    stopSequences: stopWords,
-    temperature,
-  };
-
-  const fm = fromModelId(modelName as Models, {
+  const fm = new BedrockRuntimeClient({
     region: bedrockConfig.region,
     credentials: bedrockConfig.credentials,
-    ...query,
   });
 
   const responses: string[] = [];
@@ -1376,29 +1366,32 @@ export async function call_bedrock(
       if (should_cancel && should_cancel()) throw new UserForcedPrematureExit();
 
       // Grab the response
-      let response: string;
-      if (
-        modelName.startsWith("anthropic") ||
-        modelName.startsWith("mistral") ||
-        modelName.startsWith("meta")
-      ) {
-        const chat_history: ChatHistory = construct_openai_chat_history(
-          prompt,
-          params?.chat_history,
-          params?.system_msg,
-        );
+      const chat_history: ChatHistory = construct_openai_chat_history(
+        prompt,
+        params?.chat_history,
+      );
 
-        response = (
-          await fm.chat(to_bedrock_chat_history(chat_history), {
-            modelArgs: { ...(params as Map<string, any>) },
-          })
-        ).message;
-      } else {
-        response = await fm.generate(prompt, {
-          modelArgs: { ...(params as Map<string, any>) },
-        });
+      if (modelName.includes("titan")) {
+        params = { textGenerationConfig: { ...params } };
+      }
+
+      const response = (
+        await fm.send(
+          new ConverseCommand({
+            modelId: modelName,
+            messages: to_bedrock_chat_history(chat_history),
+            inferenceConfig: {
+              stopSequences: stopWords,
+            },
+            additionalModelRequestFields: { ...params },
+          }),
+        )
+      ).output?.message?.content;
+      if (response === undefined) {
+        throw new Error("Bedrock API returned an empty response.");
+      } else {
+        responses.push(response[0].text ?? "");
       }
-      responses.push(response);
     }
   } catch (error: any) {
     throw new Error(error?.message ?? error.toString());

chainforge/react-server/src/store.tsx

@@ -198,6 +198,13 @@ export const initLLMProviderMenu: (LLMSpec | LLMGroup)[] = [
         base_model: "br.ai21.j2",
         temp: 0.9,
       },
+      {
+        name: "AI21 Jamba Instruct",
+        emoji: "🐍",
+        model: NativeLLM.Bedrock_Jamba_Instruct,
+        base_model: "br.ai21.jamba",
+        temp: 0.9,
+      },
       {
         name: "Amazon Titan",
         emoji: "🏛️",