jorel
Version:
A unified wrapper for working with LLMs from multiple providers, including streams, images, documents & automatic tool use.
390 lines (389 loc) • 16.5 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleVertexAiProvider = exports.VertexAiHarmCategory = exports.VertexAiHarmBlockThreshold = void 0;
const vertexai_1 = require("@google-cloud/vertexai");
Object.defineProperty(exports, "VertexAiHarmBlockThreshold", { enumerable: true, get: function () { return vertexai_1.HarmBlockThreshold; } });
Object.defineProperty(exports, "VertexAiHarmCategory", { enumerable: true, get: function () { return vertexai_1.HarmCategory; } });
const zod_1 = require("zod");
const providers_1 = require("../../providers");
const shared_1 = require("../../shared");
const convert_llm_message_1 = require("./convert-llm-message");
const defaultSafetySettings = [
{
category: vertexai_1.HarmCategory.HARM_CATEGORY_UNSPECIFIED,
threshold: vertexai_1.HarmBlockThreshold.BLOCK_ONLY_HIGH,
},
{
category: vertexai_1.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold: vertexai_1.HarmBlockThreshold.BLOCK_ONLY_HIGH,
},
{
category: vertexai_1.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold: vertexai_1.HarmBlockThreshold.BLOCK_ONLY_HIGH,
},
{
category: vertexai_1.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
threshold: vertexai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
{
category: vertexai_1.HarmCategory.HARM_CATEGORY_HARASSMENT,
threshold: vertexai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
];
/** Provides access to GoogleVertexAi and other compatible services */
class GoogleVertexAiProvider {
constructor({ project, location, keyFilename, safetySettings, name } = {}) {
/** @internal */
this.safetySettings = defaultSafetySettings;
this.name = name || GoogleVertexAiProvider.defaultName;
const config = {
project: project || process.env.GCP_PROJECT,
location: location || process.env.GCP_LOCATION,
keyFilename: keyFilename || process.env.GOOGLE_APPLICATION_CREDENTIALS,
};
if (!config.project)
throw new Error("[GoogleVertexAiProvider] Missing GCP project. Either pass it as config.project or set the GCP_PROJECT environment variable");
if (!config.location)
throw new Error("[GoogleVertexAiProvider] Missing GCP location. Either pass it as config.location or set the GCP_LOCATION environment variable");
this.client = new vertexai_1.VertexAI({
googleAuthOptions: {
projectId: config.project,
keyFilename: config.keyFilename,
},
location: config.location,
project: config.project,
});
if (safetySettings) {
this.safetySettings = safetySettings;
}
}
async generateResponse(model, messages, config = {}) {
const start = Date.now();
const { chatMessages, systemMessage } = await (0, convert_llm_message_1.convertLlmMessagesToVertexAiMessages)(messages);
const generativeModel = this.client.getGenerativeModel({
model,
});
const temperature = config.temperature ?? undefined;
const maxTokens = config.maxTokens ?? undefined;
let response;
try {
// Note: Google Vertex AI SDK doesn't support AbortSignal directly
// Check for cancellation before making the request
if (config.abortSignal?.aborted) {
throw new shared_1.JorElAbortError("Request was aborted");
}
response = (await generativeModel.generateContent({
contents: chatMessages,
systemInstruction: systemMessage,
tools: config.tools?.asLlmFunctions?.map((f) => {
const functionDeclarations = [
{
name: f.function.name,
description: f.function.description,
parameters: f.function.parameters,
},
];
return { functionDeclarations };
}),
generationConfig: {
temperature,
maxOutputTokens: maxTokens,
responseMimeType: config.json ? "application/json" : "text/plain",
responseSchema: config.json && typeof config.json !== "boolean"
? config.json instanceof zod_1.ZodObject
? (0, shared_1.zodSchemaToJsonSchema)(config.json)
: config.json
: undefined,
},
toolConfig: (0, providers_1.toolChoiceToVertexAi)(config.tools?.hasTools ?? false, config.toolChoice),
safetySettings: this.safetySettings,
})).response;
}
catch (error) {
if (error instanceof shared_1.JorElAbortError) {
throw error;
}
if (error instanceof vertexai_1.ClientError) {
throw new Error(`[GoogleVertexAiProvider] Error generating content: ${error.message}`);
}
if (error instanceof vertexai_1.GoogleApiError) {
throw new Error(`[GoogleVertexAiProvider] Error generating content: ${error.message}, code: ${error.code}, status: ${error.status}, details: ${error.errorDetails}`);
}
throw error;
}
const inputTokens = response.usageMetadata?.promptTokenCount;
const outputTokens = response.usageMetadata?.candidatesTokenCount;
const responseContent = response.candidates && response.candidates.length > 0
? response.candidates[0].content
: { role: "model", parts: [{ text: "" }] };
const reasoningContent = null;
const content = responseContent.parts
.filter((p) => !!p.text)
.map((p) => p.text)
.join("")
.trim();
const toolCalls = responseContent.parts
.filter((p) => p.functionCall)
.map((p) => {
const functionCall = p.functionCall;
return {
id: (0, shared_1.generateUniqueId)(),
request: {
id: (0, shared_1.generateRandomId)(),
function: {
name: functionCall.name,
arguments: functionCall.args,
},
},
approvalState: "noApprovalRequired",
executionState: "pending",
result: null,
error: null,
};
});
const durationMs = Date.now() - start;
const provider = this.name;
return {
...(0, providers_1.generateAssistantMessage)(content, reasoningContent, toolCalls),
meta: {
model,
provider,
temperature,
durationMs,
inputTokens,
outputTokens,
},
};
}
async *generateResponseStream(model, messages, config = {}) {
const start = Date.now();
const provider = this.name;
const { chatMessages, systemMessage } = await (0, convert_llm_message_1.convertLlmMessagesToVertexAiMessages)(messages);
const generativeModel = this.client.getGenerativeModel({
model,
});
const temperature = config.temperature ?? undefined;
const maxTokens = config.maxTokens ?? undefined;
// Note: Google Vertex AI SDK doesn't support AbortSignal directly
// Check for cancellation before making the request
if (config.abortSignal?.aborted) {
yield {
type: "response",
role: "assistant",
content: "",
reasoningContent: null,
meta: {
model,
provider,
temperature,
durationMs: 0,
inputTokens: undefined,
outputTokens: undefined,
},
stopReason: "userCancelled",
};
return;
}
let response;
try {
response = await generativeModel.generateContentStream({
contents: chatMessages,
systemInstruction: systemMessage,
tools: config.tools?.asLlmFunctions?.map((f) => {
const functionDeclarations = [
{
name: f.function.name,
description: f.function.description,
parameters: f.function.parameters,
},
];
return { functionDeclarations };
}),
generationConfig: {
temperature,
maxOutputTokens: maxTokens,
responseMimeType: config.json ? "application/json" : "text/plain",
responseSchema: config.json && typeof config.json !== "boolean"
? config.json instanceof zod_1.ZodObject
? (0, shared_1.zodSchemaToJsonSchema)(config.json)
: config.json
: undefined,
},
toolConfig: (0, providers_1.toolChoiceToVertexAi)(config.tools?.hasTools ?? false, config.toolChoice),
safetySettings: this.safetySettings,
});
}
catch (error) {
const isAbort = error instanceof Error &&
(error.message.toLowerCase().includes("aborted") ||
error.name === "AbortError" ||
error instanceof shared_1.JorElAbortError);
const stopReason = isAbort ? "userCancelled" : "generationError";
yield {
type: "response",
role: "assistant",
content: "",
reasoningContent: null,
meta: {
model,
provider,
temperature,
durationMs: 0,
inputTokens: undefined,
outputTokens: undefined,
},
stopReason,
error: stopReason === "generationError"
? {
message: error instanceof Error ? error.message : String(error),
type: "unknown",
}
: undefined,
};
return;
}
const _toolCalls = [];
let streamedContent = "";
let error;
let aborted = false;
try {
for await (const res of response.stream) {
// Check for cancellation during streaming
if (config.abortSignal?.aborted) {
aborted = true;
break;
}
const content = res.candidates && res.candidates.length > 0
? res.candidates[0].content
: { role: "model", parts: [{ text: "" }] };
if (content && content.parts && content.parts.length > 0) {
// Handle function calls in the stream
const functionCalls = content.parts.filter((p) => p.functionCall);
for (const part of functionCalls) {
if (part.functionCall) {
_toolCalls.push({
function: {
name: part.functionCall.name,
arguments: part.functionCall.args,
},
});
}
}
// Handle text content
const textContent = content.parts.map((part) => ("text" in part ? part.text : "")).join("");
if (textContent.length > 0) {
streamedContent += textContent;
const chunkId = (0, shared_1.generateUniqueId)();
yield { type: "chunk", content: textContent, chunkId };
}
}
}
}
catch (e) {
error = {
message: e instanceof Error ? e.message : String(e),
type: "unknown",
};
}
const durationMs = Date.now() - start;
// Determine stop reason and error message
const stopReason = config.abortSignal?.aborted ? "userCancelled" : error ? "generationError" : "completed";
// Log non-abort errors
if (error && stopReason === "generationError") {
config.logger?.error("GoogleVertexAiProvider", `Stream error: ${error.message}`);
}
// Try to get final response for token counts, but use streamed content as fallback
let finalContent = streamedContent;
let inputTokens;
let outputTokens;
if (!error && !aborted && !config.abortSignal?.aborted) {
try {
const r = await response.response;
const rawContent = r.candidates && r.candidates.length > 0 ? r.candidates[0].content : { role: "model", parts: [{ text: "" }] };
finalContent = rawContent.parts.map((p) => p.text).join("");
inputTokens = r.usageMetadata?.promptTokenCount;
outputTokens = r.usageMetadata?.candidatesTokenCount;
}
catch {
// Use streamed content if final response fails
}
}
const reasoningContent = null;
const toolCalls = _toolCalls.map((call) => {
return {
id: (0, shared_1.generateUniqueId)(),
request: {
id: (0, shared_1.generateRandomId)(),
function: {
name: call.function.name,
arguments: call.function.arguments,
},
},
approvalState: config.tools?.getTool(call.function.name)?.requiresConfirmation
? "requiresApproval"
: "noApprovalRequired",
executionState: "pending",
result: null,
error: null,
};
});
const meta = {
model,
provider,
temperature,
durationMs,
inputTokens,
outputTokens,
};
if (toolCalls.length > 0) {
yield {
type: "response",
role: "assistant_with_tools",
content: finalContent,
reasoningContent,
toolCalls,
meta,
stopReason,
error: stopReason === "generationError" ? error : undefined,
};
}
else {
yield {
type: "response",
role: "assistant",
content: finalContent,
reasoningContent,
meta,
stopReason,
error: stopReason === "generationError" ? error : undefined,
};
}
}
async getAvailableModels() {
return [];
}
async countTokens(model, contents) {
const generativeModel = this.client.getGenerativeModel({
model: model,
safetySettings: this.safetySettings,
});
const response = await generativeModel.countTokens({
contents,
});
const inputTokens = response.totalTokens;
const characterCount = contents.reduce((acc, content) => {
return acc + content.parts.reduce((acc, part) => acc + ("text" in part ? part?.text?.length || 0 : 0), 0);
}, 0);
return {
model,
inputTokens,
characterCount,
};
}
// eslint-disable-next-line @typescript-eslint/no-unused-vars
async createEmbedding(model, text, abortSignal) {
throw new Error("Embeddings are not yet supported for Vertex AI");
}
}
exports.GoogleVertexAiProvider = GoogleVertexAiProvider;
GoogleVertexAiProvider.defaultName = "google-vertex-ai";