1import * as bedrock from "@distilled.cloud/aws/bedrock-agent";
2import * as Effect from "effect/Effect";
3import * as Schedule from "effect/Schedule";
4import * as Stream from "effect/Stream";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import { createPhysicalName } from "../../PhysicalName.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { createInternalTags, diffTags, hasAlchemyTags } from "../../Tags.ts";
11
12/**
13 * The vector-store and embedding configuration of a knowledge base. Passed
14 * through to the Bedrock API unchanged — see the AWS SDK
15 * `KnowledgeBaseConfiguration` shape. The common case is a `VECTOR` type with
16 * `vectorKnowledgeBaseConfiguration.embeddingModelArn` set.
17 */
18export type KnowledgeBaseConfiguration = bedrock.KnowledgeBaseConfiguration;
19
20/**
21 * The backing vector store (OpenSearch Serverless, Pinecone, RDS/pgvector,
22 * S3 Vectors, ...). Passed through to the Bedrock API unchanged — see the AWS
23 * SDK `StorageConfiguration` shape. Omit for a managed vector store.
24 */
25export type StorageConfiguration = bedrock.StorageConfiguration;
26
27export interface KnowledgeBaseProps {
28 /**
29 * Name of the knowledge base (1-100 characters). If omitted, a
30 * deterministic physical name is generated from the app, stage, and
31 * logical ID. Changing the name triggers a replacement.
32 */
33 name?: string;
34 /**
35 * A description of the knowledge base.
36 */
37 description?: string;
38 /**
39 * The ARN of an IAM role Bedrock assumes to access the embedding model,
40 * the vector store, and (via data sources) the source data. Must trust
41 * `bedrock.amazonaws.com`.
42 */
43 roleArn: string;
44 /**
45 * The embedding + vector-store type configuration. Changing the type
46 * triggers a replacement.
47 */
48 knowledgeBaseConfiguration: KnowledgeBaseConfiguration;
49 /**
50 * The backing vector store. Omit for a Bedrock-managed store.
51 */
52 storageConfiguration?: StorageConfiguration;
53 /**
54 * Tags to apply to the knowledge base. Merged with internal Alchemy tags.
55 */
56 tags?: Record<string, string>;
57}
58
59export interface KnowledgeBase extends Resource<
60 "AWS.Bedrock.KnowledgeBase",
61 KnowledgeBaseProps,
62 {
63 /**
64 * The unique identifier of the knowledge base.
65 */
66 knowledgeBaseId: string;
67 /**
68 * The ARN of the knowledge base.
69 */
70 knowledgeBaseArn: string;
71 /**
72 * Name of the knowledge base.
73 */
74 name: string;
75 /**
76 * The ARN of the execution role the knowledge base assumes to access the
77 * embedding model and vector store.
78 */
79 roleArn: string;
80 }
81> {}
82
83/**
84 * An Amazon Bedrock knowledge base — a managed RAG index that embeds source
85 * documents into a vector store for retrieval.
86 *
87 * `KnowledgeBase` owns the index configuration; attach one or more
88 * {@link DataSource}s (e.g. an S3 bucket) to feed it documents, then trigger
89 * ingestion. Query it at runtime with the {@link Retrieve} and
90 * {@link RetrieveAndGenerate} bindings, or attach it to an {@link Agent}.
91 *
92 * The `roleArn` must grant Bedrock access to the embedding model, the vector
93 * store, and the source data. The vector store (`storageConfiguration`) must
94 * already exist — provision an OpenSearch Serverless collection (with a
95 * vector index) or another supported store first.
96 *
97 * ### Creating Knowledge Bases
98 * **Example:** OpenSearch Serverless Backed Knowledge Base
99 * ```typescript
100 * import * as Bedrock from "alchemy/AWS/Bedrock";
101 *
102 * const kb = yield* Bedrock.KnowledgeBase("docs", {
103 * roleArn: role.roleArn,
104 * knowledgeBaseConfiguration: {
105 * type: "VECTOR",
106 * vectorKnowledgeBaseConfiguration: {
107 * embeddingModelArn:
108 * "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0",
109 * },
110 * },
111 * storageConfiguration: {
112 * type: "OPENSEARCH_SERVERLESS",
113 * opensearchServerlessConfiguration: {
114 * collectionArn: collection.arn,
115 * vectorIndexName: "bedrock-index",
116 * fieldMapping: {
117 * vectorField: "bedrock-vector",
118 * textField: "bedrock-text",
119 * metadataField: "bedrock-metadata",
120 * },
121 * },
122 * },
123 * });
124 * ```
125 *
126 * @resource
127 */
128export const KnowledgeBase = Resource<KnowledgeBase>(
129 "AWS.Bedrock.KnowledgeBase",
130);
131
132/** KB status values indicating an in-flight transition to wait out. */
133const KB_TRANSIENT = new Set(["CREATING", "UPDATING", "DELETING"]);
134
135/**
136 * A knowledge base with attached data sources rejects deletion with a
137 * `ConflictException` until they are gone (the engine deletes them first, but
138 * their teardown is eventually consistent). Bounded retry (~60s), wrapped in
139 * an explicitly-typed helper so the `Effect.retry` return type does not widen
140 * the provider layer's requirement to `unknown` (see PATTERNS §7).
141 */
142const retryWhileConflict = <A, E extends { _tag: string }, R>(
143 self: Effect.Effect<A, E, R>,
144): Effect.Effect<A, E, R> =>
145 Effect.retry(self, {
146 while: (e) => e._tag === "ConflictException",
147 schedule: Schedule.max([Schedule.fixed("3 seconds"), Schedule.recurs(20)]),
148 });
149
150export const KnowledgeBaseProvider = () =>
151 Provider.effect(
152 KnowledgeBase,
153 Effect.gen(function* () {
154 const createName = Effect.fn(function* (
155 id: string,
156 props: Pick<KnowledgeBaseProps, "name">,
157 ) {
158 return (
159 props.name ?? (yield* createPhysicalName({ id, maxLength: 100 }))
160 );
161 });
162
163 const getKbOrUndefined = Effect.fn(function* (knowledgeBaseId: string) {
164 return yield* bedrock.getKnowledgeBase({ knowledgeBaseId }).pipe(
165 Effect.map((r) => r.knowledgeBase),
166 Effect.catchTag("ResourceNotFoundException", () =>
167 Effect.succeed(undefined),
168 ),
169 );
170 });
171
172 const findByName = Effect.fn(function* (name: string) {
173 const pages = yield* bedrock.listKnowledgeBases
174 .pages({})
175 .pipe(Stream.runCollect);
176 return Array.from(pages)
177 .flatMap((page) => page.knowledgeBaseSummaries ?? [])
178 .find((s) => s.name === name)?.knowledgeBaseId;
179 });
180
181 const fetchObservedTags = Effect.fn(function* (resourceArn: string) {
182 return yield* bedrock.listTagsForResource({ resourceArn }).pipe(
183 Effect.map((r) => (r.tags ?? {}) as Record<string, string>),
184 Effect.catchTag("ResourceNotFoundException", () =>
185 Effect.succeed({} as Record<string, string>),
186 ),
187 );
188 });
189
190 const waitForSettled = Effect.fn(function* (knowledgeBaseId: string) {
191 return yield* bedrock.getKnowledgeBase({ knowledgeBaseId }).pipe(
192 Effect.map((r) => r.knowledgeBase),
193 Effect.catchTag("ResourceNotFoundException", () =>
194 Effect.succeed(undefined),
195 ),
196 Effect.repeat({
197 schedule: Schedule.fixed("5 seconds"),
198 until: (kb) => kb === undefined || !KB_TRANSIENT.has(kb.status),
199 times: 60,
200 }),
201 );
202 });
203
204 return KnowledgeBase.Provider.of({
205 stables: ["knowledgeBaseId", "knowledgeBaseArn", "name"],
206
207 list: () =>
208 Effect.gen(function* () {
209 const pages = yield* bedrock.listKnowledgeBases
210 .pages({})
211 .pipe(Stream.runCollect);
212 const summaries = Array.from(pages).flatMap(
213 (page) => page.knowledgeBaseSummaries ?? [],
214 );
215 const results = yield* Effect.forEach(
216 summaries,
217 (s) =>
218 getKbOrUndefined(s.knowledgeBaseId).pipe(
219 Effect.map((kb) =>
220 kb === undefined
221 ? undefined
222 : {
223 knowledgeBaseId: kb.knowledgeBaseId,
224 knowledgeBaseArn: kb.knowledgeBaseArn,
225 name: kb.name,
226 roleArn: kb.roleArn,
227 },
228 ),
229 ),
230 { concurrency: 5 },
231 );
232 return results.filter(
233 (r): r is KnowledgeBase["Attributes"] => r !== undefined,
234 );
235 }),
236
237 read: Effect.fn(function* ({ id, olds, output }) {
238 const kbId =
239 output?.knowledgeBaseId ??
240 (yield* findByName(
241 output?.name ??
242 (yield* createName(id, olds ?? ({} as KnowledgeBaseProps))),
243 ));
244 if (kbId === undefined) return undefined;
245 const kb = yield* getKbOrUndefined(kbId);
246 if (kb === undefined || kb.status === "DELETING") return undefined;
247 const attrs = {
248 knowledgeBaseId: kb.knowledgeBaseId,
249 knowledgeBaseArn: kb.knowledgeBaseArn,
250 name: kb.name,
251 roleArn: kb.roleArn,
252 };
253 const tags = yield* fetchObservedTags(kb.knowledgeBaseArn);
254 return (yield* hasAlchemyTags(id, tags)) ? attrs : Unowned(attrs);
255 }),
256
257 diff: Effect.fn(function* ({ id, news, olds }) {
258 if (!isResolved(news)) return undefined;
259 const oldName = yield* createName(
260 id,
261 olds ?? ({} as KnowledgeBaseProps),
262 );
263 const newName = yield* createName(
264 id,
265 news ?? ({} as KnowledgeBaseProps),
266 );
267 if (oldName !== newName) {
268 return { action: "replace" } as const;
269 }
270 if (
271 (olds?.knowledgeBaseConfiguration?.type ?? undefined) !==
272 (news?.knowledgeBaseConfiguration?.type ?? undefined)
273 ) {
274 return { action: "replace" } as const;
275 }
276 // description, role, and tags converge via update.
277 }),
278
279 reconcile: Effect.fn(function* ({
280 id,
281 news = {} as KnowledgeBaseProps,
282 output,
283 session,
284 }) {
285 const name = output?.name ?? (yield* createName(id, news));
286 const internalTags = yield* createInternalTags(id);
287 const desiredTags = { ...news.tags, ...internalTags };
288
289 // 1. OBSERVE
290 let kb = output?.knowledgeBaseId
291 ? yield* getKbOrUndefined(output.knowledgeBaseId)
292 : undefined;
293 if (kb === undefined) {
294 const foundId = yield* findByName(name);
295 if (foundId !== undefined) {
296 kb = yield* getKbOrUndefined(foundId);
297 }
298 }
299
300 if (kb === undefined) {
301 // 2. ENSURE
302 const created = yield* bedrock.createKnowledgeBase({
303 name,
304 description: news.description,
305 roleArn: news.roleArn,
306 knowledgeBaseConfiguration: news.knowledgeBaseConfiguration,
307 storageConfiguration: news.storageConfiguration,
308 tags: desiredTags,
309 });
310 kb = created.knowledgeBase;
311 kb = (yield* waitForSettled(kb.knowledgeBaseId)) ?? kb;
312 } else {
313 // 3. SYNC — converge mutable settings.
314 kb = (yield* waitForSettled(kb.knowledgeBaseId)) ?? kb;
315 const drifted =
316 kb.description !== news.description ||
317 kb.roleArn !== news.roleArn ||
318 JSON.stringify(kb.knowledgeBaseConfiguration) !==
319 JSON.stringify(news.knowledgeBaseConfiguration) ||
320 JSON.stringify(kb.storageConfiguration ?? null) !==
321 JSON.stringify(news.storageConfiguration ?? null);
322 if (drifted) {
323 yield* bedrock.updateKnowledgeBase({
324 knowledgeBaseId: kb.knowledgeBaseId,
325 name,
326 description: news.description,
327 roleArn: news.roleArn,
328 knowledgeBaseConfiguration: news.knowledgeBaseConfiguration,
329 storageConfiguration: news.storageConfiguration,
330 });
331 kb = (yield* waitForSettled(kb.knowledgeBaseId)) ?? kb;
332 }
333 }
334
335 const knowledgeBaseId = kb.knowledgeBaseId;
336 const knowledgeBaseArn = kb.knowledgeBaseArn;
337
338 // 3b. SYNC TAGS against observed cloud tags.
339 const observedTags = yield* fetchObservedTags(knowledgeBaseArn);
340 const { upsert, removed } = diffTags(observedTags, desiredTags);
341 if (upsert.length > 0) {
342 yield* bedrock.tagResource({
343 resourceArn: knowledgeBaseArn,
344 tags: Object.fromEntries(
345 upsert.map(({ Key, Value }) => [Key, Value]),
346 ),
347 });
348 }
349 if (removed.length > 0) {
350 yield* bedrock.untagResource({
351 resourceArn: knowledgeBaseArn,
352 tagKeys: removed,
353 });
354 }
355
356 yield* session.note(knowledgeBaseArn);
357 return {
358 knowledgeBaseId,
359 knowledgeBaseArn,
360 name,
361 roleArn: news.roleArn,
362 };
363 }),
364
365 delete: Effect.fn(function* ({ output }) {
366 yield* retryWhileConflict(
367 bedrock.deleteKnowledgeBase({
368 knowledgeBaseId: output.knowledgeBaseId,
369 }),
370 ).pipe(
371 Effect.catchTag("ResourceNotFoundException", () => Effect.void),
372 );
373 }),
374 });
375 }),
376 );