Merge pull request #9 from co-cddo/s3_vectors

S3_vectors

Author: gillespied

pyproject.toml

@@ -1,14 +1,15 @@
 [project]
 name = "gds-idea-cdk-constructs"
-version = "0.3.0"
+version = "0.4.0"
 description = "A repo for commonly used constructs in the team."
 readme = "README.md"
 authors = [
-    { name = "David Gillespie", email = "david.gillespie@digital.cabinet-office.gov.uk" }
+    { name = "David Gillespie", email = "david.gillespie@digital.cabinet-office.gov.uk" },
+    { name = "Jose Orjales", email = "jose.orjales@digital.cabinet-office.gov.uk" },
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "aws-cdk-lib>=2.180.0",
+    "aws-cdk-lib>=2.243.0",
     "boto3>=1.35.0",
 ]
 

src/gds_idea_cdk_constructs/knowledge_base/__init__.py

@@ -0,0 +1,19 @@
+"""Reusable CDK construct for Bedrock Knowledge Bases."""
+
+from .props import (
+    ChunkingConfig,
+    ChunkingStrategy,
+    EmbeddingModel,
+    KnowledgeBaseProps,
+    StorageType,
+)
+from .stack import KnowledgeBase
+
+__all__ = [
+    "KnowledgeBase",
+    "KnowledgeBaseProps",
+    "ChunkingConfig",
+    "ChunkingStrategy",
+    "StorageType",
+    "EmbeddingModel",
+]

src/gds_idea_cdk_constructs/knowledge_base/_storage_strategies.py

@@ -0,0 +1,217 @@
+"""Storage strategy pattern for knowledge base vector backends
+
+This module defines the ``IStorageStrategy`` ABC and concrete implementations
+for each supported :class:`~.props.StorageType`.  The stack selects a strategy
+via :data:`STORAGE_STRATEGY_MAP` — consumers never interact with strategy
+classes directly.
+
+To add a new storage backend:
+    1. Implement :class:`IStorageStrategy`.
+    2. Add a :class:`~.props.StorageType` enum member.
+    3. Register the class in :data:`STORAGE_STRATEGY_MAP`.
+"""
+
+from abc import ABC, abstractmethod
+
+from aws_cdk import (
+    CfnResource,
+    RemovalPolicy,
+    aws_bedrock as bedrock,
+    aws_iam as iam,
+    aws_s3vectors as s3vectors,
+)
+from constructs import Construct
+
+from .props import KnowledgeBaseProps, StorageType
+
+
+class IStorageStrategy(ABC):
+    """Interface for knowledge base vector storage backends.
+
+    Each implementation creates the backend-specific infrastructure (e.g.,
+    S3 Vector Bucket + Index) and exposes the information the Bedrock
+    ``CfnKnowledgeBase`` needs to reference it.
+
+    Args:
+        scope: The CDK construct scope.
+        kb_props: Knowledge base configuration properties.
+    """
+
+    def __init__(self, scope: Construct, kb_props: KnowledgeBaseProps) -> None:
+        self.scope = scope
+        self.kb_props = kb_props
+
+    @abstractmethod
+    def create_storage_resources(
+        self, app_prefix: str, env_name: str, retain: bool
+    ) -> None:
+        """Provision backend-specific storage resources.
+
+        Args:
+            app_prefix: Naming prefix derived from the application name.
+            env_name: Deployment environment name (e.g., ``"development"``).
+            retain: Whether to apply RETAIN removal policy.
+        """
+
+    @abstractmethod
+    def get_storage_configuration(
+        self,
+    ) -> bedrock.CfnKnowledgeBase.StorageConfigurationProperty:
+        """Return the Bedrock storage configuration property.
+
+        Returns:
+            A ``StorageConfigurationProperty`` that tells Bedrock where
+            to store and retrieve vectors.
+        """
+
+    @abstractmethod
+    def get_iam_policy_statements(self) -> list[iam.PolicyStatement]:
+        """Return IAM policy statements the KB service role needs.
+
+        Returns:
+            A list of ``PolicyStatement`` objects granting the KB role
+            access to the storage backend.
+        """
+
+    @abstractmethod
+    def get_resource_dependencies(self) -> list[CfnResource]:
+        """Return CloudFormation resources the KB must depend on.
+
+        Returns:
+            A list of ``CfnResource`` objects that must be created before
+            the Knowledge Base.
+        """
+
+
+class S3VectorsStorageStrategy(IStorageStrategy):
+    """S3 Vectors storage backend — serverless vector storage built on S3.
+
+    Creates:
+        - ``CfnVectorBucket`` for vector storage
+        - ``CfnIndex`` for the vector index (float32, configurable dimensions
+          and distance metric)
+    """
+
+    def __init__(self, scope: Construct, kb_props: KnowledgeBaseProps) -> None:
+        super().__init__(scope, kb_props)
+        self._vector_bucket: s3vectors.CfnVectorBucket | None = None
+        self._vector_index: s3vectors.CfnIndex | None = None
+
+    def create_storage_resources(
+        self, app_prefix: str, env_name: str, retain: bool
+    ) -> None:
+        """Create S3 Vector Bucket and Index.
+
+        Args:
+            app_prefix: Naming prefix derived from the application name.
+            env_name: Deployment environment name.
+            retain: Whether to apply RETAIN removal policy.
+        """
+        removal_policy = RemovalPolicy.RETAIN if retain else RemovalPolicy.DESTROY
+
+        self._vector_bucket = s3vectors.CfnVectorBucket(
+            self.scope,
+            "VectorBucket",
+            vector_bucket_name=f"{app_prefix}-vectors-{env_name}",
+        )
+        self._vector_bucket.apply_removal_policy(removal_policy)
+
+        self._vector_index = s3vectors.CfnIndex(
+            self.scope,
+            "VectorIndex",
+            index_name=f"{app_prefix}-index-{env_name}",
+            vector_bucket_name=self._vector_bucket.vector_bucket_name,
+            data_type="float32",
+            dimension=self.kb_props.resolved_embedding_dimensions(),
+            distance_metric=self.kb_props.distance_metric,
+            metadata_configuration=s3vectors.CfnIndex.MetadataConfigurationProperty(
+                non_filterable_metadata_keys=["content"],
+            ),
+        )
+        self._vector_index.add_dependency(self._vector_bucket)
+
+    def get_storage_configuration(
+        self,
+    ) -> bedrock.CfnKnowledgeBase.StorageConfigurationProperty:
+        """Return S3 Vectors storage configuration for Bedrock.
+
+        Returns:
+            A ``StorageConfigurationProperty`` pointing to the S3 Vector
+            Bucket and Index.
+
+        Raises:
+            RuntimeError: If called before ``create_storage_resources``.
+        """
+        if self._vector_bucket is None or self._vector_index is None:
+            raise RuntimeError(
+                "create_storage_resources() must be called before "
+                "get_storage_configuration()"
+            )
+
+        return bedrock.CfnKnowledgeBase.StorageConfigurationProperty(
+            type="S3_VECTORS",
+            s3_vectors_configuration=bedrock.CfnKnowledgeBase.S3VectorsConfigurationProperty(
+                vector_bucket_arn=self._vector_bucket.attr_vector_bucket_arn,
+                index_name=self._vector_index.index_name,
+            ),
+        )
+
+    def get_iam_policy_statements(self) -> list[iam.PolicyStatement]:
+        """Return S3 Vectors IAM permissions for the KB service role.
+
+        Returns:
+            A list containing a single ``PolicyStatement`` granting
+            read/write access to the vector bucket and index.
+
+        Raises:
+            RuntimeError: If called before ``create_storage_resources``.
+        """
+        if self._vector_bucket is None or self._vector_index is None:
+            raise RuntimeError(
+                "create_storage_resources() must be called before "
+                "get_iam_policy_statements()"
+            )
+
+        return [
+            iam.PolicyStatement(
+                sid="S3VectorsAccess",
+                actions=[
+                    "s3vectors:CreateIndex",
+                    "s3vectors:GetIndex",
+                    "s3vectors:PutVectors",
+                    "s3vectors:GetVectors",
+                    "s3vectors:DeleteVectors",
+                    "s3vectors:QueryVectors",
+                    "s3vectors:ListVectors",
+                    "s3vectors:GetVectorBucket",
+                    "s3vectors:ListVectorBuckets",
+                ],
+                resources=[
+                    self._vector_bucket.attr_vector_bucket_arn,
+                    self._vector_index.attr_index_arn,
+                ],
+            )
+        ]
+
+    def get_resource_dependencies(self) -> list[CfnResource]:
+        """Return the CfnIndex as a dependency for the Knowledge Base.
+
+        Returns:
+            A list containing the ``CfnIndex`` (which transitively depends
+            on the ``CfnVectorBucket``).
+
+        Raises:
+            RuntimeError: If called before ``create_storage_resources``.
+        """
+        if self._vector_index is None:
+            raise RuntimeError(
+                "create_storage_resources() must be called before "
+                "get_resource_dependencies()"
+            )
+
+        return [self._vector_index]
+
+
+STORAGE_STRATEGY_MAP: dict[StorageType, type[IStorageStrategy]] = {
+    StorageType.S3_VECTORS: S3VectorsStorageStrategy,
+}

src/gds_idea_cdk_constructs/knowledge_base/lambda_handlers/__init__.py

@@ -0,0 +1,63 @@
+"""SQS-triggered Lambda handler for Bedrock Knowledge Base sync
+
+Invoked by an SQS queue that receives S3 event notifications.  The queue
+is configured with a batching window (default 20 minutes) so that many
+S3 object events are collected into a single Lambda invocation.  This
+handler then fires **one** ``StartIngestionJob`` call regardless of how
+many files changed, avoiding API throttling.
+
+Environment variables:
+    KNOWLEDGE_BASE_ID: Bedrock Knowledge Base ID.
+    DATA_SOURCE_ID: Bedrock Data Source ID.
+"""
+
+import os
+
+import boto3
+
+KNOWLEDGE_BASE_ID = os.environ.get("KNOWLEDGE_BASE_ID", "")
+DATA_SOURCE_ID = os.environ.get("DATA_SOURCE_ID", "")
+
+
+def handler(event, context):
+    """Start a Bedrock KB ingestion job.
+
+    The SQS batch may contain many S3 event records — we intentionally
+    ignore the individual records and trigger a full data-source sync.
+    Bedrock handles incremental ingestion internally (only new/changed
+    files are re-embedded).
+    """
+    if not KNOWLEDGE_BASE_ID or not DATA_SOURCE_ID:
+        print(
+            "ERROR: KNOWLEDGE_BASE_ID and DATA_SOURCE_ID must be set. "
+            f"KNOWLEDGE_BASE_ID={KNOWLEDGE_BASE_ID!r}, "
+            f"DATA_SOURCE_ID={DATA_SOURCE_ID!r}"
+        )
+        return {"statusCode": 400, "body": "Missing environment variables"}
+
+    client = boto3.client("bedrock-agent")
+
+    num_records = len(event.get("Records", []))
+    print(
+        f"Starting ingestion job for KB={KNOWLEDGE_BASE_ID}, "
+        f"DataSource={DATA_SOURCE_ID} "
+        f"(triggered by {num_records} SQS message(s))"
+    )
+
+    try:
+        response = client.start_ingestion_job(
+            knowledgeBaseId=KNOWLEDGE_BASE_ID,
+            dataSourceId=DATA_SOURCE_ID,
+        )
+        ingestion_job = response.get("ingestionJob", {})
+        job_id = ingestion_job.get("ingestionJobId", "unknown")
+        status = ingestion_job.get("status", "unknown")
+        print(f"Ingestion job started: jobId={job_id}, status={status}")
+
+        return {"statusCode": 200, "body": f"Ingestion job {job_id} started"}
+
+    except Exception as e:
+        # Log but do NOT re-raise — a raised exception causes SQS to retry,
+        # which would start duplicate ingestion jobs.
+        print(f"ERROR starting ingestion job: {e}")
+        return {"statusCode": 500, "body": str(e)}