1import * as bedrock from "@distilled.cloud/aws/bedrock-agent";
2import * as Effect from "effect/Effect";
3import * as Schedule from "effect/Schedule";
4import { isResolved } from "../../Diff.ts";
5import { createPhysicalName } from "../../PhysicalName.ts";
6import * as Provider from "../../Provider.ts";
7import { Resource } from "../../Resource.ts";
8
9/**
10 * The connector configuration for a data source. Passed through to the
11 * Bedrock API unchanged — see the AWS SDK `DataSourceConfiguration` shape.
12 * The common case is a `S3` type with `s3Configuration.bucketArn` set.
13 */
14export type DataSourceConfiguration = bedrock.DataSourceConfiguration;
15
16/**
17 * Controls document chunking, parsing, and custom transformation before
18 * embedding. Passed through unchanged — see the AWS SDK
19 * `VectorIngestionConfiguration` shape.
20 */
21export type VectorIngestionConfiguration = bedrock.VectorIngestionConfiguration;
22
23/**
24 * What happens to the vector embeddings when the data source is deleted.
25 * `DELETE` removes them; `RETAIN` leaves them in the vector store.
26 */
27export type DataDeletionPolicy = "RETAIN" | "DELETE";
28
29export interface DataSourceProps {
30 /**
31 * The id of the {@link KnowledgeBase} this data source feeds. Accepts a
32 * knowledge base's `knowledgeBaseId` output. Changing the knowledge base
33 * triggers a replacement.
34 */
35 knowledgeBaseId: string;
36 /**
37 * Name of the data source (1-100 characters). If omitted, a deterministic
38 * physical name is generated from the app, stage, and logical ID. Changing
39 * the name triggers a replacement.
40 */
41 name?: string;
42 /**
43 * A description of the data source.
44 */
45 description?: string;
46 /**
47 * The connector configuration (S3 bucket, web crawler, Confluence, ...).
48 */
49 dataSourceConfiguration: DataSourceConfiguration;
50 /**
51 * What happens to vector embeddings when the data source is deleted.
52 * @default "RETAIN"
53 */
54 dataDeletionPolicy?: DataDeletionPolicy;
55 /**
56 * Document chunking / parsing configuration applied before embedding.
57 */
58 vectorIngestionConfiguration?: VectorIngestionConfiguration;
59}
60
61export interface DataSource extends Resource<
62 "AWS.Bedrock.DataSource",
63 DataSourceProps,
64 {
65 /**
66 * The unique identifier of the knowledge base the data source belongs to.
67 */
68 knowledgeBaseId: string;
69 /**
70 * The unique identifier of the data source.
71 */
72 dataSourceId: string;
73 /**
74 * Name of the data source.
75 */
76 name: string;
77 }
78> {}
79
80/**
81 * A data source attached to an Amazon Bedrock {@link KnowledgeBase} — the
82 * origin of the documents the knowledge base embeds and indexes.
83 *
84 * The most common source is an S3 bucket. After the data source is created,
85 * start an ingestion job (`bedrock-agent:StartIngestionJob`) to crawl the
86 * source, chunk + embed the documents, and write them to the vector store.
87 * Ingestion is not part of the desired-state lifecycle — trigger it whenever
88 * the underlying documents change.
89 *
90 * ### Creating Data Sources
91 * **Example:** S3 Data Source
92 * ```typescript
93 * import * as Bedrock from "alchemy/AWS/Bedrock";
94 *
95 * const source = yield* Bedrock.DataSource("docs-bucket", {
96 * knowledgeBaseId: kb.knowledgeBaseId,
97 * dataSourceConfiguration: {
98 * type: "S3",
99 * s3Configuration: { bucketArn: bucket.bucketArn },
100 * },
101 * dataDeletionPolicy: "DELETE",
102 * });
103 * ```
104 *
105 * @resource
106 */
107export const DataSource = Resource<DataSource>("AWS.Bedrock.DataSource");
108
109/** Data-source status values indicating an in-flight transition. */
110const DS_TRANSIENT = new Set(["CREATING", "UPDATING", "DELETING"]);
111
112/**
113 * A data source rejects deletion with a `ConflictException` while the parent
114 * knowledge base is busy. Bounded retry (~60s), wrapped in an explicitly-typed
115 * helper so the `Effect.retry` return type does not widen the provider layer's
116 * requirement to `unknown` (see PATTERNS §7).
117 */
118const retryWhileConflict = <A, E extends { _tag: string }, R>(
119 self: Effect.Effect<A, E, R>,
120): Effect.Effect<A, E, R> =>
121 Effect.retry(self, {
122 while: (e) => e._tag === "ConflictException",
123 schedule: Schedule.max([Schedule.fixed("3 seconds"), Schedule.recurs(20)]),
124 });
125
126export const DataSourceProvider = () =>
127 Provider.effect(
128 DataSource,
129 Effect.gen(function* () {
130 const createName = Effect.fn(function* (
131 id: string,
132 props: Pick<DataSourceProps, "name">,
133 ) {
134 return (
135 props.name ?? (yield* createPhysicalName({ id, maxLength: 100 }))
136 );
137 });
138
139 const getDsOrUndefined = Effect.fn(function* (
140 knowledgeBaseId: string,
141 dataSourceId: string,
142 ) {
143 return yield* bedrock
144 .getDataSource({ knowledgeBaseId, dataSourceId })
145 .pipe(
146 Effect.map((r) => r.dataSource),
147 Effect.catchTag("ResourceNotFoundException", () =>
148 Effect.succeed(undefined),
149 ),
150 );
151 });
152
153 const waitForSettled = Effect.fn(function* (
154 knowledgeBaseId: string,
155 dataSourceId: string,
156 ) {
157 return yield* bedrock
158 .getDataSource({ knowledgeBaseId, dataSourceId })
159 .pipe(
160 Effect.map((r) => r.dataSource),
161 Effect.catchTag("ResourceNotFoundException", () =>
162 Effect.succeed(undefined),
163 ),
164 Effect.repeat({
165 schedule: Schedule.fixed("3 seconds"),
166 until: (ds) => ds === undefined || !DS_TRANSIENT.has(ds.status),
167 times: 40,
168 }),
169 );
170 });
171
172 return DataSource.Provider.of({
173 stables: ["knowledgeBaseId", "dataSourceId", "name"],
174
175 // Data sources are scoped to a parent knowledge base — enumeration
176 // requires the KB id, so this returns empty (the engine keys off state).
177 list: () => Effect.succeed([]),
178
179 read: Effect.fn(function* ({ olds, output }) {
180 const knowledgeBaseId =
181 output?.knowledgeBaseId ?? olds?.knowledgeBaseId;
182 const dataSourceId = output?.dataSourceId;
183 if (knowledgeBaseId === undefined || dataSourceId === undefined) {
184 return undefined;
185 }
186 const ds = yield* getDsOrUndefined(knowledgeBaseId, dataSourceId);
187 if (ds === undefined || ds.status === "DELETING") return undefined;
188 return {
189 knowledgeBaseId: ds.knowledgeBaseId,
190 dataSourceId: ds.dataSourceId,
191 name: ds.name,
192 };
193 }),
194
195 diff: Effect.fn(function* ({ id, news, olds }) {
196 if (!isResolved(news)) return undefined;
197 if (
198 (olds?.knowledgeBaseId ?? undefined) !==
199 (news?.knowledgeBaseId ?? undefined)
200 ) {
201 return { action: "replace" } as const;
202 }
203 const oldName = yield* createName(
204 id,
205 olds ?? ({} as DataSourceProps),
206 );
207 const newName = yield* createName(
208 id,
209 news ?? ({} as DataSourceProps),
210 );
211 if (oldName !== newName) {
212 return { action: "replace" } as const;
213 }
214 // description, config, and deletion policy converge via update.
215 }),
216
217 reconcile: Effect.fn(function* ({
218 id,
219 news = {} as DataSourceProps,
220 output,
221 session,
222 }) {
223 const name = output?.name ?? (yield* createName(id, news));
224 const knowledgeBaseId = news.knowledgeBaseId;
225
226 // 1. OBSERVE
227 let ds = output?.dataSourceId
228 ? yield* getDsOrUndefined(knowledgeBaseId, output.dataSourceId)
229 : undefined;
230
231 if (ds === undefined) {
232 // 2. ENSURE
233 const created = yield* bedrock.createDataSource({
234 knowledgeBaseId,
235 name,
236 description: news.description,
237 dataSourceConfiguration: news.dataSourceConfiguration,
238 dataDeletionPolicy: news.dataDeletionPolicy,
239 vectorIngestionConfiguration: news.vectorIngestionConfiguration,
240 });
241 ds = created.dataSource;
242 ds =
243 (yield* waitForSettled(knowledgeBaseId, ds.dataSourceId)) ?? ds;
244 } else {
245 // 3. SYNC
246 ds =
247 (yield* waitForSettled(knowledgeBaseId, ds.dataSourceId)) ?? ds;
248 const drifted =
249 ds.description !== news.description ||
250 JSON.stringify(ds.dataSourceConfiguration) !==
251 JSON.stringify(news.dataSourceConfiguration) ||
252 (ds.dataDeletionPolicy ?? undefined) !==
253 (news.dataDeletionPolicy ?? undefined) ||
254 JSON.stringify(ds.vectorIngestionConfiguration ?? null) !==
255 JSON.stringify(news.vectorIngestionConfiguration ?? null);
256 if (drifted) {
257 yield* bedrock.updateDataSource({
258 knowledgeBaseId,
259 dataSourceId: ds.dataSourceId,
260 name,
261 description: news.description,
262 dataSourceConfiguration: news.dataSourceConfiguration,
263 dataDeletionPolicy: news.dataDeletionPolicy,
264 vectorIngestionConfiguration: news.vectorIngestionConfiguration,
265 });
266 ds =
267 (yield* waitForSettled(knowledgeBaseId, ds.dataSourceId)) ?? ds;
268 }
269 }
270
271 yield* session.note(ds.dataSourceId);
272 return {
273 knowledgeBaseId,
274 dataSourceId: ds.dataSourceId,
275 name,
276 };
277 }),
278
279 delete: Effect.fn(function* ({ output }) {
280 yield* retryWhileConflict(
281 bedrock.deleteDataSource({
282 knowledgeBaseId: output.knowledgeBaseId,
283 dataSourceId: output.dataSourceId,
284 }),
285 ).pipe(
286 Effect.catchTag("ResourceNotFoundException", () => Effect.void),
287 );
288 }),
289 });
290 }),
291 );