1import * as Cloudflare from "@/Cloudflare/index.ts";
2import type { RuntimeContext } from "@/RuntimeContext.ts";
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 { Gateway } from "./Gateway.ts";
11
12// `@cf/meta/llama-3.1-8b-instruct` was deprecated by Cloudflare on
13// 2026-05-30 (the API answers 410), so use the supported fast 3.3 model.
14const MODEL = "@cf/meta/llama-3.3-70b-instruct-fp8-fast";
15// Tool tests need a model that reliably honors `tool_choice: "required"`.
16// The previous choice (`@cf/moonshotai/kimi-k2.6`) is a reasoning model
17// that intermittently emitted reasoning/text and *no* tool call, flaking
18// the tool assertions. Llama 3.3 70B is non-reasoning and reliably returns
19// an OpenAI-shaped `tool_calls` array under `tool_choice: "required"`.
20const TOOL_MODEL = "@cf/meta/llama-3.3-70b-instruct-fp8-fast";
21
22const GetWeather = Tool.make("get_weather", {
23 description:
24 "Get the current weather for a city. Always call this tool when the user asks about the weather.",
25 parameters: Schema.Struct({
26 city: Schema.String,
27 }),
28 success: Schema.Struct({
29 city: Schema.String,
30 temperatureF: Schema.Number,
31 condition: Schema.String,
32 }),
33});
34
35const WeatherToolkit = Toolkit.make(GetWeather);
36
37const WeatherToolkitLayer = WeatherToolkit.toLayer({
38 get_weather: ({ city }) =>
39 Effect.succeed({
40 city,
41 temperatureF: 72,
42 condition: "sunny",
43 }),
44});
45
46export default class LanguageModelTestWorker extends Cloudflare.Worker<LanguageModelTestWorker>()(
47 "LanguageModelTestWorker",
48 {
49 main: import.meta.url,
50 },
51 Effect.gen(function* () {
52 const aiGateway = yield* Cloudflare.AI.QueryGateway(Gateway);
53
54 const languageModel = aiGateway.model({
55 model: MODEL,
56 parameters: { temperature: 0.7, maxTokens: 1024 },
57 });
58 const toolLanguageModel = aiGateway.model({
59 model: TOOL_MODEL,
60 parameters: { temperature: 0.2, maxTokens: 1024 },
61 });
62 const dumpRawStream = (
63 model: string,
64 messages: ReadonlyArray<{ role: string; content: string }>,
65 streamOptions: Record<string, unknown> | undefined,
66 ) =>
67 Effect.gen(function* () {
68 const ai = yield* aiGateway.raw;
69 const gatewayId = yield* aiGateway.id;
70 const response = yield* Effect.tryPromise({
71 try: () =>
72 ai.run(
73 model as keyof AiModels,
74 {
75 messages,
76 stream: true,
77 ...streamOptions,
78 } as unknown as AiModels[keyof AiModels]["inputs"],
79 { gateway: { id: gatewayId }, returnRawResponse: true },
80 ),
81 catch: (e) => e,
82 });
83 if (response.body == null)
84 return HttpServerResponse.text("no body", { status: 500 });
85 const body = Stream.fromReadableStream<Uint8Array, never>({
86 evaluate: () => response.body!,
87 onError: () => undefined as never,
88 });
89 return HttpServerResponse.stream(body, {
90 headers: { "content-type": "text/event-stream" },
91 });
92 }).pipe(Effect.orDie);
93
94 return {
95 fetch: Effect.gen(function* () {
96 const request = yield* HttpServerRequest;
97 const ctx = yield* Effect.context<RuntimeContext>();
98 const url = new URL(request.url, "http://worker");
99 const prompt =
100 url.searchParams.get("prompt") ??
101 "Say the single word 'pong' and nothing else.";
102
103 if (url.pathname === "/generate") {
104 const response = yield* AiLanguageModel.generateText({ prompt }).pipe(
105 Effect.orDie,
106 );
107 return yield* HttpServerResponse.json({
108 text: response.text,
109 finishReason: response.finishReason,
110 usage: {
111 inputTokens: response.usage.inputTokens.total,
112 outputTokens: response.usage.outputTokens.total,
113 },
114 });
115 }
116
117 if (url.pathname === "/raw-stream") {
118 const model = url.searchParams.get("model") ?? MODEL;
119 const includeUsage = url.searchParams.get("include_usage") === "1";
120 return yield* dumpRawStream(
121 model,
122 [{ role: "user", content: prompt }],
123 includeUsage
124 ? { stream_options: { include_usage: true } }
125 : undefined,
126 );
127 }
128
129 if (url.pathname === "/raw-tool-stream") {
130 // Dump the raw Workers AI SSE for a streamed tool call so we can
131 // see exactly how the model encodes tool_calls in the stream.
132 return yield* dumpRawStream(
133 TOOL_MODEL,
134 [{ role: "user", content: prompt }],
135 {
136 tools: [
137 {
138 type: "function",
139 function: {
140 name: "get_weather",
141 description: "Get the current weather for a city.",
142 parameters: {
143 type: "object",
144 properties: { city: { type: "string" } },
145 required: ["city"],
146 },
147 },
148 },
149 ],
150 tool_choice: "required",
151 },
152 );
153 }
154
155 if (url.pathname === "/test-stream") {
156 // Synthetic stream: 5 chunks with 200ms gaps. If the client sees
157 // them at staggered timestamps, worker→edge streaming works.
158 // If they all land at once, output is buffered downstream.
159 const encoder = new TextEncoder();
160 const body = Stream.range(0, 5).pipe(
161 Stream.mapEffect((i) =>
162 Effect.sleep(Duration.millis(200)).pipe(
163 Effect.as(encoder.encode(`data: chunk-${i}\n\n`)),
164 ),
165 ),
166 );
167 return HttpServerResponse.stream(body, {
168 headers: { "content-type": "text/event-stream" },
169 });
170 }
171
172 if (url.pathname === "/tool") {
173 const response = yield* AiLanguageModel.generateText({
174 prompt,
175 toolkit: WeatherToolkit,
176 toolChoice: "required",
177 }).pipe(
178 Effect.provide(WeatherToolkitLayer),
179 Effect.provide(toolLanguageModel),
180 Effect.orDie,
181 );
182 return yield* HttpServerResponse.json({
183 text: response.text,
184 finishReason: response.finishReason,
185 toolCalls: response.toolCalls.map((call) => ({
186 id: call.id,
187 name: call.name,
188 params: call.params,
189 })),
190 toolResults: response.toolResults.map((result) => ({
191 id: result.id,
192 name: result.name,
193 result: result.result,
194 isFailure: result.isFailure,
195 })),
196 });
197 }
198
199 if (url.pathname === "/tool-stream") {
200 const encoder = new TextEncoder();
201 const body = AiLanguageModel.streamText({
202 prompt,
203 toolkit: WeatherToolkit,
204 toolChoice: "required",
205 }).pipe(
206 Stream.map((part) =>
207 encoder.encode(`data: ${JSON.stringify(part)}\n\n`),
208 ),
209 Stream.provide(WeatherToolkitLayer),
210 Stream.provide(toolLanguageModel),
211 Stream.provideContext(ctx),
212 );
213 return HttpServerResponse.stream(body, {
214 headers: { "content-type": "text/event-stream" },
215 });
216 }
217
218 if (url.pathname === "/stream") {
219 const encoder = new TextEncoder();
220 const body = AiLanguageModel.streamText({ prompt }).pipe(
221 Stream.map((part) =>
222 encoder.encode(`data: ${JSON.stringify(part)}\n\n`),
223 ),
224 Stream.provide(languageModel),
225 Stream.provideContext(ctx),
226 );
227 return HttpServerResponse.stream(body, {
228 headers: { "content-type": "text/event-stream" },
229 });
230 }
231
232 return HttpServerResponse.text("ok");
233 }).pipe(Effect.provide(languageModel)),
234 };
235 }).pipe(Effect.provide(Cloudflare.AI.QueryGatewayBinding)),
236) {}