1import type * as runtime from "@cloudflare/workers-types";
2import * as aisearch from "@distilled.cloud/cloudflare/aisearch";
3import * as Effect from "effect/Effect";
4import * as Stream from "effect/Stream";
5import type * as HttpClient from "effect/http/HttpClient";
6import type { RuntimeContext } from "../../RuntimeContext.ts";
7import type { Credentials } from "../Credentials.ts";
8import { SearchError, type QuerySearchClient } from "./QuerySearch.ts";
9import type { QuerySearchNamespaceClient } from "./QuerySearchNamespace.ts";
10
11// Shared HTTP scaffolding for the AI Search `*Local` binding layers. NOT
12// re-exported from `index.ts` — only the per-level `*Local` layers are public.
13//
14// The `ai_search` / `ai_search_namespace` Worker bindings proxy the same AI
15// Search REST API that distilled wraps, so the whole data plane
16// (`search` / `chatCompletions` / `info` / `stats` / `list`) is reachable over
17// HTTP with the current credentials — no Worker host, no native binding.
18//
19// The one wrinkle: distilled decodes responses to **camelCase**, while the
20// binding client contract types come from `@cloudflare/workers-types` and are
21// the **snake_case** wire shape. These adapters translate both directions
22// (mirroring the D1 `*Local` shim), passing user-controlled maps
23// (`item.metadata`, retrieval `filters`) through untouched.
24
25/**
26 * Injectable auth shared by a future Http (scoped-token) impl and the Local
27 * (current-credentials) impl. `authorize` discharges the
28 * `Credentials | HttpClient` requirement of a distilled op; `accountId` is the
29 * Cloudflare account the ops run against.
30 */
31export interface SearchAuth {
32 authorize: <A, E>(
33 eff: Effect.Effect<A, E, Credentials | HttpClient.HttpClient>,
34 ) => Effect.Effect<A, E>;
35 accountId: string;
36}
37
38const u = <T>(v: T | null | undefined): T | undefined =>
39 v == null ? undefined : v;
40
41const run = <A, E>(
42 auth: SearchAuth,
43 eff: Effect.Effect<A, E, Credentials | HttpClient.HttpClient>,
44): Effect.Effect<A, SearchError, RuntimeContext> =>
45 auth.authorize(eff).pipe(
46 Effect.mapError((cause) => {
47 const message =
48 cause instanceof Error
49 ? cause.message
50 : typeof (cause as { message?: unknown } | undefined)?.message ===
51 "string"
52 ? (cause as { message: string }).message
53 : "AI Search HTTP error";
54 return new SearchError({ message, cause });
55 }),
56 );
57
58// ── response mappers (distilled camelCase -> runtime snake_case) ─────────────
59
60type ChunkIn = aisearch.SearchNamespaceInstanceResponse["chunks"][number];
61
62const mapChunk = (
63 c: ChunkIn,
64): runtime.AiSearchSearchResponse["chunks"][number] => ({
65 id: c.id,
66 type: c.type,
67 score: c.score,
68 text: c.text,
69 item: {
70 key: c.item?.key ?? "",
71 timestamp: u(c.item?.timestamp),
72 metadata: u(c.item?.metadata) as Record<string, unknown> | undefined,
73 },
74 scoring_details: c.scoringDetails
75 ? {
76 keyword_score: u(c.scoringDetails.keywordScore),
77 vector_score: u(c.scoringDetails.vectorScore),
78 keyword_rank: u(c.scoringDetails.keywordRank),
79 vector_rank: u(c.scoringDetails.vectorRank),
80 reranking_score: u(c.scoringDetails.rerankingScore),
81 fusion_method: u(c.scoringDetails.fusionMethod) as
82 | "rrf"
83 | "max"
84 | undefined,
85 }
86 : undefined,
87});
88
89const mapSearch = (
90 r: aisearch.SearchNamespaceInstanceResponse,
91): runtime.AiSearchSearchResponse => ({
92 search_query: r.searchQuery ?? "",
93 chunks: r.chunks.map(mapChunk),
94});
95
96const mapChat = (
97 r: aisearch.ChatCompletionsNamespaceInstanceResponse,
98): runtime.AiSearchChatCompletionsResponse => ({
99 id: u(r.id),
100 object: u(r.object),
101 model: u(r.model),
102 choices: r.choices.map((choice) => ({
103 index: u(choice.index),
104 message: {
105 role: choice.message.role as
106 | "system"
107 | "developer"
108 | "user"
109 | "assistant"
110 | "tool",
111 content:
112 typeof choice.message.content === "string"
113 ? choice.message.content
114 : choice.message.content == null
115 ? null
116 : JSON.stringify(choice.message.content),
117 },
118 })),
119 chunks: r.chunks.map(mapChunk),
120});
121
122const mapStats = (
123 r: aisearch.StatsNamespaceInstanceResponse,
124): runtime.AiSearchStatsResponse => ({
125 queued: u(r.queued),
126 running: u(r.running),
127 completed: u(r.completed),
128 error: u(r.error),
129 skipped: u(r.skipped),
130 outdated: u(r.outdated),
131 last_activity: u(r.lastActivity),
132 engine: r.engine
133 ? {
134 vectorize: r.engine.vectorize
135 ? {
136 vectorsCount: r.engine.vectorize.vectorsCount,
137 dimensions: r.engine.vectorize.dimensions,
138 }
139 : undefined,
140 r2: r.engine.r2
141 ? {
142 payloadSizeBytes: r.engine.r2.payloadSizeBytes,
143 metadataSizeBytes: r.engine.r2.metadataSizeBytes,
144 objectCount: r.engine.r2.objectCount,
145 }
146 : undefined,
147 }
148 : undefined,
149});
150
151const mapInfo = (
152 r: aisearch.ReadNamespaceInstanceResponse,
153): runtime.AiSearchInstanceInfo => ({
154 id: r.id,
155 type: u(r.type) as runtime.AiSearchInstanceInfo["type"],
156 source: u(r.source),
157 source_params: u(r.sourceParams),
158 paused: u(r.paused),
159 status: u(r.status),
160 namespace: u(r.namespace),
161 created_at: u(r.createdAt),
162 modified_at: u(r.modifiedAt),
163 token_id: u(r.tokenId),
164 ai_gateway_id: u(r.aiGatewayId),
165 rewrite_query: u(r.rewriteQuery),
166 reranking: u(r.reranking),
167 embedding_model: u(r.embeddingModel),
168 ai_search_model: u(r.aiSearchModel),
169 rewrite_model: u(r.rewriteModel),
170 reranking_model: u(r.rerankingModel),
171 hybrid_search_enabled: u(r.hybridSearchEnabled),
172 index_method: r.indexMethod
173 ? { vector: r.indexMethod.vector, keyword: r.indexMethod.keyword }
174 : undefined,
175 fusion_method: u(r.fusionMethod) as "max" | "rrf" | undefined,
176 indexing_options: r.indexingOptions
177 ? {
178 keyword_tokenizer: u(r.indexingOptions.keywordTokenizer) as
179 | "porter"
180 | "trigram"
181 | undefined,
182 }
183 : undefined,
184 retrieval_options: r.retrievalOptions
185 ? {
186 keyword_match_mode: u(r.retrievalOptions.keywordMatchMode) as
187 | "and"
188 | "or"
189 | undefined,
190 boost_by: u(r.retrievalOptions.boostBy)?.map((b) => ({
191 field: b.field,
192 direction: u(b.direction) as
193 | "asc"
194 | "desc"
195 | "exists"
196 | "not_exists"
197 | undefined,
198 })),
199 }
200 : undefined,
201 chunk_size: u(r.chunkSize),
202 chunk_overlap: u(r.chunkOverlap),
203 score_threshold: u(r.scoreThreshold),
204 max_num_results: u(r.maxNumResults),
205 cache: u(r.cache),
206 cache_threshold: u(
207 r.cacheThreshold,
208 ) as runtime.AiSearchInstanceInfo["cache_threshold"],
209 custom_metadata: u(r.customMetadata)?.map((m) => ({
210 field_name: m.fieldName,
211 data_type: m.dataType as "text" | "number" | "boolean" | "datetime",
212 })),
213 sync_interval: u(
214 r.syncInterval,
215 ) as runtime.AiSearchInstanceInfo["sync_interval"],
216 metadata: u(r.metadata) as Record<string, unknown> | undefined,
217});
218
219const mapList = (
220 result: aisearch.ReadNamespaceInstanceResponse[],
221): runtime.AiSearchListResponse => ({
222 result: result.map(mapInfo),
223});
224
225const mapMulti = (
226 r: aisearch.SearchNamespaceResponse,
227): runtime.AiSearchMultiSearchResponse => ({
228 search_query: r.searchQuery ?? "",
229 chunks: r.chunks.map((c) => ({ ...mapChunk(c), instance_id: c.instanceId })),
230 errors: u(r.errors)?.map((e) => ({
231 instance_id: e.instanceId,
232 message: e.message,
233 })),
234});
235
236// ── request mappers (runtime snake_case -> distilled camelCase) ──────────────
237
238const mapMessages = (messages: runtime.AiSearchMessage[]) =>
239 messages.map((m) => ({
240 role: m.role as "system" | "developer" | "user" | "assistant" | "tool",
241 content:
242 typeof m.content === "string" || m.content == null
243 ? m.content
244 : (m.content as unknown[]).map((part) => {
245 const p = part as Record<string, unknown>;
246 if ("image_url" in p || "imageUrl" in p) {
247 const img = (p.image_url ?? p.imageUrl) as { url: string };
248 return { type: "image_url" as const, imageUrl: { url: img.url } };
249 }
250 return { type: "text" as const, text: p.text as string };
251 }),
252 }));
253
254const mapOptions = (o: runtime.AiSearchOptions | undefined) => {
255 if (!o) return undefined;
256 return {
257 cache: o.cache
258 ? { enabled: o.cache.enabled, cacheThreshold: o.cache.cache_threshold }
259 : undefined,
260 queryRewrite: o.query_rewrite
261 ? {
262 enabled: o.query_rewrite.enabled,
263 model: o.query_rewrite.model,
264 rewritePrompt: o.query_rewrite.rewrite_prompt,
265 }
266 : undefined,
267 reranking: o.reranking
268 ? {
269 enabled: o.reranking.enabled,
270 model: o.reranking.model,
271 matchThreshold: o.reranking.match_threshold,
272 }
273 : undefined,
274 retrieval: o.retrieval
275 ? {
276 retrievalType: o.retrieval.retrieval_type,
277 fusionMethod: o.retrieval.fusion_method,
278 keywordMatchMode: o.retrieval.keyword_match_mode,
279 matchThreshold: o.retrieval.match_threshold,
280 maxNumResults: o.retrieval.max_num_results,
281 contextExpansion: o.retrieval.context_expansion,
282 returnOnFailure: o.retrieval.return_on_failure,
283 // `filters` is a Vectorize metadata filter over user-defined fields —
284 // pass through untouched.
285 filters: o.retrieval.filters as Record<string, unknown> | undefined,
286 boostBy: o.retrieval.boost_by,
287 }
288 : undefined,
289 };
290};
291
292// ── client builders ──────────────────────────────────────────────────────────
293
294/** Effect resolving `{ name (namespace), id (instanceId) }` at apply time. */
295export type InstanceRef = Effect.Effect<{ name: string; id: string }>;
296
297/** Effect resolving the namespace `name` at apply time. */
298export type NamespaceRef = Effect.Effect<string>;
299
300const dieRaw = (kind: string): Effect.Effect<never> =>
301 Effect.die(
302 new Error(
303 `The AI Search ${kind} *Local binding runs over the HTTP API; the raw native runtime binding is only available inside a deployed Worker.`,
304 ),
305 );
306
307/**
308 * Build a single-instance {@link QuerySearchClient} over the HTTP API. `ref`
309 * resolves the `{ namespace, instanceId }` at apply time.
310 */
311export const makeLocalSearchClient = (
312 auth: SearchAuth,
313 ref: InstanceRef,
314): QuerySearchClient => {
315 const withRef = <A, E>(
316 fn: (r: {
317 name: string;
318 id: string;
319 }) => Effect.Effect<A, E, Credentials | HttpClient.HttpClient>,
320 ) => Effect.flatMap(ref, (r) => run(auth, fn(r)));
321
322 return {
323 raw: dieRaw("instance"),
324 search: (params) =>
325 withRef((r) =>
326 aisearch.searchNamespaceInstance({
327 accountId: auth.accountId,
328 name: r.name,
329 id: r.id,
330 ...("query" in params && params.query !== undefined
331 ? { query: params.query }
332 : { messages: mapMessages(params.messages ?? []) }),
333 aiSearchOptions: mapOptions(params.ai_search_options),
334 } as aisearch.SearchNamespaceInstanceRequest),
335 ).pipe(Effect.map(mapSearch)),
336 chatCompletions: (params) =>
337 withRef((r) =>
338 aisearch.chatCompletionsNamespaceInstance({
339 accountId: auth.accountId,
340 name: r.name,
341 id: r.id,
342 messages: mapMessages(params.messages),
343 aiSearchOptions: mapOptions(params.ai_search_options),
344 } as aisearch.ChatCompletionsNamespaceInstanceRequest),
345 ).pipe(Effect.map(mapChat)),
346 info: () =>
347 withRef((r) =>
348 aisearch.readNamespaceInstance({
349 accountId: auth.accountId,
350 name: r.name,
351 id: r.id,
352 }),
353 ).pipe(Effect.map(mapInfo)),
354 stats: () =>
355 withRef((r) =>
356 aisearch.statsNamespaceInstance({
357 accountId: auth.accountId,
358 name: r.name,
359 id: r.id,
360 }),
361 ).pipe(Effect.map(mapStats)),
362 } satisfies QuerySearchClient;
363};
364
365/**
366 * Build a namespace {@link QuerySearchNamespaceClient} over the HTTP API.
367 * `.get(instanceName)` scopes a single-instance client to `(namespace,
368 * instanceName)`.
369 */
370export const makeLocalSearchNamespaceClient = (
371 auth: SearchAuth,
372 ref: NamespaceRef,
373): QuerySearchNamespaceClient => ({
374 raw: dieRaw("namespace"),
375 get: (instanceName) =>
376 makeLocalSearchClient(
377 auth,
378 Effect.map(ref, (name) => ({ name, id: instanceName })),
379 ),
380 list: (params) =>
381 Effect.flatMap(ref, (name) =>
382 run(
383 auth,
384 Stream.runCollect(
385 aisearch.listNamespaceInstances.pages({
386 accountId: auth.accountId,
387 name,
388 perPage: params?.per_page,
389 orderBy: params?.order_by,
390 orderByDirection: params?.order_by_direction,
391 search: params?.search,
392 }),
393 ),
394 ),
395 ).pipe(
396 Effect.map((chunk) =>
397 mapList(
398 Array.from(chunk).flatMap(
399 (page) =>
400 (page.result ??
401 []) as unknown as aisearch.ReadNamespaceInstanceResponse[],
402 ),
403 ),
404 ),
405 ),
406 search: (params) =>
407 Effect.flatMap(ref, (name) =>
408 run(
409 auth,
410 aisearch.searchNamespace({
411 accountId: auth.accountId,
412 name,
413 aiSearchOptions: {
414 instanceIds: params.ai_search_options.instance_ids,
415 ...mapOptions(params.ai_search_options),
416 },
417 ...("query" in params && params.query !== undefined
418 ? { query: params.query }
419 : { messages: mapMessages(params.messages ?? []) }),
420 } as aisearch.SearchNamespaceRequest),
421 ),
422 ).pipe(Effect.map(mapMulti)),
423});