複数インスタンスを横断して検索する

// The shared instance that every tenant can search.
const GENERAL_INSTANCE = "general-knowledge";
// Each tenant also has its own instance, holding only that tenant's content.
const tenantInstance = (tenantId) => `tenant-${tenantId}`;

// Seed content, so each instance returns something on the first query. In a
// real app you would provision instances and their content ahead of time.
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) {
		// Identify the tenant and read the query.
		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 });
		}

		// One-time setup for this demo: create both instances and seed each.
		await ensureInstance(
			env,
			GENERAL_INSTANCE,
			"support-hours.md",
			GENERAL_DOC,
		);
		await ensureInstance(env, tenantInstance(tenantId), "plan.md", TENANT_DOC);

		// Search the shared instance and this tenant's instance in one call.
		// AI Search fans out to both, merges the results, and tags each chunk
		// with the instance_id it came from. You can pass 1 to 10 instance IDs.
		const results = await env.AI_SEARCH.search({
			query,
			ai_search_options: {
				instance_ids: [GENERAL_INSTANCE, tenantInstance(tenantId)],
			},
		});

		// Use the instance_id tag to separate shared results from tenant results.
		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,
			})),
			// If one instance fails, the other still returns. Failures land here.
			errors: results.errors,
		});
	},
};

// Create an instance with built-in storage and seed one document. Safe to call
// on every request: create() throws if the instance already exists, so the
// try/catch skips setup once it is in place.
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 {
		// The instance already exists and has been seeded.
	}
}

export interface Env {
	AI_SEARCH: AiSearchNamespace;
}

// The shared instance that every tenant can search.
const GENERAL_INSTANCE = "general-knowledge";
// Each tenant also has its own instance, holding only that tenant's content.
const tenantInstance = (tenantId: string) => `tenant-${tenantId}`;

// Seed content, so each instance returns something on the first query. In a
// real app you would provision instances and their content ahead of time.
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> {
		// Identify the tenant and read the query.
		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 });
		}

		// One-time setup for this demo: create both instances and seed each.
		await ensureInstance(
			env,
			GENERAL_INSTANCE,
			"support-hours.md",
			GENERAL_DOC,
		);
		await ensureInstance(env, tenantInstance(tenantId), "plan.md", TENANT_DOC);

		// Search the shared instance and this tenant's instance in one call.
		// AI Search fans out to both, merges the results, and tags each chunk
		// with the instance_id it came from. You can pass 1 to 10 instance IDs.
		const results = await env.AI_SEARCH.search({
			query,
			ai_search_options: {
				instance_ids: [GENERAL_INSTANCE, tenantInstance(tenantId)],
			},
		});

		// Use the instance_id tag to separate shared results from tenant results.
		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,
			})),
			// If one instance fails, the other still returns. Failures land here.
			errors: results.errors,
		});
	},
} satisfies ExportedHandler<Env>;

// Create an instance with built-in storage and seed one document. Safe to call
// on every request: create() throws if the instance already exists, so the
// try/catch skips setup once it is in place.
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 {
		// The instance already exists and has been seeded.
	}
}