AI SDK

import { createAISearchNamespace } from "ai-search-provider";
import { generateText, streamText } from "ai";

const INSTANCE_NAME = "knowledge-base";

const SAMPLE_DOC = `# Caching on Cloudflare
Cloudflare caches static assets at the edge. Use Cache Rules to control what is
cached, set an Edge Cache TTL to control how long objects stay in cache, and
purge the cache after a deploy.`;

// Create the instance with hybrid search, or update it if it already exists.
async function ensureInstance(env) {
	// index_method with both vector and keyword enables hybrid search.
	const hybrid = { index_method: { vector: true, keyword: true } };
	try {
		await env.AI_SEARCH.create({ id: INSTANCE_NAME, ...hybrid });
	} catch {
		await env.AI_SEARCH.get(INSTANCE_NAME).update(hybrid);
	}
}

export default {
	async fetch(request, env) {
		const url = new URL(request.url);
		// createAISearchNamespace adapts the binding to the AI SDK provider API.
		const aiSearch = createAISearchNamespace({ binding: env.AI_SEARCH });

		// Visit /setup once to create the instance and index a document.
		if (url.pathname === "/setup") {
			await ensureInstance(env);

			const instance = aiSearch.get(INSTANCE_NAME);

			// upload() queues the file and returns immediately. Indexing runs in
			// the background, so poll the item's status until it is searchable.
			const { id, key } = await instance.items.upload("caching.md", SAMPLE_DOC);

			let info = await instance.items.get(id).info();
			while (info.status === "queued" || info.status === "running") {
				await new Promise((resolve) => setTimeout(resolve, 2_000));
				info = await instance.items.get(id).info();
			}

			return Response.json({ key, status: info.status });
		}

		// Query the instance (see the next step).
		return new Response("Visit /setup first, then query with ?q=");
	},
};

import { createAISearchNamespace } from "ai-search-provider";
import { generateText, streamText } from "ai";

interface Env {
	AI_SEARCH: AiSearchNamespace;
}

const INSTANCE_NAME = "knowledge-base";

const SAMPLE_DOC = `# Caching on Cloudflare
Cloudflare caches static assets at the edge. Use Cache Rules to control what is
cached, set an Edge Cache TTL to control how long objects stay in cache, and
purge the cache after a deploy.`;

// Create the instance with hybrid search, or update it if it already exists.
async function ensureInstance(env: Env) {
	// index_method with both vector and keyword enables hybrid search.
	const hybrid = { index_method: { vector: true, keyword: true } };
	try {
		await env.AI_SEARCH.create({ id: INSTANCE_NAME, ...hybrid });
	} catch {
		await env.AI_SEARCH.get(INSTANCE_NAME).update(hybrid);
	}
}

export default {
	async fetch(request, env): Promise<Response> {
		const url = new URL(request.url);
		// createAISearchNamespace adapts the binding to the AI SDK provider API.
		const aiSearch = createAISearchNamespace({ binding: env.AI_SEARCH });

		// Visit /setup once to create the instance and index a document.
		if (url.pathname === "/setup") {
			await ensureInstance(env);

			const instance = aiSearch.get(INSTANCE_NAME);

			// upload() queues the file and returns immediately. Indexing runs in
			// the background, so poll the item's status until it is searchable.
			const { id, key } = await instance.items.upload("caching.md", SAMPLE_DOC);

			let info = await instance.items.get(id).info();
			while (info.status === "queued" || info.status === "running") {
				await new Promise((resolve) => setTimeout(resolve, 2_000));
				info = await instance.items.get(id).info();
			}

			return Response.json({ key, status: info.status });
		}

		// Query the instance (see the next step).
		return new Response("Visit /setup first, then query with ?q=");
	},
} satisfies ExportedHandler<Env>;

const query = url.searchParams.get("q") ?? "How does caching work?";

// chat() returns an AI SDK model that retrieves matching chunks and generates
// a grounded answer in one call, so there is no separate search step.
const { text, sources } = await generateText({
	model: aiSearch.get(INSTANCE_NAME).chat({
		ai_search_options: {
			// "hybrid" ranks results from both the vector and keyword indexes.
			retrieval: { retrieval_type: "hybrid", max_num_results: 5 },
		},
	}),
	messages: [{ role: "user", content: query }],
});

// `sources` holds the retrieved chunks, so you can cite them alongside `text`.
return Response.json({ text, sources });

const query = url.searchParams.get("q") ?? "How does caching work?";

// chat() returns an AI SDK model that retrieves matching chunks and generates
// a grounded answer in one call, so there is no separate search step.
const { text, sources } = await generateText({
	model: aiSearch.get(INSTANCE_NAME).chat({
		ai_search_options: {
			// "hybrid" ranks results from both the vector and keyword indexes.
			retrieval: { retrieval_type: "hybrid", max_num_results: 5 },
		},
	}),
	messages: [{ role: "user", content: query }],
});

// `sources` holds the retrieved chunks, so you can cite them alongside `text`.
return Response.json({ text, sources });

// streamText returns right away; tokens stream in as they are generated.
const result = streamText({
	model: aiSearch.get(INSTANCE_NAME).chat(),
	messages: [{ role: "user", content: query }],
});

// toTextStreamResponse() streams the generated text only.
return result.toTextStreamResponse();

// streamText returns right away; tokens stream in as they are generated.
const result = streamText({
	model: aiSearch.get(INSTANCE_NAME).chat(),
	messages: [{ role: "user", content: query }],
});

// toTextStreamResponse() streams the generated text only.
return result.toTextStreamResponse();

// sendSources forwards each retrieved chunk as a source-url part. They arrive
// before the first text part, so you can render citations as the answer streams.
return result.toUIMessageStreamResponse({ sendSources: true });

// sendSources forwards each retrieved chunk as a source-url part. They arrive
// before the first text part, so you can render citations as the answer streams.
return result.toUIMessageStreamResponse({ sendSources: true });

import { createWorkersAI } from "workers-ai-provider";
import { generateText, tool, stepCountIs } from "ai";
import { z } from "zod";

const instance = aiSearch.get(INSTANCE_NAME);
// The tool-caller must support function calling. Use a dedicated model here
// rather than the AI Search chat model, which retrieves on every request.
const workersai = createWorkersAI({ binding: env.AI });

const { text } = await generateText({
	model: workersai("@cf/zai-org/glm-5.2"),
	messages: [{ role: "user", content: query }],
	tools: {
		search_knowledge_base: tool({
			description: "Search the indexed knowledge base for relevant content.",
			inputSchema: z.object({
				query: z.string().describe("The search query"),
			}),
			// The model decides when to call this; it searches the instance.
			execute: async ({ query }) =>
				instance.search({
					query,
					ai_search_options: { retrieval: { max_num_results: 5 } },
				}),
		}),
	},
	// Cap the tool-call loop so the model cannot invoke tools indefinitely.
	stopWhen: stepCountIs(5),
});

import { createWorkersAI } from "workers-ai-provider";
import { generateText, tool, stepCountIs } from "ai";
import { z } from "zod";

const instance = aiSearch.get(INSTANCE_NAME);
// The tool-caller must support function calling. Use a dedicated model here
// rather than the AI Search chat model, which retrieves on every request.
const workersai = createWorkersAI({ binding: env.AI });

const { text } = await generateText({
	model: workersai("@cf/zai-org/glm-5.2"),
	messages: [{ role: "user", content: query }],
	tools: {
		search_knowledge_base: tool({
			description: "Search the indexed knowledge base for relevant content.",
			inputSchema: z.object({
				query: z.string().describe("The search query"),
			}),
			// The model decides when to call this; it searches the instance.
			execute: async ({ query }) =>
				instance.search({
					query,
					ai_search_options: { retrieval: { max_num_results: 5 } },
				}),
		}),
	},
	// Cap the tool-call loop so the model cannot invoke tools indefinitely.
	stopWhen: stepCountIs(5),
});