ai-functions
Version:
Core AI primitives for building intelligent applications
192 lines • 7.98 kB
JavaScript
/**
* Google GenAI (Gemini) Adapter
*
* Google doesn't have a native batch API like OpenAI/Anthropic, so this
* adapter fakes batch processing via concurrent direct calls and tracks the
* job state locally (see `LocalJobStore` in `./provider.js`).
*
* For true async batch processing, consider Google Cloud Batch with Vertex AI.
*
* @see https://ai.google.dev/gemini-api/docs
*
* @packageDocumentation
*/
import { LocalJobStore, processConcurrently, registerBatchAdapter, registerFlexAdapter, tryParseJson, } from './provider.js';
// ============================================================================
// Google GenAI client configuration
// ============================================================================
let googleApiKey;
let googleBaseUrl = 'https://generativelanguage.googleapis.com/v1beta';
// AI Gateway configuration (optional - for routing through Cloudflare AI Gateway)
let gatewayUrl;
let gatewayToken;
/** Configure the Google GenAI client. */
export function configureGoogleGenAI(options) {
if (options.apiKey)
googleApiKey = options.apiKey;
if (options.baseUrl)
googleBaseUrl = options.baseUrl;
if (options.gatewayUrl)
gatewayUrl = options.gatewayUrl;
if (options.gatewayToken)
gatewayToken = options.gatewayToken;
}
function getConfig() {
const gwUrl = gatewayUrl || process.env['AI_GATEWAY_URL'];
const gwToken = gatewayToken || process.env['AI_GATEWAY_TOKEN'];
if (gwUrl && gwToken) {
return {
apiKey: '',
baseUrl: googleBaseUrl,
gatewayUrl: gwUrl,
gatewayToken: gwToken,
};
}
const key = googleApiKey || process.env['GOOGLE_API_KEY'] || process.env['GEMINI_API_KEY'];
if (!key) {
throw new Error('Google API key not configured. Set GOOGLE_API_KEY or GEMINI_API_KEY, or use AI_GATEWAY_URL and AI_GATEWAY_TOKEN');
}
return { apiKey: key, baseUrl: googleBaseUrl };
}
// ============================================================================
// Local job tracking
// ============================================================================
const jobs = new LocalJobStore('google_batch');
// ============================================================================
// Google GenAI batch adapter (BatchProvider port)
// ============================================================================
const googleAdapter = {
async submit(items, options) {
const model = options.model || 'gemini-2.0-flash';
const { id, state } = jobs.create(items, options);
// Drive the job state machine in the background; `waitForCompletion`
// will poll the in-memory state.
const completion = (async () => {
state.status = 'in_progress';
const results = await processConcurrently(items, (item) => processGoogleItem(item, model), {
concurrency: 10,
onWaveComplete: (partial) => {
state.results = partial;
},
});
state.results = results;
state.status = results.every((r) => r.status === 'completed') ? 'completed' : 'failed';
state.completedAt = new Date();
return results;
})();
const job = {
id,
provider: 'google',
status: 'pending',
totalItems: items.length,
completedItems: 0,
failedItems: 0,
createdAt: state.createdAt,
...(options.webhookUrl !== undefined && { webhookUrl: options.webhookUrl }),
};
return { job, completion };
},
async getStatus(batchId) {
return jobs.snapshot(batchId, 'google');
},
async cancel(batchId) {
if (jobs.has(batchId)) {
jobs.get(batchId).status = 'cancelled';
}
},
async getResults(batchId) {
return jobs.get(batchId).results;
},
async waitForCompletion(batchId, pollInterval = 1000) {
return jobs.waitForCompletion(batchId, pollInterval);
},
};
// ============================================================================
// Google GenAI flex adapter (FlexAdapter port)
// ============================================================================
const googleFlexAdapter = {
async submitFlex(items, options) {
const model = options.model || 'gemini-2.0-flash';
return processConcurrently(items, (item) => processGoogleItem(item, model), {
concurrency: 10,
});
},
};
// ============================================================================
// Per-item processing
// ============================================================================
async function processGoogleItem(item, model) {
const config = getConfig();
if (config.gatewayUrl && config.gatewayToken) {
return processGoogleItemViaGateway(item, config, model);
}
const modelName = model.startsWith('models/') ? model : `models/${model}`;
const url = `${config.baseUrl}/${modelName}:generateContent?key=${config.apiKey}`;
const response = await fetch(url, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(buildGeminiBody(item)),
});
if (!response.ok) {
const error = await response.text();
throw new Error(`Google GenAI API error: ${response.status} ${error}`);
}
return parseGeminiResponse(item, (await response.json()));
}
/**
* Process a Google GenAI item via Cloudflare AI Gateway.
* Gateway URL format: {gateway_url}/google-ai-studio/v1beta/models/{model}:generateContent
*/
async function processGoogleItemViaGateway(item, config, model) {
const modelName = model.startsWith('models/') ? model.replace('models/', '') : model;
const url = `${config.gatewayUrl}/google-ai-studio/v1beta/models/${modelName}:generateContent`;
const response = await fetch(url, {
method: 'POST',
headers: {
'cf-aig-authorization': `Bearer ${config.gatewayToken}`,
'Content-Type': 'application/json',
},
body: JSON.stringify(buildGeminiBody(item)),
});
if (!response.ok) {
const error = await response.text();
throw new Error(`Google GenAI via Gateway error: ${response.status} ${error}`);
}
return parseGeminiResponse(item, (await response.json()));
}
function buildGeminiBody(item) {
// Gemini handles system instructions as part of the user message.
const userText = item.options?.system
? `System instruction: ${item.options.system}\n\nUser request: ${item.prompt}`
: item.prompt;
const contents = [{ role: 'user', parts: [{ text: userText }] }];
const generationConfig = {
maxOutputTokens: item.options?.maxTokens || 8192,
...(item.options?.temperature !== undefined && { temperature: item.options.temperature }),
...(item.schema && { responseMimeType: 'application/json' }),
};
return { contents, generationConfig };
}
function parseGeminiResponse(item, data) {
const content = data.candidates?.[0]?.content?.parts?.[0]?.text;
return {
id: item.id,
customId: item.id,
status: 'completed',
result: tryParseJson(content, !!item.schema),
...(data.usageMetadata && {
usage: {
promptTokens: data.usageMetadata.promptTokenCount,
completionTokens: data.usageMetadata.candidatesTokenCount,
totalTokens: data.usageMetadata.totalTokenCount,
},
}),
};
}
// ============================================================================
// Register adapters
// ============================================================================
registerBatchAdapter('google', googleAdapter);
registerFlexAdapter('google', googleFlexAdapter);
export { googleAdapter, googleFlexAdapter };
//# sourceMappingURL=google.js.map