1import * as AWS from "@/AWS";
2import { DataSource, KnowledgeBase } from "@/AWS/Bedrock";
3import * as S3 from "@/AWS/S3";
4import * as Test from "@/Test/Alchemy";
5import * as bedrock from "@distilled.cloud/aws/bedrock-agent";
6import { expect } from "alchemy-test";
7import * as Data from "effect/Data";
8import * as Effect from "effect/Effect";
9import * as Schedule from "effect/Schedule";
10import * as HttpClient from "effect/http/HttpClient";
11import BedrockKbTestFunctionLive, { BedrockKbTestFunction } from "./kb-handler";
12
13const { test } = Test.make({ providers: AWS.providers() });
14
15// The KnowledgeBase + DataSource live lifecycle needs a pre-existing vector
16// store. Provisioning an OpenSearch Serverless collection is cheap, but the
17// vector INDEX must be created through the collection's data-plane API
18// (SigV4-signed HTTP to the collection endpoint) — there is no AWS control
19// -plane call for it — which is out of scope for this suite. Supply a ready
20// collection + index (and a KB service role that can read the embedding model,
21// the collection, and S3) via env vars to run this against real infra.
22//
23// AWS_TEST_SLOW=1
24// BEDROCK_KB_ROLE_ARN=arn:aws:iam::...:role/...
25// BEDROCK_KB_COLLECTION_ARN=arn:aws:aoss:...:collection/...
26// BEDROCK_KB_INDEX_NAME=bedrock-index
27// BEDROCK_KB_EMBEDDING_MODEL_ARN=arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0
28const roleArn = process.env.BEDROCK_KB_ROLE_ARN;
29const collectionArn = process.env.BEDROCK_KB_COLLECTION_ARN;
30const indexName = process.env.BEDROCK_KB_INDEX_NAME ?? "bedrock-index";
31const embeddingModelArn = process.env.BEDROCK_KB_EMBEDDING_MODEL_ARN;
32
33const gated =
34 !process.env.AWS_TEST_SLOW ||
35 !roleArn ||
36 !collectionArn ||
37 !embeddingModelArn;
38
39const findKb = (knowledgeBaseId: string) =>
40 bedrock.getKnowledgeBase({ knowledgeBaseId }).pipe(
41 Effect.map((r) => r.knowledgeBase),
42 Effect.catchTag("ResourceNotFoundException", () =>
43 Effect.succeed(undefined),
44 ),
45 );
46
47class StillExists extends Data.TaggedError("StillExists")<{
48 readonly id: string;
49}> {}
50
51const assertKbGone = (knowledgeBaseId: string) =>
52 findKb(knowledgeBaseId).pipe(
53 Effect.flatMap((kb) =>
54 kb === undefined || kb.status === "DELETING"
55 ? Effect.void
56 : Effect.fail(new StillExists({ id: knowledgeBaseId })),
57 ),
58 Effect.retry({
59 while: (e) => e._tag === "StillExists",
60 schedule: Schedule.max([Schedule.exponential(2000), Schedule.recurs(10)]),
61 }),
62 );
63
64test.provider.skipIf(gated)(
65 "create knowledge base + S3 data source, then delete",
66 (stack) =>
67 Effect.gen(function* () {
68 yield* stack.destroy();
69
70 const result = yield* stack.deploy(
71 Effect.gen(function* () {
72 const bucket = yield* S3.Bucket("kb-docs", { forceDestroy: true });
73 const kb = yield* KnowledgeBase("TestKb", {
74 roleArn: roleArn!,
75 knowledgeBaseConfiguration: {
76 type: "VECTOR",
77 vectorKnowledgeBaseConfiguration: {
78 embeddingModelArn: embeddingModelArn!,
79 },
80 },
81 storageConfiguration: {
82 type: "OPENSEARCH_SERVERLESS",
83 opensearchServerlessConfiguration: {
84 collectionArn: collectionArn!,
85 vectorIndexName: indexName,
86 fieldMapping: {
87 vectorField: "bedrock-vector",
88 textField: "bedrock-text",
89 metadataField: "bedrock-metadata",
90 },
91 },
92 },
93 tags: { Environment: "test" },
94 });
95 const source = yield* DataSource("TestSource", {
96 knowledgeBaseId: kb.knowledgeBaseId,
97 dataSourceConfiguration: {
98 type: "S3",
99 s3Configuration: { bucketArn: bucket.bucketArn },
100 },
101 dataDeletionPolicy: "DELETE",
102 });
103 return {
104 knowledgeBaseId: kb.knowledgeBaseId,
105 knowledgeBaseArn: kb.knowledgeBaseArn,
106 dataSourceId: source.dataSourceId,
107 };
108 }),
109 );
110
111 const kb = yield* findKb(result.knowledgeBaseId);
112 expect(kb?.status).toBe("ACTIVE");
113
114 const ds = yield* bedrock
115 .getDataSource({
116 knowledgeBaseId: result.knowledgeBaseId,
117 dataSourceId: result.dataSourceId,
118 })
119 .pipe(Effect.map((r) => r.dataSource));
120 expect(ds.status).toBe("AVAILABLE");
121
122 yield* stack.destroy();
123 yield* assertKbGone(result.knowledgeBaseId);
124 }),
125 {
126 tags: ["provider:aws", "provider:aws:bedrock", "provider:aws:s3", "live"],
127 timeout: 600_000,
128 },
129);
130
131// Runtime coverage for the KB-scoped bindings (document ingestion, ingestion
132// jobs, retrieve, streaming RAG) through a deployed Lambda — gated behind the
133// same vector-store env prerequisites as the lifecycle test above.
134test.provider.skipIf(gated)(
135 "KB-scoped bindings work from a deployed Lambda",
136 (stack) =>
137 Effect.gen(function* () {
138 yield* stack.destroy();
139
140 const { functionUrl } = yield* stack.deploy(
141 Effect.gen(function* () {
142 return yield* BedrockKbTestFunction;
143 }).pipe(Effect.provide(BedrockKbTestFunctionLive)),
144 );
145 const baseUrl = functionUrl!.replace(/\/+$/, "");
146
147 // Fresh function URLs take a few seconds to start serving 200s.
148 yield* HttpClient.get(`${baseUrl}/ping`).pipe(
149 Effect.flatMap((r) =>
150 r.status === 200
151 ? Effect.succeed(r)
152 : Effect.fail(new Error(`not ready: ${r.status}`)),
153 ),
154 Effect.retry({
155 schedule: Schedule.max([
156 Schedule.fixed("2 seconds"),
157 Schedule.recurs(75),
158 ]),
159 }),
160 );
161
162 const getJson = (path: string) =>
163 HttpClient.get(`${baseUrl}${path}`).pipe(Effect.flatMap((r) => r.json));
164 const postJson = (path: string) =>
165 HttpClient.post(`${baseUrl}${path}`).pipe(
166 Effect.flatMap((r) => r.json),
167 );
168
169 // 1. Direct document ingestion into the CUSTOM data source.
170 const ingested = (yield* postJson("/ingest")) as { status?: string };
171 expect(ingested.status).toBeDefined();
172
173 // 2. Poll the document to INDEXED.
174 const docStatus = (yield* getJson("/doc-status").pipe(
175 Effect.repeat({
176 schedule: Schedule.spaced("5 seconds"),
177 until: (r): boolean =>
178 (r as { status?: string }).status === "INDEXED" ||
179 (r as { status?: string }).status === "FAILED",
180 times: 36,
181 }),
182 )) as { status?: string };
183 expect(docStatus.status).toBe("INDEXED");
184
185 // 3. The document is tracked.
186 const docs = (yield* getJson("/docs")) as { count: number };
187 expect(docs.count).toBeGreaterThan(0);
188
189 // 4. Semantic retrieval finds the ingested passage.
190 const retrieved = (yield* getJson("/retrieve").pipe(
191 Effect.repeat({
192 schedule: Schedule.spaced("5 seconds"),
193 until: (r): boolean => (r as { count: number }).count > 0,
194 times: 12,
195 }),
196 )) as { count: number; passages: string[] };
197 expect(retrieved.count).toBeGreaterThan(0);
198
199 // 5. Streaming RAG produces a grounded, non-empty answer.
200 const rag = (yield* getJson("/rag-stream")) as {
201 events: number;
202 text: string;
203 };
204 expect(rag.events).toBeGreaterThan(0);
205 expect(rag.text.trim().length).toBeGreaterThan(0);
206
207 // 6. Ingestion-job lifecycle on the S3 data source.
208 const sync = (yield* postJson("/sync")) as { jobId: string };
209 expect(sync.jobId).toBeDefined();
210 const job = (yield* getJson(`/job?id=${sync.jobId}`).pipe(
211 Effect.repeat({
212 schedule: Schedule.spaced("5 seconds"),
213 until: (r): boolean =>
214 (r as { status: string }).status === "COMPLETE" ||
215 (r as { status: string }).status === "FAILED",
216 times: 36,
217 }),
218 )) as { status: string };
219 expect(job.status).toBe("COMPLETE");
220 const jobs = (yield* getJson("/jobs")) as { count: number };
221 expect(jobs.count).toBeGreaterThan(0);
222
223 // 7. Delete the document again.
224 const deleted = (yield* postJson("/delete-doc")) as { status?: string };
225 expect(deleted.status).toBeDefined();
226
227 yield* stack.destroy();
228 }),
229 {
230 tags: [
231 "provider:aws",
232 "provider:aws:bedrock",
233 "provider:aws:lambda",
234 "provider:aws:s3",
235 "live",
236 ],
237 timeout: 600_000,
238 },
239);