Files
shockbot/utils/vertex.ts
T
Colin McDonnell c43ed65c3b Add Vertex AI routing support (#753)
* add Vertex AI routing support

* include Vertex smokes in action CI
2026-05-20 16:49:08 +00:00

85 lines
3.0 KiB
TypeScript

import { randomUUID } from "node:crypto";
import { mkdirSync, rmSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { VERTEX_MODEL_ID_ENV } from "../models.ts";
export const VERTEX_SERVICE_ACCOUNT_JSON_ENV = "VERTEX_SERVICE_ACCOUNT_JSON";
export const GOOGLE_APPLICATION_CREDENTIALS_ENV = "GOOGLE_APPLICATION_CREDENTIALS";
export const GOOGLE_CLOUD_PROJECT_ENV = "GOOGLE_CLOUD_PROJECT";
export const VERTEX_LOCATION_ENV = "VERTEX_LOCATION";
export type VertexCredentials = {
credentialsPath: string;
secretDir: string;
};
function hasEnvVar(name: string): boolean {
const value = process.env[name];
return typeof value === "string" && value.length > 0;
}
export function isVertexRoute(model: string | undefined): boolean {
const vertexId = process.env[VERTEX_MODEL_ID_ENV]?.trim();
return model !== undefined && vertexId !== undefined && vertexId === model;
}
export function readProjectIdFromVertexServiceAccountJson(): string | undefined {
const blob = process.env[VERTEX_SERVICE_ACCOUNT_JSON_ENV];
if (!blob) return undefined;
try {
const parsed: unknown = JSON.parse(blob);
if (!parsed || typeof parsed !== "object" || !("project_id" in parsed)) {
return undefined;
}
const projectId = parsed.project_id;
return typeof projectId === "string" && projectId.length > 0 ? projectId : undefined;
} catch {
return undefined;
}
}
function createSecretDir(): string {
const base = process.env.PULLFROG_SECRET_HOME || process.env.HOME || homedir();
const secretDir = join(base, ".pullfrog", "secrets", randomUUID());
mkdirSync(secretDir, { recursive: true, mode: 0o700 });
return secretDir;
}
export function materializeVertexCredentials(params: {
model: string | undefined;
}): VertexCredentials | undefined {
if (!isVertexRoute(params.model)) return undefined;
const blob = process.env[VERTEX_SERVICE_ACCOUNT_JSON_ENV];
if (!blob) return undefined;
const secretDir = createSecretDir();
const credentialsPath = join(secretDir, "vertex-sa.json");
writeFileSync(credentialsPath, blob, { mode: 0o600 });
process.env[GOOGLE_APPLICATION_CREDENTIALS_ENV] = credentialsPath;
const projectId = readProjectIdFromVertexServiceAccountJson();
if (projectId && !hasEnvVar(GOOGLE_CLOUD_PROJECT_ENV)) {
process.env[GOOGLE_CLOUD_PROJECT_ENV] = projectId;
}
return { credentialsPath, secretDir };
}
export function cleanupVertexCredentials(credentials: VertexCredentials | undefined): void {
if (!credentials) return;
rmSync(credentials.secretDir, { recursive: true, force: true });
}
export function applyClaudeVertexEnv(env: Record<string, string | undefined>): void {
env.CLAUDE_CODE_USE_VERTEX = "1";
env.ANTHROPIC_VERTEX_PROJECT_ID ??= env[GOOGLE_CLOUD_PROJECT_ENV];
env.CLOUD_ML_REGION ??= env[VERTEX_LOCATION_ENV];
}
export function resolveVertexOpenCodeModel(model: string | undefined): string | undefined {
return isVertexRoute(model) && model ? `google-vertex/${model}` : undefined;
}