jorel
Version:
The easiest way to use LLMs, including streams, images, documents, tools and various agent scenarios.
216 lines (215 loc) • 9.42 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleGenerativeAIProvider = void 0;
const genai_1 = require("@google/genai");
const zod_1 = require("zod");
const __1 = require("..");
const shared_1 = require("../../shared");
const convert_llm_message_1 = require("./convert-llm-message");
class GoogleGenerativeAIProvider {
constructor(options = {}) {
this.name = options.name || GoogleGenerativeAIProvider.defaultName;
const apiKey = options.apiKey || process.env.GOOGLE_AI_API_KEY;
if (!apiKey) {
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 genai_1.GoogleGenAI({ apiKey });
this.safetySettings = options.safetySettings;
}
async generateResponse(model, messages, config = {}) {
const start = Date.now();
try {
const { contents, systemInstruction } = (0, convert_llm_message_1.convertLlmMessagesToGoogleGenerativeAiMessages)(messages);
const requestConfig = this.prepareGenerationConfig(config);
// Add system instruction to config if present
if (systemInstruction) {
requestConfig.systemInstruction = systemInstruction;
}
const result = await this.client.models.generateContent({
model,
contents,
config: requestConfig,
});
const content = result.text ?? "";
const functionCalls = result.functionCalls ?? [];
const toolCalls = functionCalls.length > 0
? 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,
}))
: undefined;
const durationMs = Date.now() - start;
return {
...(0, __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 { contents, systemInstruction } = (0, convert_llm_message_1.convertLlmMessagesToGoogleGenerativeAiMessages)(messages);
const requestConfig = this.prepareGenerationConfig(config);
// Add system instruction to config if present
if (systemInstruction) {
requestConfig.systemInstruction = systemInstruction;
}
const streamResult = await this.client.models.generateContentStream({
model,
contents,
config: requestConfig,
});
let fullContent = "";
const toolCalls = [];
for await (const chunk of streamResult) {
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);
}
}
}
if (chunkText) {
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 __1.initialGoogleGenAiModels;
}
async createEmbedding(model, text) {
const result = await this.client.models.embedContent({
model,
contents: [{ role: "user", parts: [{ text }] }],
});
if (!result.embeddings || result.embeddings.length === 0) {
throw new Error("No embedding returned");
}
return result.embeddings[0].values ?? [];
}
// Helper method for preparing request configuration
prepareGenerationConfig(config) {
const requestConfig = {
safetySettings: this.safetySettings,
};
// Add generation config
if (config.temperature !== undefined || config.maxTokens !== undefined || config.json) {
requestConfig.temperature = config.temperature ?? undefined;
requestConfig.maxOutputTokens = config.maxTokens ?? undefined;
if (config.json) {
requestConfig.responseMimeType = "application/json";
if (typeof config.json !== "boolean") {
requestConfig.responseSchema =
config.json instanceof zod_1.ZodObject ? (0, shared_1.zodSchemaToJsonSchema)(config.json) : config.json;
}
}
}
// Add tools
if (config.tools?.asLlmFunctions?.length) {
requestConfig.tools = [
{
functionDeclarations: config.tools.asLlmFunctions.map((f) => ({
name: f.function.name,
description: f.function.description,
parameters: f.function.parameters, // TODO: Improve types
})),
},
];
}
// Add tool config
if (config.tools?.hasTools && config.toolChoice) {
let mode = genai_1.FunctionCallingConfigMode.AUTO;
if (config.toolChoice === "none") {
mode = genai_1.FunctionCallingConfigMode.NONE;
}
else if (config.toolChoice === "required") {
mode = genai_1.FunctionCallingConfigMode.ANY;
}
requestConfig.toolConfig = {
functionCallingConfig: {
mode,
},
};
}
return requestConfig;
}
}
exports.GoogleGenerativeAIProvider = GoogleGenerativeAIProvider;
GoogleGenerativeAIProvider.defaultName = "google-generative-ai";