1import * as Duration from "effect/Duration";
2import * as Effect from "effect/Effect";
3import * as Layer from "effect/Layer";
4import type * as Redacted from "effect/Redacted";
5import * as Schema from "effect/Schema";
6import * as Stream from "effect/Stream";
7import { AiError, LanguageModel, Response, Tool } from "effect/ai";
8import { toCodecOpenAI } from "effect/ai/OpenAiStructuredOutput";
9import * as Sse from "effect/encoding/Sse";
10import * as FetchHttpClient from "effect/http/FetchHttpClient";
11import * as HttpClient from "effect/http/HttpClient";
12import * as HttpClientRequest from "effect/http/HttpClientRequest";
13import type * as HttpClientResponse from "effect/http/HttpClientResponse";
14import type { RuntimeContext } from "../RuntimeContext.ts";
15
16/** Runtime-only connection obtained from QueryAIGateway, never an account API key. */
17export interface LanguageModelClient {
18 /** OpenAI-compatible base URL ending in /v1. */
19 readonly chatBaseUrl: Effect.Effect<string, never, RuntimeContext>;
20 /** Redacted credential with ai_gateway:invoke scope. */
21 readonly token: Effect.Effect<
22 Redacted.Redacted<string>,
23 never,
24 RuntimeContext
25 >;
26}
27
28/** Configuration for Neon's OpenAI-compatible Chat Completions endpoint. */
29export interface LanguageModelOptions {
30 /** Model ID from Neon's catalog that supports Chat Completions. */
31 readonly model: string;
32 /** Model-dependent generation settings, sent without silently rewriting them. */
33 readonly parameters?: {
34 /** Sampling temperature. */
35 readonly temperature?: number;
36 /** Maximum output tokens; mutually exclusive with maxCompletionTokens. */
37 readonly maxTokens?: number;
38 /** Completion budget for models using max_completion_tokens. */
39 readonly maxCompletionTokens?: number;
40 /** Nucleus sampling probability. */
41 readonly topP?: number;
42 /** Provider-supported deterministic sampling seed. */
43 readonly seed?: number;
44 /** Repetition penalty based on token frequency. */
45 readonly frequencyPenalty?: number;
46 /** Repetition penalty based on token presence. */
47 readonly presencePenalty?: number;
48 };
49}
50
51/**
52 * Effect AI over Neon's documented `/v1/chat/completions` route. Supports text,
53 * image inputs, function tools, JSON-schema output and cancellable SSE streams.
54 * Tool and structured-output support depends on the selected model. Responses-only
55 * models (including Codex), native Anthropic thinking/cache controls, audio,
56 * provider-executed tools and image generation are not supported by this adapter.
57 * Use the bound client's dialect URLs with a native SDK for those APIs.
58 *
59 * ### Generate text
60 * **Example:** A request-scoped Effect model
61 * ```typescript
62 * const ai = yield* Neon.QueryAIGateway(gateway);
63 * const model = ai.model({ model: "gpt-5-mini" });
64 * // Inside the Function or Worker request handler:
65 * const reply = yield* LanguageModel.generateText({ prompt: "Say hello." }).pipe(
66 * Effect.provide(model),
67 * );
68 * ```
69 *
70 * ### Generate structured output
71 * **Example:** Decode the response with an Effect Schema
72 * ```typescript
73 * const reply = yield* LanguageModel.generateObject({
74 * prompt: "Return a short greeting.",
75 * schema: Schema.Struct({ greeting: Schema.String }),
76 * }).pipe(Effect.provide(ai.model({ model: "gpt-5-mini" })));
77 * ```
78 *
79 * @layer
80 * @provides effect/ai/LanguageModel
81 * @product AI Gateway
82 * @category AI Gateway
83 */
84export const makeLanguageModelLayer = (
85 options: LanguageModelOptions & { readonly client: LanguageModelClient },
86): Layer.Layer<LanguageModel.LanguageModel, never, RuntimeContext> =>
87 Layer.effect(LanguageModel.LanguageModel, makeLanguageModel(options)).pipe(
88 Layer.provide(FetchHttpClient.layer),
89 );
90
91/** Construct with an injected Effect HTTP client for custom transports and protocol tests. */
92export const makeLanguageModel = ({
93 client,
94 model,
95 parameters,
96}: LanguageModelOptions & {
97 readonly client: LanguageModelClient;
98}): Effect.Effect<
99 LanguageModel.LanguageModel,
100 never,
101 RuntimeContext | HttpClient.HttpClient
102> =>
103 Effect.gen(function* () {
104 const http = yield* HttpClient.HttpClient;
105 const baseUrl = yield* client.chatBaseUrl;
106 const token = yield* client.token;
107 const send = (options: LanguageModel.ProviderOptions, stream: boolean) =>
108 Effect.gen(function* () {
109 const body = yield* requestBody(options, model, parameters, stream);
110 const request = yield* HttpClientRequest.bodyJson(body)(
111 HttpClientRequest.post(
112 `${baseUrl.replace(/\/$/, "")}/chat/completions`,
113 ).pipe(
114 HttpClientRequest.bearerToken(token),
115 HttpClientRequest.accept(
116 stream ? "text/event-stream" : "application/json",
117 ),
118 ),
119 ).pipe(
120 Effect.mapError(() =>
121 invalidRequest("Request is not JSON serializable"),
122 ),
123 );
124 const response = yield* http.execute(request).pipe(
125 Effect.mapError(() =>
126 error(
127 new AiError.NetworkError({
128 reason: "TransportError",
129 request: {
130 method: "POST",
131 url: "/v1/chat/completions",
132 urlParams: [],
133 headers: {},
134 },
135 description: "Neon AI Gateway transport failed",
136 }),
137 ),
138 ),
139 );
140 if (response.status < 200 || response.status >= 300) {
141 return yield* httpError(response);
142 }
143 return response;
144 });
145 return yield* LanguageModel.make({
146 codecTransformer: toCodecOpenAI,
147 generateText: (options) =>
148 Effect.gen(function* () {
149 const response = yield* send(options, false);
150 const raw = yield* response.json.pipe(
151 Effect.mapError(() => invalidOutput("Expected a JSON completion")),
152 );
153 const decoded = yield* decode(Completion, raw);
154 const choice = decoded.choices[0];
155 if (!choice) return yield* invalidOutput("Completion has no choices");
156 const parts: Array<Response.PartEncoded> = [];
157 const message = choice.message;
158 if (message.refusal)
159 return yield* error(
160 new AiError.ContentPolicyError({
161 description: "Model refused the request",
162 }),
163 );
164 if (message.reasoning_content)
165 parts.push({ type: "reasoning", text: message.reasoning_content });
166 if (typeof message.content === "string") {
167 if (message.content)
168 parts.push({ type: "text", text: message.content });
169 } else if (message.content) {
170 for (const block of message.content)
171 parts.push({ type: "text", text: block.text });
172 }
173 for (const call of message.tool_calls ?? []) {
174 parts.push({
175 type: "tool-call",
176 id: call.id,
177 name: call.function.name,
178 params: yield* parseArguments(call.function.arguments),
179 });
180 }
181 parts.push({
182 type: "finish",
183 reason: finishReason(choice.finish_reason),
184 usage: usage(decoded.usage),
185 });
186 return parts;
187 }),
188 streamText: (options) =>
189 Stream.unwrap(
190 Effect.gen(function* () {
191 const response = yield* send(options, true);
192 return completionStream(response);
193 }),
194 ),
195 });
196 });
197
198const error = (reason: AiError.AiError["reason"]) =>
199 AiError.make({
200 module: "Neon.LanguageModel",
201 method: "chatCompletions",
202 reason,
203 });
204const invalidOutput = (description: string) =>
205 error(new AiError.InvalidOutputError({ description }));
206const invalidRequest = (description: string) =>
207 error(new AiError.InvalidRequestError({ description }));
208const decode = <S extends Schema.ConstraintDecoder<unknown>>(
209 schema: S,
210 value: unknown,
211) =>
212 Schema.decodeUnknownEffect(schema)(value).pipe(
213 Effect.mapError(() =>
214 invalidOutput("Invalid Chat Completions response shape"),
215 ),
216 );
217const parseArguments = (value: string) =>
218 Schema.decodeUnknownEffect(
219 Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown)),
220 )(value).pipe(
221 Effect.mapError(() =>
222 invalidOutput("Invalid function-call JSON arguments"),
223 ),
224 );
225
226const TokenCount = Schema.Int.check(Schema.isGreaterThanOrEqualTo(0));
227const Usage = Schema.Struct({
228 prompt_tokens: Schema.optional(TokenCount),
229 completion_tokens: Schema.optional(TokenCount),
230 prompt_tokens_details: Schema.optional(
231 Schema.NullOr(
232 Schema.Struct({ cached_tokens: Schema.optional(TokenCount) }),
233 ),
234 ),
235 completion_tokens_details: Schema.optional(
236 Schema.NullOr(
237 Schema.Struct({ reasoning_tokens: Schema.optional(TokenCount) }),
238 ),
239 ),
240});
241const Content = Schema.NullOr(
242 Schema.Union([
243 Schema.String,
244 Schema.Array(
245 Schema.Struct({ type: Schema.Literal("text"), text: Schema.String }),
246 ),
247 ]),
248);
249const Completion = Schema.Struct({
250 choices: Schema.Array(
251 Schema.Struct({
252 finish_reason: Schema.NullOr(Schema.String),
253 message: Schema.Struct({
254 content: Schema.optional(Content),
255 reasoning_content: Schema.optional(Schema.NullOr(Schema.String)),
256 refusal: Schema.optional(Schema.NullOr(Schema.String)),
257 tool_calls: Schema.optional(
258 Schema.Array(
259 Schema.Struct({
260 id: Schema.String,
261 type: Schema.Literal("function"),
262 function: Schema.Struct({
263 name: Schema.String,
264 arguments: Schema.String,
265 }),
266 }),
267 ),
268 ),
269 }),
270 }),
271 ),
272 usage: Schema.optional(Schema.NullOr(Usage)),
273});
274const Chunk = Schema.Struct({
275 choices: Schema.Array(
276 Schema.Struct({
277 index: TokenCount,
278 finish_reason: Schema.optional(Schema.NullOr(Schema.String)),
279 delta: Schema.Struct({
280 content: Schema.optional(Schema.NullOr(Schema.String)),
281 reasoning_content: Schema.optional(Schema.NullOr(Schema.String)),
282 refusal: Schema.optional(Schema.NullOr(Schema.String)),
283 tool_calls: Schema.optional(
284 Schema.Array(
285 Schema.Struct({
286 index: TokenCount,
287 id: Schema.optional(Schema.String),
288 type: Schema.optional(Schema.Literal("function")),
289 function: Schema.optional(
290 Schema.Struct({
291 name: Schema.optional(Schema.String),
292 arguments: Schema.optional(Schema.String),
293 }),
294 ),
295 }),
296 ),
297 ),
298 }),
299 }),
300 ),
301 usage: Schema.optional(Schema.NullOr(Usage)),
302});
303const ErrorBody = Schema.Struct({
304 error: Schema.Struct({
305 code: Schema.optional(
306 Schema.Union([Schema.String, Schema.Number, Schema.Null]),
307 ),
308 message: Schema.optional(Schema.String),
309 }),
310});
311
312const httpError = (
313 response: HttpClientResponse.HttpClientResponse,
314): Effect.Effect<never, AiError.AiError> =>
315 Effect.gen(function* () {
316 // Never retain upstream bodies, headers or request objects in public errors.
317 const body = yield* response.json.pipe(
318 Effect.flatMap(Schema.decodeUnknownEffect(ErrorBody)),
319 Effect.catch(() => Effect.succeed(undefined)),
320 );
321 const status = response.status;
322 if (
323 status === 403 &&
324 body?.error.message === "ai gateway not enabled for account"
325 )
326 return yield* error(new AiError.QuotaExhaustedError({}));
327 if (status === 401 || status === 403)
328 return yield* error(
329 new AiError.AuthenticationError({
330 kind: status === 401 ? "InvalidKey" : "InsufficientPermissions",
331 }),
332 );
333 if (
334 status === 402 ||
335 body?.error.code === "REQUEST_LIMIT_EXCEEDED" ||
336 body?.error.code === "insufficient_quota"
337 )
338 return yield* error(new AiError.QuotaExhaustedError({}));
339 if (status === 429) {
340 const seconds = Number(response.headers["retry-after"]);
341 return yield* error(
342 new AiError.RateLimitError({
343 ...(Number.isFinite(seconds) && seconds >= 0
344 ? { retryAfter: Duration.seconds(seconds) }
345 : {}),
346 }),
347 );
348 }
349 if (status >= 500)
350 return yield* error(
351 new AiError.InternalProviderError({
352 description: `Neon AI Gateway returned HTTP ${status}`,
353 }),
354 );
355 return yield* invalidRequest(
356 `Neon AI Gateway rejected the request (HTTP ${status}); check the model, endpoint and parameters`,
357 );
358 });
359
360const finishReason = (
361 reason: string | null | undefined,
362): Response.FinishReason => {
363 switch (reason) {
364 case "stop":
365 return "stop";
366 case "length":
367 return "length";
368 case "tool_calls":
369 return "tool-calls";
370 case "content_filter":
371 return "content-filter";
372 case null:
373 case undefined:
374 return "unknown";
375 default:
376 return "other";
377 }
378};
379const usage = (value: typeof Usage.Type | null | undefined) => {
380 const input = value?.prompt_tokens;
381 const cached = value?.prompt_tokens_details?.cached_tokens;
382 const output = value?.completion_tokens;
383 const reasoning = value?.completion_tokens_details?.reasoning_tokens;
384 return new Response.Usage({
385 inputTokens: {
386 total: input,
387 uncached:
388 input === undefined ? undefined : Math.max(0, input - (cached ?? 0)),
389 cacheRead: cached,
390 cacheWrite: undefined,
391 },
392 outputTokens: {
393 total: output,
394 text:
395 output === undefined
396 ? undefined
397 : Math.max(0, output - (reasoning ?? 0)),
398 reasoning,
399 },
400 });
401};
402
403const requestBody = (
404 options: LanguageModel.ProviderOptions,
405 model: string,
406 parameters: LanguageModelOptions["parameters"],
407 stream: boolean,
408) =>
409 Effect.gen(function* () {
410 if (
411 parameters?.maxTokens !== undefined &&
412 parameters.maxCompletionTokens !== undefined
413 )
414 return yield* invalidRequest(
415 "Set maxTokens or maxCompletionTokens, not both",
416 );
417 if (options.tools.some(Tool.isProviderDefined))
418 return yield* invalidRequest(
419 "Provider-executed tools require a native provider API",
420 );
421 const messages: Array<Record<string, unknown>> = [];
422 for (const message of options.prompt.content) {
423 if (message.role === "system") {
424 messages.push({ role: "system", content: message.content });
425 } else if (message.role === "user") {
426 const content: Array<unknown> = [];
427 for (const part of message.content) {
428 if (part.type === "text")
429 content.push({ type: "text", text: part.text });
430 else if (
431 part.type === "file" &&
432 part.mediaType.startsWith("image/")
433 ) {
434 const url = yield* Effect.sync(() => {
435 if (part.data instanceof URL) return part.data.toString();
436 if (part.data instanceof Uint8Array) {
437 let binary = "";
438 for (const byte of part.data)
439 binary += String.fromCharCode(byte);
440 return `data:${part.mediaType};base64,${btoa(binary)}`;
441 }
442 return /^(data:|https?:)/.test(part.data)
443 ? part.data
444 : `data:${part.mediaType};base64,${part.data}`;
445 });
446 content.push({ type: "image_url", image_url: { url } });
447 } else
448 return yield* invalidRequest(
449 "Only text and image user inputs are supported",
450 );
451 }
452 messages.push({ role: "user", content });
453 } else if (message.role === "assistant") {
454 const text: Array<string> = [];
455 const toolCalls: Array<unknown> = [];
456 for (const part of message.content) {
457 if (part.type === "text") text.push(part.text);
458 else if (part.type === "tool-call")
459 toolCalls.push({
460 id: part.id,
461 type: "function",
462 function: {
463 name: part.name,
464 arguments: yield* stringify(part.params),
465 },
466 });
467 else if (part.type !== "reasoning")
468 return yield* invalidRequest(
469 "Unsupported assistant content for Chat Completions",
470 );
471 }
472 messages.push({
473 role: "assistant",
474 content: text.join("") || null,
475 ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
476 });
477 } else {
478 for (const part of message.content) {
479 if (part.type !== "tool-result")
480 return yield* invalidRequest(
481 "Unsupported tool message for Chat Completions",
482 );
483 messages.push({
484 role: "tool",
485 tool_call_id: part.id,
486 content:
487 typeof part.result === "string"
488 ? part.result
489 : yield* stringify(part.result),
490 });
491 }
492 }
493 }
494 const choice = options.toolChoice;
495 const selected =
496 typeof choice === "object" && "oneOf" in choice
497 ? options.tools.filter((tool) => choice.oneOf.includes(tool.name))
498 : options.tools;
499 const tools = yield* Effect.try({
500 try: () =>
501 selected.map((tool) => ({
502 type: "function",
503 function: {
504 name: tool.name,
505 description: Tool.getDescription(tool),
506 parameters: Tool.getJsonSchema(tool),
507 },
508 })),
509 catch: () =>
510 invalidRequest("Tool schema cannot be represented as JSON Schema"),
511 });
512 const responseFormat = options.responseFormat;
513 const jsonSchema =
514 responseFormat.type === "json"
515 ? yield* Effect.try({
516 try: () => toCodecOpenAI(responseFormat.schema).jsonSchema,
517 catch: () =>
518 error(
519 new AiError.UnsupportedSchemaError({
520 description:
521 "Schema is not supported by OpenAI structured output",
522 }),
523 ),
524 })
525 : undefined;
526 return {
527 model,
528 messages,
529 stream,
530 ...(stream ? { stream_options: { include_usage: true } } : {}),
531 ...(tools.length
532 ? {
533 tools,
534 tool_choice:
535 typeof choice === "object"
536 ? "tool" in choice
537 ? { type: "function", function: { name: choice.tool } }
538 : (choice.mode ?? "auto")
539 : choice,
540 }
541 : {}),
542 ...(jsonSchema
543 ? {
544 response_format: {
545 type: "json_schema",
546 json_schema: {
547 name:
548 responseFormat.type === "json"
549 ? responseFormat.objectName
550 : "response",
551 schema: jsonSchema,
552 strict: true,
553 },
554 },
555 }
556 : {}),
557 temperature: parameters?.temperature,
558 max_tokens: parameters?.maxTokens,
559 max_completion_tokens: parameters?.maxCompletionTokens,
560 top_p: parameters?.topP,
561 seed: parameters?.seed,
562 frequency_penalty: parameters?.frequencyPenalty,
563 presence_penalty: parameters?.presencePenalty,
564 };
565 });
566const stringify = (value: unknown) =>
567 Effect.try({
568 try: () => JSON.stringify(value),
569 catch: () => invalidRequest("Prompt contains a non-JSON value"),
570 });
571
572const completionStream = (
573 response: HttpClientResponse.HttpClientResponse,
574): Stream.Stream<Response.StreamPartEncoded, AiError.AiError> =>
575 Stream.unwrap(
576 Effect.sync(() => {
577 let done = false;
578 let reason: string | undefined;
579 let tokenUsage: typeof Usage.Type | null | undefined;
580 let textStarted = false;
581 let reasoningStarted = false;
582 const calls = new Map<
583 number,
584 { id: string; name: string; arguments: string; started: boolean }
585 >();
586 const chunks = response.stream.pipe(
587 Stream.mapError(() =>
588 invalidOutput("Chat Completions stream transport failed"),
589 ),
590 Stream.decodeText(),
591 Stream.pipeThroughChannel(Sse.decode<AiError.AiError, unknown>()),
592 Stream.mapError(() =>
593 invalidOutput("Invalid or interrupted SSE stream"),
594 ),
595 Stream.takeUntil((event) => event.data === "[DONE]"),
596 Stream.mapEffect((event) =>
597 Effect.gen(function* () {
598 const parts: Array<Response.StreamPartEncoded> = [];
599 if (event.data === "[DONE]") {
600 done = true;
601 return parts;
602 }
603 if (!event.data) return parts;
604 if (event.event === "error")
605 return yield* error(
606 new AiError.InternalProviderError({
607 description: "Neon AI Gateway stream reported an error",
608 }),
609 );
610 const chunk = yield* Schema.decodeUnknownEffect(
611 Schema.fromJsonString(Chunk),
612 )(event.data).pipe(
613 Effect.mapError(() =>
614 invalidOutput("Invalid Chat Completions stream chunk"),
615 ),
616 );
617 if (chunk.usage != null) tokenUsage = chunk.usage;
618 for (const choice of chunk.choices) {
619 if (choice.index !== 0) continue;
620 if (choice.finish_reason != null) reason = choice.finish_reason;
621 const delta = choice.delta;
622 if (delta.refusal)
623 return yield* error(
624 new AiError.ContentPolicyError({
625 description: "Model refused the request",
626 }),
627 );
628 if (delta.reasoning_content) {
629 if (!reasoningStarted) {
630 reasoningStarted = true;
631 parts.push({ type: "reasoning-start", id: "reasoning" });
632 }
633 parts.push({
634 type: "reasoning-delta",
635 id: "reasoning",
636 delta: delta.reasoning_content,
637 });
638 }
639 if (delta.content) {
640 if (!textStarted) {
641 textStarted = true;
642 parts.push({ type: "text-start", id: "text" });
643 }
644 parts.push({
645 type: "text-delta",
646 id: "text",
647 delta: delta.content,
648 });
649 }
650 for (const deltaCall of delta.tool_calls ?? []) {
651 const call = calls.get(deltaCall.index) ?? {
652 id: "",
653 name: "",
654 arguments: "",
655 started: false,
656 };
657 if (deltaCall.id) {
658 if (call.id && call.id !== deltaCall.id)
659 return yield* invalidOutput(
660 "Tool call ID changed within a stream",
661 );
662 call.id = deltaCall.id;
663 }
664 if (deltaCall.function?.name) {
665 if (call.started && call.name !== deltaCall.function.name)
666 return yield* invalidOutput(
667 "Tool name changed within a stream",
668 );
669 call.name = deltaCall.function.name;
670 }
671 const args = deltaCall.function?.arguments ?? "";
672 call.arguments += args;
673 if (!call.started && call.id && call.name) {
674 call.started = true;
675 parts.push({
676 type: "tool-params-start",
677 id: call.id,
678 name: call.name,
679 });
680 if (call.arguments)
681 parts.push({
682 type: "tool-params-delta",
683 id: call.id,
684 delta: call.arguments,
685 });
686 } else if (call.started && args)
687 parts.push({
688 type: "tool-params-delta",
689 id: call.id,
690 delta: args,
691 });
692 calls.set(deltaCall.index, call);
693 }
694 }
695 return parts;
696 }),
697 ),
698 Stream.flatMap(Stream.fromIterable),
699 );
700 const finish = Stream.unwrap(
701 Effect.gen(function* () {
702 if (!done || !reason)
703 return yield* invalidOutput(
704 "Chat Completions stream ended before its finish reason and [DONE]",
705 );
706 const parts: Array<Response.StreamPartEncoded> = [];
707 if (reasoningStarted)
708 parts.push({ type: "reasoning-end", id: "reasoning" });
709 if (textStarted) parts.push({ type: "text-end", id: "text" });
710 for (const call of calls.values()) {
711 if (!call.started)
712 return yield* invalidOutput("Incomplete streamed tool identity");
713 parts.push({ type: "tool-params-end", id: call.id });
714 if (reason !== "length" && reason !== "content_filter") {
715 parts.push({
716 type: "tool-call",
717 id: call.id,
718 name: call.name,
719 params: yield* parseArguments(call.arguments),
720 });
721 }
722 }
723 parts.push({
724 type: "finish",
725 reason: finishReason(reason),
726 usage: usage(tokenUsage),
727 });
728 return Stream.fromIterable(parts);
729 }),
730 );
731 return Stream.concat(chunks, finish);
732 }),
733 );