Agents SDK

import { AIChatAgent } from "@cloudflare/ai-chat";
import { routeAgentRequest } from "agents";
import { createWorkersAI } from "workers-ai-provider";
import { streamText, convertToModelMessages, tool, stepCountIs } from "ai";
import { z } from "zod";

const INSTANCE_ID = "knowledge-base";

const SEED_DOC = `# Getting started
AI Search indexes your content so an agent can retrieve it at query time.`;

// AIChatAgent stores the conversation history and calls onChatMessage() once
// for each user message.
export class SearchAgent extends AIChatAgent {
	// Guard so the one-time instance setup runs only once per running agent.
	ready = false;

	// Create the agent's instance with hybrid search enabled, then seed it so
	// the first query has content. create() throws if the instance already
	// exists, so the try/catch makes this idempotent.
	async ensureInstance() {
		if (this.ready) return;
		try {
			// index_method with both vector and keyword enables hybrid search.
			await this.env.AI_SEARCH.create({
				id: INSTANCE_ID,
				index_method: { vector: true, keyword: true },
			});
			// upload() queues the file; indexing runs in the background. Poll the
			// item status until it is searchable so the first query has content.
			const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
			const { id } = await instance.items.upload(
				"getting-started.md",
				SEED_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();
			}
		} catch {
			// Instance already exists.
		}
		this.ready = true;
	}

	// Runs on every chat message: make sure the instance exists, then stream a
	// tool-using response.
	async onChatMessage() {
		await this.ensureInstance();

		// The agent exposes AI Search's search() as a tool instead of using AI
		// Search's own chat model, so this model must support function calling to
		// decide when to search and drive the tool calls.
		const workersai = createWorkersAI({ binding: this.env.AI });

		const result = streamText({
			model: workersai("@cf/zai-org/glm-5.2"),
			system:
				"You are a support assistant. Use search_knowledge_base to find " +
				"relevant content before answering, and cite what you use.",
			// this.messages is the stored chat history; convert it to the format
			// the model expects.
			messages: await convertToModelMessages(this.messages),
			tools: {
				search_knowledge_base: tool({
					description: "Search the knowledge base for relevant content.",
					inputSchema: z.object({
						query: z.string().describe("The user's question or search terms"),
					}),
					// Hybrid search runs by default because the instance indexes
					// both vectors and keywords.
					execute: async ({ query }) => {
						const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
						return await instance.search({
							query,
							ai_search_options: { retrieval: { max_num_results: 5 } },
						});
					},
				}),
				save_resolution: tool({
					description:
						"Save a resolved answer to the knowledge base for reuse.",
					inputSchema: z.object({
						title: z.string().describe("Short descriptive title"),
						content: z.string().describe("The resolution to save"),
					}),
					// upload() returns as soon as the file is queued; indexing then
					// finishes in the background.
					execute: async ({ title, content }) => {
						const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
						const item = await instance.items.upload(`${title}.md`, content);
						return { key: item.key, status: item.status };
					},
				}),
			},
			// Cap the tool-call loop so the agent cannot run tools indefinitely.
			stopWhen: stepCountIs(5),
		});

		// Stream the reply, including any tool activity, back to the client.
		return result.toUIMessageStreamResponse();
	}
}

export default {
	async fetch(request, env) {
		// Route the request to the matching agent instance, keyed by the URL.
		return (
			(await routeAgentRequest(request, env)) ||
			new Response("Not found", { status: 404 })
		);
	},
};

import { AIChatAgent } from "@cloudflare/ai-chat";
import { routeAgentRequest } from "agents";
import { createWorkersAI } from "workers-ai-provider";
import { streamText, convertToModelMessages, tool, stepCountIs } from "ai";
import { z } from "zod";

const INSTANCE_ID = "knowledge-base";

const SEED_DOC = `# Getting started
AI Search indexes your content so an agent can retrieve it at query time.`;

// AIChatAgent stores the conversation history and calls onChatMessage() once
// for each user message.
export class SearchAgent extends AIChatAgent {
	// Guard so the one-time instance setup runs only once per running agent.
	private ready = false;

	// Create the agent's instance with hybrid search enabled, then seed it so
	// the first query has content. create() throws if the instance already
	// exists, so the try/catch makes this idempotent.
	private async ensureInstance() {
		if (this.ready) return;
		try {
			// index_method with both vector and keyword enables hybrid search.
			await this.env.AI_SEARCH.create({
				id: INSTANCE_ID,
				index_method: { vector: true, keyword: true },
			});
			// upload() queues the file; indexing runs in the background. Poll the
			// item status until it is searchable so the first query has content.
			const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
			const { id } = await instance.items.upload(
				"getting-started.md",
				SEED_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();
			}
		} catch {
			// Instance already exists.
		}
		this.ready = true;
	}

	// Runs on every chat message: make sure the instance exists, then stream a
	// tool-using response.
	async onChatMessage() {
		await this.ensureInstance();

		// The agent exposes AI Search's search() as a tool instead of using AI
		// Search's own chat model, so this model must support function calling to
		// decide when to search and drive the tool calls.
		const workersai = createWorkersAI({ binding: this.env.AI });

		const result = streamText({
			model: workersai("@cf/zai-org/glm-5.2"),
			system:
				"You are a support assistant. Use search_knowledge_base to find " +
				"relevant content before answering, and cite what you use.",
			// this.messages is the stored chat history; convert it to the format
			// the model expects.
			messages: await convertToModelMessages(this.messages),
			tools: {
				search_knowledge_base: tool({
					description: "Search the knowledge base for relevant content.",
					inputSchema: z.object({
						query: z.string().describe("The user's question or search terms"),
					}),
					// Hybrid search runs by default because the instance indexes
					// both vectors and keywords.
					execute: async ({ query }) => {
						const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
						return await instance.search({
							query,
							ai_search_options: { retrieval: { max_num_results: 5 } },
						});
					},
				}),
				save_resolution: tool({
					description:
						"Save a resolved answer to the knowledge base for reuse.",
					inputSchema: z.object({
						title: z.string().describe("Short descriptive title"),
						content: z.string().describe("The resolution to save"),
					}),
					// upload() returns as soon as the file is queued; indexing then
					// finishes in the background.
					execute: async ({ title, content }) => {
						const instance = this.env.AI_SEARCH.get(INSTANCE_ID);
						const item = await instance.items.upload(`${title}.md`, content);
						return { key: item.key, status: item.status };
					},
				}),
			},
			// Cap the tool-call loop so the agent cannot run tools indefinitely.
			stopWhen: stepCountIs(5),
		});

		// Stream the reply, including any tool activity, back to the client.
		return result.toUIMessageStreamResponse();
	}
}

export default {
	async fetch(request: Request, env: Env) {
		// Route the request to the matching agent instance, keyed by the URL.
		return (
			(await routeAgentRequest(request, env)) ||
			new Response("Not found", { status: 404 })
		);
	},
} satisfies ExportedHandler<Env>;