1import * as Bedrock from "@/AWS/Bedrock";
2import * as Lambda from "@/AWS/Lambda";
3import * as Duration from "effect/Duration";
4import * as Effect from "effect/Effect";
5import * as Schema from "effect/Schema";
6import * as Stream from "effect/Stream";
7import { LanguageModel as AiLanguageModel, Tool, Toolkit } from "effect/ai";
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, "language-model-handler.ts");
13
14// Nova Micro via the us cross-region inference profile: the cheapest
15// on-demand conversational model in the testing account, with full
16// Converse tool-use + streaming support.
17const MODEL = "us.amazon.nova-micro-v1:0";
18// A second bound model to exercise the per-call `modelId` override.
19const LITE_MODEL = "us.amazon.nova-lite-v1:0";
20
21const GetWeather = Tool.make("get_weather", {
22 description:
23 "Get the current weather for a city. Always call this tool when the user asks about the weather.",
24 parameters: Schema.Struct({
25 city: Schema.String,
26 }),
27 success: Schema.Struct({
28 city: Schema.String,
29 temperatureF: Schema.Number,
30 condition: Schema.String,
31 }),
32});
33
34const WeatherToolkit = Toolkit.make(GetWeather);
35
36const WeatherToolkitLayer = WeatherToolkit.toLayer({
37 get_weather: ({ city }) =>
38 Effect.succeed({
39 city,
40 temperatureF: 72,
41 condition: "sunny",
42 }),
43});
44
45const toSse = (parts: Iterable<unknown>): string =>
46 [...parts].map((part) => `data: ${JSON.stringify(part)}\n\n`).join("");
47
48export class BedrockLanguageModelFunction extends Lambda.Function<Lambda.Function>()(
49 "BedrockLanguageModelFunction",
50) {}
51
52export default BedrockLanguageModelFunction.make(
53 {
54 main,
55 functionUrl: true,
56 // Model inference regularly exceeds Lambda's 3s default timeout.
57 timeout: Duration.seconds(120),
58 },
59 Effect.gen(function* () {
60 const model = yield* Bedrock.LanguageModel([MODEL, LITE_MODEL], {
61 parameters: { maxTokens: 1024, temperature: 0.2 },
62 });
63
64 return {
65 fetch: Effect.gen(function* () {
66 const request = yield* HttpServerRequest;
67 const url = new URL(request.originalUrl);
68 const pathname = url.pathname;
69 const prompt =
70 url.searchParams.get("prompt") ??
71 "Say the single word 'pong' and nothing else.";
72
73 // Cheap readiness route — no Bedrock call.
74 if (pathname === "/ping") {
75 return yield* HttpServerResponse.json({ ok: true });
76 }
77
78 if (pathname === "/generate") {
79 const response = yield* AiLanguageModel.generateText({ prompt });
80 return yield* HttpServerResponse.json({
81 text: response.text,
82 finishReason: response.finishReason,
83 usage: {
84 inputTokens: response.usage.inputTokens.total,
85 outputTokens: response.usage.outputTokens.total,
86 },
87 });
88 }
89
90 if (pathname === "/generate-short") {
91 // Runtime override: clamp the same bound model to a tiny budget.
92 const response = yield* AiLanguageModel.generateText({
93 prompt,
94 }).pipe(
95 Bedrock.withModelParameters({ maxTokens: 8, temperature: 0 }),
96 );
97 return yield* HttpServerResponse.json({
98 text: response.text,
99 finishReason: response.finishReason,
100 outputTokens: response.usage.outputTokens.total,
101 });
102 }
103
104 if (pathname === "/generate-lite") {
105 // Runtime override: route this call to the second bound model.
106 const response = yield* AiLanguageModel.generateText({
107 prompt,
108 }).pipe(Bedrock.withModelParameters({ modelId: LITE_MODEL }));
109 return yield* HttpServerResponse.json({
110 text: response.text,
111 finishReason: response.finishReason,
112 });
113 }
114
115 if (pathname === "/stream") {
116 // Collected server-side: Lambda function URLs buffer responses by
117 // default, and the tests assert on the part sequence, not
118 // incremental delivery.
119 const parts = yield* Stream.runCollect(
120 AiLanguageModel.streamText({ prompt }),
121 );
122 return HttpServerResponse.text(toSse(parts), {
123 headers: { "content-type": "text/event-stream" },
124 });
125 }
126
127 if (pathname === "/tool") {
128 const response = yield* AiLanguageModel.generateText({
129 prompt,
130 toolkit: WeatherToolkit,
131 toolChoice: "required",
132 }).pipe(Effect.provide(WeatherToolkitLayer));
133 return yield* HttpServerResponse.json({
134 text: response.text,
135 finishReason: response.finishReason,
136 toolCalls: response.toolCalls.map((call) => ({
137 id: call.id,
138 name: call.name,
139 params: call.params,
140 })),
141 toolResults: response.toolResults.map((result) => ({
142 id: result.id,
143 name: result.name,
144 result: result.result,
145 isFailure: result.isFailure,
146 })),
147 });
148 }
149
150 if (pathname === "/tool-stream") {
151 const parts = yield* Stream.runCollect(
152 AiLanguageModel.streamText({
153 prompt,
154 toolkit: WeatherToolkit,
155 toolChoice: "required",
156 }).pipe(Stream.provide(WeatherToolkitLayer)),
157 );
158 return HttpServerResponse.text(toSse(parts), {
159 headers: { "content-type": "text/event-stream" },
160 });
161 }
162
163 return yield* HttpServerResponse.json(
164 { error: "Not found", pathname },
165 { status: 404 },
166 );
167 }).pipe(
168 // Surface adapter/model failures as a JSON 500 so live-test runs can
169 // read the failure without digging through CloudWatch.
170 Effect.catchTag("AiError", (error) =>
171 HttpServerResponse.json(
172 { error: String(error.message) },
173 { status: 500 },
174 ),
175 ),
176 Effect.provide(model),
177 Effect.orDie,
178 ),
179 };
180 }).pipe(Effect.provide(Bedrock.LanguageModelHttp)),
181);