streaming chat, switch to conversation API

app.py

@@ -34,34 +34,43 @@ def process_values_js():
 def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temperature, max_tokens, model: str, region):
     try:
         llm = LLM.create_llm(model)
-        body = llm.generate_body(message, history, system_prompt, temperature, max_tokens)
+        messages = llm.generate_body(message, history)
 
         config = Config(
-            read_timeout=600,
-            connect_timeout=30,
-            retries={
+            read_timeout = 600,
+            connect_timeout = 30,
+            retries = {
                 'max_attempts': 10,
                 'mode': 'adaptive'
             }
         )
 
         sess = boto3.Session(
-            aws_access_key_id=aws_access,
-            aws_secret_access_key=aws_secret,
-            aws_session_token=aws_token,
-            region_name=region)
+            aws_access_key_id = aws_access,
+            aws_secret_access_key = aws_secret,
+            aws_session_token = aws_token,
+            region_name = region)
         br = sess.client(service_name="bedrock-runtime", config = config)
 
-        response = br.invoke_model(body=body, modelId=f"{model}",
-                                accept="application/json", contentType="application/json")
-        response_body = json.loads(response.get('body').read())
-        result = llm.read_response(response_body)
+        response = br.converse_stream(
+            modelId = model,
+            messages = messages,
+            system = [{"text": system_prompt}],
+            inferenceConfig = {
+                "temperature": temperature,
+                "maxTokens": max_tokens,
+            }
+        )
+        response_stream = response.get('stream')
+        
+        partial_response = ""
+        for chunk in llm.read_response(response_stream):
+            partial_response += chunk
+            yield partial_response
 
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
 
-    return result
-
 def import_history(history, file):
     with open(file.name, mode="rb") as f:
         content = f.read()

llm.py

@@ -7,211 +7,175 @@ from doc2json import process_docx
 import fitz
 from PIL import Image
 import io
+import boto3
+from botocore.config import Config
+import re
 
 # constants
 log_to_console = False
-
-def process_pdf_img(pdf_fn: str):
-    pdf = fitz.open(pdf_fn)
-    message_parts = []
-
-    for page in pdf.pages():
-        # Create a transformation matrix for rendering at the calculated scale
-        mat = fitz.Matrix(0.6, 0.6)
-        
-        # Render the page to a pixmap
-        pix = page.get_pixmap(matrix=mat, alpha=False)
-        
-        # Convert pixmap to PIL Image
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        
-        # Convert PIL Image to bytes
-        img_byte_arr = io.BytesIO()
-        img.save(img_byte_arr, format='PNG')
-        img_byte_arr = img_byte_arr.getvalue()
-        
-        # Encode image to base64
-        base64_encoded = base64.b64encode(img_byte_arr).decode('utf-8')
-        
-        # Append the message part
-        message_parts.append({
-            "type": "text",
-            "text": f"Page {page.number} of file '{pdf_fn}'"
-        })
-        message_parts.append({"type": "image", "source": {"type": "base64",
-            "media_type": "image/png",
-            "data": base64_encoded}})
-
-    pdf.close()
-
-    return message_parts
-
-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')}
-
-def encode_file(fn: str) -> list:
-    user_msg_parts = []
-
-    if fn.endswith(".docx"):
-        user_msg_parts.append({"type": "text", "text": process_docx(fn)})
-    elif fn.endswith(".pdf"):
-        user_msg_parts.extend(process_pdf_img(fn))
-    else:
-        with open(fn, mode="rb") as f:
-            content = f.read()
-
-        isImage = False
-        if isinstance(content, bytes):
-            try:
-                # try to add as image
-                content = encode_image(content)
-                isImage = True
-            except:
-                # not an image, try text
-                content = content.decode('utf-8', 'replace')
-        else:
-            content = str(content)
-
-        if isImage:
-            user_msg_parts.append({"type": "image",
-                                "source": content})
-        else:
-            fname = os.path.basename(fn)
-            user_msg_parts.append({"type": "text", "text": f"``` {fname}\n{content}\n```"})
-
-    return user_msg_parts
+use_document_message_type = False  # AWS document message type usage
 
 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
 
+class LLM:
     @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}")
+        return LLM()
+
+    def generate_body(self, message, history):
+        messages = []
+
+        # AWS API requires strict user, assi, user, ... sequence
+        lastTypeHuman = False
 
-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:
-                if type(human) is tuple:
-                    user_msg_parts.extend(encode_file(human[0]))
+                if lastTypeHuman:
+                    last_msg = messages.pop()
+                    user_msg_parts = last_msg["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 isinstance(human, tuple):
+                    user_msg_parts.extend(self._process_file(human[0]))
+                else:
+                    user_msg_parts.extend([{"text": human}])
 
-        if message['text']:
-            user_msg_parts.append({"type": "text", "text": message['text']})
-        if message['files']:
-            for file in message['files']:
-                user_msg_parts.extend(encode_file(file['path']))
-        history_claude_format.append({"role": "user", "content": user_msg_parts})
+                messages.append({"role": "user", "content": user_msg_parts})
+                lastTypeHuman = True
+            if assi:
+                messages.append({"role": "assistant", "content": [{"text": assi}]})
+                lastTypeHuman = False
+        
         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
-            })
+        if message.text:
+            user_msg_parts.append({"text": message.text})
+        if message.files:
+            for file in message.files:
+                user_msg_parts.extend(self._process_file(file.path))
         
-        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 human:
-                if type(human) is tuple:
-                    prompt += f"[INST] {encode_file(human[0])} [/INST]"
-                else:
-                    prompt += f"[INST] {human} [/INST]"
+        if user_msg_parts:
+            messages.append({"role": "user", "content": user_msg_parts})
+        
+        return messages
 
-            if assi is not None:
-                prompt += f"{assi}</s>"
+    def _process_file(self, file_path):
+        if use_document_message_type and self._is_supported_document_type(file_path):
+            return [self._create_document_message(file_path)]
+        else:
+            return self._encode_file(file_path)
 
-        if message['text'] or message['files']:
-            prompt += "[INST] "
+    def _is_supported_document_type(self, file_path):
+        supported_extensions = ['.pdf', '.csv', '.doc', '.docx', '.xls', '.xlsx', '.html', '.txt', '.md']
+        return os.path.splitext(file_path)[1].lower() in supported_extensions
 
-            if message['text']:
-                prompt += message['text']
-            if message['files']:
-                for file in message['files']:
-                    prompt += f"{encode_file(file['path'])}\n"
-            
-            prompt += " [/INST]"
+    def _create_document_message(self, file_path):
+        with open(file_path, 'rb') as file:
+            file_content = file.read()
+        
+        file_name = re.sub(r'[^a-zA-Z0-9\s\-\(\)\[\]]', '', os.path.basename(file_path))[:200].strip() or "unnamed_file"
+        file_extension = os.path.splitext(file_path)[1][1:]  # Remove the dot
+
+        return {
+            "document": {
+                "name": file_name,
+                "format": file_extension,
+                "source": {
+                    "bytes": file_content
+                }
+            }
+        }
+
+    def _encode_file(self, fn: str) -> list:
+        if fn.endswith(".docx"):
+            return [{"text": process_docx(fn)}]
+        elif fn.endswith(".pdf"):
+            return self._process_pdf_img(fn)
+        else:
+            with open(fn, mode="rb") as f:
+                content = f.read()
+
+            if isinstance(content, bytes):
+                try:
+                    # try to add as image
+                    image_data = self._encode_image(content)
+                    return [{"image": image_data}]
+                except:
+                    # not an image, try text
+                    content = content.decode('utf-8', 'replace')
+            else:
+                content = str(content)
 
-        if log_to_console:
-            print(f"br_prompt: {str(prompt)}")
+            fname = os.path.basename(fn)
+            return [{"text": f"``` {fname}\n{content}\n```"}]
 
-        body = json.dumps({
-            "prompt": prompt,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-        })
+    def _process_pdf_img(self, pdf_fn: str):
+        pdf = fitz.open(pdf_fn)
+        message_parts = []
 
-        return body
+        for page in pdf.pages():
+            # Create a transformation matrix for rendering at the calculated scale
+            mat = fitz.Matrix(0.6, 0.6)
+            
+            # Render the page to a pixmap
+            pix = page.get_pixmap(matrix=mat, alpha=False)
+            
+            # Convert pixmap to PIL Image
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            
+            # Convert PIL Image to bytes
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+            
+            # Append the message parts
+            message_parts.append({"text": f"Page {page.number} of file '{pdf_fn}'"})
+            message_parts.append({"image": {
+                "format": "png",
+                "source": {"bytes": img_byte_arr}
+            }})
+
+        pdf.close()
+
+        return message_parts
+
+    def _encode_image(self, image_data):
+        # 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")
 
-    @staticmethod
-    def read_response(response_body) -> Type[str]:
-        return response_body.get('outputs')[0].get('text')
+        return {
+                "format": image_type,
+                "source": {"bytes": image_data}
+            }
+
+    def read_response(self, response_stream):
+        for event in response_stream:
+            if 'contentBlockDelta' in event:
+                yield event['contentBlockDelta']['delta']['text']
+            if 'messageStop' in event:
+                if log_to_console:
+                    print(f"\nStop reason: {event['messageStop']['stopReason']}")
+            if 'metadata' in event:
+                metadata = event['metadata']
+                if 'usage' in metadata and log_to_console:
+                    print("\nToken usage:")
+                    print(f"Input tokens: {metadata['usage']['inputTokens']}")
+                    print(f"Output tokens: {metadata['usage']['outputTokens']}")
+                    print(f"Total tokens: {metadata['usage']['totalTokens']}")