Merge pull request #2289 from giselles-ai/feat/add-ai-gateway-headers-to-github-ingest

Add stripe-customer-id header to GitHub ingest flow

Author: shige

apps/studio.giselles.ai/lib/vector-stores/document/ingest/ingest-document.ts

@@ -11,6 +11,7 @@ import { and, eq, lt, or } from "drizzle-orm";
 import { db, documentVectorStoreSources } from "@/db";
 import type { TeamWithSubscription } from "@/services/teams";
 import { fetchTeamByDbId } from "@/services/teams/fetch-team";
+import { buildAiGatewayHeaders } from "../../shared/ai-gateway-headers";
 import {
 	deleteDocumentEmbeddingsByProfiles,
 	getDocumentVectorStoreSource,
@@ -117,30 +118,6 @@ async function resolveEmbeddingTelemetryContext(
 	}
 }
 
-function buildAiGatewayHeaders(
-	telemetryContext: DocumentEmbeddingTelemetryContext | null,
-): Record<string, string> | undefined {
-	const headers: Record<string, string> = {
-		"http-referer":
-			process.env.AI_GATEWAY_HTTP_REFERER ?? "https://giselles.ai",
-		"x-title": process.env.AI_GATEWAY_X_TITLE ?? "Giselle",
-	};
-
-	const stripeCustomerId = telemetryContext?.team.activeCustomerId ?? undefined;
-	const teamPlan = telemetryContext?.team.plan;
-	if (stripeCustomerId !== undefined) {
-		headers["stripe-customer-id"] = stripeCustomerId;
-		headers["stripe-restricted-access-key"] =
-			process.env.STRIPE_AI_GATEWAY_RESTRICTED_ACCESS_KEY ?? "";
-	} else if (teamPlan === "pro" || teamPlan === "team") {
-		console.warn(
-			`Stripe customer ID not found for document ingest (team: ${telemetryContext?.team.id})`,
-		);
-	}
-
-	return headers;
-}
-
 /**
  * Ingest a document source by extracting, chunking, and embedding text content
  * This function:
@@ -289,7 +266,9 @@ export async function ingestDocument(
 		const insertedProfileIds: EmbeddingProfileId[] = [];
 
 		// Build AI Gateway headers for billing attribution
-		const aiGatewayHeaders = buildAiGatewayHeaders(telemetryContext);
+		const aiGatewayHeaders = buildAiGatewayHeaders(
+			telemetryContext?.team ?? null,
+		);
 
 		try {
 			for (const embeddingProfileId of embeddingProfileIds) {

apps/studio.giselles.ai/lib/vector-stores/github/ingest/blobs/ingest-github-blobs.ts

@@ -24,6 +24,7 @@ export async function ingestGitHubBlobs(params: {
 	teamDbId: number;
 	embeddingProfileId: EmbeddingProfileId;
 	embeddingComplete?: EmbeddingCompleteCallback;
+	headers?: Record<string, string>;
 }): Promise<void> {
 	const { repositoryIndexDbId, isInitialIngest } = await getRepositoryIndexInfo(
 		params.source,
@@ -55,6 +56,7 @@ export async function ingestGitHubBlobs(params: {
 		}),
 		embeddingProfileId: params.embeddingProfileId,
 		embeddingComplete: params.embeddingComplete,
+		headers: params.headers,
 	});
 
 	const result = await ingest();

apps/studio.giselles.ai/lib/vector-stores/github/ingest/issues/ingest-github-issues.ts

@@ -21,6 +21,7 @@ export async function ingestGitHubIssues(params: {
 	teamDbId: number;
 	embeddingProfileId: EmbeddingProfileId;
 	embeddingComplete?: EmbeddingCompleteCallback;
+	headers?: Record<string, string>;
 }): Promise<void> {
 	const { repositoryIndexDbId } = await getRepositoryIndexInfo(
 		params.source,
@@ -80,6 +81,7 @@ export async function ingestGitHubIssues(params: {
 		},
 		embeddingProfileId: params.embeddingProfileId,
 		embeddingComplete: params.embeddingComplete,
+		headers: params.headers,
 	});
 
 	const result = await ingest();

apps/studio.giselles.ai/lib/vector-stores/github/ingest/process-repository.ts

@@ -15,6 +15,7 @@ import {
 } from "@/db";
 import type { TeamWithSubscription } from "@/services/teams";
 import { fetchTeamByDbId } from "@/services/teams/fetch-team";
+import { buildAiGatewayHeaders } from "../../shared/ai-gateway-headers";
 import type { RepositoryWithStatuses } from "../types";
 import {
 	createBlobMetadata,
@@ -76,12 +77,15 @@ const CONTENT_PROCESSORS: Record<
 			},
 		});
 
+		const headers = buildAiGatewayHeaders(team);
+
 		await ingestGitHubBlobs({
 			octokitClient: octokit,
 			source: { owner, repo, commitSha: commit.sha },
 			teamDbId,
 			embeddingProfileId,
 			embeddingComplete,
+			headers,
 		});
 
 		return {
@@ -108,12 +112,15 @@ const CONTENT_PROCESSORS: Record<
 			},
 		});
 
+		const headers = buildAiGatewayHeaders(team);
+
 		await ingestGitHubPullRequests({
 			githubAuthConfig: buildGitHubAuthConfig(installationId),
 			source: { owner, repo },
 			teamDbId,
 			embeddingProfileId,
 			embeddingComplete,
+			headers,
 		});
 
 		const lastPrNumber = await getLastIngestedPrNumber(
@@ -144,12 +151,15 @@ const CONTENT_PROCESSORS: Record<
 			},
 		});
 
+		const headers = buildAiGatewayHeaders(team);
+
 		await ingestGitHubIssues({
 			githubAuthConfig: buildGitHubAuthConfig(installationId),
 			source: { owner, repo },
 			teamDbId,
 			embeddingProfileId,
 			embeddingComplete,
+			headers,
 		});
 
 		const lastIssueNumber = await getLastIngestedIssueNumber(

apps/studio.giselles.ai/lib/vector-stores/github/ingest/pull-requests/ingest-github-pull-requests.ts

@@ -26,6 +26,7 @@ export async function ingestGitHubPullRequests(params: {
 	teamDbId: number;
 	embeddingProfileId: EmbeddingProfileId;
 	embeddingComplete?: EmbeddingCompleteCallback;
+	headers?: Record<string, string>;
 }): Promise<void> {
 	const { repositoryIndexDbId } = await getRepositoryIndexInfo(
 		params.source,
@@ -59,6 +60,7 @@ export async function ingestGitHubPullRequests(params: {
 		}),
 		embeddingProfileId: params.embeddingProfileId,
 		embeddingComplete: params.embeddingComplete,
+		headers: params.headers,
 	});
 
 	const result = await ingest();

apps/studio.giselles.ai/lib/vector-stores/shared/ai-gateway-headers.ts

@@ -0,0 +1,29 @@
+import type { TeamWithSubscription } from "@/services/teams";
+
+/**
+ * Build AI Gateway headers for billing attribution.
+ * This function creates headers required for AI Gateway requests,
+ * including Stripe customer ID for usage tracking.
+ */
+export function buildAiGatewayHeaders(
+	team: TeamWithSubscription | null,
+): Record<string, string> {
+	const headers: Record<string, string> = {
+		"http-referer":
+			process.env.AI_GATEWAY_HTTP_REFERER ?? "https://giselles.ai",
+		"x-title": process.env.AI_GATEWAY_X_TITLE ?? "Giselle",
+	};
+
+	const stripeCustomerId = team?.activeCustomerId ?? undefined;
+	if (stripeCustomerId !== undefined) {
+		headers["stripe-customer-id"] = stripeCustomerId;
+		headers["stripe-restricted-access-key"] =
+			process.env.STRIPE_AI_GATEWAY_RESTRICTED_ACCESS_KEY ?? "";
+	} else if (team?.plan === "pro" || team?.plan === "team") {
+		console.warn(
+			`Stripe customer ID not found for vector store ingest (team: ${team?.id})`,
+		);
+	}
+
+	return headers;
+}

packages/rag/src/ingest/pipeline.ts

@@ -45,6 +45,9 @@ export interface IngestPipelineOptions<
 	parallelLimit?: number;
 	onProgress?: (progress: IngestProgress) => void;
 	onError?: (error: IngestError) => void;
+
+	// Optional headers for AI Gateway requests
+	headers?: Record<string, string>;
 }
 
 const DEFAULT_MAX_BATCH_SIZE = 100;
@@ -113,6 +116,7 @@ export function createPipeline<
 		{
 			embeddingComplete,
 			transport: useGateway ? "gateway" : "provider",
+			headers: useGateway ? options.headers : undefined,
 		},
 	);