singularity-forge/packages/google-gemini-cli-provider/src/index.js

99 lines
No EOL
4.2 KiB
JavaScript

/**
* Google Gemini CLI transport helper.
*
* Purpose: keep the Gemini CLI Core auth and content-generator wiring in a
* dedicated workspace package so provider code can depend on one small helper
* instead of embedding the upstream integration inline.
*
* Consumer: `@singularity-forge/ai` Google Gemini provider, plus SF-side
* background catalog discovery.
*/
import { AuthType, CodeAssistServer, getOauthClient, makeFakeConfig, setupUser, } from "@google/gemini-cli-core";
import { createContentGenerator, createContentGeneratorConfig, } from "@google/gemini-cli-core/dist/src/core/contentGenerator.js";
/**
* Create a Gemini CLI Core content generator for a model.
*
* Purpose: centralize the Code Assist setup and OAuth bootstrap logic in a
* reusable package so SF's Gemini provider can stay focused on stream shaping.
*
* Consumer: the Google Gemini provider in pi-ai.
*/
export async function createGeminiCliContentGenerator(options) {
const cwd = options.cwd ?? process.cwd();
const config = makeFakeConfig({
model: options.modelId,
cwd,
targetDir: options.targetDir ?? cwd,
});
const generatorConfig = await createContentGeneratorConfig(config, AuthType.LOGIN_WITH_GOOGLE);
return createContentGenerator(generatorConfig, config);
}
/**
* Discover the active gemini-cli account: tier, project, and every model the
* account has access to (with per-model usage fraction and reset time).
*
* Best-effort: any failure (OAuth expired, no project, network) returns null
* silently so callers can downgrade gracefully.
*
* Consumer: SF-side background catalog cache, usage UI, capacity diagnostics.
*/
export async function snapshotGeminiCliAccount(cwd) {
try {
const config = makeFakeConfig({ cwd: cwd ?? process.cwd() });
const authClient = await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, config);
const userData = await setupUser(authClient, config);
const projectId = userData?.projectId;
if (!projectId || typeof projectId !== "string")
return null;
const server = new CodeAssistServer(authClient, projectId, { headers: {} });
const data = await server.retrieveUserQuota({ project: projectId });
// Dedup buckets per modelId, keeping the WORST quota (lowest
// remainingFraction). Code Assist sometimes returns multiple buckets
// for the same model when more than one quota window applies; the
// pessimistic choice is what every consumer (UI, capacity diagnostics,
// model picker) actually wants to surface.
const byModel = new Map();
for (const b of data?.buckets ?? []) {
const modelId = typeof b.modelId === "string" ? b.modelId : "";
if (!modelId)
continue;
const remainingFraction = typeof b.remainingFraction === "number" ? b.remainingFraction : 1;
const bucket = {
modelId,
usedFraction: 1 - remainingFraction,
remainingFraction,
resetTime: typeof b.resetTime === "string" ? b.resetTime : undefined,
};
const existing = byModel.get(modelId);
if (!existing || bucket.remainingFraction < existing.remainingFraction) {
byModel.set(modelId, bucket);
}
}
const models = Array.from(byModel.values()).sort((a, b) => a.modelId.localeCompare(b.modelId));
if (models.length === 0)
return null;
return {
projectId,
userTierId: typeof userData?.userTier === "string" ? userData.userTier : undefined,
userTierName: userData?.userTierName,
paidTier: userData?.paidTier
? { id: userData.paidTier.id, name: userData.paidTier.name }
: undefined,
models,
};
}
catch {
return null;
}
}
/**
* Convenience wrapper: just the model IDs the active gemini-cli account has
* access to. Returns null on failure (same contract as snapshotGeminiCliAccount).
*/
export async function discoverGeminiCliModels(cwd) {
const snap = await snapshotGeminiCliAccount(cwd);
if (!snap)
return null;
return snap.models.map((m) => m.modelId);
}
//# sourceMappingURL=index.js.map