feat: intention management

source/infrastructure/lib/api/api-stack.ts

@@ -449,6 +449,9 @@ export class ApiConstruct extends Construct {
         handler: "intention.lambda_handler",
         environment: {
           INTENTION_TABLE_NAME: props.chatStackOutputs.intentionTableName,
+          INDEX_TABLE_NAME: props.sharedConstructOutputs.indexTable.tableName,
+          CHATBOT_TABLE_NAME: props.sharedConstructOutputs.chatbotTable.tableName,
+          MODEL_TABLE_NAME: props.sharedConstructOutputs.modelTable.tableName,
           S3_BUCKET: s3Bucket.bucketName,
         },
         layers: [apiLambdaOnlineSourceLayer],
@@ -529,6 +532,8 @@ export class ApiConstruct extends Construct {
         proxy: true,
       });
       const apiResourceIntentionManagement = api.root.addResource("intention");
+      const indexScan = apiResourceIntentionManagement.addResource("index-used-scan")
+      indexScan.addMethod("POST", lambdaIntentionIntegration, this.genMethodOption(api, auth, null));
       // apiResourceIntentionManagement.addMethod("DELETE", lambdaIntentionIntegration, this.genMethodOption(api, auth, null));
       const presignedUrl = apiResourceIntentionManagement.addResource("execution-presigned-url");
       presignedUrl.addMethod("POST", lambdaIntentionIntegration, {...

source/lambda/intention/constant.py

@@ -1,11 +1,42 @@
+from enum import Enum, unique
+
+
 DEFAULT_MAX_ITEMS = 50
 DEFAULT_SIZE = 50
 HTTPS_PORT_NUMBER = "443"
 DEFAULT_CONTENT_TYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
 ROOT_RESOURCE = "/intention"
 PRESIGNED_URL_RESOURCE = f"{ROOT_RESOURCE}/execution-presigned-url"
 EXECUTION_RESOURCE = f"{ROOT_RESOURCE}/executions"
+INDEX_USED_SCAN_RESOURCE = f"{ROOT_RESOURCE}/index-used-scan"
 DOWNLOAD_RESOURCE = f"{ROOT_RESOURCE}/download-template"
 SECRET_NAME = "opensearch-master-user"
 AOS_INDEX = "aics_intention_index"
-BULK_SIZE = 100000000
+BULK_SIZE = 100000000
+
+@unique
+class IndexType(Enum):
+    QD = "qd"
+    QQ = "qq"
+    INTENTION = "intention"
+
+@unique
+class KBType(Enum):
+    AOS = "aos"
+
+
+@unique
+class Status(Enum):
+    ACTIVE = "active"
+    INACTIVE = "inactive"
+
+ModelDimensionMap = {
+    "amazon.titan-embed-text-v1": 1536,
+    "cohere.embed-english-v3": 1024,
+    "amazon.titan-embed-text-v2:0": 1024
+}
+
+@unique
+class ModelType(Enum):
+    EMBEDDING = "embedding_and_rerank"
+    LLM = "llm"

source/lambda/intention/ddb_utils.py

@@ -0,0 +1,124 @@
+from datetime import datetime, timezone
+from typing import Any
+
+from constant import IndexType, KBType, ModelType, Status
+from embeddings import get_embedding_info
+
+def create_item_if_not_exist(ddb_table, item_key: dict, body: str):
+    response = ddb_table.get_item(Key=item_key)
+    item = response.get("Item")
+    if not item:
+        ddb_table.put_item(Item=body)
+        return False, item
+    return True, item
+
+def check_item_exist(ddb_table, item_key: dict):
+    response = ddb_table.get_item(Key=item_key)
+    item = response.get("Item")
+    return True, item if item else False, None
+
+
+def initiate_model(
+    model_table, group_name, model_id, embedding_endpoint, create_time=None
+):
+    if not create_time:
+        create_time = str(datetime.now(timezone.utc))
+    embedding_info = dict(get_embedding_info(embedding_endpoint))
+    embedding_info["ModelEndpoint"] = embedding_endpoint
+    create_item_if_not_exist(
+        model_table,
+        {"groupName": group_name, "modelId": model_id},
+        {
+            "groupName": group_name,
+            "modelId": model_id,
+            "modelType": ModelType.EMBEDDING.value,
+            "parameter": embedding_info,
+            "createTime": create_time,
+            "updateTime": create_time,
+            "status": Status.ACTIVE.value,
+        },
+    )
+    return embedding_info["ModelType"]
+
+
+def initiate_index(
+    index_table,
+    group_name,
+    index_id,
+    model_id,
+    index_type,
+    tag,
+    create_time=None,
+    description="",
+):
+    if not create_time:
+        create_time = str(datetime.now(timezone.utc))
+
+    db_body = {
+        "groupName": group_name,
+        "indexId": index_id,
+        "indexType": index_type,
+        "kbType": KBType.AOS.value,
+        "modelIds": {"embedding": model_id},
+        "tag": tag,
+        "createTime": create_time,
+        "status": Status.ACTIVE.value,
+    }
+
+    create_item_if_not_exist(
+        index_table, {"groupName": group_name, "indexId": index_id}, db_body
+    )
+
+
+def initiate_chatbot(
+    chatbot_table, group_name, chatbot_id, index_id, index_type, tag, create_time=None
+):
+    if not create_time:
+        create_time = str(datetime.now(timezone.utc))
+    is_existed, item = create_item_if_not_exist(
+        chatbot_table,
+        {"groupName": group_name, "chatbotId": chatbot_id},
+        {
+            "groupName": group_name,
+            "chatbotId": chatbot_id,
+            "languages": ["zh"],
+            "indexIds": {index_type: {"count": 1, "value": {tag: index_id}}},
+            "createTime": create_time,
+            "updateTime": create_time,
+            "status": Status.ACTIVE.value,
+        },
+    )
+
+    if is_existed:
+        index_id_dict = item.get("indexIds", {})
+        append_index = True
+        if index_type in index_id_dict:
+            # Append it with the same index type
+            for key in index_id_dict[index_type]["value"].keys():
+                if key == tag:
+                    append_index = False
+                    break
+
+            if append_index:
+                item["indexIds"][index_type]["value"][tag] = index_id
+                item["indexIds"][index_type]["count"] = len(
+                    item["indexIds"][index_type]["value"]
+                )
+                chatbot_table.put_item(Item=item)
+        else:
+            # Add a new index type
+            item["indexIds"][index_type] = {"count": 1, "value": {tag: index_id}}
+            chatbot_table.put_item(Item=item)
+
+
+def is_chatbot_existed(ddb_table, group_name: str, chatbot_id: str):
+    response = ddb_table.get_item(
+        Key={
+            "groupName": group_name,
+            "chatbotId": chatbot_id,
+        },
+    )
+    item = response.get("Item")
+    if not item:
+        return False
+    return True

source/lambda/intention/embeddings.py

@@ -0,0 +1,56 @@
+def get_embedding_info(embedding_endpoint_name):
+    """
+    Get the embedding info from the endpoint name
+    """
+    # Get the embedding info from the endpoint name
+    if "bge-large-zh" in embedding_endpoint_name:
+        embeddings_model_provider = "BAAI"
+        embeddings_model_name = "bge-large-zh-v1-5"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "bge-large-zh"
+
+    elif "bge-large-en" in embedding_endpoint_name:
+        embeddings_model_provider = "BAAI"
+        embeddings_model_name = "bge-large-en-v1-5"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "bge-large-en"
+
+    elif "bge-m3" in embedding_endpoint_name:
+        embeddings_model_provider = "BAAI"
+        embeddings_model_name = "bge-m3"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "m3"
+
+    elif "embedding" in embedding_endpoint_name:
+        embeddings_model_provider = "Netease"
+        embeddings_model_name = "bce_embedding_model.tar.gz"
+        embeddings_model_dimensions = 768
+        embeddings_model_type = "bce"
+
+    elif "cohere" in embedding_endpoint_name:
+        embeddings_model_provider = "Cohere"
+        embeddings_model_name = "cohere.embed-english-v3"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "bedrock"
+    elif "titan-embed-text-v1" in embedding_endpoint_name:
+        embeddings_model_provider = "Titan"
+        embeddings_model_name = "amazon.titan-embed-text-v1"
+        embeddings_model_dimensions = 1536
+        embeddings_model_type = "bedrock"
+    elif "titan-embed-text-v2" in embedding_endpoint_name:
+        embeddings_model_provider = "Titan"
+        embeddings_model_name = "amazon.titan-embed-text-v2:0"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "bedrock"
+    else:
+        embeddings_model_provider = "Not Found"
+        embeddings_model_name = "Not Found"
+        embeddings_model_dimensions = 1024
+        embeddings_model_type = "Not Found"
+
+    return {
+        "ModelProvider": embeddings_model_provider,
+        "ModelName": embeddings_model_name,
+        "ModelDimension": embeddings_model_dimensions,
+        "ModelType": embeddings_model_type,
+    }