jorel
Version:
The easiest way to use LLMs, including streams, images, documents, tools and various agent scenarios.
286 lines (285 loc) • 12.2 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 || "google-vertex-ai";
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 {
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 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 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, toolCalls),
meta: {
model,
provider,
temperature,
durationMs,
inputTokens,
outputTokens,
},
};
}
async *generateResponseStream(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;
const 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,
});
const durationMs = Date.now() - start;
const _toolCalls = [];
for await (const res of response.stream) {
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) {
yield { type: "chunk", content: textContent };
}
}
}
const r = await response.response;
const rawContent = r.candidates && r.candidates.length > 0 ? r.candidates[0].content : { role: "model", parts: [{ text: "" }] };
const inputTokens = r.usageMetadata?.promptTokenCount;
const outputTokens = r.usageMetadata?.candidatesTokenCount;
const content = rawContent.parts.map((p) => p.text).join("");
const provider = this.name;
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,
toolCalls,
meta,
};
}
else {
yield {
type: "response",
role: "assistant",
content,
meta,
};
}
}
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) {
throw new Error("Embeddings are not yet supported for Vertex AI");
}
}
exports.GoogleVertexAiProvider = GoogleVertexAiProvider;