1import * as Bedrock from "@/AWS/Bedrock";
2import * as Lambda from "@/AWS/Lambda";
3import * as S3 from "@/AWS/S3";
4import * as Duration from "effect/Duration";
5import * as Effect from "effect/Effect";
6import * as Layer from "effect/Layer";
7import * as Stream from "effect/Stream";
8import { HttpServerRequest } from "effect/http/HttpServerRequest";
9import * as HttpServerResponse from "effect/http/HttpServerResponse";
10import path from "pathe";
11
12const main = path.resolve(import.meta.dirname, "kb-handler.ts");
13
14// Same env-gated vector-store prerequisites as KnowledgeBase.slow.test.ts —
15// this fixture only ever deploys from that (skipIf-gated) suite.
16const roleArn = process.env.BEDROCK_KB_ROLE_ARN ?? "";
17const collectionArn = process.env.BEDROCK_KB_COLLECTION_ARN ?? "";
18const indexName = process.env.BEDROCK_KB_INDEX_NAME ?? "bedrock-index";
19const embeddingModelArn = process.env.BEDROCK_KB_EMBEDDING_MODEL_ARN ?? "";
20
21// Nova Micro via the us cross-region inference profile — the cheapest
22// on-demand generation model in the testing account.
23const MODEL = "us.amazon.nova-micro-v1:0";
24
25const DOC_ID = "welcome-doc";
26const DOC_TEXT =
27 "Alchemy is an Infrastructure-as-Effects framework that deploys cloud resources from typed Effect programs.";
28
29export class BedrockKbTestFunction extends Lambda.Function<Lambda.Function>()(
30 "BedrockKbTestFunction",
31) {}
32
33export default BedrockKbTestFunction.make(
34 {
35 main,
36 functionUrl: true,
37 // Retrieval + generation regularly exceed Lambda's 3s default timeout.
38 timeout: Duration.seconds(60),
39 },
40 Effect.gen(function* () {
41 const bucket = yield* S3.Bucket("kb-docs", { forceDestroy: true });
42 const kb = yield* Bedrock.KnowledgeBase("BindingsKb", {
43 roleArn,
44 knowledgeBaseConfiguration: {
45 type: "VECTOR",
46 vectorKnowledgeBaseConfiguration: {
47 embeddingModelArn,
48 },
49 },
50 storageConfiguration: {
51 type: "OPENSEARCH_SERVERLESS",
52 opensearchServerlessConfiguration: {
53 collectionArn,
54 vectorIndexName: indexName,
55 fieldMapping: {
56 vectorField: "bedrock-vector",
57 textField: "bedrock-text",
58 metadataField: "bedrock-metadata",
59 },
60 },
61 },
62 });
63 // S3-backed source exercises the ingestion-job bindings; CUSTOM source
64 // exercises direct document ingestion.
65 const s3Source = yield* Bedrock.DataSource("BindingsS3Source", {
66 knowledgeBaseId: kb.knowledgeBaseId,
67 dataSourceConfiguration: {
68 type: "S3",
69 s3Configuration: { bucketArn: bucket.bucketArn },
70 },
71 dataDeletionPolicy: "DELETE",
72 });
73 const customSource = yield* Bedrock.DataSource("BindingsCustomSource", {
74 knowledgeBaseId: kb.knowledgeBaseId,
75 dataSourceConfiguration: { type: "CUSTOM" },
76 dataDeletionPolicy: "DELETE",
77 });
78
79 const knowledgeBaseId = yield* kb.knowledgeBaseId;
80
81 const ingestDocuments =
82 yield* Bedrock.IngestKnowledgeBaseDocuments(customSource);
83 const getDocuments = yield* Bedrock.GetKnowledgeBaseDocuments(customSource);
84 const listDocuments =
85 yield* Bedrock.ListKnowledgeBaseDocuments(customSource);
86 const deleteDocuments =
87 yield* Bedrock.DeleteKnowledgeBaseDocuments(customSource);
88 const startIngestionJob = yield* Bedrock.StartIngestionJob(s3Source);
89 const getIngestionJob = yield* Bedrock.GetIngestionJob(s3Source);
90 const listIngestionJobs = yield* Bedrock.ListIngestionJobs(s3Source);
91 const retrieve = yield* Bedrock.Retrieve(kb);
92 const ragStream = yield* Bedrock.RetrieveAndGenerateStream(kb, MODEL);
93
94 return {
95 fetch: Effect.gen(function* () {
96 const request = yield* HttpServerRequest;
97 const url = new URL(request.originalUrl);
98 const pathname = url.pathname;
99
100 // Cheap readiness route — no Bedrock call.
101 if (request.method === "GET" && pathname === "/ping") {
102 return yield* HttpServerResponse.json({ ok: true });
103 }
104
105 if (request.method === "POST" && pathname === "/ingest") {
106 const { documentDetails } = yield* ingestDocuments({
107 documents: [
108 {
109 content: {
110 dataSourceType: "CUSTOM",
111 custom: {
112 customDocumentIdentifier: { id: DOC_ID },
113 sourceType: "IN_LINE",
114 inlineContent: {
115 type: "TEXT",
116 textContent: { data: DOC_TEXT },
117 },
118 },
119 },
120 },
121 ],
122 });
123 return yield* HttpServerResponse.json({
124 status: documentDetails?.[0]?.status,
125 });
126 }
127
128 if (request.method === "GET" && pathname === "/doc-status") {
129 const { documentDetails } = yield* getDocuments({
130 documentIdentifiers: [
131 { dataSourceType: "CUSTOM", custom: { id: DOC_ID } },
132 ],
133 });
134 return yield* HttpServerResponse.json({
135 status: documentDetails?.[0]?.status,
136 });
137 }
138
139 if (request.method === "GET" && pathname === "/docs") {
140 const { documentDetails } = yield* listDocuments({});
141 return yield* HttpServerResponse.json({
142 count: documentDetails?.length ?? 0,
143 });
144 }
145
146 if (request.method === "POST" && pathname === "/delete-doc") {
147 const { documentDetails } = yield* deleteDocuments({
148 documentIdentifiers: [
149 { dataSourceType: "CUSTOM", custom: { id: DOC_ID } },
150 ],
151 });
152 return yield* HttpServerResponse.json({
153 status: documentDetails?.[0]?.status,
154 });
155 }
156
157 if (request.method === "POST" && pathname === "/sync") {
158 const { ingestionJob } = yield* startIngestionJob({});
159 return yield* HttpServerResponse.json({
160 jobId: ingestionJob.ingestionJobId,
161 status: ingestionJob.status,
162 });
163 }
164
165 if (request.method === "GET" && pathname === "/job") {
166 const jobId = url.searchParams.get("id") ?? "";
167 const { ingestionJob } = yield* getIngestionJob({
168 ingestionJobId: jobId,
169 });
170 return yield* HttpServerResponse.json({
171 status: ingestionJob.status,
172 });
173 }
174
175 if (request.method === "GET" && pathname === "/jobs") {
176 const { ingestionJobSummaries } = yield* listIngestionJobs({});
177 return yield* HttpServerResponse.json({
178 count: ingestionJobSummaries.length,
179 });
180 }
181
182 if (request.method === "GET" && pathname === "/retrieve") {
183 const result = yield* retrieve({
184 retrievalQuery: { text: "What is Alchemy?" },
185 retrievalConfiguration: {
186 vectorSearchConfiguration: { numberOfResults: 3 },
187 },
188 });
189 return yield* HttpServerResponse.json({
190 count: result.retrievalResults.length,
191 passages: result.retrievalResults.map((r) => r.content?.text),
192 });
193 }
194
195 if (request.method === "GET" && pathname === "/rag-stream") {
196 // knowledgeBaseId is an Accessor<string> at init scope — yield it
197 // again at request time to resolve the concrete value.
198 const kbId = yield* knowledgeBaseId;
199 const result = yield* ragStream({
200 input: { text: "What is Alchemy?" },
201 retrieveAndGenerateConfiguration: {
202 type: "KNOWLEDGE_BASE",
203 knowledgeBaseConfiguration: {
204 knowledgeBaseId: kbId,
205 modelArn: MODEL,
206 },
207 },
208 });
209 const events = yield* Stream.runCollect(result.stream);
210 const text = Array.from(events)
211 .map((event) => event.output?.text ?? "")
212 .join("");
213 return yield* HttpServerResponse.json({
214 sessionId: result.sessionId,
215 events: events.length,
216 text,
217 });
218 }
219
220 return yield* HttpServerResponse.json(
221 { error: "Not found", method: request.method, pathname },
222 { status: 404 },
223 );
224 }).pipe(Effect.orDie),
225 };
226 }).pipe(
227 Effect.provide(
228 Layer.mergeAll(
229 Bedrock.DeleteKnowledgeBaseDocumentsHttp,
230 Bedrock.GetIngestionJobHttp,
231 Bedrock.GetKnowledgeBaseDocumentsHttp,
232 Bedrock.IngestKnowledgeBaseDocumentsHttp,
233 Bedrock.ListIngestionJobsHttp,
234 Bedrock.ListKnowledgeBaseDocumentsHttp,
235 Bedrock.RetrieveHttp,
236 Bedrock.RetrieveAndGenerateStreamHttp,
237 Bedrock.StartIngestionJobHttp,
238 ),
239 ),
240 ),
241);