diff --git a/infra/lib/compute-stack.ts b/infra/lib/compute-stack.ts index 9ff5d48..2f4a186 100644 --- a/infra/lib/compute-stack.ts +++ b/infra/lib/compute-stack.ts @@ -12,6 +12,7 @@ import * as lambdaEventSources from 'aws-cdk-lib/aws-lambda-event-sources'; import * as bedrock from 'aws-cdk-lib/aws-bedrock'; import * as opensearchserverless from 'aws-cdk-lib/aws-opensearchserverless'; import * as logs from 'aws-cdk-lib/aws-logs'; +import * as cr from 'aws-cdk-lib/custom-resources'; import { Construct } from 'constructs'; export interface ComputeStackProps extends cdk.StackProps { @@ -90,8 +91,28 @@ export class ComputeStack extends cdk.Stack { resources: [`arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v2:0`], })); - // OpenSearch Serverless data access policy - new opensearchserverless.CfnAccessPolicy(this, 'OssDataAccessPolicy', { + // Lambda to pre-create the vector index (retries until AOSS access policy propagates) + const indexCreatorFn = new lambda.Function(this, 'OssIndexCreator', { + runtime: lambda.Runtime.PYTHON_3_12, + handler: 'app.handler', + code: lambda.Code.fromAsset('../lambdas/oss-index-creator', { + bundling: { + image: lambda.Runtime.PYTHON_3_12.bundlingImage, + command: [ + 'bash', '-c', + 'pip install -r requirements.txt -t /asset-output && cp -au . /asset-output', + ], + }, + }), + timeout: cdk.Duration.minutes(6), + }); + + indexCreatorFn.addToRolePolicy(new iam.PolicyStatement({ + actions: ['aoss:APIAccessAll'], + resources: [ossCollection.attrArn], + })); + + const ossDataAccessPolicy = new opensearchserverless.CfnAccessPolicy(this, 'OssDataAccessPolicy', { name: 'proposal-system-kb-access', type: 'data', policy: JSON.stringify([{ @@ -99,11 +120,27 @@ export class ComputeStack extends cdk.Stack { { ResourceType: 'collection', Resource: ['collection/proposal-system-kb'], Permission: ['aoss:*'] }, { ResourceType: 'index', Resource: ['index/proposal-system-kb/*'], Permission: ['aoss:*'] }, ], - Principal: [kbRole.roleArn, `arn:aws:iam::${this.account}:root`], + Principal: [kbRole.roleArn, indexCreatorFn.role!.roleArn], }]), }); + ossDataAccessPolicy.addDependency(ossCollection); + + const indexProvider = new cr.Provider(this, 'OssIndexProvider', { + onEventHandler: indexCreatorFn, + }); + + const ossIndex = new cdk.CustomResource(this, 'OssIndex', { + serviceToken: indexProvider.serviceToken, + properties: { + Endpoint: ossCollection.attrCollectionEndpoint, + IndexName: 'proposal-system-index', + VectorField: 'embedding', + TextField: 'text', + MetadataField: 'metadata', + }, + }); + ossIndex.node.addDependency(ossDataAccessPolicy); - // Bedrock Knowledge Base const knowledgeBase = new bedrock.CfnKnowledgeBase(this, 'KnowledgeBase', { name: 'proposal-system-kb', roleArn: kbRole.roleArn, @@ -126,6 +163,7 @@ export class ComputeStack extends cdk.Stack { }, }, }); + knowledgeBase.node.addDependency(ossIndex); // KB Data Source (S3 library bucket) const dataSource = new bedrock.CfnDataSource(this, 'KbDataSource', { diff --git a/lambdas/oss-index-creator/app.py b/lambdas/oss-index-creator/app.py new file mode 100644 index 0000000..ec23316 --- /dev/null +++ b/lambdas/oss-index-creator/app.py @@ -0,0 +1,90 @@ +import json +import time +import boto3 +from opensearchpy import OpenSearch, RequestsHttpConnection +from requests_aws4auth import AWS4Auth + + +def handler(event, context): + print(f"Event: {json.dumps(event, default=str)}") + + if event["RequestType"] == "Delete": + print("Delete request, returning") + return {"PhysicalResourceId": event.get("PhysicalResourceId", "none")} + + props = event["ResourceProperties"] + endpoint = props["Endpoint"].replace("https://", "") + index_name = props["IndexName"] + vector_field = props["VectorField"] + text_field = props["TextField"] + metadata_field = props["MetadataField"] + + print(f"Endpoint: {endpoint}") + print(f"Index: {index_name}") + + session = boto3.Session() + credentials = session.get_credentials().get_frozen_credentials() + region = session.region_name + print(f"Region: {region}") + + awsauth = AWS4Auth( + credentials.access_key, + credentials.secret_key, + region, + "aoss", + session_token=credentials.token, + ) + + client = OpenSearch( + hosts=[{"host": endpoint, "port": 443}], + http_auth=awsauth, + use_ssl=True, + verify_certs=True, + connection_class=RequestsHttpConnection, + timeout=30, + ) + + index_body = { + "settings": { + "index": {"knn": True, "knn.algo_param.ef_search": 512} + }, + "mappings": { + "properties": { + vector_field: { + "type": "knn_vector", + "dimension": 1024, + "method": { + "engine": "faiss", + "name": "hnsw", + "space_type": "l2", + }, + }, + text_field: {"type": "text"}, + metadata_field: {"type": "text"}, + } + }, + } + + for attempt in range(30): + try: + print(f"Attempt {attempt}: creating index...") + response = client.indices.create(index=index_name, body=index_body) + print(f"Index created successfully: {response}") + # Verify the index exists + exists = client.indices.exists(index=index_name) + print(f"Index exists check: {exists}") + return {"PhysicalResourceId": index_name} + except Exception as e: + error_str = str(e) + print(f"Attempt {attempt} error: {error_str}") + if "resource_already_exists_exception" in error_str: + print("Index already exists, returning success") + return {"PhysicalResourceId": index_name} + if "403" in error_str and attempt < 29: + print(f"403 error, retrying in 10s (attempt {attempt}/29)") + time.sleep(10) + continue + print(f"Fatal error on attempt {attempt}: {error_str}") + raise + + raise Exception("Timeout waiting for AOSS access policy propagation") diff --git a/lambdas/oss-index-creator/requirements.txt b/lambdas/oss-index-creator/requirements.txt new file mode 100644 index 0000000..c0ac3c5 --- /dev/null +++ b/lambdas/oss-index-creator/requirements.txt @@ -0,0 +1,3 @@ +opensearch-py>=2.4.0 +requests-aws4auth>=1.2.0 +requests>=2.31.0