1import * as Neon from "@/Neon";
2import * as Config from "effect/Config";
3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";
5import * as Schema from "effect/Schema";
6import * as Stream from "effect/Stream";
7import { LanguageModel, Tool, Toolkit } from "effect/ai";
8import { HttpServerRequest } from "effect/http/HttpServerRequest";
9import * as HttpServerResponse from "effect/http/HttpServerResponse";
10import { languageModelGateway } from "./language-model-resources.ts";
11
12const Sum = Tool.make("sum", {
13 description: "Add two numbers.",
14 parameters: Schema.Struct({ a: Schema.Number, b: Schema.Number }),
15 success: Schema.Number,
16});
17const Tools = Toolkit.make(Sum);
18const handlers = Tools.toLayer({ sum: ({ a, b }) => Effect.succeed(a + b) });
19
20export const languageModelHandler = (source = languageModelGateway) =>
21 Effect.gen(function* () {
22 const gateway = yield* Neon.QueryAIGateway(source);
23 const model = Layer.unwrap(
24 Config.String("AI_MODEL").pipe(
25 Effect.map((model) => gateway.model({ model })),
26 ),
27 );
28 return {
29 fetch: Effect.gen(function* () {
30 const request = yield* HttpServerRequest;
31 if (request.url.startsWith("/object")) {
32 const response = yield* LanguageModel.generateObject({
33 prompt: "Return a JSON object with greeting equal to hello.",
34 schema: Schema.Struct({ greeting: Schema.String }),
35 });
36 return yield* HttpServerResponse.json(response.value);
37 }
38 if (request.url.startsWith("/tool")) {
39 const response = yield* LanguageModel.generateText({
40 prompt: "Use sum to add 2 and 3.",
41 toolkit: Tools,
42 toolChoice: { tool: "sum" },
43 }).pipe(Effect.provide(handlers));
44 return yield* HttpServerResponse.json({
45 results: response.toolResults.map((part) => part.result),
46 });
47 }
48 if (request.url.startsWith("/stream")) {
49 const parts = yield* LanguageModel.streamText({
50 prompt: "Say hello.",
51 }).pipe(Stream.runCollect);
52 return yield* HttpServerResponse.json({
53 text: parts
54 .filter((part) => part.type === "text-delta")
55 .map((part) => part.delta)
56 .join(""),
57 finished: parts.some((part) => part.type === "finish"),
58 });
59 }
60 const response = yield* LanguageModel.generateText({
61 prompt: "Say hello.",
62 });
63 return yield* HttpServerResponse.json({ text: response.text });
64 }).pipe(
65 Effect.provide(model),
66 Effect.catchTag("AiError", (error) =>
67 HttpServerResponse.json(
68 { reason: error.reason._tag },
69 { status: 502 },
70 ),
71 ),
72 Effect.orDie,
73 ),
74 };
75 });