Skip to content

Commit f10451d

Browse files
keremnalbantKerem Nalbantkrokoko
authored
feat: add additional props for opensearch serverless and bedrock (#460)
* feat(construct): add name prop to knowledge base * feat(construct): add collectionName and standbyReplicas props to VectorCollection --------- Co-authored-by: Kerem Nalbant <[email protected]> Co-authored-by: Alain Krok <[email protected]>
1 parent 1136f18 commit f10451d

10 files changed

+166
-31
lines changed

apidocs/classes/opensearchserverless.VectorCollection.md

+11-1
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ It includes all policies.
2929
- [dataAccessPolicy](opensearchserverless.VectorCollection.md#dataaccesspolicy)
3030
- [dataAccessPolicyDocument](opensearchserverless.VectorCollection.md#dataaccesspolicydocument)
3131
- [node](opensearchserverless.VectorCollection.md#node)
32+
- [standbyReplicas](opensearchserverless.VectorCollection.md#standbyreplicas)
3233

3334
### Methods
3435

@@ -40,14 +41,15 @@ It includes all policies.
4041

4142
### constructor
4243

43-
**new VectorCollection**(`scope`, `id`): [`VectorCollection`](opensearchserverless.VectorCollection.md)
44+
**new VectorCollection**(`scope`, `id`, `props?`): [`VectorCollection`](opensearchserverless.VectorCollection.md)
4445

4546
#### Parameters
4647

4748
| Name | Type |
4849
| :------ | :------ |
4950
| `scope` | `Construct` |
5051
| `id` | `string` |
52+
| `props?` | [`VectorCollectionProps`](../interfaces/opensearchserverless.VectorCollectionProps.md) |
5153

5254
#### Returns
5355

@@ -117,6 +119,14 @@ The tree node.
117119

118120
Construct.node
119121

122+
___
123+
124+
### standbyReplicas
125+
126+
**standbyReplicas**: [`VectorCollectionStandbyReplicas`](../enums/opensearchserverless.VectorCollectionStandbyReplicas.md)
127+
128+
Indicates whether to use standby replicas for the collection.
129+
120130
## Methods
121131

122132
### grantDataAccess
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
[@cdklabs/generative-ai-cdk-constructs](../README.md) / [opensearchserverless](../modules/opensearchserverless.md) / VectorCollectionStandbyReplicas
2+
3+
# Enumeration: VectorCollectionStandbyReplicas
4+
5+
[opensearchserverless](../modules/opensearchserverless.md).VectorCollectionStandbyReplicas
6+
7+
## Table of contents
8+
9+
### Enumeration Members
10+
11+
- [DISABLED](opensearchserverless.VectorCollectionStandbyReplicas.md#disabled)
12+
- [ENABLED](opensearchserverless.VectorCollectionStandbyReplicas.md#enabled)
13+
14+
## Enumeration Members
15+
16+
### DISABLED
17+
18+
**DISABLED** = ``"DISABLED"``
19+
20+
___
21+
22+
### ENABLED
23+
24+
**ENABLED** = ``"ENABLED"``

apidocs/interfaces/bedrock.KnowledgeBaseProps.md

+9
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ Properties for a knowledge base
1515
- [indexName](bedrock.KnowledgeBaseProps.md#indexname)
1616
- [instruction](bedrock.KnowledgeBaseProps.md#instruction)
1717
- [knowledgeBaseState](bedrock.KnowledgeBaseProps.md#knowledgebasestate)
18+
- [name](bedrock.KnowledgeBaseProps.md#name)
1819
- [tags](bedrock.KnowledgeBaseProps.md#tags)
1920
- [vectorField](bedrock.KnowledgeBaseProps.md#vectorfield)
2021
- [vectorIndex](bedrock.KnowledgeBaseProps.md#vectorindex)
@@ -85,6 +86,14 @@ Specifies whether to use the knowledge base or not when sending an InvokeAgent r
8586

8687
___
8788

89+
### name
90+
91+
`Optional` `Readonly` **name**: `string`
92+
93+
The name of the knowledge base.
94+
95+
___
96+
8897
### tags
8998

9099
`Optional` `Readonly` **tags**: `Record`\<`string`, `string`\>
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
[@cdklabs/generative-ai-cdk-constructs](../README.md) / [opensearchserverless](../modules/opensearchserverless.md) / VectorCollectionProps
2+
3+
# Interface: VectorCollectionProps
4+
5+
[opensearchserverless](../modules/opensearchserverless.md).VectorCollectionProps
6+
7+
## Table of contents
8+
9+
### Properties
10+
11+
- [collectionName](opensearchserverless.VectorCollectionProps.md#collectionname)
12+
- [standbyReplicas](opensearchserverless.VectorCollectionProps.md#standbyreplicas)
13+
14+
## Properties
15+
16+
### collectionName
17+
18+
`Readonly` **collectionName**: `string`
19+
20+
The name of the collection.
21+
22+
___
23+
24+
### standbyReplicas
25+
26+
`Optional` `Readonly` **standbyReplicas**: [`VectorCollectionStandbyReplicas`](../enums/opensearchserverless.VectorCollectionStandbyReplicas.md)
27+
28+
Indicates whether to use standby replicas for the collection.
29+
30+
**`Default`**
31+
32+
```ts
33+
ENABLED
34+
```

apidocs/modules/opensearchserverless.md

+8
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,14 @@
44

55
## Table of contents
66

7+
### Enumerations
8+
9+
- [VectorCollectionStandbyReplicas](../enums/opensearchserverless.VectorCollectionStandbyReplicas.md)
10+
711
### Classes
812

913
- [VectorCollection](../classes/opensearchserverless.VectorCollection.md)
14+
15+
### Interfaces
16+
17+
- [VectorCollectionProps](../interfaces/opensearchserverless.VectorCollectionProps.md)

src/cdk-lib/bedrock/knowledge-base.ts

+22-17
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,11 @@ export interface KnowledgeBaseProps {
9999
*/
100100
readonly embeddingsModel: BedrockFoundationModel;
101101

102+
/**
103+
* The name of the knowledge base.
104+
*/
105+
readonly name?: string;
106+
102107
/**
103108
* The description of the knowledge base.
104109
*
@@ -241,7 +246,7 @@ export class KnowledgeBase extends Construct {
241246
const textField = 'AMAZON_BEDROCK_TEXT_CHUNK';
242247
const metadataField = 'AMAZON_BEDROCK_METADATA';
243248

244-
this.description= props.description ?? 'CDK deployed Knowledge base'; // even though this prop is optional, if no value is provided it will fail to deploy
249+
this.description = props.description ?? 'CDK deployed Knowledge base'; // even though this prop is optional, if no value is provided it will fail to deploy
245250
this.knowledgeBaseState = props.knowledgeBaseState ?? 'ENABLED';
246251

247252

@@ -251,10 +256,11 @@ export class KnowledgeBase extends Construct {
251256
validateIndexParameters(props.vectorIndex, indexName, vectorField);
252257
}
253258

254-
this.name = generatePhysicalNameV2(
259+
this.name = props.name ?? generatePhysicalNameV2(
255260
this,
256261
'KB',
257262
{ maxLength: 32 });
263+
258264
const roleName = generatePhysicalNameV2(
259265
this,
260266
'AmazonBedrockExecutionRoleForKnowledgeBase',
@@ -422,7 +428,7 @@ export class KnowledgeBase extends Construct {
422428
tags: props.tags,
423429
});
424430

425-
this.knowledgeBaseInstance=knowledgeBase;
431+
this.knowledgeBaseInstance = knowledgeBase;
426432

427433
const kbCRPolicy = new iam.Policy(this, 'KBCRPolicy', {
428434
// roles: [crProvider.role],
@@ -583,8 +589,7 @@ export class KnowledgeBase extends Construct {
583589
* @returns The instance of VectorCollection, VectorStoreType.
584590
* @internal This is an internal core function and should not be called directly.
585591
*/
586-
private handleOpenSearchDefaultVectorCollection():
587-
{
592+
private handleOpenSearchDefaultVectorCollection(): {
588593
vectorStore: VectorCollection;
589594
vectorStoreType: VectorStoreType;
590595
} {
@@ -606,7 +611,7 @@ export class KnowledgeBase extends Construct {
606611
knowledgeBaseId: this.knowledgeBaseId,
607612
knowledgeBaseState: this.knowledgeBaseState,
608613
};
609-
agent.knowledgeBases=[agentKnowledgeBaseProperty];
614+
agent.knowledgeBases = [agentKnowledgeBaseProperty];
610615
}
611616
}
612617

@@ -636,15 +641,15 @@ function validateVectorIndex(
636641
) {
637642
if (!(vectorStore instanceof VectorCollection) && vectorIndex) {
638643
throw new Error('If vectorStore is not of type VectorCollection, vectorIndex should not be provided ' +
639-
'in KnowledgeBase construct.');
644+
'in KnowledgeBase construct.');
640645
}
641646
if (!(vectorStore instanceof VectorCollection) && indexName) {
642647
throw new Error('If vectorStore is not of type VectorCollection, indexName should not be provided ' +
643-
'in KnowledgeBase construct.');
648+
'in KnowledgeBase construct.');
644649
}
645650
if (!(vectorStore instanceof VectorCollection) && vectorField) {
646651
throw new Error('If vectorStore is not of type VectorCollection, vectorField should not be provided ' +
647-
'in KnowledgeBase construct.');
652+
'in KnowledgeBase construct.');
648653
}
649654
}
650655

@@ -667,19 +672,19 @@ function validateIndexParameters(
667672
if (vectorIndex.indexName !== 'bedrock-knowledge-base-default-index') {
668673
if (vectorIndex.indexName !== indexName) {
669674
throw new Error('Default value of indexName is `bedrock-knowledge-base-default-index`.' +
670-
' If you create VectorIndex manually and assign vectorIndex to value other than' +
671-
' `bedrock-knowledge-base-default-index` then you must provide the same value in KnowledgeBase construct.' +
672-
' If you created VectorIndex manually and set it to `bedrock-knowledge-base-default-index`' +
673-
' then do not assign indexName in KnowledgeBase construct.');
675+
' If you create VectorIndex manually and assign vectorIndex to value other than' +
676+
' `bedrock-knowledge-base-default-index` then you must provide the same value in KnowledgeBase construct.' +
677+
' If you created VectorIndex manually and set it to `bedrock-knowledge-base-default-index`' +
678+
' then do not assign indexName in KnowledgeBase construct.');
674679
}
675680
}
676681
if (vectorIndex.vectorField !== 'bedrock-knowledge-base-default-vector') {
677682
if (vectorIndex.vectorField !== vectorField) {
678683
throw new Error('Default value of vectorField is `bedrock-knowledge-base-default-vector`.' +
679-
' If you create VectorIndex manually and assign vectorField to value other than' +
680-
' `bedrock-knowledge-base-default-field` then you must provide the same value in KnowledgeBase construct.' +
681-
' If you created VectorIndex manually and set it to `bedrock-knowledge-base-default-vector`' +
682-
' then do not assign vectorField in KnowledgeBase construct.');
684+
' If you create VectorIndex manually and assign vectorField to value other than' +
685+
' `bedrock-knowledge-base-default-field` then you must provide the same value in KnowledgeBase construct.' +
686+
' If you created VectorIndex manually and set it to `bedrock-knowledge-base-default-vector`' +
687+
' then do not assign vectorField in KnowledgeBase construct.');
683688
}
684689
}
685690
}

src/cdk-lib/opensearchserverless/vector-collection.ts

+29-2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,25 @@ import * as oss from 'aws-cdk-lib/aws-opensearchserverless';
1616
import { Construct } from 'constructs';
1717
import { generatePhysicalNameV2 } from '../../common/helpers/utils';
1818

19+
export enum VectorCollectionStandbyReplicas {
20+
ENABLED = 'ENABLED',
21+
DISABLED = 'DISABLED',
22+
}
23+
24+
export interface VectorCollectionProps {
25+
/**
26+
* The name of the collection.
27+
*/
28+
readonly collectionName: string;
29+
30+
/**
31+
* Indicates whether to use standby replicas for the collection.
32+
*
33+
* @default ENABLED
34+
*/
35+
readonly standbyReplicas?: VectorCollectionStandbyReplicas;
36+
}
37+
1938
/**
2039
* Deploys an OpenSearch Serverless Collection to be used as a vector store.
2140
*
@@ -27,6 +46,11 @@ export class VectorCollection extends Construct {
2746
*/
2847
public collectionName: string;
2948

49+
/**
50+
* Indicates whether to use standby replicas for the collection.
51+
*/
52+
public standbyReplicas: VectorCollectionStandbyReplicas;
53+
3054
/**
3155
* The ID of the collection.
3256
*/
@@ -52,14 +76,16 @@ export class VectorCollection extends Construct {
5276
*/
5377
private dataAccessPolicyDocument: any[] = [];
5478

55-
constructor(scope: Construct, id: string) {
79+
constructor(scope: Construct, id: string, props?: VectorCollectionProps) {
5680
super(scope, id);
5781

58-
this.collectionName = generatePhysicalNameV2(
82+
this.collectionName = props?.collectionName ?? generatePhysicalNameV2(
5983
this,
6084
'VectorStore',
6185
{ maxLength: 32, lower: true });
6286

87+
this.standbyReplicas = props?.standbyReplicas ?? VectorCollectionStandbyReplicas.ENABLED;
88+
6389
const encryptionPolicyName = generatePhysicalNameV2(this,
6490
'EncryptionPolicy',
6591
{ maxLength: 32, lower: true });
@@ -103,6 +129,7 @@ export class VectorCollection extends Construct {
103129
const collection = new oss.CfnCollection(this, 'VectorCollection', {
104130
name: this.collectionName,
105131
type: 'VECTORSEARCH',
132+
standbyReplicas: this.standbyReplicas,
106133
});
107134

108135
this.collectionArn = collection.attrArn;

test/cdk-lib/bedrock/__snapshots__/agent.test.ts.snap

+1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

test/cdk-lib/bedrock/knowledge-base.test.ts

+10-8
Original file line numberDiff line numberDiff line change
@@ -96,20 +96,22 @@ describe('KnowledgeBase', () => {
9696
});
9797

9898
test('Should correctly initialize with custom props', () => {
99+
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
99100
const vectorStore = new AmazonAuroraDefaultVectorStore(stack, 'AuroraDefaultVectorStore6', {
100-
embeddingsModelVectorDimension: BedrockFoundationModel.TITAN_EMBED_TEXT_V1.vectorDimensions!,
101+
embeddingsModelVectorDimension: model.vectorDimensions!,
101102
});
102-
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
103103
const knowledgeBase = new KnowledgeBase(stack, 'AuroraDefaultKnowledgeBase', {
104104
embeddingsModel: model,
105105
vectorStore: vectorStore,
106106
instruction: 'Test instruction',
107+
name: 'TestKnowledgeBase',
107108
});
108109

109110
expect(knowledgeBase.instruction).toBe('Test instruction');
110111
expect(knowledgeBase.name).toBeDefined();
111112
expect(knowledgeBase.role).toBeDefined();
112113
expect(knowledgeBase.vectorStore).toBe(vectorStore);
114+
expect(knowledgeBase.name).toBe('TestKnowledgeBase');
113115
});
114116

115117
test('Should correctly initialize role with necessary permissions', () => {
@@ -129,10 +131,10 @@ describe('KnowledgeBase', () => {
129131
});
130132

131133
test('Should throw error when vectorStore is not VectorCollection and indexName is provided', () => {
134+
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
132135
const vectorStore = new AmazonAuroraDefaultVectorStore(stack, 'AmazonAuroraDefaultVectorStore6', {
133-
embeddingsModelVectorDimension: BedrockFoundationModel.TITAN_EMBED_TEXT_V1.vectorDimensions!,
136+
embeddingsModelVectorDimension: model.vectorDimensions!,
134137
});
135-
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
136138

137139
expect(() => {
138140
new KnowledgeBase(stack, 'AuroraDefaultKnowledgeBase6', {
@@ -144,10 +146,10 @@ describe('KnowledgeBase', () => {
144146
});
145147

146148
test('Should throw error when vectorStore is not VectorCollection and vectorField is provided', () => {
149+
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
147150
const vectorStore = new AmazonAuroraDefaultVectorStore(stack, 'AmazonAuroraDefaultVectorStore5', {
148-
embeddingsModelVectorDimension: BedrockFoundationModel.TITAN_EMBED_TEXT_V1.vectorDimensions!,
151+
embeddingsModelVectorDimension: model.vectorDimensions!,
149152
});
150-
const model = BedrockFoundationModel.TITAN_EMBED_TEXT_V1;
151153

152154
expect(() => {
153155
new KnowledgeBase(stack, 'AuroraDefaultKnowledgeBase5', {
@@ -187,10 +189,10 @@ describe('KnowledgeBase', () => {
187189
});
188190

189191
test('Should correctly initialize with AmazonAuroraDefaultVectorStore and custom embeddingsModel', () => {
192+
const model = BedrockFoundationModel.COHERE_EMBED_ENGLISH_V3;
190193
const vectorStore = new AmazonAuroraDefaultVectorStore(stack, 'AuroraDefaultVectorStore2', {
191-
embeddingsModelVectorDimension: BedrockFoundationModel.COHERE_EMBED_ENGLISH_V3.vectorDimensions!,
194+
embeddingsModelVectorDimension: model.vectorDimensions!,
192195
});
193-
const model = BedrockFoundationModel.COHERE_EMBED_ENGLISH_V3;
194196

195197
const knowledgeBase = new KnowledgeBase(stack, 'AuroraDefaultKnowledgeBase2', {
196198
embeddingsModel: model,

0 commit comments

Comments
 (0)