diff --git a/bin/config.ts b/bin/config.ts index ca8b5d021..72496dd29 100644 --- a/bin/config.ts +++ b/bin/config.ts @@ -39,6 +39,9 @@ export function getConfig(): SystemConfig { createIndex: false, enterprise: false, }, + knowledgeBase: { + enabled: false, + }, }, embeddingsModels: [ { diff --git a/cli/magic-config.ts b/cli/magic-config.ts index 4584199e1..07151c2e2 100644 --- a/cli/magic-config.ts +++ b/cli/magic-config.ts @@ -156,6 +156,8 @@ const embeddingModels = [ ? config.llms?.sagemaker.length > 0 : false; options.huggingfaceApiSecretArn = config.llms?.huggingfaceApiSecretArn; + options.enableSagemakerModelsSchedule = + config.llms?.sagemakerSchedule?.enabled; options.enableSagemakerModelsSchedule = config.llms?.sagemakerSchedule?.enabled; options.timezonePicker = config.llms?.sagemakerSchedule?.timezonePicker; @@ -194,6 +196,7 @@ const embeddingModels = [ (m) => m.default )[0].name; options.kendraExternal = config.rag.engines.kendra.external; + options.kbExternal = config.rag.engines.knowledgeBase?.external ?? []; options.kendraEnterprise = config.rag.engines.kendra.enterprise; // Advanced settings @@ -586,6 +589,7 @@ async function processCreateOptions(options: any): Promise { { message: "Aurora", name: "aurora" }, { message: "OpenSearch", name: "opensearch" }, { message: "Kendra (managed)", name: "kendra" }, + { message: "Bedrock KnowldgeBase", name: "knowledgeBase" }, ], validate(choices: any) { return (this as any).skipped || choices.length > 0 @@ -694,6 +698,82 @@ async function processCreateOptions(options: any): Promise { }); newKendra = kendraInstance.newKendra; } + + // Knowledge Bases + let newKB = + answers.enableRag && answers.ragsToEnable.includes("knowledgeBase"); + const kbExternal: any[] = []; + const existingKBIndices = Array.from(options.kbExternal || []); + while (newKB === true) { + const existingIndex: any = existingKBIndices.pop(); + const kbQ = [ + { + type: "input", + name: "name", + message: "KnowledgeBase source name", + validate(v: string) { + return RegExp(/^\w[\w-_]*\w$/).test(v); + }, + initial: existingIndex?.name, + }, + { + type: "autocomplete", + limit: 8, + name: "region", + choices: ["us-east-1", "us-west-2"], + message: `Region of the Bedrock Knowledge Base index${ + existingIndex?.region ? " (" + existingIndex?.region + ")" : "" + }`, + initial: ["us-east-1", "us-west-2"].indexOf(existingIndex?.region), + }, + { + type: "input", + name: "roleArn", + message: + "Cross account role Arn to assume to call the Bedrock KnowledgeBase, leave empty if not needed", + validate: (v: string) => { + const valid = iamRoleRegExp.test(v); + return v.length === 0 || valid; + }, + initial: existingIndex?.roleArn ?? "", + }, + { + type: "input", + name: "knowledgeBaseId", + message: "Bedrock KnowledgeBase ID", + validate(v: string) { + return /[A-Z0-9]{10}/.test(v); + }, + initial: existingIndex?.knowledgeBaseId, + }, + { + type: "confirm", + name: "enabled", + message: "Enable this knowledge base", + initial: existingIndex?.enabled ?? true, + }, + { + type: "confirm", + name: "newKB", + message: "Do you want to add another Bedrock KnowledgeBase source", + initial: false, + }, + ]; + const kbInstance: any = await enquirer.prompt(kbQ); + const ext = (({ enabled, name, roleArn, knowledgeBaseId, region }) => ({ + enabled, + name, + roleArn, + knowledgeBaseId, + region, + }))(kbInstance); + if (ext.roleArn === "") ext.roleArn = undefined; + kbExternal.push({ + ...ext, + }); + newKB = kbInstance.newKB; + } + const modelsPrompts = [ { type: "select", @@ -1078,6 +1158,10 @@ async function processCreateOptions(options: any): Promise { external: [{}], enterprise: false, }, + knowledgeBase: { + enabled: false, + external: [{}], + }, }, embeddingsModels: [{}], crossEncoderModels: [{}], @@ -1107,6 +1191,9 @@ async function processCreateOptions(options: any): Promise { config.rag.engines.kendra.createIndex || kendraExternal.length > 0; config.rag.engines.kendra.external = [...kendraExternal]; config.rag.engines.kendra.enterprise = answers.kendraEnterprise; + config.rag.engines.knowledgeBase.enabled = + config.rag.engines.knowledgeBase.external.length > 0; + config.rag.engines.knowledgeBase.external = [...kbExternal]; console.log("\n✨ This is the chosen configuration:\n"); console.log(JSON.stringify(config, undefined, 2)); diff --git a/lib/chatbot-api/functions/api-handler/index.py b/lib/chatbot-api/functions/api-handler/index.py index 8a7884e6e..e9c67b0d6 100644 --- a/lib/chatbot-api/functions/api-handler/index.py +++ b/lib/chatbot-api/functions/api-handler/index.py @@ -15,6 +15,7 @@ from routes.documents import router as documents_router from routes.kendra import router as kendra_router from routes.user_feedback import router as user_feedback_router +from routes.bedrock_kb import router as bedrock_kb_router tracer = Tracer() logger = Logger() @@ -32,6 +33,7 @@ app.include_router(documents_router) app.include_router(kendra_router) app.include_router(user_feedback_router) +app.include_router(bedrock_kb_router) @logger.inject_lambda_context( diff --git a/lib/chatbot-api/functions/api-handler/routes/bedrock_kb.py b/lib/chatbot-api/functions/api-handler/routes/bedrock_kb.py new file mode 100644 index 000000000..9d74893d7 --- /dev/null +++ b/lib/chatbot-api/functions/api-handler/routes/bedrock_kb.py @@ -0,0 +1,21 @@ +import genai_core.parameters +import genai_core.bedrock_kb +from pydantic import BaseModel +from aws_lambda_powertools import Logger, Tracer +from aws_lambda_powertools.event_handler.appsync import Router + +tracer = Tracer() +router = Router() +logger = Logger() + + +class KendraDataSynchRequest(BaseModel): + workspaceId: str + + +@router.resolver(field_name="listBedrockKnowledgeBases") +@tracer.capture_method +def list_bedrock_kbs(): + indexes = genai_core.bedrock_kb.list_bedrock_kbs() + + return indexes diff --git a/lib/chatbot-api/functions/api-handler/routes/rag.py b/lib/chatbot-api/functions/api-handler/routes/rag.py index 72a239bcf..a7eb4da6a 100644 --- a/lib/chatbot-api/functions/api-handler/routes/rag.py +++ b/lib/chatbot-api/functions/api-handler/routes/rag.py @@ -35,6 +35,11 @@ def engines(): "name": "Amazon Kendra", "enabled": engines.get("kendra", {}).get("enabled", False) == True, }, + { + "id": "bedrock_kb", + "name": "Bedrock Knowledge Bases", + "enabled": engines.get("knowledgeBase", {}).get("enabled", False) == True, + }, ] return ret_value diff --git a/lib/chatbot-api/functions/api-handler/routes/workspaces.py b/lib/chatbot-api/functions/api-handler/routes/workspaces.py index 1c71a24fe..937567730 100644 --- a/lib/chatbot-api/functions/api-handler/routes/workspaces.py +++ b/lib/chatbot-api/functions/api-handler/routes/workspaces.py @@ -1,6 +1,7 @@ import re import genai_core.types import genai_core.kendra +import genai_core.bedrock_kb import genai_core.parameters import genai_core.workspaces from pydantic import BaseModel @@ -55,6 +56,13 @@ class CreateWorkspaceKendraRequest(BaseModel): useAllData: bool +class CreateWorkspaceBedrockKBRequest(BaseModel): + kind: str + name: str + knowledgeBaseId: str + hybridSearch: bool + + @router.resolver(field_name="listWorkspaces") @tracer.capture_method def list_workspaces(): @@ -115,6 +123,16 @@ def create_kendra_workspace(input: dict): return ret_value +@router.resolver(field_name="createBedrockKBWorkspace") +@tracer.capture_method +def create_bedrock_kb_workspace(input: dict): + config = genai_core.parameters.get_config() + + request = CreateWorkspaceBedrockKBRequest(**input) + ret_value = _create_workspace_bedrock_kb(request, config) + return ret_value + + def _create_workspace_aurora(request: CreateWorkspaceAuroraRequest, config: dict): workspace_name = request.name.strip() embedding_models = config["rag"]["embeddingsModels"] @@ -291,6 +309,39 @@ def _create_workspace_kendra(request: CreateWorkspaceKendraRequest, config: dict ) +def _create_workspace_bedrock_kb( + request: CreateWorkspaceBedrockKBRequest, config: dict +): + workspace_name = request.name.strip() + kbs = genai_core.bedrock_kb.list_bedrock_kbs() + + workspace_name_match = name_regex.match(workspace_name) + workspace_name_is_match = bool(workspace_name_match) + if ( + len(workspace_name) == 0 + or len(workspace_name) > 100 + or not workspace_name_is_match + ): + raise genai_core.types.CommonError("Invalid workspace name") + + knowledge_base = None + for current in kbs: + if current["id"] == request.knowledgeBaseId: + knowledge_base = current + break + + if knowledge_base is None: + raise genai_core.types.CommonError("Knowledge Base id not found") + + return _convert_workspace( + genai_core.workspaces.create_workspace_bedrock_kb( + workspace_name=workspace_name, + knowledge_base=knowledge_base, + hybrid_search=request.hybridSearch, + ) + ) + + def _convert_workspace(workspace: dict): kendra_index_external = workspace.get("kendra_index_external") @@ -320,6 +371,8 @@ def _convert_workspace(workspace: dict): "kendraIndexId": workspace.get("kendra_index_id"), "kendraIndexExternal": kendra_index_external, "kendraUseAllData": workspace.get("kendra_use_all_data", kendra_index_external), + "knowledgeBaseId": workspace.get("knowledge_base_id"), + "knowledgeBaseExternal": workspace.get("knowledge_base_external"), "createdAt": workspace.get("created_at"), "updatedAt": workspace.get("updated_at"), } diff --git a/lib/chatbot-api/rest-api.ts b/lib/chatbot-api/rest-api.ts index 758813876..634d72bf6 100644 --- a/lib/chatbot-api/rest-api.ts +++ b/lib/chatbot-api/rest-api.ts @@ -188,6 +188,33 @@ export class ApiResolvers extends Construct { ); } + if (props.config.rag.engines.knowledgeBase.enabled) { + for (const item of props.config.rag.engines.knowledgeBase.external || + []) { + if (item.roleArn) { + apiHandler.addToRolePolicy( + new iam.PolicyStatement({ + actions: ["sts:AssumeRole"], + resources: [item.roleArn], + }) + ); + } else { + apiHandler.addToRolePolicy( + new iam.PolicyStatement({ + actions: ["bedrock:Retrieve"], + resources: [ + `arn:${cdk.Aws.PARTITION}:bedrock:${ + item.region ?? cdk.Aws.REGION + }:${cdk.Aws.ACCOUNT_ID}:knowledge-base/${ + item.knowledgeBaseId + }`, + ], + }) + ); + } + } + } + for (const item of props.config.rag.engines.kendra.external ?? []) { if (item.roleArn) { apiHandler.addToRolePolicy( diff --git a/lib/chatbot-api/schema/schema.graphql b/lib/chatbot-api/schema/schema.graphql index f5a8f42cc..4601bfdb3 100644 --- a/lib/chatbot-api/schema/schema.graphql +++ b/lib/chatbot-api/schema/schema.graphql @@ -23,6 +23,13 @@ input CreateWorkspaceKendraInput { useAllData: Boolean! } +input CreateWorkspaceBedrockKBInput { + name: String! + kind: String! + knowledgeBaseId: String! + hybridSearch: Boolean! +} + input CreateWorkspaceOpenSearchInput { name: String! kind: String! @@ -150,6 +157,12 @@ type KendraIndex @aws_cognito_user_pools { external: Boolean! } +type BedrockKB @aws_cognito_user_pools { + id: String! + name: String! + external: Boolean! +} + input ListDocumentsInput { workspaceId: String! documentType: String! @@ -302,6 +315,8 @@ type Workspace @aws_cognito_user_pools { kendraIndexId: String kendraIndexExternal: Boolean kendraUseAllData: Boolean + knowledgeBaseId: String + knowledgeBaseExternal: Boolean createdAt: AWSDateTime! updatedAt: AWSDateTime! } @@ -315,6 +330,8 @@ type Channel @aws_iam @aws_cognito_user_pools { type Mutation { createKendraWorkspace(input: CreateWorkspaceKendraInput!): Workspace! @aws_cognito_user_pools + createBedrockKBWorkspace(input: CreateWorkspaceBedrockKBInput!): Workspace! + @aws_cognito_user_pools createOpenSearchWorkspace(input: CreateWorkspaceOpenSearchInput!): Workspace! @aws_cognito_user_pools createAuroraWorkspace(input: CreateWorkspaceAuroraInput!): Workspace! @@ -358,6 +375,7 @@ type Query { @aws_cognito_user_pools getSession(id: String!): Session @aws_cognito_user_pools listKendraIndexes: [KendraIndex!]! @aws_cognito_user_pools + listBedrockKnowledgeBases: [BedrockKB!]! @aws_cognito_user_pools isKendraDataSynching(workspaceId: String!): Boolean @aws_cognito_user_pools listDocuments(input: ListDocumentsInput!): DocumentsResult! @aws_cognito_user_pools diff --git a/lib/model-interfaces/langchain/functions/request-handler/adapters/shared/meta/llama3_instruct.py b/lib/model-interfaces/langchain/functions/request-handler/adapters/shared/meta/llama3_instruct.py index f400b1119..3acf38943 100644 --- a/lib/model-interfaces/langchain/functions/request-handler/adapters/shared/meta/llama3_instruct.py +++ b/lib/model-interfaces/langchain/functions/request-handler/adapters/shared/meta/llama3_instruct.py @@ -14,7 +14,7 @@ You are an helpful assistant that provides concise answers to user questions with as little sentences as possible and at maximum 3 sentences. You do not repeat yourself. You avoid bulleted list or emojis.{EOD}{{chat_history}}{USER_HEADER} -{{input}}{EOD}{ASSISTANT_HEADER}""" +Context: {{input}}{EOD}{ASSISTANT_HEADER}""" Llama3QAPrompt = f"""{BEGIN_OF_TEXT}{SYSTEM_HEADER} diff --git a/lib/model-interfaces/langchain/index.ts b/lib/model-interfaces/langchain/index.ts index 0057283e6..e213606a9 100644 --- a/lib/model-interfaces/langchain/index.ts +++ b/lib/model-interfaces/langchain/index.ts @@ -182,6 +182,31 @@ export class LangChainInterface extends Construct { } } + if (props.config.rag.engines.knowledgeBase.enabled) { + for (const item of props.config.rag.engines.knowledgeBase.external || + []) { + if (item.roleArn) { + requestHandler.addToRolePolicy( + new iam.PolicyStatement({ + actions: ["sts:AssumeRole"], + resources: [item.roleArn], + }) + ); + } else { + requestHandler.addToRolePolicy( + new iam.PolicyStatement({ + actions: ["bedrock:Retrieve"], + resources: [ + `arn:${cdk.Aws.PARTITION}:bedrock:${ + item.region ?? cdk.Aws.REGION + }:${cdk.Aws.ACCOUNT_ID}:knowledge-base/${item.knowledgeBaseId}`, + ], + }) + ); + } + } + } + props.sessionsTable.grantReadWriteData(requestHandler); props.messagesTopic.grantPublish(requestHandler); props.shared.apiKeysSecret.grantRead(requestHandler); diff --git a/lib/rag-engines/workspaces/functions/delete-workspace-workflow/delete/index.py b/lib/rag-engines/workspaces/functions/delete-workspace-workflow/delete/index.py index a0a2abb18..59b473d25 100644 --- a/lib/rag-engines/workspaces/functions/delete-workspace-workflow/delete/index.py +++ b/lib/rag-engines/workspaces/functions/delete-workspace-workflow/delete/index.py @@ -3,6 +3,7 @@ import genai_core.aurora.delete import genai_core.opensearch.delete import genai_core.kendra.delete +import genai_core.bedrock_kb.delete from aws_lambda_powertools import Logger from aws_lambda_powertools.utilities.typing import LambdaContext @@ -17,10 +18,12 @@ def lambda_handler(event, context: LambdaContext): raise genai_core.types.CommonError("Workspace not found") if workspace["engine"] == "aurora": - genai_core.aurora.delete.delete_aurora_workspace(workspace) + genai_core.aurora.delete.delete_workspace(workspace) elif workspace["engine"] == "opensearch": - genai_core.opensearch.delete.delete_open_search_workspace(workspace) + genai_core.opensearch.delete.delete_workspace(workspace) elif workspace["engine"] == "kendra": - genai_core.kendra.delete.delete_kendra_workspace(workspace) + genai_core.kendra.delete.delete_workspace(workspace) + elif workspace["engine"] == "bedrock_kb": + genai_core.bedrock_kb.delete.delete_workspace(workspace) else: raise genai_core.types.CommonError("Workspace engine not supported") diff --git a/lib/shared/layers/python-sdk/python/genai_core/aurora/delete.py b/lib/shared/layers/python-sdk/python/genai_core/aurora/delete.py index 6df615e13..732eccbd9 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/aurora/delete.py +++ b/lib/shared/layers/python-sdk/python/genai_core/aurora/delete.py @@ -19,7 +19,7 @@ dynamodb = boto3.resource("dynamodb") -def delete_aurora_workspace(workspace: dict): +def delete_workspace(workspace: dict): workspace_id = workspace["workspace_id"] genai_core.utils.delete_files_with_prefix.delete_files_with_prefix( UPLOAD_BUCKET_NAME, workspace_id diff --git a/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/__init__.py b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/__init__.py new file mode 100644 index 000000000..b81303d67 --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/__init__.py @@ -0,0 +1,2 @@ +from .query import query_workspace_bedrock_kb +from .indexes import list_bedrock_kbs diff --git a/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/client.py b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/client.py new file mode 100644 index 000000000..0d4f2e2f7 --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/client.py @@ -0,0 +1,46 @@ +import os +import boto3 +import genai_core.types +import genai_core.parameters + +sts_client = boto3.client("sts") + + +def get_kb_runtime_client_for_id(knowledge_base_id: str): + + config = genai_core.parameters.get_config() + kb_config = config.get("rag", {}).get("engines", {}).get("knowledgeBase", {}) + external = kb_config.get("external", []) + + for kb in external: + current_id = kb.get("knowledgeBaseId", "") + current_name = kb.get("name", "") + region_name = kb.get("region") + role_arn = kb.get("roleArn") + + if not current_id or not current_name: + continue + + if current_id == knowledge_base_id: + config_data = {"service_name": "bedrock-agent-runtime"} + if region_name: + config_data["region_name"] = region_name + + if role_arn: + assumed_role_object = sts_client.assume_role( + RoleArn=role_arn, + RoleSessionName="AssumedRoleSession", + ) + + credentials = assumed_role_object["Credentials"] + config_data["aws_access_key_id"] = credentials["AccessKeyId"] + config_data["aws_secret_access_key"] = credentials["SecretAccessKey"] + config_data["aws_session_token"] = credentials["SessionToken"] + + client = boto3.client(**config_data) + + return client + + raise genai_core.types.CommonError( + f"Could not find Amazon Bedrock KnowledgeBase ID {knowledge_base_id}" + ) diff --git a/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/delete.py b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/delete.py new file mode 100644 index 000000000..e94310b7e --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/delete.py @@ -0,0 +1,24 @@ +import os +import boto3 +import genai_core.utils.delete_files_with_prefix + + +WORKSPACES_TABLE_NAME = os.environ["WORKSPACES_TABLE_NAME"] +DOCUMENTS_TABLE_NAME = os.environ.get("DOCUMENTS_TABLE_NAME") + + +WORKSPACE_OBJECT_TYPE = "workspace" + +dynamodb = boto3.resource("dynamodb") + + +def delete_workspace(workspace: dict): + workspace_id = workspace["workspace_id"] + + workspaces_table = dynamodb.Table(WORKSPACES_TABLE_NAME) + + response = workspaces_table.delete_item( + Key={"workspace_id": workspace_id, "object_type": WORKSPACE_OBJECT_TYPE}, + ) + + print(f"Delete Item succeeded: {response}") diff --git a/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/indexes.py b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/indexes.py new file mode 100644 index 000000000..466a88bb5 --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/indexes.py @@ -0,0 +1,27 @@ +import os +import genai_core.parameters + + +def list_bedrock_kbs(): + config = genai_core.parameters.get_config() + kb_config = config.get("rag", {}).get("engines", {}).get("knowledgeBase", {}) + external = kb_config.get("external", {}) + + ret_value = [] + + for kb in external: + current_id = kb.get("knowledgeBaseId", "") + current_name = kb.get("name", "") + + if not current_id or not current_name: + continue + + ret_value.append( + { + "id": current_id, + "name": current_name, + "external": True, + } + ) + + return ret_value diff --git a/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/query.py b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/query.py new file mode 100644 index 000000000..4d62e30e5 --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/bedrock_kb/query.py @@ -0,0 +1,75 @@ +import os +import re +import genai_core.types +from typing import List +from .client import get_kb_runtime_client_for_id + +s3_pattern = re.compile(r"(s3-|s3\.)?(.*)\.amazonaws\.com") + + +def query_workspace_bedrock_kb( + workspace_id: str, workspace: dict, query: str, limit: int, full_response: bool +): + knowledge_base_id = workspace.get("knowledge_base_id") + search_type = "HYBRID" if workspace.get("hybrid_search", False) else "SEMANTIC" + + if not knowledge_base_id: + raise genai_core.types.CommonError( + f"Could not find Amazon Bedrock KnowledgeBase ID for workspace {workspace_id}" + ) + + client = get_kb_runtime_client_for_id(knowledge_base_id) + limit = max(1, min(100, limit)) + + result = client.retrieve( + knowledgeBaseId=knowledge_base_id, + retrievalQuery={"text": query}, + retrievalConfiguration={ + "vectorSearchConfiguration": { + "numberOfResults": limit, + "overrideSearchType": search_type, + } + }, + ) + + items = result["retrievalResults"] + items = _convert_records("bedrock_kb", workspace_id, items) + + ret_value = { + "engine": "bedrock_kb", + "items": items, + } + + return ret_value + + +def _convert_records(source: str, workspace_id: str, records: List[dict]): + converted_records = [] + _id = 0 + for record in records: + + path = record.get("location", {}).get("s3Location", {}).get("uri", "") + content = record.get("content", {}).get("text", "") + score = record.get("score", 0) + + converted = { + "chunk_id": str(_id), + "workspace_id": workspace_id, + "document_id": "", + "document_sub_id": None, + "document_type": "object", + "document_sub_type": None, + "path": path, + "language": None, + "title": "", + "content": content, + "content_complement": None, + "metadata": None, + "sources": [source], + "vector_search_score": score, + "score": None, + } + _id += 1 + converted_records.append(converted) + + return converted_records diff --git a/lib/shared/layers/python-sdk/python/genai_core/kendra/delete.py b/lib/shared/layers/python-sdk/python/genai_core/kendra/delete.py index 130b54ea6..158326897 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/kendra/delete.py +++ b/lib/shared/layers/python-sdk/python/genai_core/kendra/delete.py @@ -17,7 +17,7 @@ dynamodb = boto3.resource("dynamodb") -def delete_kendra_workspace(workspace: dict): +def delete_workspace(workspace: dict): workspace_id = workspace["workspace_id"] genai_core.utils.delete_files_with_prefix.delete_files_with_prefix( UPLOAD_BUCKET_NAME, workspace_id diff --git a/lib/shared/layers/python-sdk/python/genai_core/opensearch/delete.py b/lib/shared/layers/python-sdk/python/genai_core/opensearch/delete.py index 954bc3a42..1e4217ee6 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/opensearch/delete.py +++ b/lib/shared/layers/python-sdk/python/genai_core/opensearch/delete.py @@ -18,7 +18,7 @@ dynamodb = boto3.resource("dynamodb") -def delete_open_search_workspace(workspace: dict): +def delete_workspace(workspace: dict): workspace_id = workspace["workspace_id"] index_name = workspace_id.replace("-", "") @@ -66,7 +66,7 @@ def delete_open_search_workspace(workspace: dict): "document_id": item["document_id"], } ) - + print(f"Deleted {len(items_to_delete)} items.") response = workspaces_table.delete_item( diff --git a/lib/shared/layers/python-sdk/python/genai_core/opensearch/query.py b/lib/shared/layers/python-sdk/python/genai_core/opensearch/query.py index 05d359012..96aafbff0 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/opensearch/query.py +++ b/lib/shared/layers/python-sdk/python/genai_core/opensearch/query.py @@ -128,7 +128,14 @@ def query_workspace_open_search( ret_items = list(filter(lambda val: val["score"] > threshold, unique_items))[ :limit ] - if len(ret_items) < limit: + if len(ret_items) < limit and len(unique_items) > len(ret_items): + unique_items = list( + filter( + lambda record: record["chunk_id"] + not in [r["chunk_id"] for r in ret_items], + unique_items, + ) + ) unique_items = sorted( unique_items, key=lambda x: x["vector_search_score"] or -1, reverse=True ) @@ -209,4 +216,4 @@ def keyword_query(client, index_name: str, text: str, size: int = 25): ret_value = response["hits"]["hits"] ret_value = ret_value if ret_value is not None else [] - return ret_value \ No newline at end of file + return ret_value diff --git a/lib/shared/layers/python-sdk/python/genai_core/semantic_search.py b/lib/shared/layers/python-sdk/python/genai_core/semantic_search.py index 495e7f6df..a81fb0094 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/semantic_search.py +++ b/lib/shared/layers/python-sdk/python/genai_core/semantic_search.py @@ -4,6 +4,7 @@ from genai_core.aurora import query_workspace_aurora from genai_core.opensearch import query_workspace_open_search from genai_core.kendra import query_workspace_kendra +from genai_core.bedrock_kb import query_workspace_bedrock_kb def semantic_search( @@ -29,6 +30,10 @@ def semantic_search( return query_workspace_kendra( workspace_id, workspace, query, limit, full_response ) + elif workspace["engine"] == "bedrock_kb": + return query_workspace_bedrock_kb( + workspace_id, workspace, query, limit, full_response + ) raise genai_core.types.CommonError( "Semantic search is not supported for this workspace" diff --git a/lib/shared/layers/python-sdk/python/genai_core/types.py b/lib/shared/layers/python-sdk/python/genai_core/types.py index e0d303b55..aa0f0b7b9 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/types.py +++ b/lib/shared/layers/python-sdk/python/genai_core/types.py @@ -27,6 +27,12 @@ class Workspace(BaseModel): engine: str +class WorkspaceStatus(Enum): + SUBMITTED = "submitted" + READY = "ready" + CREATING = "creating" + + class Provider(Enum): BEDROCK = "bedrock" OPENAI = "openai" diff --git a/lib/shared/layers/python-sdk/python/genai_core/workspaces.py b/lib/shared/layers/python-sdk/python/genai_core/workspaces.py index 7652158f4..34ee1b9c9 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/workspaces.py +++ b/lib/shared/layers/python-sdk/python/genai_core/workspaces.py @@ -4,6 +4,7 @@ import boto3 import genai_core.embeddings from datetime import datetime +from .types import WorkspaceStatus from genai_core.types import Task dynamodb = boto3.resource("dynamodb") @@ -120,7 +121,7 @@ def create_workspace_aurora( "object_type": WORKSPACE_OBJECT_TYPE, "format_version": 1, "name": workspace_name, - "engine": "aurora", + "engine": WorkspaceStatus.SUBMITTED.value, "status": "submitted", "embeddings_model_provider": embeddings_model_provider, "embeddings_model_name": embeddings_model_name, @@ -188,7 +189,7 @@ def create_workspace_open_search( "format_version": 1, "name": workspace_name, "engine": "opensearch", - "status": "submitted", + "status": WorkspaceStatus.SUBMITTED, "embeddings_model_provider": embeddings_model_provider, "embeddings_model_name": embeddings_model_name, "embeddings_model_dimensions": embeddings_model_dimensions, @@ -240,7 +241,7 @@ def create_workspace_kendra( "format_version": 1, "name": workspace_name, "engine": "kendra", - "status": "submitted", + "status": WorkspaceStatus.SUBMITTED.value, "kendra_index_id": kendra_index_id, "kendra_index_external": kendra_index_external, "kendra_use_all_data": use_all_data, @@ -268,6 +269,37 @@ def create_workspace_kendra( return item +def create_workspace_bedrock_kb( + workspace_name: str, knowledge_base: dict, hybrid_search: bool +): + workspace_id = str(uuid.uuid4()) + timestamp = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ") + knowledge_base_id = knowledge_base["id"] + external = knowledge_base["external"] + + item = { + "workspace_id": workspace_id, + "object_type": WORKSPACE_OBJECT_TYPE, + "format_version": 1, + "name": workspace_name, + "engine": "bedrock_kb", + "status": WorkspaceStatus.READY.value, + "knowledge_base_id": knowledge_base_id, + "knowledge_base_external": external, + "hybrid_search": hybrid_search, + "documents": 0, + "vectors": 0, + "size_in_bytes": 0, + "created_at": timestamp, + "updated_at": timestamp, + } + + response = table.put_item(Item=item) + print(response) + + return item + + def delete_workspace(workspace_id: str): response = table.get_item( Key={"workspace_id": workspace_id, "object_type": WORKSPACE_OBJECT_TYPE} @@ -290,4 +322,4 @@ def delete_workspace(workspace_id: str): ), ) - print(response) \ No newline at end of file + print(response) diff --git a/lib/shared/types.ts b/lib/shared/types.ts index c9501ff64..e41a051c4 100644 --- a/lib/shared/types.ts +++ b/lib/shared/types.ts @@ -143,6 +143,15 @@ export interface SystemConfig { }[]; enterprise?: boolean; }; + knowledgeBase: { + enabled: boolean; + external?: { + name: string; + knowledgeBaseId: string; + region?: SupportedRegion; + roleArn?: string; + }[]; + }; }; embeddingsModels: { provider: ModelProvider; diff --git a/lib/user-interface/react-app/package-lock.json b/lib/user-interface/react-app/package-lock.json index ce364c816..674e6e47f 100644 --- a/lib/user-interface/react-app/package-lock.json +++ b/lib/user-interface/react-app/package-lock.json @@ -19741,4 +19741,4 @@ } } } -} +} \ No newline at end of file diff --git a/lib/user-interface/react-app/public/images/welcome/chat-modes.png b/lib/user-interface/react-app/public/images/welcome/chat-modes.png new file mode 100644 index 000000000..930495e23 Binary files /dev/null and b/lib/user-interface/react-app/public/images/welcome/chat-modes.png differ diff --git a/lib/user-interface/react-app/src/common/api-client/api-client.ts b/lib/user-interface/react-app/src/common/api-client/api-client.ts index 61fcbcae0..582a5e7d5 100644 --- a/lib/user-interface/react-app/src/common/api-client/api-client.ts +++ b/lib/user-interface/react-app/src/common/api-client/api-client.ts @@ -10,19 +10,21 @@ import { SemanticSearchClient } from "./semantic-search-client"; import { DocumentsClient } from "./documents-client"; import { KendraClient } from "./kendra-client"; import { UserFeedbackClient } from "./user-feedback-client"; +import { BedrockKBClient } from "./kb-client"; export class ApiClient { - private _healthClient: HealthClient | undefined; - private _ragEnginesClient: RagEnginesClient | undefined; - private _embeddingsClient: EmbeddingsClient | undefined; - private _crossEncodersClient: CrossEncodersClient | undefined; - private _modelsClient: ModelsClient | undefined; - private _workspacesClient: WorkspacesClient | undefined; - private _sessionsClient: SessionsClient | undefined; - private _semanticSearchClient: SemanticSearchClient | undefined; - private _documentsClient: DocumentsClient | undefined; - private _kendraClient: KendraClient | undefined; - private _userFeedbackClient: UserFeedbackClient | undefined; + private _healthClient?: HealthClient; + private _ragEnginesClient?: RagEnginesClient; + private _embeddingsClient?: EmbeddingsClient; + private _crossEncodersClient?: CrossEncodersClient; + private _modelsClient?: ModelsClient; + private _workspacesClient?: WorkspacesClient; + private _sessionsClient?: SessionsClient; + private _semanticSearchClient?: SemanticSearchClient; + private _documentsClient?: DocumentsClient; + private _kendraClient?: KendraClient; + private _userFeedbackClient?: UserFeedbackClient; + private _bedrockKBClient?: BedrockKBClient; public get health() { if (!this._healthClient) { @@ -104,6 +106,14 @@ export class ApiClient { return this._kendraClient; } + public get bedrockKB() { + if (!this._bedrockKBClient) { + this._bedrockKBClient = new BedrockKBClient(); + } + + return this._bedrockKBClient; + } + public get userFeedback() { if (!this._userFeedbackClient) { this._userFeedbackClient = new UserFeedbackClient(); diff --git a/lib/user-interface/react-app/src/common/api-client/kb-client.ts b/lib/user-interface/react-app/src/common/api-client/kb-client.ts new file mode 100644 index 000000000..75c48a748 --- /dev/null +++ b/lib/user-interface/react-app/src/common/api-client/kb-client.ts @@ -0,0 +1,17 @@ +import { API } from "aws-amplify"; +import { GraphQLQuery, GraphQLResult } from "@aws-amplify/api"; +import { listBedrockKnowledgeBases } from "../../graphql/queries"; +import { ListBedrockKnowledgeBasesQuery } from "../../API"; + +export class BedrockKBClient { + async listKnowledgeBases(): Promise< + GraphQLResult> + > { + const result = await API.graphql< + GraphQLQuery + >({ + query: listBedrockKnowledgeBases, + }); + return result; + } +} diff --git a/lib/user-interface/react-app/src/common/api-client/workspaces-client.ts b/lib/user-interface/react-app/src/common/api-client/workspaces-client.ts index 58124abc8..7f72e82d5 100644 --- a/lib/user-interface/react-app/src/common/api-client/workspaces-client.ts +++ b/lib/user-interface/react-app/src/common/api-client/workspaces-client.ts @@ -5,6 +5,7 @@ import { createAuroraWorkspace, createKendraWorkspace, createOpenSearchWorkspace, + createBedrockKBWorkspace, deleteWorkspace, } from "../../graphql/mutations"; import { @@ -14,6 +15,7 @@ import { CreateKendraWorkspaceMutation, CreateOpenSearchWorkspaceMutation, DeleteWorkspaceMutation, + CreateBedrockKBWorkspaceMutation, } from "../../API"; export class WorkspacesClient { @@ -109,4 +111,18 @@ export class WorkspacesClient { }); return result; } + + async createBedrockKBWorkspace(params: { + name: string; + knowledgeBaseId: string; + hybridSearch: boolean; + }): Promise>> { + const result = API.graphql>({ + query: createBedrockKBWorkspace, + variables: { + input: { ...params, kind: "bedrock_kb" }, + }, + }); + return result; + } } diff --git a/lib/user-interface/react-app/src/common/constants.ts b/lib/user-interface/react-app/src/common/constants.ts index be302ac8f..23da9d32f 100644 --- a/lib/user-interface/react-app/src/common/constants.ts +++ b/lib/user-interface/react-app/src/common/constants.ts @@ -37,6 +37,7 @@ export abstract class Labels { aurora: "Aurora Serverless v2 (pgvector)", opensearch: "OpenSearch Serverless", kendra: "Kendra", + bedrock_kb: "Bedrock Knowledge Base", }; static statusTypeMap: Record = { diff --git a/lib/user-interface/react-app/src/common/types.ts b/lib/user-interface/react-app/src/common/types.ts index b67839dd6..5396ce5f7 100644 --- a/lib/user-interface/react-app/src/common/types.ts +++ b/lib/user-interface/react-app/src/common/types.ts @@ -95,3 +95,9 @@ export interface KendraWorkspaceCreateInput { kendraIndex: SelectProps.Option | null; useAllData: boolean; } + +export interface BedrockKBWorkspaceCreateInput { + name: string; + knowledgeBaseId: SelectProps.Option | null; + hybridSearch: boolean; +} diff --git a/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx b/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx index 05bc9d8bc..84dc8c1dc 100644 --- a/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx +++ b/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx @@ -112,62 +112,64 @@ export default function ChatMessage(props: ChatMessageProps) { container: "jsonContainer", }} /> - {props.message.metadata.documents && ( - <> -
- - Copied to clipboard - - } - > -
- { + {props.message.metadata.documents && + (props.message.metadata.documents as RagDocument[]).length > + 0 && ( + <> +
+ + Copied to clipboard + + } + > +
+ { return { id: `${i}`, label: p.metadata.path?.split("/").at(-1) ?? p.metadata.title ?? p.metadata.document_id.slice(-8), + href: p.metadata.path, content: ( - <> -