jorel
Version:
The easiest way to use LLMs, including streams, images, documents, tools and various agent scenarios.
230 lines (229 loc) • 9.7 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleGenerativeAIProvider = void 0;
const generative_ai_1 = require("@google/generative-ai");
const zod_1 = require("zod");
const providers_1 = require("../../providers");
const shared_1 = require("../../shared");
const convert_llm_message_1 = require("./convert-llm-message");
class GoogleGenerativeAIProvider {
constructor({ apiKey, safetySettings, name } = {}) {
this.name = name || "google-generative-ai";
const key = apiKey || process.env.GOOGLE_AI_API_KEY;
if (!key) {
throw new Error("[GoogleGenerativeAIProvider] Missing API key. Either pass it as config.apiKey or set the GOOGLE_AI_API_KEY environment variable");
}
this.client = new generative_ai_1.GoogleGenerativeAI(key);
this.safetySettings = safetySettings;
}
async generateResponse(model, messages, config = {}) {
const start = Date.now();
try {
const generativeModel = this.client.getGenerativeModel({
model,
});
const { contents, systemInstruction } = (0, convert_llm_message_1.convertLlmMessagesToGoogleGenerativeAiMessages)(messages);
const tools = this.prepareTools(config);
const toolConfig = this.prepareToolConfig(config);
const generationConfig = this.prepareGenerationConfig(config);
const safetySettings = this.safetySettings;
const result = await generativeModel.generateContent({
contents,
systemInstruction,
generationConfig,
tools,
toolConfig,
safetySettings,
});
const response = result.response;
const content = response.text();
const functionCalls = response.functionCalls();
const toolCalls = functionCalls?.map((functionCall) => ({
id: (0, shared_1.generateUniqueId)(),
request: {
id: (0, shared_1.generateRandomId)(),
function: {
name: functionCall.name,
arguments: functionCall.args,
},
},
approvalState: config.tools?.getTool(functionCall.name)?.requiresConfirmation
? "requiresApproval"
: "noApprovalRequired",
executionState: "pending",
result: null,
error: null,
}));
const durationMs = Date.now() - start;
return {
...(0, providers_1.generateAssistantMessage)(content, toolCalls),
meta: {
model,
provider: this.name,
temperature: config.temperature ?? undefined,
durationMs,
inputTokens: undefined,
outputTokens: undefined,
},
};
}
catch (error) {
throw new Error(`[GoogleGenerativeAIProvider] Error generating content: ${error}`);
}
}
async *generateResponseStream(model, messages, config = {}) {
const start = Date.now();
try {
const tools = this.prepareTools(config);
const toolConfig = this.prepareToolConfig(config);
const generationConfig = this.prepareGenerationConfig(config);
const safetySettings = this.safetySettings;
const generativeModel = this.client.getGenerativeModel({
model,
generationConfig,
tools,
toolConfig,
safetySettings,
});
const { contents, systemInstruction } = (0, convert_llm_message_1.convertLlmMessagesToGoogleGenerativeAiMessages)(messages);
const result = await generativeModel.generateContentStream({
contents,
generationConfig,
systemInstruction,
tools,
toolConfig,
safetySettings,
});
let fullContent = "";
const toolCalls = [];
for await (const chunk of result.stream) {
const chunkText = chunk.text();
fullContent += chunkText;
// Check for function calls in each chunk
const functionCalls = chunk.functionCalls?.();
if (functionCalls && functionCalls.length > 0) {
// Process new function calls that haven't been seen before
for (const functionCall of functionCalls) {
// Check if this function call is already in our toolCalls array
const existingToolCall = toolCalls.find((tc) => tc.request.function.name === functionCall.name &&
JSON.stringify(tc.request.function.arguments) === JSON.stringify(functionCall.args));
if (!existingToolCall) {
const newToolCall = {
id: (0, shared_1.generateUniqueId)(),
request: {
id: (0, shared_1.generateRandomId)(),
function: {
name: functionCall.name,
arguments: functionCall.args,
},
},
approvalState: config.tools?.getTool(functionCall.name)?.requiresConfirmation
? "requiresApproval"
: "noApprovalRequired",
executionState: "pending",
result: null,
error: null,
};
toolCalls.push(newToolCall);
}
}
}
yield { type: "chunk", content: chunkText };
}
const durationMs = Date.now() - start;
const meta = {
model,
provider: this.name,
temperature: config.temperature ?? undefined,
durationMs,
inputTokens: undefined,
outputTokens: undefined,
};
// If we have tool calls, yield a response with tools
if (toolCalls.length > 0) {
yield {
type: "response",
role: "assistant_with_tools",
content: fullContent,
toolCalls,
meta,
};
}
else {
yield {
type: "response",
role: "assistant",
content: fullContent,
meta,
};
}
}
catch (error) {
throw new Error(`[GoogleGenerativeAIProvider] Error generating content stream: ${error}`);
}
}
async getAvailableModels() {
return providers_1.defaultGoogleGenAiModels;
}
async createEmbedding(model, text) {
const generativeModel = this.client.getGenerativeModel({ model });
const result = await generativeModel.embedContent(text);
if (!result.embedding) {
throw new Error("No embedding returned");
}
return result.embedding.values;
}
// Helper methods for preparing request parameters
prepareTools(config) {
if (!config.tools?.asLlmFunctions?.length) {
return undefined;
}
return config.tools.asLlmFunctions.map((f) => {
const functionDeclaration = {
name: f.function.name,
description: f.function.description,
parameters: f.function.parameters,
};
return { functionDeclarations: [functionDeclaration] };
});
}
prepareToolConfig(config) {
if (!config.tools?.hasTools) {
return undefined;
}
function toolChoiceToFunctionCallingMode(toolChoice) {
if (!toolChoice || toolChoice === "auto") {
return generative_ai_1.FunctionCallingMode.AUTO;
}
if (toolChoice === "none") {
return generative_ai_1.FunctionCallingMode.NONE;
}
if (toolChoice === "required") {
return generative_ai_1.FunctionCallingMode.ANY;
}
return generative_ai_1.FunctionCallingMode.ANY;
}
return {
functionCallingConfig: {
mode: toolChoiceToFunctionCallingMode(config.toolChoice),
allowedFunctionNames: undefined,
},
};
}
prepareGenerationConfig(config) {
const generationConfig = {
temperature: config.temperature ?? undefined,
maxOutputTokens: config.maxTokens ?? undefined,
responseSchema: undefined,
};
if (config.json) {
generationConfig.responseMimeType = "application/json";
if (typeof config.json !== "boolean") {
generationConfig.responseSchema =
config.json instanceof zod_1.ZodObject ? (0, shared_1.zodSchemaToJsonSchema)(config.json) : config.json;
}
}
return generationConfig;
}
}
exports.GoogleGenerativeAIProvider = GoogleGenerativeAIProvider;