跨多个实例搜索

// 所有租户都可以搜索的共享实例。
const GENERAL_INSTANCE = "general-knowledge";
// 每个租户也有其自身的实例，仅保存该租户的内容。
const tenantInstance = (tenantId) => `tenant-${tenantId}`;

// 种子内容，以便每个实例在首次查询时都返回内容。在实际应用中，你会提前预置实例及其内容。
const GENERAL_DOC = `# Support hours
Support is available Monday to Friday, 9am to 5pm UTC for all customers.`;
const TENANT_DOC = `# Your plan
This account is on the Enterprise plan with a dedicated success manager.`;

export default {
	async fetch(request, env) {
		// 识别租户并读取查询。
		const tenantId = request.headers.get("x-tenant-id");
		if (!tenantId) {
			return new Response("Missing x-tenant-id header", { status: 400 });
		}
		const query = new URL(request.url).searchParams.get("q");
		if (!query) {
			return new Response("Add a ?q= query parameter", { status: 400 });
		}

		// 本 Demo 的一次性设置：创建两个实例并为各自填充种子数据。
		await ensureInstance(
			env,
			GENERAL_INSTANCE,
			"support-hours.md",
			GENERAL_DOC,
		);
		await ensureInstance(env, tenantInstance(tenantId), "plan.md", TENANT_DOC);

		// 在一次调用中搜索共享实例和该租户的实例。
		// AI Search 会扇出 (fan out) 到两者、合并结果，并用其来自的 instance_id 标记每个 chunk。你可以传递 1 到 10 个实例 ID。
		const results = await env.AI_SEARCH.search({
			query,
			ai_search_options: {
				instance_ids: [GENERAL_INSTANCE, tenantInstance(tenantId)],
			},
		});

		// 使用 instance_id 标记将共享结果与租户结果区分开来。
		const fromGeneral = results.chunks.filter(
			(chunk) => chunk.instance_id === GENERAL_INSTANCE,
		);
		const fromTenant = results.chunks.filter(
			(chunk) => chunk.instance_id === tenantInstance(tenantId),
		);

		return Response.json({
			query: results.search_query,
			general: fromGeneral.map((chunk) => ({
				key: chunk.item.key,
				text: chunk.text,
			})),
			tenant: fromTenant.map((chunk) => ({
				key: chunk.item.key,
				text: chunk.text,
			})),
			// 如果某个实例失败，另一个实例仍会返回。失败会落在此处。
			errors: results.errors,
		});
	},
};

// 创建带有内置存储的实例并填充一份文档种子。每次请求调用都是安全的：如果实例已存在则 create() 会抛出异常，因此 try/catch 会在就位后跳过设置。
async function ensureInstance(env, id, key, content) {
	try {
		await env.AI_SEARCH.create({ id });
		await env.AI_SEARCH.get(id).items.uploadAndPoll(key, content, {
			timeoutMs: 60_000,
		});
	} catch {
		// 实例已存在且已被填充种子数据。
	}
}

export interface Env {
	AI_SEARCH: AiSearchNamespace;
}

// 所有租户都可以搜索的共享实例。
const GENERAL_INSTANCE = "general-knowledge";
// 每个租户也有其自身的实例，仅保存该租户的内容。
const tenantInstance = (tenantId: string) => `tenant-${tenantId}`;

// 种子内容，以便每个实例在首次查询时都返回内容。在实际应用中，你会提前预置实例及其内容。
const GENERAL_DOC = `# Support hours
Support is available Monday to Friday, 9am to 5pm UTC for all customers.`;
const TENANT_DOC = `# Your plan
This account is on the Enterprise plan with a dedicated success manager.`;

export default {
	async fetch(request, env): Promise<Response> {
		// 识别租户并读取查询。
		const tenantId = request.headers.get("x-tenant-id");
		if (!tenantId) {
			return new Response("Missing x-tenant-id header", { status: 400 });
		}
		const query = new URL(request.url).searchParams.get("q");
		if (!query) {
			return new Response("Add a ?q= query parameter", { status: 400 });
		}

		// 本 Demo 的一次性设置：创建两个实例并为各自填充种子数据。
		await ensureInstance(
			env,
			GENERAL_INSTANCE,
			"support-hours.md",
			GENERAL_DOC,
		);
		await ensureInstance(env, tenantInstance(tenantId), "plan.md", TENANT_DOC);

		// 在一次调用中搜索共享实例和该租户的实例。
		// AI Search 会扇出 (fan out) 到两者、合并结果，并用其来自的 instance_id 标记每个 chunk。你可以传递 1 到 10 个实例 ID。
		const results = await env.AI_SEARCH.search({
			query,
			ai_search_options: {
				instance_ids: [GENERAL_INSTANCE, tenantInstance(tenantId)],
			},
		});

		// 使用 instance_id 标记将共享结果与租户结果区分开来。
		const fromGeneral = results.chunks.filter(
			(chunk) => chunk.instance_id === GENERAL_INSTANCE,
		);
		const fromTenant = results.chunks.filter(
			(chunk) => chunk.instance_id === tenantInstance(tenantId),
		);

		return Response.json({
			query: results.search_query,
			general: fromGeneral.map((chunk) => ({
				key: chunk.item.key,
				text: chunk.text,
			})),
			tenant: fromTenant.map((chunk) => ({
				key: chunk.item.key,
				text: chunk.text,
			})),
			// 如果某个实例失败，另一个实例仍会返回。失败会落在此处。
			errors: results.errors,
		});
	},
} satisfies ExportedHandler<Env>;

// 创建带有内置存储的实例并填充一份文档种子。每次请求调用都是安全的：如果实例已存在则 create() 会抛出异常，因此 try/catch 会在就位后跳过设置。
async function ensureInstance(
	env: Env,
	id: string,
	key: string,
	content: string,
) {
	try {
		await env.AI_SEARCH.create({ id });
		await env.AI_SEARCH.get(id).items.uploadAndPoll(key, content, {
			timeoutMs: 60_000,
		});
	} catch {
		// 实例已存在且已被填充种子数据。
	}
}