Add support for cohere/titan embedding models

bin/config.ts

@@ -52,6 +52,22 @@ export function getConfig(): SystemConfig {
           provider: "bedrock",
           name: "amazon.titan-embed-text-v1",
           dimensions: 1536,
+        },
+        //Support for inputImage is not yet implemented for amazon.titan-embed-image-v1
+        {
+          provider: "bedrock",
+          name: "amazon.titan-embed-image-v1",
+          dimensions: 1024,
+        },
+        {
+          provider: "bedrock",
+          name: "cohere.embed-english-v3",
+          dimensions: 1024,
+        },
+        {
+          provider: "bedrock",
+          name: "cohere.embed-multilingual-v3",
+          dimensions: 1024,
           default: true,
         },
         {
@@ -71,4 +87,4 @@ export function getConfig(): SystemConfig {
   };
 }
 
-export const config: SystemConfig = getConfig();
+export const config: SystemConfig = getConfig();

cli/magic-config.ts

@@ -33,6 +33,22 @@ const embeddingModels = [
     name: "amazon.titan-embed-text-v1",
     dimensions: 1536,
   },
+  //Support for inputImage is not yet implemented for amazon.titan-embed-image-v1
+  {
+    provider: "bedrock",
+    name: "amazon.titan-embed-image-v1",
+    dimensions: 1024,
+  },
+  {
+    provider: "bedrock",
+    name: "cohere.embed-english-v3",
+    dimensions: 1024,
+  },
+  {
+    provider: "bedrock",
+    name: "cohere.embed-multilingual-v3",
+    dimensions: 1024,
+  },
   {
     provider: "openai",
     name: "text-embedding-ada-002",
@@ -322,15 +338,22 @@ async function processCreateOptions(options: any): Promise<void> {
   }
   const modelsPrompts = [
     {
-      type: "select",
+      type: "select", 
       name: "defaultEmbedding",
-      message: "Which is the default embedding model",
-      choices: embeddingModels.map((m) => ({ name: m.name, value: m })),
-      initial: options.defaultEmbedding || undefined,
-      skip(): boolean {
-        return !(this as any).state.answers.enableRag;
+      message: "Select a default embedding model",
+      choices: embeddingModels.map(m => ({name: m.name, value: m})),
+      initial: options.defaultEmbedding,
+      validate(value: string) {
+        if ((this as any).state.answers.enableRag) {
+          return value ? true : 'Select a default embedding model'; 
+        }
+
+        return true;
       },
-    },
+      skip() {
+        return !answers.enableRag; 
+      }
+    }
   ];
   const models: any = await enquirer.prompt(modelsPrompts);
 
@@ -410,4 +433,4 @@ async function processCreateOptions(options: any): Promise<void> {
   ).create
     ? createConfig(config)
     : console.log("Skipping");
-}
+}

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

@@ -1,20 +1,22 @@
 import genai_core.types
 import genai_core.parameters
 import genai_core.embeddings
-from typing import List
+from typing import List, Optional
 from pydantic import BaseModel
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler.appsync import Router
+from genai_core.types import CommonError, Provider, Task
 
 tracer = Tracer()
 router = Router()
 logger = Logger()
 
 
 class EmbeddingsRequest(BaseModel):
-    provider: str
+    provider: Provider
     model: str
     passages: List[str]
+    task: Optional[Task] = Task.STORE
 
 
 @router.resolver(field_name="listEmbeddingModels")
@@ -34,13 +36,13 @@ def embeddings(input: dict):
     )
 
     if selected_model is None:
-        raise genai_core.types.CommonError("Model not found")
+        raise CommonError("Model not found")
 
     ret_value = genai_core.embeddings.generate_embeddings(
-        selected_model, request.passages
+        selected_model, request.passages, request.task
     )
 
     return [
         {"vector": v, "passage": request.passages[idx]}
         for idx, v in enumerate(ret_value)
-    ]
+    ]

lib/chatbot-api/schema/schema.graphql

@@ -41,6 +41,7 @@ input CalculateEmbeddingsInput {
   provider: String!
   model: String!
   passages: [String]!
+  task: String!
 }
 
 type CrawlerProperties @aws_cognito_user_pools {
@@ -362,4 +363,4 @@ schema {
   query: Query
   mutation: Mutation
   subscription: Subscription
-}
+}

lib/shared/layers/python-sdk/python/genai_core/aurora/query.py

@@ -7,6 +7,7 @@
 from genai_core.aurora.connection import AuroraConnection
 from genai_core.aurora.utils import convert_types
 from aws_lambda_powertools import Logger
+from genai_core.types import CommonError, Task
 
 logger = Logger()
 
@@ -35,17 +36,17 @@ def query_workspace_aurora(
     )
 
     if selected_model is None:
-        raise genai_core.types.CommonError("Embeddings model not found")
+        raise CommonError("Embeddings model not found")
 
     cross_encoder_model = genai_core.cross_encoder.get_cross_encoder_model(
         cross_encoder_model_provider, cross_encoder_model_name
     )
 
     if cross_encoder_model is None:
-        raise genai_core.types.CommonError("Cross encoder model not found")
+        raise CommonError("Cross encoder model not found")
 
     query_embeddings = genai_core.embeddings.generate_embeddings(
-        selected_model, [query]
+        selected_model, [query], Task.RETRIEVE.value
     )[0]
 
     language_name, detected_languages = genai_core.utils.comprehend.get_query_language(
@@ -291,7 +292,7 @@ def _convert_records(source: str, records: List[dict]):
             converted["keyword_search_score"] = record[12]
             converted["vector_search_score"] = None
         else:
-            raise genai_core.types.CommonError("Unknown source")
+            raise CommonError("Unknown source")
 
         converted_records.append(converted)
 

lib/shared/layers/python-sdk/python/genai_core/chunks.py

@@ -1,11 +1,11 @@
 import os
 import uuid
 import boto3
-import genai_core.types
 import genai_core.documents
 import genai_core.embeddings
 import genai_core.aurora.chunks
 import genai_core.opensearch.chunks
+from genai_core.types import CommonError,Task
 from typing import List, Optional
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
@@ -37,10 +37,10 @@ def add_chunks(
     )
 
     if embeddings_model is None:
-        raise genai_core.types.CommonError("Embeddings model not found")
+        raise CommonError("Embeddings model not found")
 
     chunk_embeddings = genai_core.embeddings.generate_embeddings(
-        embeddings_model, chunks
+        embeddings_model, chunks, Task.STORE.value
     )
     chunk_ids = [uuid.uuid4() for _ in chunks]
 
@@ -77,7 +77,7 @@ def add_chunks(
             replace=replace,
         )
     else:
-        raise genai_core.types.CommonError("Engine not supported")
+        raise CommonError("Engine not supported")
 
     added_vectors = result["added_vectors"]
     genai_core.documents.set_document_vectors(
@@ -100,7 +100,7 @@ def split_content(workspace: dict, content: str):
 
         return text_data
 
-    raise genai_core.types.CommonError("Chunking strategy not supported")
+    raise CommonError("Chunking strategy not supported")
 
 
 def store_chunks_on_s3(

lib/shared/layers/python-sdk/python/genai_core/embeddings.py

@@ -4,7 +4,7 @@
 import random
 import botocore
 import numpy as np
-import genai_core.types
+from genai_core.types import EmbeddingsModel, CommonError, Provider, Task
 import genai_core.clients
 import genai_core.parameters
 from typing import List, Optional
@@ -13,22 +13,22 @@
 
 
 def generate_embeddings(
-    model: genai_core.types.EmbeddingsModel, input: List[str], batch_size: int = 50
+    model: EmbeddingsModel, input: List[str], task: str = "store", batch_size: int = 50
 ) -> List[List[float]]:
     input = list(map(lambda x: x[:10000], 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 == "openai":
+        if model.provider == Provider.OPENAI.value:
             ret_value.extend(_generate_embeddings_openai(model, batch))
-        elif model.provider == "bedrock":
-            ret_value.extend(_generate_embeddings_bedrock(model, batch))
-        elif model.provider == "sagemaker":
+        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 genai_core.types.CommonError(f"Unknown provider")
+            raise CommonError(f"Unknown provider: {model.provider}")
 
     return ret_value
 
@@ -44,25 +44,25 @@ def get_embeddings_models():
 
 
 def get_embeddings_model(
-    provider: str, name: str
-) -> Optional[genai_core.types.EmbeddingsModel]:
+    provider: Provider, name: str
+) -> Optional[EmbeddingsModel]:
     config = genai_core.parameters.get_config()
     models = config["rag"]["embeddingsModels"]
 
     for model in models:
-        if model["provider"] == provider and model["name"] == name:
-            return genai_core.types.EmbeddingsModel(**model)
+        if model["provider"] == provider.value and model["name"] == name:
+            return EmbeddingsModel(**model)
 
     return None
 
 
 def _generate_embeddings_openai(
-    model: genai_core.types.EmbeddingsModel, input: List[str]
+    model: EmbeddingsModel, input: List[str]
 ):
     openai = genai_core.clients.get_openai_client()
 
     if not openai:
-        raise genai_core.types.CommonError(
+        raise CommonError(
             "OpenAI API is not available. Please set OPENAI_API_KEY."
         )
 
@@ -73,13 +73,23 @@ def _generate_embeddings_openai(
 
 
 def _generate_embeddings_bedrock(
-    model: genai_core.types.EmbeddingsModel, input: List[str]
+    model: EmbeddingsModel, input: List[str], task: Task
 ):
     bedrock = genai_core.clients.get_bedrock_client()
 
     if not bedrock:
-        raise genai_core.types.CommonError("Bedrock is not enabled.")
+        raise CommonError("Bedrock is not enabled.")
 
+    model_provider = model.name.split(".")[0]
+    if model_provider == Provider.AMAZON.value:
+        return _generate_embeddings_amazon(model, input, bedrock)
+    elif model_provider == Provider.COHERE.value:
+        return _generate_embeddings_cohere(model, input, task, bedrock)
+    else:
+        raise CommonError(f"Unknown embeddings provider \"{model_provider}\"")
+
+
+def _generate_embeddings_amazon(model: EmbeddingsModel, input: List[str], bedrock):
     ret_value = []
     for value in input:
         body = json.dumps({"inputText": value})
@@ -97,12 +107,25 @@ def _generate_embeddings_bedrock(
     ret_value = np.array(ret_value)
     ret_value = ret_value / np.linalg.norm(ret_value, axis=1, keepdims=True)
     ret_value = ret_value.tolist()
-
     return ret_value
 
 
+def _generate_embeddings_cohere(model: EmbeddingsModel, input: List[str], task: Task, bedrock):
+    input_type = "search_query" if task == Task.RETRIEVE else "search_document"
+    body = json.dumps({"texts": input, "input_type": input_type})
+    response = bedrock.invoke_model(
+        body=body,
+        modelId=model.name,
+        accept="application/json",
+        contentType="application/json",
+    )
+    response_body = json.loads(response.get("body").read())
+    embeddings = response_body.get("embeddings")
+
+    return embeddings
+
 def _generate_embeddings_sagemaker(
-    model: genai_core.types.EmbeddingsModel, input: List[str]
+    model: EmbeddingsModel, input: List[str]
 ):
     client = genai_core.clients.get_sagemaker_client()
 
@@ -132,4 +155,4 @@ def _generate_embeddings_sagemaker(
                 continue
             else:
                 # If the exception was due to another reason, raise it.
-                raise error
+                raise error