Claude-3 and Mistral support

app.py

@@ -5,10 +5,9 @@ import boto3
 
 from doc2json import process_docx
 from settings_mgr import generate_download_settings_js, generate_upload_settings_js
+from llm import LLM
 
 dump_controls = False
-log_to_console = False
-
 
 def add_text(history, text):
     history = history + [(text, None)]
@@ -59,26 +58,10 @@ def process_values_js():
     }
     """
 
-def bot(message, history, aws_access, aws_secret, aws_token, temperature, max_tokens, model, region):
+def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temperature, max_tokens, model: str, region):
     try:
-        prompt = "\n\n"
-        for human, assi in history:
-            if prompt is not None:
-                prompt += f"Human: {human}\n\n"
-            if assi is not None:
-                prompt += f"Assistant: {assi}\n\n"
-        if message:
-            prompt += f"Human: {message}\n\n"
-        prompt += f"Assistant:"
-
-        if log_to_console:
-            print(f"br_prompt: {str(prompt)}")
-
-        body = json.dumps({
-            "prompt": prompt,
-            "max_tokens_to_sample": max_tokens,
-            "temperature": temperature,
-        })
+        llm = LLM.create_llm(model)
+        body = llm.generate_body(message, history, system_prompt, temperature, max_tokens)
 
         sess = boto3.Session(
             aws_access_key_id=aws_access,
@@ -87,21 +70,18 @@ def bot(message, history, aws_access, aws_secret, aws_token, temperature, max_to
             region_name=region)
         br = sess.client(service_name="bedrock-runtime")
 
-        response = br.invoke_model(body=body, modelId=f"anthropic.{model}",
+        response = br.invoke_model(body=body, modelId=f"{model}",
                                 accept="application/json", contentType="application/json")
         response_body = json.loads(response.get('body').read())
-        br_result = response_body.get('completion')
+        br_result = llm.read_response(response_body)
 
         history[-1][1] = br_result
-        if log_to_console:
-            print(f"br_result: {str(history)}")
 
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
 
     return "", history
 
-
 def import_history(history, file):
     with open(file.name, mode="rb") as f:
         content = f.read()
@@ -117,16 +97,18 @@ def import_history(history, file):
     return history
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Amazon™️ Bedrock™️ Chat™️ (Nils' Version™️) feat. Anthropic™️ Claude-2™️")
+    gr.Markdown("# Amazon™️ Bedrock™️ Chat™️ (Nils' Version™️) feat. Mistral™️ AI & Anthropic™️ Claude™️")
 
     with gr.Accordion("Settings"):
         aws_access = gr.Textbox(label="AWS Access Key", elem_id="aws_access")
         aws_secret = gr.Textbox(label="AWS Secret Key", elem_id="aws_secret")
         aws_token = gr.Textbox(label="AWS Session Token", elem_id="aws_token")
-        model = gr.Dropdown(label="Model", value="claude-v2:1", allow_custom_value=True, elem_id="model",
-                            choices=["claude-v2:1", "claude-v2"])
-        region = gr.Dropdown(label="Region", value="eu-central-1", allow_custom_value=True, elem_id="region",
-                            choices=["eu-central-1", "us-east-1", "us-west-1"])
+        model = gr.Dropdown(label="Model", value="anthropic.claude-3-sonnet-20240229-v1:0", allow_custom_value=True, elem_id="model",
+                            choices=["anthropic.claude-3-sonnet-20240229-v1:0", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-v2:1", "anthropic.claude-v2",
+                                     "mistral.mistral-7b-instruct-v0:2", "mistral.mixtral-8x7b-instruct-v0:1", "mistral.mistral-large-2402-v1:0"])
+        system_prompt = gr.TextArea("You are a helpful AI.", label="System Prompt", lines=3, max_lines=250, elem_id="system_prompt")  
+        region = gr.Dropdown(label="Region", value="eu-west-3", allow_custom_value=True, elem_id="region",
+                            choices=["eu-central-1", "eu-west-3", "us-east-1", "us-west-1"])
         temp = gr.Slider(0, 1, label="Temperature", elem_id="temp", value=1)
         max_tokens = gr.Slider(1, 200000, label="Max. Tokens", elem_id="max_tokens", value=4000)
         save_button = gr.Button("Save Settings")  
@@ -136,7 +118,7 @@ with gr.Blocks() as demo:
 
         load_button.click(load_settings, js="""  
             () => {  
-                let elems = ['#aws_access textarea', '#aws_secret textarea', '#aws_token textarea', '#temp input', '#max_tokens input', '#model', '#region'];
+                let elems = ['#aws_access textarea', '#aws_secret textarea', '#aws_token textarea', '#system_prompt textarea', '#temp input', '#max_tokens input', '#model', '#region'];
                 elems.forEach(elem => {
                     let item = document.querySelector(elem);
                     let event = new InputEvent('input', { bubbles: true });
@@ -146,11 +128,12 @@ with gr.Blocks() as demo:
             }  
         """)
 
-        save_button.click(save_settings, [aws_access, aws_secret, aws_token, temp, max_tokens, model, region], js="""  
-            (acc, sec, tok, prompt, temp, ntok, model, region) => {  
+        save_button.click(save_settings, [aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region], js="""  
+            (acc, sec, tok, system_prompt, temp, ntok, model, region) => {  
                 localStorage.setItem('aws_access', acc);  
                 localStorage.setItem('aws_secret', sec);  
                 localStorage.setItem('aws_token', tok);  
+                localStorage.setItem('system_prompt', system_prompt);
                 localStorage.setItem('temp', document.querySelector('#temp input').value);  
                 localStorage.setItem('max_tokens', document.querySelector('#max_tokens input').value);  
                 localStorage.setItem('model', model);  
@@ -161,11 +144,12 @@ with gr.Blocks() as demo:
         control_ids = [('aws_access', '#aws_access textarea'),
                        ('aws_secret', '#aws_secret textarea'),
                        ('aws_token', '#aws_token textarea'),
+                       ('system_prompt', '#system_prompt textarea'),
                        ('temp', '#temp input'),
                        ('max_tokens', '#max_tokens input'),
                        ('model', '#model'),
                        ('region', '#region')]
-        controls = [aws_access, aws_secret, aws_token, temp, max_tokens, model, region]
+        controls = [aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region]
 
         dl_settings_button.click(None, controls, js=generate_download_settings_js("amz_chat_settings.bin", control_ids))
         ul_settings_button.click(None, None, None, js=generate_upload_settings_js(control_ids))
@@ -187,7 +171,7 @@ with gr.Blocks() as demo:
         )
         submit_btn = gr.Button("🚀 Send", scale=0)
         submit_click = submit_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [txt, chatbot, aws_access, aws_secret, aws_token, temp, max_tokens, model, region], [txt, chatbot],
+            bot, [txt, chatbot, aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region], [txt, chatbot],
         )
         submit_click.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
 
@@ -256,7 +240,7 @@ with gr.Blocks() as demo:
         import_button.upload(import_history, inputs=[chatbot, import_button], outputs=[chatbot])
 
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-        bot, [txt, chatbot, aws_access, aws_secret, aws_token, temp, max_tokens, model, region], [txt, chatbot],
+        bot, [txt, chatbot, aws_access, aws_secret, aws_token, system_prompt, temp, max_tokens, model, region], [txt, chatbot],
     )
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
     file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False, postprocess=False)

llm.py

@@ -0,0 +1,134 @@
+from abc import ABC, abstractmethod
+from typing import Type, TypeVar
+import base64
+import json
+
+# constants
+image_embed_prefix = "🖼️🆙 "
+log_to_console = False
+
+def encode_image(image_data):
+    """Generates a prefix for image base64 data in the required format for the
+    four known image formats: png, jpeg, gif, and webp.
+
+    Args:
+    image_data: The image data, encoded in base64.
+
+    Returns:
+    An object encoding the image
+    """
+
+    # Get the first few bytes of the image data.
+    magic_number = image_data[:4]
+  
+    # Check the magic number to determine the image type.
+    if magic_number.startswith(b'\x89PNG'):
+        image_type = 'png'
+    elif magic_number.startswith(b'\xFF\xD8'):
+        image_type = 'jpeg'
+    elif magic_number.startswith(b'GIF89a'):
+        image_type = 'gif'
+    elif magic_number.startswith(b'RIFF'):
+        if image_data[8:12] == b'WEBP':
+            image_type = 'webp'
+        else:
+            # Unknown image type.
+            raise Exception("Unknown image type")
+    else:
+        # Unknown image type.
+        raise Exception("Unknown image type")
+
+    return {"type": "base64",
+            "media_type": "image/" + image_type,
+            "data": base64.b64encode(image_data).decode('utf-8')}
+
+LLMClass = TypeVar('LLMClass', bound='LLM')
+class LLM(ABC):
+    @abstractmethod
+    def generate_body(message, history, system_prompt, temperature, max_tokens):
+        pass
+
+    @abstractmethod
+    def read_response(message, history, system_prompt, temperature, max_tokens):
+        pass
+
+    @staticmethod
+    def create_llm(model: str) -> Type[LLMClass]:
+        if model.startswith("anthropic.claude"):
+            return Claude()
+        elif model.startswith("mistral."):
+            return Mistral()
+        else:
+            raise ValueError(f"Unsupported model: {model}")
+
+class Claude(LLM):
+    @staticmethod
+    def generate_body(message, history, system_prompt, temperature, max_tokens):
+        history_claude_format = []
+        user_msg_parts = []
+        for human, assi in history:
+            if human is not None:
+                if human.startswith(image_embed_prefix):
+                    with open(human.lstrip(image_embed_prefix), mode="rb") as f:
+                        content = f.read()
+                    user_msg_parts.append({"type": "image",
+                                            "source": encode_image(content)})
+                else:
+                    user_msg_parts.append({"type": "text", "text": human})
+
+            if assi is not None:
+                if user_msg_parts:
+                    history_claude_format.append({"role": "user", "content": user_msg_parts})
+                    user_msg_parts = []
+
+                history_claude_format.append({"role": "assistant", "content": assi})
+
+        if message:
+            user_msg_parts.append({"type": "text", "text": human})
+            
+        if user_msg_parts:
+            history_claude_format.append({"role": "user", "content": user_msg_parts})
+
+        if log_to_console:
+            print(f"br_prompt: {str(history_claude_format)}")
+
+        body = json.dumps({
+                "anthropic_version": "bedrock-2023-05-31",
+                "system": system_prompt,
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "messages": history_claude_format
+            })
+        
+        return body
+
+    @staticmethod
+    def read_response(response_body) -> Type[str]:
+        return response_body.get('content')[0].get('text')
+
+class Mistral(LLM):
+    @staticmethod
+    def generate_body(message, history, system_prompt, temperature, max_tokens):
+        prompt = "<s>"
+        for human, assi in history:
+            if prompt is not None:
+                prompt += f"[INST] {human} [/INST]\n"
+            if assi is not None:
+                prompt += f"{assi}</s>\n"
+        if message:
+            prompt += f"[INST] {message} [/INST]"
+
+        if log_to_console:
+            print(f"br_prompt: {str(prompt)}")
+
+        body = json.dumps({
+            "prompt": prompt,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        })
+
+        return body
+
+    @staticmethod
+    def read_response(response_body) -> Type[str]:
+        return response_body.get('outputs')[0].get('text')

requirements.txt

@@ -1,4 +1,4 @@
-gradio
-langchain
-boto3
+gradio
+langchain
+boto3>1.34.54
 lxml