1import * as Lambda from "@/AWS/Lambda";
2import * as S3Vectors from "@/AWS/S3Vectors";
3import * as Duration from "effect/Duration";
4import * as Effect from "effect/Effect";
5import * as Layer from "effect/Layer";
6import { HttpServerRequest } from "effect/http/HttpServerRequest";
7import * as HttpServerResponse from "effect/http/HttpServerResponse";
8import path from "pathe";
9
10const main = path.resolve(import.meta.dirname, "vectors-handler.ts");
11
12export class VectorsTestFunction extends Lambda.Function<Lambda.Function>()(
13 "S3VectorsTestFunction",
14) {}
15
16export default VectorsTestFunction.make(
17 {
18 main,
19 functionUrl: true,
20 timeout: Duration.seconds(30),
21 },
22 Effect.gen(function* () {
23 const bucket = yield* S3Vectors.VectorBucket("VBucket", {});
24 const index = yield* S3Vectors.Index("VIndex", {
25 vectorBucketName: bucket.vectorBucketName,
26 dimension: 4,
27 distanceMetric: "cosine",
28 });
29 // Least-privilege split bindings for writes and reads, plus the
30 // ReadWrite client to prove the composed level end-to-end.
31 const writer = yield* S3Vectors.VectorsWrite(index);
32 const reader = yield* S3Vectors.VectorsRead(index);
33 const vectors = yield* S3Vectors.Vectors(index);
34
35 return {
36 fetch: Effect.gen(function* () {
37 const request = yield* HttpServerRequest;
38 const url = new URL(request.originalUrl);
39 const pathname = url.pathname;
40
41 if (request.method === "GET" && pathname === "/ping") {
42 return yield* HttpServerResponse.json({ ok: true });
43 }
44
45 if (request.method === "GET" && pathname === "/put") {
46 yield* writer.put({
47 vectors: [
48 { key: "a", data: { float32: [1, 0, 0, 0] } },
49 { key: "b", data: { float32: [0, 1, 0, 0] } },
50 { key: "c", data: { float32: [0.9, 0.1, 0, 0] } },
51 ],
52 });
53 return yield* HttpServerResponse.json({ put: 3 });
54 }
55
56 if (request.method === "GET" && pathname === "/query") {
57 const result = yield* reader.query({
58 topK: 2,
59 queryVector: { float32: [1, 0, 0, 0] },
60 returnDistance: true,
61 });
62 return yield* HttpServerResponse.json({
63 keys: result.vectors.map((v) => v.key),
64 distanceMetric: result.distanceMetric,
65 });
66 }
67
68 if (request.method === "GET" && pathname === "/get") {
69 const result = yield* reader.get({
70 keys: ["a"],
71 returnData: true,
72 });
73 return yield* HttpServerResponse.json({
74 keys: result.vectors.map((v) => v.key),
75 });
76 }
77
78 if (request.method === "GET" && pathname === "/list") {
79 const result = yield* reader.list({});
80 return yield* HttpServerResponse.json({
81 keys: result.vectors.map((v) => v.key),
82 });
83 }
84
85 if (request.method === "GET" && pathname === "/delete") {
86 yield* writer.delete({ keys: ["b"] });
87 return yield* HttpServerResponse.json({ deleted: "b" });
88 }
89
90 if (request.method === "GET" && pathname === "/rw") {
91 // Round-trip through the ReadWrite client.
92 yield* vectors.put({
93 vectors: [{ key: "rw", data: { float32: [0, 0, 1, 0] } }],
94 });
95 const result = yield* vectors.get({ keys: ["rw"] });
96 return yield* HttpServerResponse.json({
97 keys: result.vectors.map((v) => v.key),
98 });
99 }
100
101 return yield* HttpServerResponse.json(
102 { error: "Not found", method: request.method, pathname },
103 { status: 404 },
104 );
105 }).pipe(Effect.orDie),
106 };
107 }).pipe(
108 Effect.provide(
109 Layer.mergeAll(
110 S3Vectors.VectorsHttp,
111 S3Vectors.VectorsReadHttp,
112 S3Vectors.VectorsWriteHttp,
113 ),
114 ),
115 ),
116);