1import * as Neon from "alchemy/Neon";
2import type { RuntimeContext } from "alchemy/RuntimeContext";
3import * as Config from "effect/Config";
4import * as Effect from "effect/Effect";
5import * as Schema from "effect/Schema";
6import type * as Scope from "effect/Scope";
7import * as Stream from "effect/Stream";
8import * as LanguageModel from "effect/ai/LanguageModel";
9import * as Sse from "effect/encoding/Sse";
10import * as HttpServerRequest from "effect/http/HttpServerRequest";
11import * as HttpServerResponse from "effect/http/HttpServerResponse";
12import { resources } from "../resources.ts";
13import { gateway } from "./resources.ts";
14
15/** Authenticated, opt-in text streaming backed by the bound Effect model. */
16export const chat = (
17 ai: Pick<Neon.QueryAIGatewayClient, "model">,
18 environment: Readonly<Record<string, string | undefined>>,
19) =>
20 Effect.gen(function* () {
21 const request = yield* HttpServerRequest.HttpServerRequest;
22 if (
23 !environment.NEON_EXAMPLE_API_KEY ||
24 request.headers.authorization !==
25 `Bearer ${environment.NEON_EXAMPLE_API_KEY}`
26 ) {
27 return HttpServerResponse.text("Unauthorized", { status: 401 });
28 }
29 const body = yield* request.json.pipe(
30 Effect.flatMap(
31 Schema.decodeUnknownEffect(Schema.Struct({ prompt: Schema.String })),
32 ),
33 Effect.catch(() => Effect.succeed(undefined)),
34 );
35 if (!body)
36 return HttpServerResponse.text("A JSON prompt is required", {
37 status: 400,
38 });
39 if (body.prompt.trim().length === 0 || body.prompt.length > 4000) {
40 return HttpServerResponse.text(
41 "Expected a nonempty prompt of at most 4000 characters",
42 { status: 400 },
43 );
44 }
45 if (
46 environment.NEON_AI_ALLOW_PAID !== "true" ||
47 !environment.NEON_AI_MODEL
48 ) {
49 return HttpServerResponse.text(
50 "Explicit paid inference opt-in and a configured model are required",
51 { status: 503 },
52 );
53 }
54
55 const native = yield* Neon.FunctionRequest;
56 const model = ai.model({
57 model: environment.NEON_AI_MODEL,
58 parameters: { maxTokens: 128 },
59 });
60 // The response body is consumed after the handler returns, in the transferred request scope.
61 const context = yield* Effect.context<RuntimeContext | Scope.Scope>();
62 const events = LanguageModel.streamText({ prompt: body.prompt }).pipe(
63 Stream.flatMap((part): Stream.Stream<Record<string, unknown>> => {
64 switch (part.type) {
65 case "text-start":
66 case "text-end":
67 return Stream.succeed({ type: part.type, id: part.id });
68 case "text-delta":
69 return Stream.succeed({
70 type: "text-delta",
71 id: part.id,
72 delta: part.delta,
73 });
74 case "finish":
75 return Stream.make(
76 { type: "finish-step" },
77 { type: "finish", finishReason: part.reason },
78 );
79 default:
80 return Stream.empty;
81 }
82 }),
83 Stream.provide(model),
84 Stream.provideContext(context),
85 Stream.catchTag("AiError", () =>
86 Stream.succeed({
87 type: "error",
88 errorText:
89 "Gateway rejected the request. Check model access, paid-plan entitlement and prepaid credits. No credits were purchased.",
90 }),
91 ),
92 );
93 const stream = Stream.concat(
94 Stream.make({ type: "start" }, { type: "start-step" }),
95 events,
96 ).pipe(
97 Stream.map((event) => JSON.stringify(event)),
98 Stream.concat(Stream.succeed("[DONE]")),
99 Stream.map((data): Sse.Event => ({
100 _tag: "Event",
101 event: "message",
102 id: undefined,
103 data,
104 })),
105 Stream.pipeThroughChannel(Sse.encode<never, unknown>()),
106 Stream.encodeText,
107 Stream.interruptWhen(
108 Effect.callback<void>((resume) => {
109 const abort = () => resume(Effect.void);
110 if (native.signal.aborted) abort();
111 else native.signal.addEventListener("abort", abort, { once: true });
112 return Effect.sync(() =>
113 native.signal.removeEventListener("abort", abort),
114 );
115 }),
116 ),
117 );
118 return HttpServerResponse.stream(stream, {
119 headers: {
120 "content-type": "text/event-stream",
121 "cache-control": "no-cache",
122 "x-vercel-ai-ui-message-stream": "v1",
123 },
124 });
125 });
126
127export default class EffectApi extends Neon.Function<EffectApi>()(
128 "EffectAI",
129 Effect.gen(function* () {
130 return {
131 branch: (yield* resources).branch,
132 main: import.meta.url,
133 env: {
134 NEON_EXAMPLE_API_KEY: yield* Config.Redacted("NEON_EXAMPLE_API_KEY"),
135 NEON_AI_MODEL: yield* Config.String("NEON_AI_MODEL"),
136 NEON_AI_ALLOW_PAID: yield* Config.String("NEON_AI_ALLOW_PAID").pipe(
137 Config.withDefault("false"),
138 ),
139 },
140 };
141 }),
142 Effect.gen(function* () {
143 const ai = yield* Neon.QueryAIGateway(gateway);
144 const environment = yield* Neon.FunctionEnvironment;
145 return { fetch: chat(ai, environment) };
146 }).pipe(Effect.provide(Neon.QueryAIGatewayHttp)),
147) {}