feat: add token limit chunks to embedding models (#670)

* feat: add token limit chunks to embedding models

* fix: linting, dry build and cdk-lib issues

---------

Co-authored-by: Maryam Khidir <mkhidir@amazon.de>

Author: maryamkhidir

.github/workflows/build.yaml

@@ -7,9 +7,13 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
           node-version: "20"
+      - name: Install latest CDK CLI
+        run: |
+          npm install -g aws-cdk@latest
+          cdk --version
       - name: Formatting
         run: |
           npm ci
@@ -23,7 +27,7 @@ jobs:
           npm audit
           npm run build
           npm run test
-          npx cdk synth
+          cdk synth
       - name: PyTests
         # Suppression of pip audit failure until langchain is upgraded.
         run: |

.python-version

@@ -6,7 +6,7 @@ The following Python packages may be included in this product:
 - cfnresponse==1.1.2
 - opensearch-py==2.3.1
 - openai==0.28.0
-- requests==2.32.0
+- requests==2.32.4
 - huggingface-hub
 - hf-transfer
 - aws_xray_sdk==2.12.1
@@ -817,7 +817,7 @@ Agreement.
 
 The following Python packages may be included in this product:
 
-- urllib3<2
+- urllib3==2.5.0
 
 These packages each contain the following license and notice below:
 

NOTICE

@@ -27,6 +27,16 @@ This blueprint deploys the complete AWS GenAI LLM Chatbot solution in your AWS a
 - AWS CLI configured with credentials
 - Node.js 18+ and npm
 - Python 3.8+
+- AWS CDK CLI version compatible with aws-cdk-lib 2.206.0 or later
+  ```bash
+  # Install or update the CDK CLI globally
+  npm install -g aws-cdk@latest
+  
+  # Verify the installed version
+  cdk --version
+  ```
+
+> **Important**: The CDK CLI version must be compatible with the aws-cdk-lib version used in this project (currently 2.206.0). If you encounter a "Cloud assembly schema version mismatch" error during deployment, update your CDK CLI to the latest version using the command above.
 
 ### Deployment
 

README.md

@@ -72,7 +72,7 @@ def add_user_to_group(cognito, username, group_name, user_pool_id):
 
 def handler(event, context):
     print(f"Event received: {event}")
-    
+
     # Handle different trigger types with different event structures
     if "request" in event and "userAttributes" in event["request"]:
         # POST_AUTHENTICATION trigger
@@ -88,18 +88,30 @@ def handler(event, context):
         new_group = user_attributes.get("custom:chatbot_role")
         user_pool_id = event["userPoolId"]
         trigger_type = "PRE_AUTHENTICATION"
-    elif "request" in event and "userAttributes" in event["request"] and "validationData" in event["request"]:
+    elif (
+        "request" in event
+        and "userAttributes" in event["request"]
+        and "validationData" in event["request"]
+    ):
         # POST_CONFIRMATION trigger
         user_attributes = event["request"]["userAttributes"]
         username = user_attributes.get("sub") or user_attributes.get("username")
         new_group = user_attributes.get("custom:chatbot_role")
         user_pool_id = event["userPoolId"]
         trigger_type = "POST_CONFIRMATION"
-    elif "request" in event and "userAttributes" in event["request"] and "validationData" not in event["request"]:
+    elif (
+        "request" in event
+        and "userAttributes" in event["request"]
+        and "validationData" not in event["request"]
+    ):
         # PRE_SIGN_UP trigger
         user_attributes = event["request"]["userAttributes"]
         # For Pre sign-up, username might be in different fields
-        username = user_attributes.get("sub") or user_attributes.get("username") or user_attributes.get("email")
+        username = (
+            user_attributes.get("sub")
+            or user_attributes.get("username")
+            or user_attributes.get("email")
+        )
         new_group = user_attributes.get("custom:chatbot_role")
         user_pool_id = event["userPoolId"]
         trigger_type = "PRE_SIGN_UP"
@@ -115,7 +127,7 @@ def handler(event, context):
 
     # Get default group from environment variable or use 'user' as fallback
     default_group = os.environ.get("DEFAULT_USER_GROUP", "user")
-    
+
     # If no custom:chatbot_role is provided, use default group
     if not new_group:
         new_group = default_group
@@ -125,18 +137,22 @@ def handler(event, context):
     if trigger_type == "PRE_SIGN_UP":
         print("Pre sign-up trigger - user will be created after this trigger completes")
         print(f"Will assign user to group: {new_group}")
-        print("Note: Group assignment will happen in a separate trigger (POST_CONFIRMATION)")
-        
+        print(
+            "Note: Group assignment will happen in a separate \
+            trigger (POST_CONFIRMATION)"
+        )
+
         # For Pre sign-up, we can only validate or modify the sign-up request
         # We cannot assign groups yet as the user doesn't exist
-        # The group assignment will need to happen in POST_CONFIRMATION or PRE_AUTHENTICATION
-        
+        # The group assignment will need to happen in
+        # POST_CONFIRMATION or PRE_AUTHENTICATION
+
         # You might want to add the group information to the user attributes
         # so it can be used later in POST_CONFIRMATION
         if "custom:chatbot_role" not in user_attributes:
             user_attributes["custom:chatbot_role"] = new_group
             print(f"Added custom:chatbot_role attribute: {new_group}")
-        
+
         return event
 
     # For other triggers (PRE_AUTHENTICATION, POST_AUTHENTICATION, POST_CONFIRMATION)

lib/authentication/lambda/addFederatedUserToUserGroup/index.py

@@ -14,9 +14,8 @@
 
 @router.resolver(field_name="listModels")
 @tracer.capture_method
-@permissions.approved_roles([
-    permissions.ADMIN_ROLE,
-    permissions.WORKSPACES_MANAGER_ROLE
-])
+@permissions.approved_roles(
+    [permissions.ADMIN_ROLE, permissions.WORKSPACES_MANAGER_ROLE]
+)
 def models() -> list[dict[str, Any]]:
     return genai_core.models.list_models()

lib/chatbot-api/functions/api-handler/routes/models.py

@@ -37,7 +37,7 @@ def _forward(self, model_inputs, **kwargs):
                 input_ids=input_ids.to(self.model.device),
                 attention_mask=attention_mask.to(self.model.device),
                 return_dict_in_generate=True,
-                **kwargs
+                **kwargs,
             )
 
         return {"input_ids": input_ids, "outputs": outputs}

lib/sagemaker-model/hf-custom-script-model/samples/pipeline/inference.py

@@ -1,18 +1,18 @@
-boto3==1.28.57
+boto3==1.40.13
 aws-lambda-powertools==2.25.1
 numpy==1.26.4
 cfnresponse==1.1.2
 aws_requests_auth==0.4.3
 requests-aws4auth==1.2.3
 langchain==0.3.7
 langchain-community==0.3.3
-opensearch-py==2.3.1
-psycopg2-binary==2.9.7
+opensearch-py==3.0.0
+psycopg2-binary==2.9.10
 pgvector==0.2.2
-urllib3<2
+urllib3==2.5.0
 openai==1.47.0
 beautifulsoup4==4.12.2
-requests==2.32.2
+requests==2.32.4
 attrs==23.1.0
 feedparser==6.0.11
 PyJWT==2.9.0

lib/shared/file-import-batch-job/requirements.txt

@@ -12,9 +12,9 @@ langchain-text-splitters==0.3.5
 opensearch-py==2.4.2
 psycopg2-binary==2.9.7
 pgvector==0.2.2
-urllib3<2
+urllib3==2.5.0
 beautifulsoup4==4.12.2
-requests==2.32.0
+requests==2.32.4
 attrs==23.1.0
 feedparser==6.0.11
 defusedxml==0.7.1

lib/shared/layers/common/requirements.txt

@@ -18,25 +18,85 @@
 logger = Logger()
 
 
+def get_model_token_limit(model_name):
+    # Extract provider from model name
+    model_provider = model_name.split(".")[0]
+
+    # https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html
+    # https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/embeddings
+    # https://docs.cohere.com/v2/docs/cohere-embed
+    PROVIDER_TOKEN_LIMITS = {
+        Provider.AMAZON.value: 8000,  # Amazon Titan models
+        Provider.COHERE.value: 512,  # Cohere models
+        Provider.OPENAI.value: 8191,  # OpenAI models
+        "default": 2500,  # Default fallback (2500 * 4 = 10000)
+    }
+
+    return PROVIDER_TOKEN_LIMITS.get(model_provider, PROVIDER_TOKEN_LIMITS["default"])
+
+
 def generate_embeddings(
     model: EmbeddingsModel, input: list[str], task: str = "store", batch_size: int = 50
 ) -> list[list[float]]:
-    input = [x[:10000] for x in input]
-
-    ret_value = []
-    batch_split = [input[i : i + batch_size] for i in range(0, len(input), batch_size)]
-
-    for batch in batch_split:
-        if model.provider == Provider.OPENAI.value:
-            ret_value.extend(_generate_embeddings_openai(model, batch))
-        elif model.provider == Provider.BEDROCK.value:
-            ret_value.extend(_generate_embeddings_bedrock(model, batch, task))
-        elif model.provider == Provider.SAGEMAKER.value:
-            ret_value.extend(_generate_embeddings_sagemaker(model, batch))
-        else:
-            raise CommonError(f"Unknown provider: {model.provider}")
+    try:
+        # Get model-specific token limit
+        token_limit = get_model_token_limit(model.name)
+        char_limit = min(token_limit * 4, 10000)  # Use existing 10000 char limit as max
+
+        # Chunk inputs and track mapping
+        chunked_input = []
+        chunk_mapping = []
+        current_idx = 0
+
+        for text in input:
+            # Split text into chunks if it exceeds the limit
+            if len(text) <= char_limit:
+                chunks = [text]
+            else:
+                chunks = [
+                    text[i : i + char_limit] for i in range(0, len(text), char_limit)
+                ]
+
+            # Track which chunks belong to which original input using a chunk map
+            chunk_indices = list(range(current_idx, current_idx + len(chunks)))
+            chunk_mapping.append(chunk_indices)
+            current_idx += len(chunks)
+
+            chunked_input.extend(chunks)
+
+        ret_value = []
+        batch_split = [
+            chunked_input[i : i + batch_size]
+            for i in range(0, len(chunked_input), batch_size)
+        ]
+
+        for batch in batch_split:
+            if model.provider == Provider.OPENAI.value:
+                ret_value.extend(_generate_embeddings_openai(model, batch))
+            elif model.provider == Provider.BEDROCK.value:
+                ret_value.extend(_generate_embeddings_bedrock(model, batch, task))
+            elif model.provider == Provider.SAGEMAKER.value:
+                ret_value.extend(_generate_embeddings_sagemaker(model, batch))
+            else:
+                raise CommonError(f"Unknown provider: {model.provider}")
 
-    return ret_value
+        # Combine embeddings from the same original input
+        final_embeddings = []
+        for chunks_idx in chunk_mapping:
+            if len(chunks_idx) == 1:
+                final_embeddings.append(ret_value[chunks_idx[0]])
+            else:
+                # Average the embeddings
+                chunk_embeddings = [ret_value[idx] for idx in chunks_idx]
+                avg_embedding = [
+                    sum(values) / len(values) for values in zip(*chunk_embeddings)
+                ]
+                final_embeddings.append(avg_embedding)
+
+        return final_embeddings
+    except Exception as e:
+        logger.error(f"Error in generate_embeddings: {str(e)}")
+        raise CommonError(f"Failed to generate embeddings: {str(e)}")
 
 
 def get_embeddings_models():