mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-12 02:02:13 +03:00
399 lines
16 KiB
TypeScript
399 lines
16 KiB
TypeScript
import { SignJWT, importPKCS8 } from "jose";
|
|
import { BaseExecutor, ExecuteInput } from "./base.ts";
|
|
import { PROVIDERS } from "../config/constants.ts";
|
|
|
|
interface ServiceAccount {
|
|
type: string;
|
|
project_id: string;
|
|
private_key_id: string;
|
|
private_key: string;
|
|
client_email: string;
|
|
[key: string]: unknown;
|
|
}
|
|
|
|
const TOKEN_CACHE = new Map<string, { token: string; expiresAt: number }>();
|
|
|
|
// OAuth scopes minted into the Vertex SA access token.
|
|
// - cloud-platform authorizes Vertex AI (aiplatform.googleapis.com) for chat/image execution.
|
|
// - generative-language.retriever is ADDITIONALLY required so model discovery can list the live
|
|
// catalog from generativelanguage.googleapis.com/v1beta/models — without it that listing returns
|
|
// 403 ACCESS_TOKEN_SCOPE_INSUFFICIENT and discovery silently falls back to the static ~10-model
|
|
// registry list. The extra scope is harmless for execution (cloud-platform still present) and for
|
|
// projects where it isn't needed (the mint never validates scope availability).
|
|
export const VERTEX_OAUTH_SCOPES = [
|
|
"https://www.googleapis.com/auth/cloud-platform",
|
|
"https://www.googleapis.com/auth/generative-language.retriever",
|
|
] as const;
|
|
|
|
export function parseSAFromApiKey(apiKey: string): ServiceAccount {
|
|
try {
|
|
return JSON.parse(apiKey);
|
|
} catch {
|
|
throw new Error("Vertex AI requires a valid Service Account JSON as the API key");
|
|
}
|
|
}
|
|
|
|
/**
|
|
* A Service Account credential is a JSON object (type/client_email/private_key). A Vertex AI
|
|
* Express-mode API key is an opaque non-JSON string. Distinguishing them lets the executor
|
|
* support BOTH: Service Account JSON (JWT → OAuth → project-scoped endpoint + Bearer auth) and
|
|
* Express keys (project-less publisher endpoint + x-goog-api-key auth), instead of failing every
|
|
* Express key with "requires a valid Service Account JSON".
|
|
*/
|
|
export function looksLikeServiceAccountJson(apiKey: string): boolean {
|
|
if (!apiKey || typeof apiKey !== "string") return false;
|
|
try {
|
|
const parsed = JSON.parse(apiKey);
|
|
return !!parsed && typeof parsed === "object" && !Array.isArray(parsed);
|
|
} catch {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/** True for a Vertex AI Express-mode API key (a non-empty, non-JSON, non-OAuth credential). */
|
|
export function isExpressApiKey(apiKey?: string | null): boolean {
|
|
return typeof apiKey === "string" && apiKey.trim().length > 0 && !looksLikeServiceAccountJson(apiKey);
|
|
}
|
|
|
|
export async function getAccessToken(sa: ServiceAccount): Promise<string> {
|
|
if (!sa.client_email || !sa.private_key) {
|
|
throw new Error(
|
|
"Service Account JSON is missing required fields (client_email or private_key)"
|
|
);
|
|
}
|
|
|
|
const cacheKey = sa.client_email;
|
|
const cached = TOKEN_CACHE.get(cacheKey);
|
|
|
|
// Buffer of 60 seconds
|
|
if (cached && Date.now() < cached.expiresAt - 60_000) {
|
|
return cached.token;
|
|
}
|
|
|
|
const privateKey = await importPKCS8(sa.private_key, "RS256");
|
|
const now = Math.floor(Date.now() / 1000);
|
|
|
|
const jwt = await new SignJWT({
|
|
iss: sa.client_email,
|
|
sub: sa.client_email,
|
|
aud: "https://oauth2.googleapis.com/token",
|
|
iat: now,
|
|
exp: now + 3600,
|
|
scope: VERTEX_OAUTH_SCOPES.join(" "),
|
|
})
|
|
.setProtectedHeader({ alg: "RS256", kid: sa.private_key_id })
|
|
.sign(privateKey);
|
|
|
|
const tokenRes = await fetch("https://oauth2.googleapis.com/token", {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/x-www-form-urlencoded" },
|
|
body: new URLSearchParams({
|
|
grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer",
|
|
assertion: jwt,
|
|
}),
|
|
});
|
|
|
|
if (!tokenRes.ok) {
|
|
const errorText = await tokenRes.text();
|
|
throw new Error(
|
|
`Failed to exchange JWT for Vertex access token: ${tokenRes.status} ${errorText}`
|
|
);
|
|
}
|
|
|
|
const tokenData = await tokenRes.json();
|
|
const accessToken = tokenData.access_token;
|
|
|
|
if (!accessToken) {
|
|
throw new Error("Vertex AI token exchange succeeded but no access_token found");
|
|
}
|
|
|
|
TOKEN_CACHE.set(cacheKey, {
|
|
token: accessToken,
|
|
expiresAt: (now + 3600) * 1000,
|
|
});
|
|
|
|
return accessToken;
|
|
}
|
|
|
|
const PARTNER_MODELS = new Set([
|
|
// Generic prefix, not pinned to a version: every Claude model on Vertex is an Anthropic
|
|
// partner model, never a Google-publisher one. Pinned prefixes (e.g. "claude-3-5-sonnet")
|
|
// silently break every time Anthropic ships a new generation — see issue #1985.
|
|
"claude-",
|
|
"deepseek-v3",
|
|
"deepseek-v3.2",
|
|
"deepseek-v4",
|
|
"deepseek-deepseek-r1",
|
|
"qwen3-next-80b",
|
|
"qwen3.6-",
|
|
"llama-3.1",
|
|
"mistral-",
|
|
"glm-5",
|
|
"glm-5.1",
|
|
"meta/llama",
|
|
]);
|
|
|
|
function isPartnerModel(model: string) {
|
|
const normalizedModel = model.toLowerCase();
|
|
return [...PARTNER_MODELS].some((prefix) => normalizedModel.startsWith(prefix));
|
|
}
|
|
|
|
// Anthropic models need their own branch: they use Vertex's native Anthropic Messages API
|
|
// (publishers/anthropic/.../rawPredict), not the generic OpenAI-compatible partner endpoint the
|
|
// other PARTNER_MODELS entries (DeepSeek, Qwen, Llama, Mistral, GLM) go through — the OpenAI-shaped
|
|
// endpoint 404s/"malformed argument"s for Claude models on at least some projects.
|
|
function isClaudeModel(model: string) {
|
|
return model.toLowerCase().startsWith("claude-");
|
|
}
|
|
|
|
// Defensive normalizer: target-format resolution for manually-added custom Claude models under
|
|
// "vertex"/"vertex-partner" was observed sending a Gemini-shaped body (contents/parts) to the
|
|
// Anthropic rawPredict endpoint instead of the configured "claude" format, causing a hard
|
|
// "messages: Field required" error upstream regardless of the stored per-model targetFormat. This
|
|
// converts a Gemini-shaped body to Anthropic Messages shape so the executor works either way,
|
|
// independent of that unresolved upstream resolution gap.
|
|
function toAnthropicBody(body: Record<string, unknown>): Record<string, unknown> {
|
|
const contents = body.contents as Array<{ role?: string; parts?: Array<{ text?: string }> }> | undefined;
|
|
if (!Array.isArray(contents)) return body;
|
|
|
|
const messages = contents.map((c) => ({
|
|
role: c.role === "model" ? "assistant" : "user",
|
|
content: (c.parts || []).map((p) => p.text || "").join(""),
|
|
}));
|
|
const generationConfig = body.generationConfig as { maxOutputTokens?: number } | undefined;
|
|
const systemInstruction = body.systemInstruction as { parts?: Array<{ text?: string }> } | undefined;
|
|
|
|
const converted: Record<string, unknown> = {
|
|
messages,
|
|
max_tokens: generationConfig?.maxOutputTokens || 4096,
|
|
};
|
|
if (systemInstruction?.parts?.length) {
|
|
converted.system = systemInstruction.parts.map((p) => p.text || "").join("");
|
|
}
|
|
return converted;
|
|
}
|
|
|
|
// rawPredict always returns a single complete JSON body, never real SSE framing (see buildUrl).
|
|
// When the caller actually requested a stream, synthesize a genuine Anthropic-native event
|
|
// sequence from that JSON so the existing claude-to-openai.ts (and sibling) response translators
|
|
// — which already parse real message_start/content_block_*/message_delta/message_stop events —
|
|
// can consume it correctly, instead of relying on the OpenAI-`choices`-only JSON→SSE fallback
|
|
// (open-sse/utils/jsonToSse.ts) which cannot represent Anthropic's native response shape at all.
|
|
function synthesizeClaudeSse(response: Record<string, unknown>): string {
|
|
const messageId = typeof response.id === "string" ? response.id : `msg_${Date.now()}`;
|
|
const model = typeof response.model === "string" ? response.model : "";
|
|
const usage = (response.usage as Record<string, unknown>) || {};
|
|
const stopReason = typeof response.stop_reason === "string" ? response.stop_reason : "end_turn";
|
|
const stopSequence = (response.stop_sequence as string | null | undefined) ?? null;
|
|
const content = Array.isArray(response.content) ? response.content : [];
|
|
|
|
const events: Array<{ event: string; data: Record<string, unknown> }> = [];
|
|
|
|
events.push({
|
|
event: "message_start",
|
|
data: {
|
|
type: "message_start",
|
|
message: {
|
|
id: messageId,
|
|
type: "message",
|
|
role: "assistant",
|
|
content: [],
|
|
model,
|
|
stop_reason: null,
|
|
stop_sequence: null,
|
|
usage: { input_tokens: usage.input_tokens || 0, output_tokens: 0 },
|
|
},
|
|
},
|
|
});
|
|
|
|
content.forEach((block: Record<string, unknown>, index: number) => {
|
|
if (block.type === "text") {
|
|
events.push({
|
|
event: "content_block_start",
|
|
data: { type: "content_block_start", index, content_block: { type: "text", text: "" } },
|
|
});
|
|
if (block.text) {
|
|
events.push({
|
|
event: "content_block_delta",
|
|
data: {
|
|
type: "content_block_delta",
|
|
index,
|
|
delta: { type: "text_delta", text: block.text },
|
|
},
|
|
});
|
|
}
|
|
events.push({ event: "content_block_stop", data: { type: "content_block_stop", index } });
|
|
} else if (block.type === "tool_use") {
|
|
events.push({
|
|
event: "content_block_start",
|
|
data: {
|
|
type: "content_block_start",
|
|
index,
|
|
content_block: { type: "tool_use", id: block.id, name: block.name, input: {} },
|
|
},
|
|
});
|
|
events.push({
|
|
event: "content_block_delta",
|
|
data: {
|
|
type: "content_block_delta",
|
|
index,
|
|
delta: { type: "input_json_delta", partial_json: JSON.stringify(block.input ?? {}) },
|
|
},
|
|
});
|
|
events.push({ event: "content_block_stop", data: { type: "content_block_stop", index } });
|
|
} else if (block.type === "thinking") {
|
|
events.push({
|
|
event: "content_block_start",
|
|
data: { type: "content_block_start", index, content_block: { type: "thinking", thinking: "" } },
|
|
});
|
|
if (block.thinking) {
|
|
events.push({
|
|
event: "content_block_delta",
|
|
data: {
|
|
type: "content_block_delta",
|
|
index,
|
|
delta: { type: "thinking_delta", thinking: block.thinking },
|
|
},
|
|
});
|
|
}
|
|
events.push({ event: "content_block_stop", data: { type: "content_block_stop", index } });
|
|
}
|
|
});
|
|
|
|
events.push({
|
|
event: "message_delta",
|
|
data: {
|
|
type: "message_delta",
|
|
delta: { stop_reason: stopReason, stop_sequence: stopSequence },
|
|
usage: { output_tokens: usage.output_tokens || 0 },
|
|
},
|
|
});
|
|
|
|
events.push({ event: "message_stop", data: { type: "message_stop" } });
|
|
|
|
return events.map((e) => `event: ${e.event}\ndata: ${JSON.stringify(e.data)}\n\n`).join("");
|
|
}
|
|
|
|
export class VertexExecutor extends BaseExecutor {
|
|
constructor() {
|
|
super("vertex", PROVIDERS.vertex);
|
|
}
|
|
|
|
async execute(input: ExecuteInput) {
|
|
const { credentials, log, model, stream } = input;
|
|
// Defensive: trim stray surrounding whitespace from a pasted credential.
|
|
if (typeof credentials.apiKey === "string") {
|
|
credentials.apiKey = credentials.apiKey.trim();
|
|
}
|
|
// Service Account JSON → mint a short-lived OAuth token (Bearer). An Express-mode API key is
|
|
// sent as-is via x-goog-api-key (see buildHeaders), so no token exchange is needed for it.
|
|
if (credentials.apiKey && !credentials.accessToken && looksLikeServiceAccountJson(credentials.apiKey)) {
|
|
try {
|
|
const sa = parseSAFromApiKey(credentials.apiKey);
|
|
credentials.accessToken = await getAccessToken(sa);
|
|
} catch (err: any) {
|
|
log?.error?.("VERTEX", `Failed to generate JWT token: ${err.message}`);
|
|
throw err;
|
|
}
|
|
}
|
|
if (isClaudeModel(model) && input.body && typeof input.body === "object") {
|
|
let body = input.body as Record<string, unknown>;
|
|
if (!Array.isArray(body.messages)) {
|
|
body = toAnthropicBody(body);
|
|
input.body = body;
|
|
}
|
|
// The rawPredict endpoint requires "anthropic_version" in the body (Vertex's substitute
|
|
// for the "anthropic-version" header used by Anthropic's direct API).
|
|
body.anthropic_version ??= "vertex-2023-10-16";
|
|
// Unlike Anthropic's direct API (which reads the model from the body), Vertex's
|
|
// rawPredict endpoint already encodes project/region/model in the URL and 400s with
|
|
// "model: Extra inputs are not permitted" if the translated request body still carries
|
|
// one (the openai→claude request translator copies the client's model field over).
|
|
delete body.model;
|
|
}
|
|
|
|
const result = await super.execute(input);
|
|
|
|
if (isClaudeModel(model) && stream) {
|
|
const response = result instanceof Response ? result : result?.response;
|
|
if (response?.ok) {
|
|
const contentType = response.headers.get("content-type") || "";
|
|
if (contentType.includes("application/json") && !contentType.includes("text/event-stream")) {
|
|
const jsonText = await response.text();
|
|
let newBody = jsonText;
|
|
let newContentType = contentType;
|
|
try {
|
|
newBody = synthesizeClaudeSse(JSON.parse(jsonText));
|
|
newContentType = "text/event-stream";
|
|
} catch (err) {
|
|
const message = err instanceof Error ? err.message : String(err);
|
|
log?.warn?.("VERTEX", `Failed to synthesize Claude SSE stream: ${message}`);
|
|
}
|
|
const newHeaders = new Headers(response.headers);
|
|
newHeaders.set("content-type", newContentType);
|
|
newHeaders.delete("content-length");
|
|
const newResponse = new Response(newBody, {
|
|
status: response.status,
|
|
statusText: response.statusText,
|
|
headers: newHeaders,
|
|
});
|
|
return result instanceof Response ? newResponse : { ...result, response: newResponse };
|
|
}
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
buildUrl(model: string, stream: boolean, urlIndex = 0, credentials: any = null) {
|
|
// Vertex AI Express mode: project-less v1 publisher endpoint with the API key passed as a
|
|
// ?key= query parameter (verified working contract — same as the CaptionAI GeminiClient). The
|
|
// Express key is NOT accepted as a Bearer/OAuth credential or via x-goog-api-key on this API.
|
|
if (isExpressApiKey(credentials?.apiKey) && !credentials?.accessToken) {
|
|
const expressKey = encodeURIComponent(String(credentials.apiKey).trim());
|
|
if (isPartnerModel(model)) {
|
|
// Partner (Anthropic/etc.) models are not available via Express keys; best-effort.
|
|
return `https://aiplatform.googleapis.com/v1/publishers/openapi/chat/completions?key=${expressKey}`;
|
|
}
|
|
const op = stream ? "streamGenerateContent?alt=sse&" : "generateContent?";
|
|
return `https://aiplatform.googleapis.com/v1/publishers/google/models/${model}:${op}key=${expressKey}`;
|
|
}
|
|
|
|
const region = credentials?.providerSpecificData?.region || "us-central1";
|
|
let project = "unknown-project";
|
|
|
|
if (credentials?.apiKey) {
|
|
try {
|
|
const sa = parseSAFromApiKey(credentials.apiKey);
|
|
if (sa.project_id) project = sa.project_id;
|
|
} catch {
|
|
// Ignored, handled in execute
|
|
}
|
|
}
|
|
|
|
if (isClaudeModel(model)) {
|
|
// streamRawPredict?alt=sse was verified to return a single plain JSON body (not real SSE
|
|
// framing) rather than actual chunked events, which breaks the SSE parser upstream
|
|
// ("stream ended before producing a non-ping SSE event"). rawPredict is confirmed reliable
|
|
// for both streaming and non-streaming requests; always use it here.
|
|
return `https://aiplatform.googleapis.com/v1/projects/${project}/locations/${region}/publishers/anthropic/models/${model}:rawPredict`;
|
|
}
|
|
if (isPartnerModel(model)) {
|
|
return `https://aiplatform.googleapis.com/v1/projects/${project}/locations/global/endpoints/openapi/chat/completions`;
|
|
}
|
|
return `https://aiplatform.googleapis.com/v1/projects/${project}/locations/${region}/publishers/google/models/${model}:${stream ? "streamGenerateContent?alt=sse" : "generateContent"}`;
|
|
}
|
|
|
|
buildHeaders(credentials: any, stream = true) {
|
|
const headers: Record<string, string> = { "Content-Type": "application/json" };
|
|
if (credentials.accessToken) {
|
|
headers["Authorization"] = `Bearer ${credentials.accessToken}`;
|
|
}
|
|
// Express-mode keys are carried in the ?key= query parameter (see buildUrl), not a header.
|
|
if (stream) {
|
|
headers["Accept"] = "text/event-stream";
|
|
}
|
|
return headers;
|
|
}
|
|
}
|